Fix DCP KV head mapping for GQA models (#32858)

Co-authored-by: Khoa Pham <khoa.pham@radixark.ai>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

test/registered/dcp/test_dcp_layout_unit.py

@@ -20,7 +20,9 @@ from unittest.mock import MagicMock, patch
 import torch
 
 from sglang.srt import runtime_context as rc
+from sglang.srt.configs.model_config import ModelConfig
 from sglang.srt.layers.dcp.layout import get_dcp_lens
+from sglang.srt.layers.linear import QKVParallelLinear
 from sglang.srt.mem_cache.allocator.paged import PagedTokenToKVPoolAllocator
 from sglang.srt.mem_cache.kv_cache_configurator import KVCacheConfigurator
 from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool
@@ -123,6 +125,58 @@ class TestGetDcpLens(CustomTestCase):
             real_kv_size + physical_page_size,
         )
 
+    @staticmethod
+    def _kv_head_config(*, is_draft_model: bool):
+        model_config = ModelConfig.__new__(ModelConfig)
+        model_config.hf_config = SimpleNamespace(model_type="qwen3_5_text")
+        model_config.hf_text_config = SimpleNamespace(num_key_value_heads=8)
+        model_config.is_draft_model = is_draft_model
+        return model_config
+
+    def test_model_config_uses_non_dcp_tp_size_for_kv_heads(self):
+        model_config = self._kv_head_config(is_draft_model=False)
+
+        self.assertEqual(model_config.get_num_kv_heads(16), 1)
+        self.assertEqual(model_config.get_num_kv_heads(16, dcp_size=4), 2)
+
+    def test_a_draft_keeps_kv_heads_tp_sharded_under_dcp(self):
+        """The draft pool must match what a TP-sharded draft builds; sizing it
+        with the target's dcp_size over-allocates by that factor."""
+        model_config = self._kv_head_config(is_draft_model=True)
+
+        self.assertEqual(model_config.get_num_kv_heads(16, dcp_size=4), 1)
+        self.assertEqual(model_config.get_num_kv_heads(16), 1)
+
+    def test_gqa_qkv_loader_replicates_kv_within_dcp_group(self):
+        hidden_size = 4
+        head_size = 2
+        q_weight = torch.arange(64, dtype=torch.float32).view(16, hidden_size)
+        k_weight = torch.arange(16, dtype=torch.float32).view(4, hidden_size) + 100
+        v_weight = torch.arange(16, dtype=torch.float32).view(4, hidden_size) + 200
+
+        for tp_rank in range(4):
+            layer = QKVParallelLinear(
+                hidden_size=hidden_size,
+                head_size=head_size,
+                total_num_heads=8,
+                total_num_kv_heads=2,
+                bias=False,
+                params_dtype=torch.float32,
+                tp_rank=tp_rank,
+                tp_size=4,
+                kv_tp_rank=tp_rank // 2,
+                kv_tp_size=2,
+            )
+            layer.weight_loader(layer.weight, q_weight, "q")
+            layer.weight_loader(layer.weight, k_weight, "k")
+            layer.weight_loader(layer.weight, v_weight, "v")
+
+            q, k, v = layer.weight.split([4, 2, 2], dim=0)
+            kv_start = (tp_rank // 2) * 2
+            self.assertTrue(torch.equal(q, q_weight[tp_rank * 4 : (tp_rank + 1) * 4]))
+            self.assertTrue(torch.equal(k, k_weight[kv_start : kv_start + 2]))
+            self.assertTrue(torch.equal(v, v_weight[kv_start : kv_start + 2]))
+
     def test_configurator_scales_only_the_virtual_dcp_allocator(self):
         physical_kv_size = 1024
         physical_page_size = 64

test/registered/dcp/test_qwen3p5_triton_dcp.py

@@ -2,8 +2,10 @@ import os
 import unittest
 from types import SimpleNamespace
 
+import torch
+
 from sglang.srt.utils import kill_process_tree
-from sglang.test.ci.ci_register import register_amd_ci
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.test.run_eval import run_eval
 from sglang.test.test_utils import (
     DEFAULT_URL_FOR_TEST,
@@ -18,16 +20,17 @@ register_amd_ci(
     suite="nightly-amd-accuracy-8-gpu-mi35x-qwen35-triton-dcp",
     nightly=True,
 )
+register_cuda_ci(est_time=4800, suite="nightly-4-gpu-b200", nightly=True)
 
 QWEN35_MODEL_PATH = os.environ.get("QWEN3_5_MODEL_PATH", "Qwen/Qwen3.5-397B-A17B-FP8")
 SERVER_LAUNCH_TIMEOUT = 4800
-TP_SIZE = 8
-DCP_SIZE = 2
+TP_SIZE = 4
+DCP_SIZE = 4
 GSM8K_ACCURACY_THRESHOLD = 0.90
 
 
 class TestQwen35TritonDCPGsm8k(CustomTestCase):
-    """Qwen3.5 Triton DCP (tp=8, dcp=2) full GSM8K accuracy on AMD MI35x."""
+    """Qwen3.5 Triton DCP (TP4/DCP4) full GSM8K accuracy."""
 
     @classmethod
     def setUpClass(cls):
@@ -52,8 +55,9 @@ class TestQwen35TritonDCPGsm8k(CustomTestCase):
             ),
         ]
         env = os.environ.copy()
-        env["SGLANG_USE_AITER"] = "1"
-        env["HSA_NO_SCRATCH_RECLAIM"] = "1"
+        if torch.version.hip:
+            env["SGLANG_USE_AITER"] = "1"
+            env["HSA_NO_SCRATCH_RECLAIM"] = "1"
 
         cls.process = popen_launch_server(
             cls.model,
@@ -83,7 +87,7 @@ class TestQwen35TritonDCPGsm8k(CustomTestCase):
 
         if is_in_ci():
             write_github_step_summary(
-                f"### test_a_gsm8k (qwen3.5-triton-dcp2)\n" f'{metrics["score"]=:.3f}\n'
+                f"### test_a_gsm8k (qwen3.5-triton-dcp4)\n" f'{metrics["score"]=:.3f}\n'
             )
         self.assertGreater(metrics["score"], GSM8K_ACCURACY_THRESHOLD)
 

test/registered/unit/model_executor/test_pool_configurator.py

@@ -96,7 +96,7 @@ def _make_model_runner(
     )
     mc.swa_head_dim = swa_head_dim or head_dim
     mc.swa_v_head_dim = swa_v_head_dim or v_head_dim
-    mc.get_num_kv_heads = lambda tp_size: num_kv_heads
+    mc.get_num_kv_heads = lambda tp_size, dcp_size=1: num_kv_heads
     mc.get_swa_num_kv_heads = lambda tp_size: swa_num_kv_heads or num_kv_heads
     mc.hf_config = SimpleNamespace(architectures=["LlamaForCausalLM"])
     mc.hf_config.get_text_config = lambda: mc.hf_config