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>

python/sglang/srt/configs/model_config.py

@@ -1116,14 +1116,18 @@ class ModelConfig:
             return max(per_layer)
         return self.num_attention_heads
 
-    def get_num_kv_heads(self, tensor_parallel_size) -> int:
+    def get_num_kv_heads(self, tensor_parallel_size: int, dcp_size: int = 1) -> int:
-        """Returns the number of KV heads per GPU."""
+        """Number of KV heads per GPU.
+
+        DCP ranks replicate KV, so heads shard across ``tp // dcp`` groups.
+        Drafts never join the group and ignore ``dcp_size``. With fewer heads
+        than groups, each GPU keeps one.
+        """
         total_num_kv_heads = self.get_total_num_kv_heads()
-        # If tensor parallelism is used, we divide the number of KV heads by
+        if self.is_draft_model:
-        # the tensor parallel size. We will replicate the KV heads in the
+            dcp_size = 1
-        # case where the number of KV heads is smaller than the tensor
+        kv_tensor_parallel_size = tensor_parallel_size // dcp_size
-        # parallel size so each GPU has at least one KV head.
+        return max(1, total_num_kv_heads // kv_tensor_parallel_size)
-        return max(1, total_num_kv_heads // tensor_parallel_size)
 
     def get_swa_num_kv_heads(self, tensor_parallel_size) -> int:
         """Similar to get_num_kv_heads(), but for SWA."""

python/sglang/srt/layers/attention/flashinfer_backend.py

@@ -364,7 +364,7 @@ class FlashInferAttnBackend(AttentionBackend):
             num_attention_heads=model_runner.model_config.num_attention_heads
             // get_parallel().attn_tp_size,
             num_kv_heads=model_runner.model_config.get_num_kv_heads(
-                get_parallel().attn_tp_size
+                get_parallel().attn_tp_size, get_parallel().attn_dcp_size
             ),
         )
         self.max_context_len = model_runner.model_config.context_len
@@ -1479,7 +1479,7 @@ class FlashInferIndicesUpdaterDecode:
             // get_parallel().attn_tp_size
         )
         self.num_kv_heads = model_runner.model_config.get_num_kv_heads(
-            get_parallel().attn_tp_size
+            get_parallel().attn_tp_size, get_parallel().attn_dcp_size
         )
         self.head_dim = model_runner.model_config.head_dim
         self.data_type = attn_backend.flashinfer_kv_cache_dtype
@@ -1751,7 +1751,7 @@ class FlashInferIndicesUpdaterPrefill:
             // get_parallel().attn_tp_size
         )
         self.num_kv_heads = model_runner.model_config.get_num_kv_heads(
-            get_parallel().attn_tp_size
+            get_parallel().attn_tp_size, get_parallel().attn_dcp_size
         )
         self.head_dim = model_runner.model_config.head_dim
         self.data_type = attn_backend.flashinfer_kv_cache_dtype

python/sglang/srt/layers/attention/triton_backend.py

@@ -203,7 +203,7 @@ class TritonAttnBackend(AttentionBackend):
             // get_parallel().attn_tp_size
         ) * self.dcp_size
         self.num_kv_head = model_runner.model_config.get_num_kv_heads(
-            get_parallel().attn_tp_size
+            get_parallel().attn_tp_size, get_parallel().attn_dcp_size
         )
         # The decode kernel's "// Lv" stride trick requires attn_logits.shape[-1]
         # to exactly match the layer's v_head_dim, so hybrid SWA models with

python/sglang/srt/layers/linear.py

@@ -960,6 +960,8 @@ class QKVParallelLinear(ColumnParallelLinear):
         load_presharded_attn: bool = False,
         v_head_size: Optional[int] = None,
         skip_block_quant_check: bool = False,
+        kv_tp_rank: Optional[int] = None,
+        kv_tp_size: Optional[int] = None,
     ):
         self.with_bias = bias
         self.hidden_size = hidden_size
@@ -975,12 +977,17 @@ class QKVParallelLinear(ColumnParallelLinear):
         if tp_size is None:
             tp_size = get_parallel().tp_size
         self.tp_rank, self.tp_size = tp_rank, tp_size
+        if kv_tp_rank is None:
+            kv_tp_rank = tp_rank
+        if kv_tp_size is None:
+            kv_tp_size = tp_size
+        self.kv_tp_rank, self.kv_tp_size = kv_tp_rank, kv_tp_size
         self.num_heads = divide(self.total_num_heads, tp_size)
-        if tp_size >= self.total_num_kv_heads:
+        if kv_tp_size >= self.total_num_kv_heads:
             self.num_kv_heads = 1
-            self.num_kv_head_replicas = divide(tp_size, self.total_num_kv_heads)
+            self.num_kv_head_replicas = divide(kv_tp_size, self.total_num_kv_heads)
         else:
-            self.num_kv_heads = divide(self.total_num_kv_heads, tp_size)
+            self.num_kv_heads = divide(self.total_num_kv_heads, kv_tp_size)
             self.num_kv_head_replicas = 1
         self.q_proj_shard_size = self.num_heads * self.head_size
         self.kv_proj_shard_size = self.num_kv_heads * self.head_size
@@ -1102,7 +1109,7 @@ class QKVParallelLinear(ColumnParallelLinear):
                 shard_id=shard_id,
                 shard_offset=rank_shard_offset,
                 shard_size=rank_shard_size,
-                tp_rank=self.tp_rank,
+                tp_rank=(self.tp_rank if shard_id == "q" else self.kv_tp_rank),
                 use_presharded_weights=self.use_presharded_weights,
             )
 
@@ -1156,7 +1163,7 @@ class QKVParallelLinear(ColumnParallelLinear):
             shard_id=loaded_shard_id,
             shard_offset=shard_offset,
             shard_size=shard_size,
-            tp_rank=self.tp_rank,
+            tp_rank=(self.tp_rank if loaded_shard_id == "q" else self.kv_tp_rank),
             use_presharded_weights=self.use_presharded_weights,
         )
 
@@ -1179,8 +1186,13 @@ class QKVParallelLinear(ColumnParallelLinear):
 
         if is_gguf_weight:
             output_dim = getattr(param, "output_dim", None)
-            shard_size = loaded_weight.size(output_dim) // self.tp_size
+            shard_tp_rank, shard_tp_size = (
-            start_idx = self.tp_rank * shard_size
+                (self.kv_tp_rank, self.kv_tp_size)
+                if loaded_shard_id in ("k", "v")
+                else (self.tp_rank, self.tp_size)
+            )
+            shard_size = loaded_weight.size(output_dim) // shard_tp_size
+            start_idx = shard_tp_rank * shard_size
 
             loaded_weight = loaded_weight.narrow(output_dim, start_idx, shard_size)
 
@@ -1326,7 +1338,7 @@ class QKVParallelLinear(ColumnParallelLinear):
             if loaded_shard_id == "q":
                 shard_id = self.tp_rank
             else:
-                shard_id = self.tp_rank // self.num_kv_head_replicas
+                shard_id = self.kv_tp_rank // self.num_kv_head_replicas
             start_idx = shard_id * shard_size
 
             if _is_cpu:

python/sglang/srt/mem_cache/kv_cache_configurator.py

@@ -488,7 +488,9 @@ class KVCacheConfigurator:
         bundle = init_unified_mamba_pools(
             device=self.device,
             kv_cache_dtype=self.kv_cache_dtype,
-            head_num=self.model_config.get_num_kv_heads(get_parallel().attn_tp_size),
+            head_num=self.model_config.get_num_kv_heads(
+                get_parallel().attn_tp_size, get_parallel().attn_dcp_size
+            ),
             head_dim=self.model_config.head_dim,
             page_size=self.page_size,
             start_layer=self.layer_info.start_layer,
@@ -556,7 +558,9 @@ class KVCacheConfigurator:
             enable_memory_saver=get_exec().features.enable_memory_saver,
         )
 
-        head_num = self.model_config.get_num_kv_heads(get_parallel().attn_tp_size)
+        head_num = self.model_config.get_num_kv_heads(
+            get_parallel().attn_tp_size, get_parallel().attn_dcp_size
+        )
         head_dim = self.model_config.head_dim
         if self.is_hybrid_swa_compress:
             # Asymmetric head dims between full and SWA (NPU compress path):
@@ -1082,7 +1086,9 @@ class KVCacheConfigurator:
             max_total_num_tokens,
             page_size=self.pool_page_size,
             dtype=self.kv_cache_dtype,
-            head_num=self.model_config.get_num_kv_heads(get_parallel().attn_tp_size),
+            head_num=self.model_config.get_num_kv_heads(
+                get_parallel().attn_tp_size, get_parallel().attn_dcp_size
+            ),
             head_dim=self.model_config.head_dim,
             layer_num=self.layer_info.num_effective_layers,
             device=self.device,
@@ -1120,7 +1126,9 @@ class KVCacheConfigurator:
             page_size=self.pool_page_size,
             dtype=self.kv_cache_dtype,
             post_capture_active=self.post_capture_kv_active,
-            head_num=self.model_config.get_num_kv_heads(get_parallel().attn_tp_size),
+            head_num=self.model_config.get_num_kv_heads(
+                get_parallel().attn_tp_size, get_parallel().attn_dcp_size
+            ),
             head_dim=self.model_config.head_dim,
             swa_attention_layer_ids=self.model_config.swa_attention_layer_ids,
             full_attention_layer_ids=self.model_config.full_attention_layer_ids,
@@ -1161,7 +1169,9 @@ class KVCacheConfigurator:
             max_total_num_tokens,
             page_size=self.pool_page_size,
             dtype=self.kv_cache_dtype,
-            head_num=self.model_config.get_num_kv_heads(get_parallel().attn_tp_size),
+            head_num=self.model_config.get_num_kv_heads(
+                get_parallel().attn_tp_size, get_parallel().attn_dcp_size
+            ),
             head_dim=self.model_config.head_dim,
             layer_num=self.layer_info.num_effective_layers,
             device=self.device,
@@ -1302,7 +1312,9 @@ class KVCacheConfigurator:
             page_size=self.pool_page_size,
             dtype=self.kv_cache_dtype,
             post_capture_active=self.post_capture_kv_active,
-            head_num=self.model_config.get_num_kv_heads(get_parallel().attn_tp_size),
+            head_num=self.model_config.get_num_kv_heads(
+                get_parallel().attn_tp_size, get_parallel().attn_dcp_size
+            ),
             head_dim=self.model_config.head_dim,
             swa_attention_layer_ids=swa_attention_layer_ids,
             full_attention_layer_ids=full_attention_layer_ids,
@@ -1335,7 +1347,9 @@ class KVCacheConfigurator:
                 if m3_fp8_attn_gemm_enabled(self.server_args)
                 else self.model_dtype
             ),
-            head_num=self.model_config.get_num_kv_heads(get_parallel().attn_tp_size),
+            head_num=self.model_config.get_num_kv_heads(
+                get_parallel().attn_tp_size, get_parallel().attn_dcp_size
+            ),
             head_dim=self.model_config.head_dim,
             idx_head_dim=sparse_cfg["sparse_index_dim"],
             dense_layer_ids=dense_layer_ids,
@@ -1384,7 +1398,9 @@ class KVCacheConfigurator:
             page_size=self.pool_page_size,
             size=max_total_num_tokens,
             dtype=self.kv_cache_dtype,
-            head_num=self.model_config.get_num_kv_heads(get_parallel().attn_tp_size),
+            head_num=self.model_config.get_num_kv_heads(
+                get_parallel().attn_tp_size, get_parallel().attn_dcp_size
+            ),
             head_dim=self.model_config.head_dim,
             # if draft worker, we only need 1 attention layer's kv pool
             full_attention_layer_ids=full_attention_layer_ids,
@@ -1406,7 +1422,9 @@ class KVCacheConfigurator:
             max_total_num_tokens,
             page_size=self.pool_page_size,
             dtype=self.kv_cache_dtype,
-            head_num=self.model_config.get_num_kv_heads(get_parallel().attn_tp_size),
+            head_num=self.model_config.get_num_kv_heads(
+                get_parallel().attn_tp_size, get_parallel().attn_dcp_size
+            ),
             head_dim=self.model_config.head_dim,
             v_head_dim=self.model_config.v_head_dim,
             layer_num=self.layer_info.num_effective_layers,
@@ -1439,7 +1457,9 @@ class KVCacheConfigurator:
             max_total_num_tokens,
             page_size=self.pool_page_size,
             dtype=self.kv_cache_dtype,
-            head_num=self.model_config.get_num_kv_heads(get_parallel().attn_tp_size),
+            head_num=self.model_config.get_num_kv_heads(
+                get_parallel().attn_tp_size, get_parallel().attn_dcp_size
+            ),
             head_dim=self.model_config.head_dim,
             v_head_dim=self.model_config.v_head_dim,
             layer_num=self.layer_info.num_effective_layers,

python/sglang/srt/model_executor/pool_configurator.py

@@ -196,6 +196,7 @@ class DefaultPoolConfigurator(MemoryPoolConfigurator):
 
         kv_size = torch._utils._element_size(kv_cache_dtype)
         tp_size = get_parallel().attn_tp_size
+        dcp_size = get_parallel().attn_dcp_size
 
         if kvc.use_mla_backend:
             from sglang.srt.mem_cache.kv_cache_configurator import (
@@ -289,8 +290,9 @@ class DefaultPoolConfigurator(MemoryPoolConfigurator):
             # cell_size is already a sum over heterogeneous sub-pools.
             return main_pool_bytes + indexer_bytes
         else:
+            n = model_config.get_num_kv_heads(tp_size, dcp_size)
             cell_size = (
-                model_config.get_num_kv_heads(tp_size)
+                n
                 * (model_config.head_dim + model_config.v_head_dim)
                 * effective_num_layers
                 * kv_size
@@ -299,7 +301,6 @@ class DefaultPoolConfigurator(MemoryPoolConfigurator):
             if is_float4_e2m1fn_x2(kv_cache_dtype):
                 # kv_scale_buffer
                 scale_block_size = 16
-                n = model_config.get_num_kv_heads(tp_size)
                 k = model_config.head_dim
                 cell_size = (cell_size // 2) + (
                     (n * k * effective_num_layers * 2 * kv_size) // scale_block_size
@@ -308,7 +309,6 @@ class DefaultPoolConfigurator(MemoryPoolConfigurator):
                 cell_size += n * k * 2 * kv_size
             elif self.kv_cache_dtype_str == "mxfp8":
                 scale_block_size = 32
-                n = model_config.get_num_kv_heads(tp_size)
                 cell_size += (
                     n * (model_config.head_dim + model_config.v_head_dim) * num_layers
                 ) // scale_block_size

python/sglang/srt/models/qwen3_5.py

@@ -848,15 +848,20 @@ class Qwen3_5AttentionDecoderLayer(nn.Module):
         self.hidden_size = config.hidden_size
         self.attn_tp_rank = get_parallel().attn_tp_rank
         self.attn_tp_size = get_parallel().attn_tp_size
+        # A Qwen3.5 draft is rewritten to the MTP arch (model_config._config_draft_model),
+        # so is_nextn marks it. Drafts are TP-sharded and do not replicate KV under DCP.
+        dcp_size = 1 if is_nextn else get_parallel().attn_dcp_size
+        self.kv_tp_size = self.attn_tp_size // dcp_size
+        self.kv_tp_rank = self.attn_tp_rank // dcp_size
         self.total_num_heads = config.num_attention_heads
         assert self.total_num_heads % self.attn_tp_size == 0
         self.num_heads = self.total_num_heads // self.attn_tp_size
         self.total_num_kv_heads = config.num_key_value_heads
-        if self.total_num_kv_heads >= self.attn_tp_size:
+        if self.total_num_kv_heads >= self.kv_tp_size:
-            assert self.total_num_kv_heads % self.attn_tp_size == 0
+            assert self.total_num_kv_heads % self.kv_tp_size == 0
         else:
-            assert self.attn_tp_size % self.total_num_kv_heads == 0
+            assert self.kv_tp_size % self.total_num_kv_heads == 0
-        self.num_kv_heads = max(1, self.total_num_kv_heads // self.attn_tp_size)
+        self.num_kv_heads = max(1, self.total_num_kv_heads // self.kv_tp_size)
         self.head_dim = config.head_dim or (self.hidden_size // self.num_heads)
         self.q_size = self.num_heads * self.head_dim
         self.kv_size = self.num_kv_heads * self.head_dim
@@ -886,6 +891,8 @@ class Qwen3_5AttentionDecoderLayer(nn.Module):
             dtype=torch.get_default_dtype(),
         )
 
+        # Q stays sharded across attention TP ranks; K/V are replicated within
+        # each DCP group.
         self.qkv_proj = QKVParallelLinear(
             config.hidden_size,
             self.head_dim,
@@ -895,6 +902,8 @@ class Qwen3_5AttentionDecoderLayer(nn.Module):
             quant_config=quant_config,
             tp_rank=self.attn_tp_rank,
             tp_size=self.attn_tp_size,
+            kv_tp_rank=self.kv_tp_rank,
+            kv_tp_size=self.kv_tp_size,
             prefix=add_prefix("qkv_proj", prefix),
         )
 

python/sglang/test/kits/attention_unittest/attention_methods/dense_attention.py

@@ -298,9 +298,10 @@ class TinyModelConfig:
     def get_max_num_attention_heads(self) -> int:
         return self.num_attention_heads
 
-    def get_num_kv_heads(self, tp_size: int) -> int:
+    def get_num_kv_heads(self, tp_size: int, dcp_size: int = 1) -> int:
-        assert self.num_key_value_heads % tp_size == 0
+        kv_tp_size = tp_size // dcp_size
-        return self.num_key_value_heads // tp_size
+        assert self.num_key_value_heads % kv_tp_size == 0
+        return self.num_key_value_heads // kv_tp_size
 
 
 class MockModelRunner(ModelRunner):

python/sglang/test/kits/attention_unittest/attention_methods/dual_chunk_attention.py

@@ -298,9 +298,10 @@ class TinyDualChunkModelConfig:
     def get_max_num_attention_heads(self) -> int:
         return self.num_attention_heads
 
-    def get_num_kv_heads(self, tp_size: int) -> int:
+    def get_num_kv_heads(self, tp_size: int, dcp_size: int = 1) -> int:
-        assert self.num_key_value_heads % tp_size == 0
+        kv_tp_size = tp_size // dcp_size
-        return self.num_key_value_heads // tp_size
+        assert self.num_key_value_heads % kv_tp_size == 0
+        return self.num_key_value_heads // kv_tp_size
 
 
 class DualChunkMockModelRunner(ModelRunner):

python/sglang/test/kits/attention_unittest/attention_methods/gdn_attention.py

@@ -191,9 +191,10 @@ class TinyGDNModelConfig:
     def get_max_num_attention_heads(self) -> int:
         return self.num_attention_heads
 
-    def get_num_kv_heads(self, tp_size: int) -> int:
+    def get_num_kv_heads(self, tp_size: int, dcp_size: int = 1) -> int:
-        assert self.num_key_value_heads % tp_size == 0
+        kv_tp_size = tp_size // dcp_size
-        return self.num_key_value_heads // tp_size
+        assert self.num_key_value_heads % kv_tp_size == 0
+        return self.num_key_value_heads // kv_tp_size
 
 
 class MockGDNModelRunner(ModelRunner):

python/sglang/test/kits/attention_unittest/attention_methods/kda_attention.py

@@ -196,9 +196,10 @@ class TinyKDAModelConfig:
     def get_max_num_attention_heads(self) -> int:
         return self.num_attention_heads
 
-    def get_num_kv_heads(self, tp_size: int) -> int:
+    def get_num_kv_heads(self, tp_size: int, dcp_size: int = 1) -> int:
-        assert self.num_key_value_heads % tp_size == 0
+        kv_tp_size = tp_size // dcp_size
-        return self.num_key_value_heads // tp_size
+        assert self.num_key_value_heads % kv_tp_size == 0
+        return self.num_key_value_heads // kv_tp_size
 
 
 class MockKDAModelRunner(ModelRunner):

python/sglang/test/kits/attention_unittest/attention_methods/lightning_attention.py

@@ -206,9 +206,10 @@ class TinyLightningModelConfig:
     def get_max_num_attention_heads(self) -> int:
         return self.num_attention_heads
 
-    def get_num_kv_heads(self, tp_size: int) -> int:
+    def get_num_kv_heads(self, tp_size: int, dcp_size: int = 1) -> int:
-        assert self.num_key_value_heads % tp_size == 0
+        kv_tp_size = tp_size // dcp_size
-        return self.num_key_value_heads // tp_size
+        assert self.num_key_value_heads % kv_tp_size == 0
+        return self.num_key_value_heads // kv_tp_size
 
 
 class MockLightningModelRunner(ModelRunner):

python/sglang/test/kits/attention_unittest/attention_methods/mamba2_attention.py

@@ -292,9 +292,10 @@ class TinyMamba2ModelConfig:
     def get_max_num_attention_heads(self) -> int:
         return self.num_attention_heads
 
-    def get_num_kv_heads(self, tp_size: int) -> int:
+    def get_num_kv_heads(self, tp_size: int, dcp_size: int = 1) -> int:
-        assert self.num_key_value_heads % tp_size == 0
+        kv_tp_size = tp_size // dcp_size
-        return self.num_key_value_heads // tp_size
+        assert self.num_key_value_heads % kv_tp_size == 0
+        return self.num_key_value_heads // kv_tp_size
 
 
 class MockMamba2ModelRunner(ModelRunner):

python/sglang/test/kits/attention_unittest/attention_methods/mla_attention.py

@@ -203,7 +203,7 @@ class TinyMLAModelConfig:
     def get_max_num_attention_heads(self) -> int:
         return self.num_attention_heads
 
-    def get_num_kv_heads(self, tp_size: int) -> int:
+    def get_num_kv_heads(self, tp_size: int, dcp_size: int = 1) -> int:
         return 1
 
 

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
+TP_SIZE = 4
-DCP_SIZE = 2
+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,6 +55,7 @@ class TestQwen35TritonDCPGsm8k(CustomTestCase):
             ),
         ]
         env = os.environ.copy()
+        if torch.version.hip:
             env["SGLANG_USE_AITER"] = "1"
             env["HSA_NO_SCRATCH_RECLAIM"] = "1"
 
@@ -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