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()