[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 <cursoragent@cursor.com> Co-authored-by: Shanshan Shen <467638484@qq.com>
haic0 committed
Sep 9, 2026 at 22:50 UTC
d8d53f17c231bf477aa89f25a904f9b7c54fbcb9
2 files changed
+12
tests/models/test_registry.py
+2
@@ -128,6 +128,8 @@ def test_registry_model_property(model_arch, is_mm, init_cuda, score_type):
128
# in V1.
129
# ("MLPSpeculatorPreTrainedModel", False, False),
130
("DeepseekV2ForCausalLM", True, False),
131
+ ("KimiLinearForCausalLM", True, False),
132
+ ("KimiK3ForConditionalGeneration", True, False),
133
("Qwen2VLForConditionalGeneration", True, True),
134
],
135
)
vllm/model_executor/layers/attention/mla_attention.py
+10
@@ -2537,6 +2537,16 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
2537
if num_decodes > 0:
2538
dcp_tot_seq_lens_device = None
2539
if self.dcp_world_size > 1:
2540
+ assert seq_lens is not None, (
2541
+ "MLA DCP decode requires seq_lens on CommonAttentionMetadata"
2542
+ )
2543
+ if dcp_local_seq_lens is None:
2544
+ dcp_local_seq_lens = get_dcp_local_seq_lens(
2545
+ seq_lens,
2546
+ dcp_size=self.dcp_world_size,
2547
+ dcp_rank=get_dcp_group().rank_in_group,
2548
+ cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size,
2549
+ )
2550
dcp_tot_seq_lens_device = seq_lens[:num_decodes]
2551
seq_lens = dcp_local_seq_lens
2552