minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: Khoa Pham <khoa.pham@radixark.ai>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
YAMY
2026-08-08 14:25:09 -07:00
committed by GitHub
co-authored by Khoa Pham Claude Opus 5
parent c2d90db1e3
commit a59bb931c6
17 changed files with 172 additions and 63 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/configs/model_config.py
+11 -7
View File
@@ -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."""
python/sglang/srt/layers/attention/flashinfer_backend.py
+3 -3
View File
@@ -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
+1 -1
View File
@@ -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
+20 -8
View File
@@ -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:
python/sglang/srt/mem_cache/kv_cache_configurator.py
+30 -10
View File
@@ -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
+3 -3
View File
@@ -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
+13 -4
View File
@@ -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),
)
python/sglang/test/kits/attention_unittest/attention_methods/dense_attention.py
+4 -3
View File
@@ -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):
python/sglang/test/kits/attention_unittest/attention_methods/dual_chunk_attention.py
+4 -3
View File
@@ -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):
python/sglang/test/kits/attention_unittest/attention_methods/gdn_attention.py
+4 -3
View File
@@ -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):
python/sglang/test/kits/attention_unittest/attention_methods/kda_attention.py
+4 -3
View File
@@ -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):
python/sglang/test/kits/attention_unittest/attention_methods/lightning_attention.py
+4 -3
View File
@@ -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):
python/sglang/test/kits/attention_unittest/attention_methods/mamba2_attention.py
+4 -3
View File
@@ -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):
python/sglang/test/kits/attention_unittest/attention_methods/mla_attention.py
+1 -1
View File
@@ -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