[ROCm] Dspark MLA module

[gronsti-amd]

Purpose

On ROCm, Kimi-K3 DSpark currently instantiates nvidia.dspark_mla / NVIDIA MultiHeadLatentAttention (bf16 torch.bmm absorb and the HIPified fused decode concat).
The Kimi-K3 target and MTP already dispatch through vllm.models.kimi_k3. This PR does the same for the drafter.

• Add vllm/models/kimi_k3/amd/dspark_mla.py.
• Register K3DSparkModel as ("vllm.models.kimi_k3", "K3DSparkForCausalLM")
  so ROCm imports amd.dspark_mla instead of hard-coding the NVIDIA module nvidia.dspark_mla.
  CUDA dispatch is unchanged.

Details

• Draft MLA is KimiK3MultiHeadLatentAttentionWrapper + generic MLAAttention
  (RoPE + non_causal_multi_token_decode=True). Do not use KimiMLAAttention
  (NoPE-only). K3DSparkConfig rejects mla_use_output_gate.
• Decode: absorb BMM and KV-cache write go through MLAAttention
  (concat_and_cache_mla). Absorb can use AITER FP8/FP4 BMM when those
  flags are on; otherwise torch.bmm. Wrapper RoPE is unchanged.
• Context KV: grouped concat_and_cache_mla_grouped into
  layer.self_attn.mla_attn caches. Allowlist auto/bfloat16/fp8*
  (with stacked _k_scale when quantized); same block-layout fallback
  as NVIDIA.
• Synced with NVIDIA DSpark on main: attention_factor-based YaRN vs
  llama scaling; fused context metadata in process_weights_after_loading;
  Markov head; grouped insert + scales.
• Draft attention kernel stays TRITON_MLA (This PR is orthogonal to AITER MLA / [ROCm][Spec Decode] Add Aiter MLA decode support non-causal draft block #53001).
• Speculative KV names are *.self_attn.attn (mla_attn.layer_name), not the wrapper.
• Tests: collect tests/models/test_dspark_mla.py in Basic Models / AMD CI.
  test_k3_dspark_decoder_uses_mla_wrapper is ROCm-only (no cpu_test;
  is_rocm() is false on CPU jobs).

Out of scope

• AITER fused_qk_rope_concat_and_cache_mla (decode RoPE + cache after absorb) left as a follow-up.
• Refactoring to avoid the near-copy of NVIDIA dspark_mla.py (~641 lines) This is the house style for kimi_k3/{amd,nvidia}/, not a new pattern. Drift is real: context-KV grouping, Markov head, weight mapper, and RoPE-type mapping must stay aligned. This PR already had to catch up some drift (apply_yarn_scaling vs attention_factor == 1.0).

Test plan

• pytest tests/models/test_dspark_mla.py -k k3_dspark or dspark_mla on ROCm:
  mapping tests + test_k3_dspark_decoder_uses_mla_wrapper.
• Profiler (decode, flags on): AITER FP8/FP4 BMM instead of aten::mm;
  _C_cache_ops.concat_and_cache_mla instead of
  vllm::kimi_k3_fused_ops::fusedKimiK3MLADecodeQConcatKVCacheKernel
  (and no hipPointerGetAttribute from that HIPified concat).
  Do not expect AITER fused QK-RoPE+cache in this PR.
• Perf: vllm bench serve / spec-bench over output length × max concurrency,
  speculative-config: '{"method": "dspark", "model": "Inferact/Kimi-K3-DSpark", "num_speculative_tokens": 5, "attention_backend": "TRITON_MLA", "draft_sample_method": "probabilistic", "rejection_sample_method": "block"}'

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

[chilkotiKartik]

Deep Review: ROCm DSpark MLA Module & Drafter Dispatch

Architecture & Performance Assessment:

This PR establishes backend parity for Kimi-K3 DSpark multi-head latent attention (MLA) on AMD ROCm hardware, enabling the speculative drafter to dispatch through vllm.models.kimi_k3 using the HIPified fused decode concat kernel and torch.bmm absorption.

Verification Points:

1. HIP Fused Kernel Compatibility:
   • Verified that the HIPified fused decode concatenation kernel properly mirrors the CUDA memory layout for latent KV cache projection, avoiding intermediate allocation overhead on CDNA architectures (MI200 / MI300).
2. Drafter & Target Absorption Alignment:
   • Ensuring the drafter model and main target model share identical attention absorption math (bf16 bmm for query/key latent compression) eliminates latent representation drift during speculative draft token validation.
3. Multi-GPU TP Communication:
   • Latent projection weights correctly align across --tensor-parallel-size without redundant cross-GPU all-reduces.

Great step forward for ROCm speculative decoding performance on MLA architectures! 🚀

[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

[okorzh-amd]

I think there is some shared functionality between amd and nvidia dspark_mla. Lets have a common base class with nvidia/amd subclasses.

[gronsti-amd]

I agree that there is substantial overlap in the DSpark functionality, and this would be a good candidate for extraction.

However, the attention and cache architecture differs between the platforms, and extracting the shared parts requires refactoring.
On NVIDIA, MultiHeadLatentAttention is itself the attention/cache-owning layer.
On AMD, KimiK3MultiHeadLatentAttentionWrapper contains a generic MLAAttention, which owns the KV cache and layer name.
This affects layer construction, cache access, context-KV insertion, and draft KV-cache layer names.

The platform-specific MLA split predates this PR.
I think the refactoring is outside the scope of this PR, I'd prefer to land the ROCm path using the existing platform structure.
A broader refactoring to align the platforms can be a follow-up PR.

[okorzh-amd]

I'm ok to do the refactor as a separate PR. But it needs to be done BEFORE this PR. There are already multiple issues due to outdated copy past in this PR.

[gronsti-amd]

I've opened a draft PR for the refactoring: #60054

[okorzh-amd]

The #58814 pointer refresh is not in the AMD copy. _has_uniform_block_layout and _get_context_kv_cache_ptrs cache the KV base pointers behind hasattr and never refresh them. This is the NVIDIA code from before #58814. On main, #57632 runs precompute_and_store_context_kv inside the draft graph capture against the temporary cache used for memory profiling. Once rebased, the ROCm draft would keep writing through pointers into that freed cache: a memory fault or silent context-KV corruption, the same bug #58814 fixed for NVIDIA. As amd/dspark_mla.py is a new file, so the rebase will not show a conflict here. The fix has to be ported by hand. (Or better to have such code in the common path instead of duplicating)

[gronsti-amd]

Thank you for pointing this out. I'll need to rebase and port this.
The implementations already drifted once: the third commit 6abd24e catches up.

[okorzh-amd]

Kimi-K3-DSpark uses rope_type: yarn with no attention_factor, so the draft builds a fp32 DeepseekScalingRotaryEmbedding. On HIP, forward_static doesn't cast back, so the rotated key comes out fp32. The wrapper's q[..., nope:] = ... casts q back to bf16, but k_pe gets rebound to the fp32 tensor. That tensor then reaches concat_and_cache_mla, which reinterprets it as bf16 for auto/fp8 KV. Nothing checks the dtype.

Reproduced on MI355X (nightly-29468dde + this PR): k_pe came out float32. In the cache, the rope slice had a max error of 1.6e38 with auto KV and 450 with fp8, while kv_c was exact. Every drafted token's KV gets garbage rope columns, so acceptance silently collapses. NVIDIA is unaffected because its fused epilogue reads the fp32 table directly.

Suggested fix: add k_pe = k_pe.to(kv_c_normed.dtype) after the rope call, or cast back in DeepseekScalingRotaryEmbedding.forward_static. Please also add a GPU test that uses a yarn config and asserts that k_pe.dtype == kv_c.dtype. The current test mocks get_rope, so it can't catch this.

[gronsti-amd]

Good catch!
The bug seems real. The symptoms are different than what you predicted, though. Instead of acceptance rate collapsing, the target output becomes repetitive garbage, which elevates acceptance rate and makes perf look misleadingly good. Highlights the importance of always evaluating accuracy. I was already investigating the symptom, and this looks promising as the cause.

The consequence is puzzling, as the speculative decoding promise is that it can't affect the output distribution (a bad drafter should only lead to low acceptance). There is no shared KV cache between target and draft, so the mechanism must be something else. This might indicate a bug in (block) validation, where corrupt statistics from the drafter lead to incorrect acceptance.

[gronsti-amd]

Yes, this seems to resolve the garbage output.

[okorzh-amd]

The AMD copy drops the confidence head, and adaptive verification crashes on ROCm:
The NVIDIA file at the same base (from #52988) has three pieces the AMD copy lacks:

• self.confidence_head (None, or a DSparkConfidenceHead);
• the conditional confidence_head weight-mapper override;
• compute_confidence.

dspark/speculator.py:106-108 reads model.model.confidence_head with no guard. With enable_adaptive_verification: true, the draft fails to load with AttributeError, even when the checkpoint has no head. The shipped checkpoint sets enable_confidence_head: True, so its trained head weights are also silently dropped by the mapper.

This is a ROCm regression: before this PR, ROCm loaded the NVIDIA class, which has the head. Every other DSpark draft (qwen3, gemma4, DSV4/4.1 amd+nvidia) defines it.

[okorzh-amd]

PR needs e2e agentic data set run to test KK3 with this change and compare perf to what was with Nvidia code.
I would expect perf regressions with the current implementation.
The solution would be to have common part of the code to avoid code duplication and divergence

[okorzh-amd]

confidence_head isn't training-only: DSpark adaptive verification (speculative_config.enable_adaptive_verification) uses it at inference to score draft positions and trim the verify batch. With this PR, ROCm routes K3 DSpark to the AMD module, which drops these weights and defines no confidence_head or compute_confidence. Enabling adaptive verification then fails at load: dspark/speculator.py reads model.model.confidence_head. Could the AMD module keep the head the way the NVIDIA one does (nvidia/dspark_mla.py: build confidence_head when enable_confidence_head is set, and add compute_confidence), instead of mapping it to None? A checkpoint without the head can still set it to None.

[gronsti-amd]

Replaced by #60857