[MRv2] FlashAttention PCP support for GQA on MRv2

vllm/config/model.py

@@ -1313,7 +1313,14 @@ def verify_with_parallel_config(
             )
 
         decode_context_parallel_size = parallel_config.decode_context_parallel_size
-        if decode_context_parallel_size > 1 and not self.use_mla:
+        # DCP groups span the PCP axis before TP, so DCP only splits query heads
+        # when it reaches past PCP (see pcp.maybe_all_gather_q_for_dcp). When
+        # dcp == pcp the group holds replicated Q and none of these TP-head
+        # constraints apply.
+        if (
+            decode_context_parallel_size > parallel_config.prefill_context_parallel_size
+            and not self.use_mla
+        ):
             total_num_kv_heads = self.get_total_num_kv_heads()
             if tensor_parallel_size <= total_num_kv_heads:
                 raise ValueError(

vllm/config/vllm.py

@@ -2144,10 +2144,11 @@ def _get_v2_model_runner_unsupported_features(self) -> list[str]:
         model_config = self.model_config
         speculative_config = self.speculative_config
 
-        if self.parallel_config.prefill_context_parallel_size > 1 and not (
-            model_config is not None and model_config.use_mla
-        ):
-            unsupported.append("prefill context parallelism")
+        # PCP on MRv2 is opt-in per attention backend via
+        # AttentionImplBase.supports_pcp, enforced by
+        # check_attention_cp_compatibility (vllm/v1/worker/cp_utils.py). MLA and
+        # GQA (FlashAttention) both opt in, so do not block here -- if no
+        # attention impl supports PCP the cp_utils check raises a clear error.
         if self.compilation_config.mode == CompilationMode.STOCK_TORCH_COMPILE:
             unsupported.append("stock torch.compile")