Repository navigation
[ROCm] ROCm binding on common Kimi-K3 Dspark MLA module - #60857
gronsti-amd wants to merge 11 commits into
Conversation
Keep the NVIDIA module as a re-export so the registry path is unchanged. Signed-off-by: Stig-Arne Grönroos <stig-arne.gronroos@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>
NVIDIA keeps MultiHeadLatentAttention as the cache owner. A platform override can point context-KV insert and draft layer names at another module. Signed-off-by: Stig-Arne Grönroos <stig-arne.gronroos@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>
The defaults keep today's concat and cache write. A subclass can replace the decode query and the KV rows it writes without a new custom-op argument. Signed-off-by: Stig-Arne Grönroos <stig-arne.gronroos@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>
The common module is a new file, so pre-commit lints it in full. The (list, tuple) form was already on main and fails UP038 once the file moves. Signed-off-by: Stig-Arne Grönroos <stig-arne.gronroos@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>
The file was not listed in Buildkite, so the refactor would not be collected on either platform. Signed-off-by: Stig-Arne Grönroos <stig-arne.gronroos@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>
DeepseekV32Attention already assigns a concrete RotaryEmbedding to that name. Typing the base attribute as optional made mypy reject its cache reads. Signed-off-by: Stig-Arne Grönroos <stig-arne.gronroos@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>
Replay reads the captured positions and slot-mapping allocations, so a replacement tensor freezes the pointer. Lock the NVIDIA binding class and the default hook returns. Signed-off-by: Stig-Arne Grönroos <stig-arne.gronroos@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>
ROCm keeps the wrapper factory and points the KV cache at mla_attn. The copied decoder, context-KV insert, and weight mapper stay in the common module. Signed-off-by: Stig-Arne Grönroos <stig-arne.gronroos@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>
HIP YaRN RoPE returns fp32. The query slice assignment casts q back to bf16, but k_pe stays fp32 and concat_and_cache_mla stores those bits as the KV dtype. Signed-off-by: Stig-Arne Grönroos <stig-arne.gronroos@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>
The fused decode launch applies a one-row identity table, so it is only valid when the wrapper has not already rotated. Document the duplicated YaRN setup and bind the construction test to the wrapper signature. Signed-off-by: Stig-Arne Grönroos <stig-arne.gronroos@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>
KimiLinearConfig types the latent sizes as optional. Assert them the same way the NVIDIA K3 MLA constructor does, so the factory arithmetic typechecks. Signed-off-by: Stig-Arne Grönroos <stig-arne.gronroos@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>
gronsti-amd
changed the title
|
This pull request has merge conflicts that must be resolved before it can be |
| @@ -0,0 +1,225 @@ | |||
| # SPDX-License-Identifier: Apache-2.0 | |||
| # SPDX-FileCopyrightText: Copyright contributors to the vLLM project | |||
| """ROCm binding for the Kimi-K3 DSpark MLA draft. | |||
There was a problem hiding this comment.
Overall provided solution is slower than the code it replaces.
The PR reports +0.44 ms per step in eager mode. Two things in this PR's scope cause most of it:
- The new self.rotary_emb is None fusion check (amd/mla.py:282-284) isn't needed for correctness.
- The fused decode applies a one-row cos=1, sin=0 table. On already-rotated q_pe/k_pe that does nothing.
- On MI355X I called fused_qk_rope_concat_and_cache_mla exactly as _fused_decode does. Run on pre-rotated inputs, the KV rows and the fp8 q_out are byte-identical to the unfused path.
- With VLLM_ROCM_USE_AITER=1 and fp8 KV, this check is the only thing turning off the existing fused decode for the draft. In each of 5 layers, one fused launch becomes a separate cache write plus a q concat and quant.
- Please drop the check (or move it to [ROCm] Fuse Kimi-K3 DSpark QK RoPE with the MLA cache write #60864, where the wrapper skips its own RoPE). Also drop the "double RoPE" wording at :106-108 and :279-281.
- Native YaRN costs 15 kernel launches per layer (75 per step): the forward_native elementwise chain, the slice copy-back, and the cast.
- An in-place ops.rotary_embedding(positions, q[..., nope:], k_pe, ...) with the fp32 table does it in 1. The common module already uses this op for context-KV RoPE.
- It also makes the k_pe cast unnecessary. k_pe comes out byte-identical, and q differs by at most 1 bf16 ulp on a tiny fraction of elements.
- Please re-measure in the mode production uses. The draft replays inside full CUDA graphs there, so the eager 3.41 → 3.85 ms overstates the per-step cost.
| # rotated q and k above, and this launch would apply cos=1, sin=0 on | ||
| # top of that. | ||
| fuse = ( | ||
| self._fused_qk_prep and self.rotary_emb is None and q_dcp_replicated is None |
There was a problem hiding this comment.
This check isn't needed for correctness. _fused_decode passes a one-row cos=1, sin=0 table, so on q_pe/k_pe that are already rotated it does nothing. On MI355X I called fused_qk_rope_concat_and_cache_mla exactly as _fused_decode does, on pre-rotated inputs. The KV rows and the fp8 q_out came out byte-identical to the unfused concat_and_cache_mla plus static quant. With VLLM_ROCM_USE_AITER=1 and fp8 KV, this check is the only thing turning off the existing fused decode for the draft. In each of the 5 draft layers, one launch becomes a cache write plus a q concat and quant, which is part of the reported slowdown. Could you drop it here?
| # decode-only batches; anything else falls through to MLAAttention. | ||
| attn_metadata = layer = kv_cache = slot_mapping = None | ||
| fuse = self._fused_qk_prep and q_dcp_replicated is None | ||
| # Identity-RoPE fusion is the NoPE path. A real embedding has already |
There was a problem hiding this comment.
This comment and the docstring at :106-108 describe a double RoPE that doesn't happen: cos=1, sin=0 on rotated inputs is an exact identity. The docstring also mentions a "DSpark epilogue" that only exists in #60864.
| is_sparse=False, | ||
| topk_indices_buffer=None, | ||
| ) | ||
| return KimiK3MultiHeadLatentAttentionWrapper( |
There was a problem hiding this comment.
Moving the draft onto generic MLAAttention changes the precision of its absorb matmuls under AITER:
- With VLLM_ROCM_USE_AITER=1 on gfx950, MLAAttention (mla_attention.py:717-724, :1336-1349) quantizes the bf16 draft kv_b_proj's W_UK/W_UV to MXFP4.
- Decode then runs batched_gemm_a16wfp4(prequant=True), which quantizes the activations too.
- On gfx942, or with FP4BMM off, the FP8 path runs instead.
- The NVIDIA module keeps a bf16 torch.bmm.
There was a problem hiding this comment.
Separately, this wrapper's init now runs _fused_qk_prep_supported() for the draft. That function reads attn.W_UK_T (amd/mla.py:54), which only exists after process_weights_after_loading. So construction raises AttributeError with fp8 KV + AITER MLA + FP8BMM=0 + FP4BMM=0, which is the natural workaround for the precision change. The code is older (#57640), but the draft now reaches it. getattr(attn, "W_UK_T", None) or deciding after weights load would fix it. As written, the bf16 W_UK_T branch of _fused_decode can never run.
| num_local_heads = num_heads // tp_size | ||
| scaling = qk_head_dim**-0.5 | ||
|
|
||
| rope_parameters = dict(config.rope_parameters) |
There was a problem hiding this comment.
Lines 81-103 copy the NVIDIA rope-type remap, get_rope(dtype=fp32) and the mscale^2 factor on the softmax scale, and only the docstring keeps them in sync. Deleting scaling *= mscale * mscale (:103) still passes all 36 tests. Could this be a shared helper in common/ (using yarn_get_mscale) that both bindings call? Please also add a CPU test that builds both bindings from the real YaRN rope_parameters and compares the scale, rope class, is_neox_style and cache dtype. They all match today.
| quant_config=quant_config, | ||
| prefix=f"{prefix}.q_b_proj", | ||
| ) | ||
| else: |
There was a problem hiding this comment.
This branch can't be reached: K3DSparkConfig rejects q_lora_rank=None, and the common context-KV metadata asserts it is set. Replace it with an assert?
|
|
||
| def _make_dspark_mla_attention( | ||
| *, | ||
| config: KimiLinearConfig, |
There was a problem hiding this comment.
nit: the draft config is K3DSparkConfig, not KimiLinearConfig. Same at :195, and the import at :35.
This branch has not been deployed
2 participants
Purpose
On ROCm, Kimi-K3 DSpark currently instantiates the NVIDIA draft module. This PR registers
K3DSparkModelas("vllm.models.kimi_k3", "K3DSparkForCausalLM")and, on ROCm, importsamd.dspark_mla. CUDA still importsnvidia.dspark_mla.Builds on top of #60054
Replaces #57942, which was the duplicate-and-modify approach (not touching NVIDIA).
What changes
amd/dspark_mla.pyis a binding over the common classes: a factory,build_self_attn,build_mlp, andkv_cache_layerreturningattn.mla_attn.Draft attention is
KimiK3MultiHeadLatentAttentionWrapperaround genericMLAAttention, with RoPE andnon_causal_multi_token_decode=True.KimiMLAAttentionis not used: it is NoPE-only, and would add a second wrapper.Context-KV insert, the weight mapper, the Markov head, and the confidence head stay in the common module, so the #58814 refresh and adaptive verification come along with the dispatch.
The factory duplicates the NVIDIA projection and YaRN setup.
The docstring records the settings that have to stay aligned.
HIP YaRN returns fp32. The query slice assignment casts Q back to bf16,
and rebinding
k_peused to leave it fp32.concat_and_cache_mlathen stored those bits as the KV dtype.After RoPE, the wrapper casts
k_peback withk_pe.to(kv_c_normed.dtype).On MI355 that cast removed the repetitive garbage in the target output.
Decode-only FP8 with a real embedding must not enter
_fused_decode.That path applies a one-row identity RoPE after the wrapper has already rotated.
The call site requires
rotary_emb is None.Target K3 still has
rotary_emb is None, so its launch stays. DeepSeek does not use this wrapper.Trace
The baseline draft gap still launches
fusedKimiK3MLADecodeQConcatKVCacheKernelfive times per step, 21 µs total.The trace for this PR does not use that kernel. Each draft layer instead runs native
DeepseekScalingRotaryEmbedding.forward_hip, which callsforward_native, and then an ungroupedconcat_and_cache_mla.On the steady decode shape the eager gap grows from 3.41 ms to 3.85 ms (+0.44 ms, about 13%).
This PR is slower than the NVIDIA-module baseline on ROCm. The kernels that replaced the fused decode epilogue are the native YaRN elementwise chain plus the ungrouped cache insert.
Potential follow-up
If the fusion in the follow-up (#60864) is not merged, then the following alternative should be considered:
DeepseekScalingRotaryEmbedding.forward_hipcallsforward_native.Pointing it at AITER Triton RoPE would remove the elementwise chain for this draft and for every other HIP model
that uses the same embedding.
Validation
pytest tests/models/test_dspark_mla.pyon ROCm, includingtest_k3_dspark_decoder_uses_mla_wrapper(ROCm-only) andtest_hip_yarn_k_pe_is_cast_back_to_kv_dtype(CPU; a fake fp32 RoPE, asserting the seen key is bf16).test_amd_kv_cache_layer_returns_inner_mla_attn: draft KV names end in.attn.test_fused_decode_not_taken_when_rotary_emb_is_set.Out of scope
fused_qk_rope_concat_and_cache_mla([ROCm] Fuse Kimi-K3 DSpark QK RoPE with the MLA cache write #60864).DeepseekScalingRotaryEmbedding.forward_hip.Pull Request Checklist
I used vLLM's
/pr-checklistskill. (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.