From d8d53f17c231bf477aa89f25a904f9b7c54fbcb9 Mon Sep 17 00:00:00 2001 From: haic0 <149741444+haic0@users.noreply.github.com> Date: Wed, 9 Sep 2026 22:50:24 +0800 Subject: [PATCH] [Rocm][Kimi-k3] Add pipeline_parallel support for the kimik3 model (#53664) Signed-off-by: haic0 <149741444+haic0@users.noreply.github.com> Co-authored-by: Cursor Co-authored-by: Shanshan Shen <467638484@qq.com> --- tests/models/test_registry.py | 2 ++ vllm/model_executor/layers/attention/mla_attention.py | 10 ++++++++++ 2 files changed, 12 insertions(+) diff --git a/tests/models/test_registry.py b/tests/models/test_registry.py index 3bcb0dc433..30eb80a3f8 100644 --- a/tests/models/test_registry.py +++ b/tests/models/test_registry.py @@ -128,6 +128,8 @@ def test_registry_model_property(model_arch, is_mm, init_cuda, score_type): # in V1. # ("MLPSpeculatorPreTrainedModel", False, False), ("DeepseekV2ForCausalLM", True, False), + ("KimiLinearForCausalLM", True, False), + ("KimiK3ForConditionalGeneration", True, False), ("Qwen2VLForConditionalGeneration", True, True), ], ) diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index a45f074837..daaef86898 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -2537,6 +2537,16 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]): if num_decodes > 0: dcp_tot_seq_lens_device = None if self.dcp_world_size > 1: + assert seq_lens is not None, ( + "MLA DCP decode requires seq_lens on CommonAttentionMetadata" + ) + if dcp_local_seq_lens is None: + dcp_local_seq_lens = get_dcp_local_seq_lens( + seq_lens, + dcp_size=self.dcp_world_size, + dcp_rank=get_dcp_group().rank_in_group, + cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size, + ) dcp_tot_seq_lens_device = seq_lens[:num_decodes] seq_lens = dcp_local_seq_lens