[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>

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),
     ],
 )

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