fix num_q_head && add UT

Signed-off-by: Jingchun Gao <[email protected]>

Author: Jingchun Gao

tests/distributed/test_context_parallel.py

@@ -232,6 +232,10 @@ def _compare_cp_with_tp(
     "bigcode/gpt_bigcode-santacoder": [
         CPTestSettings.detailed(),
         CPTestSettings.detailed(tp_base=2),
+        CPTestSettings.detailed(attn_backend="FLASHINFER"),
+        CPTestSettings.detailed(
+            attn_backend="FLASHINFER", cp_kv_cache_interleave_size=16
+        ),
     ],
 }
 

vllm/v1/attention/backends/flashinfer.py

@@ -478,9 +478,8 @@ def __init__(
             self.dcp_rank = 0
             self.dcp_kv_cache_interleave_size = 1
 
-        self.num_qo_heads = (
-            self.model_config.get_num_attention_heads(self.vllm_config.parallel_config)
-            * self.dcp_world_size
+        self.num_qo_heads = self.model_config.get_num_attention_heads(
+            self.vllm_config.parallel_config
         )
 
         self.num_kv_heads = self.kv_cache_spec.num_kv_heads