[Model] Fix Kimi-K3 MLA with disabled context parallelism (#50404)
Signed-off-by: Tasos Varoudis <varoudis@archtech.gr> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: Lucas Wilkinson <lwilkins@redhat.com>
Dr Tasos Varoudis committed
Aug 5, 2026 at 03:33 UTC
c416f15710bbaf3e1d843c9e08403ca19fa49427
1 file changed
+5
vllm/models/kimi_k3/nvidia/mla.py
+5
@@ -301,6 +301,11 @@ class MultiHeadLatentAttention(nn.Module, AttentionLayerBase):
301
kv_b_proj=self.kv_b_proj,
302
indexer=None,
303
)
304
+ if getattr(self.impl, "dcp_world_size", -1) < 1:
305
+ # FlashAttention requires the cp_world_size is positive and the cp_rank
306
+ # is non negative; manually set here if not set by caller (-1 is unset)
307
+ self.impl.dcp_world_size = 1
308
+ self.impl.dcp_rank = 0
309
self.q_pad_num_heads = getattr(self.impl, "q_pad_num_heads", None)
310
311
vllm_config = get_current_vllm_config()