diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index 33bd64f08..6ffda9367 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/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: - """Returns the number of KV heads per GPU.""" + def get_num_kv_heads(self, tensor_parallel_size: int, dcp_size: int = 1) -> int: + """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 - # the tensor parallel size. We will replicate the KV heads in the - # case where the number of KV heads is smaller than the tensor - # parallel size so each GPU has at least one KV head. - return max(1, total_num_kv_heads // tensor_parallel_size) + if self.is_draft_model: + dcp_size = 1 + kv_tensor_parallel_size = tensor_parallel_size // dcp_size + return max(1, total_num_kv_heads // kv_tensor_parallel_size) def get_swa_num_kv_heads(self, tensor_parallel_size) -> int: """Similar to get_num_kv_heads(), but for SWA.""" diff --git a/python/sglang/srt/layers/attention/flashinfer_backend.py b/python/sglang/srt/layers/attention/flashinfer_backend.py index 63cd3d970..324ca3f33 100644 --- a/python/sglang/srt/layers/attention/flashinfer_backend.py +++ b/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 diff --git a/python/sglang/srt/layers/attention/triton_backend.py b/python/sglang/srt/layers/attention/triton_backend.py index 9d3111e45..37a8061b1 100644 --- a/python/sglang/srt/layers/attention/triton_backend.py +++ b/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 diff --git a/python/sglang/srt/layers/linear.py b/python/sglang/srt/layers/linear.py index 67662c9a7..d42d7578f 100644 --- a/python/sglang/srt/layers/linear.py +++ b/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 - start_idx = self.tp_rank * shard_size + shard_tp_rank, shard_tp_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: diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index 485a96ee0..574760fe7 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/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, diff --git a/python/sglang/srt/model_executor/pool_configurator.py b/python/sglang/srt/model_executor/pool_configurator.py index fea221567..68bc1eddd 100644 --- a/python/sglang/srt/model_executor/pool_configurator.py +++ b/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 diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index 46c7ed45f..56ddca2dc 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/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: - assert self.total_num_kv_heads % self.attn_tp_size == 0 + if self.total_num_kv_heads >= self.kv_tp_size: + assert self.total_num_kv_heads % self.kv_tp_size == 0 else: - assert self.attn_tp_size % self.total_num_kv_heads == 0 - self.num_kv_heads = max(1, self.total_num_kv_heads // self.attn_tp_size) + assert self.kv_tp_size % self.total_num_kv_heads == 0 + 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), ) diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/dense_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/dense_attention.py index 3ed889b2f..7ca676e61 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/dense_attention.py +++ b/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: - assert self.num_key_value_heads % tp_size == 0 - return self.num_key_value_heads // tp_size + def get_num_kv_heads(self, tp_size: int, dcp_size: int = 1) -> int: + kv_tp_size = tp_size // dcp_size + assert self.num_key_value_heads % kv_tp_size == 0 + return self.num_key_value_heads // kv_tp_size class MockModelRunner(ModelRunner): diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/dual_chunk_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/dual_chunk_attention.py index 3eabff6ef..2e6831e72 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/dual_chunk_attention.py +++ b/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: - assert self.num_key_value_heads % tp_size == 0 - return self.num_key_value_heads // tp_size + def get_num_kv_heads(self, tp_size: int, dcp_size: int = 1) -> int: + kv_tp_size = tp_size // dcp_size + assert self.num_key_value_heads % kv_tp_size == 0 + return self.num_key_value_heads // kv_tp_size class DualChunkMockModelRunner(ModelRunner): diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/gdn_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/gdn_attention.py index db551e0db..50354ff4a 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/gdn_attention.py +++ b/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: - assert self.num_key_value_heads % tp_size == 0 - return self.num_key_value_heads // tp_size + def get_num_kv_heads(self, tp_size: int, dcp_size: int = 1) -> int: + kv_tp_size = tp_size // dcp_size + assert self.num_key_value_heads % kv_tp_size == 0 + return self.num_key_value_heads // kv_tp_size class MockGDNModelRunner(ModelRunner): diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/kda_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/kda_attention.py index dd53c74cb..0315c5352 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/kda_attention.py +++ b/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: - assert self.num_key_value_heads % tp_size == 0 - return self.num_key_value_heads // tp_size + def get_num_kv_heads(self, tp_size: int, dcp_size: int = 1) -> int: + kv_tp_size = tp_size // dcp_size + assert self.num_key_value_heads % kv_tp_size == 0 + return self.num_key_value_heads // kv_tp_size class MockKDAModelRunner(ModelRunner): diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/lightning_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/lightning_attention.py index a1fcf89ef..1b85b55b3 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/lightning_attention.py +++ b/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: - assert self.num_key_value_heads % tp_size == 0 - return self.num_key_value_heads // tp_size + def get_num_kv_heads(self, tp_size: int, dcp_size: int = 1) -> int: + kv_tp_size = tp_size // dcp_size + assert self.num_key_value_heads % kv_tp_size == 0 + return self.num_key_value_heads // kv_tp_size class MockLightningModelRunner(ModelRunner): diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/mamba2_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/mamba2_attention.py index 78436ba14..f46459a50 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/mamba2_attention.py +++ b/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: - assert self.num_key_value_heads % tp_size == 0 - return self.num_key_value_heads // tp_size + def get_num_kv_heads(self, tp_size: int, dcp_size: int = 1) -> int: + kv_tp_size = tp_size // dcp_size + assert self.num_key_value_heads % kv_tp_size == 0 + return self.num_key_value_heads // kv_tp_size class MockMamba2ModelRunner(ModelRunner): diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/mla_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/mla_attention.py index d2e2f04fb..21cfbdcd2 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/mla_attention.py +++ b/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 diff --git a/test/registered/dcp/test_dcp_layout_unit.py b/test/registered/dcp/test_dcp_layout_unit.py index bf445bf1e..3973443cf 100644 --- a/test/registered/dcp/test_dcp_layout_unit.py +++ b/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 diff --git a/test/registered/amd/test_qwen3p5_triton_dcp.py b/test/registered/dcp/test_qwen3p5_triton_dcp.py similarity index 83% rename from test/registered/amd/test_qwen3p5_triton_dcp.py rename to test/registered/dcp/test_qwen3p5_triton_dcp.py index 4536dd012..9a6fe4d33 100644 --- a/test/registered/amd/test_qwen3p5_triton_dcp.py +++ b/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) diff --git a/test/registered/unit/model_executor/test_pool_configurator.py b/test/registered/unit/model_executor/test_pool_configurator.py index 3f3ee96cd..b0f58df64 100644 --- a/test/registered/unit/model_executor/test_pool_configurator.py +++ b/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