From c416f15710bbaf3e1d843c9e08403ca19fa49427 Mon Sep 17 00:00:00 2001 From: Dr Tasos Varoudis Date: Wed, 5 Aug 2026 03:33:14 +0100 Subject: [PATCH] [Model] Fix Kimi-K3 MLA with disabled context parallelism (#50404) Signed-off-by: Tasos Varoudis Signed-off-by: Lucas Wilkinson Co-authored-by: Lucas Wilkinson --- vllm/models/kimi_k3/nvidia/mla.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/vllm/models/kimi_k3/nvidia/mla.py b/vllm/models/kimi_k3/nvidia/mla.py index 3334f3bbbb..2b9a20140e 100644 --- a/vllm/models/kimi_k3/nvidia/mla.py +++ b/vllm/models/kimi_k3/nvidia/mla.py @@ -301,6 +301,11 @@ class MultiHeadLatentAttention(nn.Module, AttentionLayerBase): kv_b_proj=self.kv_b_proj, indexer=None, ) + if getattr(self.impl, "dcp_world_size", -1) < 1: + # FlashAttention requires the cp_world_size is positive and the cp_rank + # is non negative; manually set here if not set by caller (-1 is unset) + self.impl.dcp_world_size = 1 + self.impl.dcp_rank = 0 self.q_pad_num_heads = getattr(self.impl, "q_pad_num_heads", None) vllm_config = get_current_vllm_config()