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

@@ -0,0 +1,96 @@
+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, register_cuda_ci
+from sglang.test.run_eval import run_eval
+from sglang.test.test_utils import (
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    is_in_ci,
+    popen_launch_server,
+    write_github_step_summary,
+)
+
+register_amd_ci(
+    est_time=4800,
+    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 = 4
+DCP_SIZE = 4
+GSM8K_ACCURACY_THRESHOLD = 0.90
+
+
+class TestQwen35TritonDCPGsm8k(CustomTestCase):
+    """Qwen3.5 Triton DCP (TP4/DCP4) full GSM8K accuracy."""
+
+    @classmethod
+    def setUpClass(cls):
+        cls.model = QWEN35_MODEL_PATH
+        cls.base_url = DEFAULT_URL_FOR_TEST
+
+        other_args = [
+            "--trust-remote-code",
+            "--tp",
+            str(TP_SIZE),
+            "--dcp-size",
+            str(DCP_SIZE),
+            "--attention-backend",
+            "triton",
+            "--context-length",
+            "1048576",
+            "--disable-radix-cache",
+            "--json-model-override-args",
+            (
+                '{"rope_scaling":{"rope_type":"yarn","factor":4.0,'
+                '"original_max_position_embeddings":262144}}'
+            ),
+        ]
+        env = os.environ.copy()
+        if torch.version.hip:
+            env["SGLANG_USE_AITER"] = "1"
+            env["HSA_NO_SCRATCH_RECLAIM"] = "1"
+
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=SERVER_LAUNCH_TIMEOUT,
+            other_args=other_args,
+            env=env,
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+    def test_a_gsm8k(self):
+        args = SimpleNamespace(
+            base_url=self.base_url,
+            model=self.model,
+            eval_name="gsm8k",
+            api="completion",
+            max_tokens=512,
+            num_examples=1319,
+            num_threads=32,
+            num_shots=5,
+        )
+        metrics = run_eval(args)
+        print(f"{metrics=}")
+
+        if is_in_ci():
+            write_github_step_summary(
+                f"### test_a_gsm8k (qwen3.5-triton-dcp4)\n" f'{metrics["score"]=:.3f}\n'
+            )
+        self.assertGreater(metrics["score"], GSM8K_ACCURACY_THRESHOLD)
+
+
+if __name__ == "__main__":
+    unittest.main()