minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 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
test/registered/dcp/test_dcp_layout_unit.py
+54
View File
@@ -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/amd/test_qwen3p5_triton_dcp.py → test/registered/dcp/test_qwen3p5_triton_dcp.py
+11 -7
View File
@@ -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
+1 -1
View File
@@ -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