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

Add KV-canary SWA + DeepSeek-V4 pool support (#26810)

Browse Source
This commit is contained in:
fzyzcjy
2026-05-31 09:55:34 +08:00
committed by GitHub
parent 9ecf314970
commit e188745ee0
10 changed files with 334 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/kv_canary/pool_patcher/adapters/dsv4.py
+45
View File
@@ -0,0 +1,45 @@
from __future__ import annotations
import torch
from sglang.srt.kv_canary.buffer_group import CanaryBufferGroup, PoolKind
from sglang.srt.kv_canary.pool_patcher.buf_info_splice import patch_buf_info_method
from sglang.srt.kv_canary.pool_patcher.buffer_alloc import alloc_canary_buf
def attach_dsv4(
*,
pool: object,
device: torch.device,
kv_token_id_vs_position_offset: int,
) -> tuple[CanaryBufferGroup, ...]:
"""Attach canary buffers to a DeepSeekV4TokenToKVPool.
TODO: only the swa_kv_pool sub-pool is wired; c4_kv_pool / c128_kv_pool /
c4_indexer_kv_pool / compress state pools are left uncovered.
"""
sub_pool = pool.swa_kv_pool
num_slots = int(sub_pool.size)
k_head = alloc_canary_buf(num_slots=num_slots, device=device)
k_tail = alloc_canary_buf(num_slots=num_slots, device=device)
group = CanaryBufferGroup(
kind=PoolKind.SWA,
k_head=k_head,
k_tail=k_tail,
v_head=None,
v_tail=None,
swa_index_lut=pool.full_to_swa_index_mapping,
kv_token_id_vs_position_offset=kv_token_id_vs_position_offset,
)
patch_buf_info_method(
pool,
method_name="get_state_buf_infos",
group=group,
has_v_half=False,
page_size=sub_pool.page_size,
)
return (group,)
python/sglang/srt/kv_canary/pool_patcher/adapters/swa.py
+71
View File
@@ -0,0 +1,71 @@
from __future__ import annotations
from typing import Optional
import torch
from sglang.srt.kv_canary.buffer_group import CanaryBufferGroup, PoolKind
from sglang.srt.kv_canary.pool_patcher.buf_info_splice import patch_buf_info_method
from sglang.srt.kv_canary.pool_patcher.buffer_alloc import alloc_canary_buf
def attach_swa(
*,
pool: object,
device: torch.device,
kv_token_id_vs_position_offset: int,
) -> tuple[CanaryBufferGroup, ...]:
full_group = _build_subpool_group(
sub_pool=pool.full_kv_pool,
kind=PoolKind.FULL,
device=device,
swa_lut=None,
kv_token_id_vs_position_offset=kv_token_id_vs_position_offset,
)
swa_group = _build_subpool_group(
sub_pool=pool.swa_kv_pool,
kind=PoolKind.SWA,
device=device,
swa_lut=pool.full_to_swa_index_mapping,
kv_token_id_vs_position_offset=kv_token_id_vs_position_offset,
)
patch_buf_info_method(
pool,
method_name="get_contiguous_buf_infos",
group=full_group,
has_v_half=True,
page_size=pool.page_size,
)
patch_buf_info_method(
pool,
method_name="get_state_buf_infos",
group=swa_group,
has_v_half=True,
page_size=pool.page_size,
)
return (full_group, swa_group)
def _build_subpool_group(
*,
sub_pool: object,
kind: PoolKind,
device: torch.device,
swa_lut: Optional[torch.Tensor],
kv_token_id_vs_position_offset: int,
) -> CanaryBufferGroup:
num_slots = int(sub_pool.k_buffer[0].shape[0])
k_head = alloc_canary_buf(num_slots=num_slots, device=device)
k_tail = alloc_canary_buf(num_slots=num_slots, device=device)
v_head = alloc_canary_buf(num_slots=num_slots, device=device)
v_tail = alloc_canary_buf(num_slots=num_slots, device=device)
return CanaryBufferGroup(
kind=kind,
k_head=k_head,
k_tail=k_tail,
v_head=v_head,
v_tail=v_tail,
swa_index_lut=swa_lut,
kv_token_id_vs_position_offset=kv_token_id_vs_position_offset,
)
python/sglang/srt/kv_canary/pool_patcher/api.py
+6
View File
@@ -7,12 +7,16 @@ import torch
from sglang.srt.kv_canary.buffer_group import CanaryBufferGroup
from sglang.srt.kv_canary.config import CanaryConfig
from sglang.srt.kv_canary.pool_patcher.adapters.dsv4 import attach_dsv4
from sglang.srt.kv_canary.pool_patcher.adapters.mha import attach_mha
from sglang.srt.kv_canary.pool_patcher.adapters.swa import attach_swa
from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
from sglang.srt.mem_cache.memory_pool import (
KVCache,
MHATokenToKVPool,
MHATokenToKVPoolFP4,
)
from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
logger = logging.getLogger(__name__)
@@ -21,6 +25,8 @@ PoolAttacher = Callable[..., tuple[CanaryBufferGroup, ...]]
_POOL_ATTACHERS: Dict[Type, PoolAttacher] = {
MHATokenToKVPool: attach_mha,
MHATokenToKVPoolFP4: attach_mha,
SWAKVPool: attach_swa,
DeepSeekV4TokenToKVPool: attach_dsv4,
}
python/sglang/test/kv_canary/consts.py
+52
View File
@@ -0,0 +1,52 @@
from __future__ import annotations
from typing import Final
# SWA e2e pool sizing for 8 reqs × ~7K prompt + 2K decode, SWA window 1024.
# FULL pool = max-total-tokens; must fit 8 × (7000 + 2048) = 72_384 to avoid preempt.
# SWA pool = max-total-tokens × ratio;
# ≥ 8 × 1024 = 8192 (else deadlock — in-flight footprint exceeds capacity);
# < 8 × 7000 = 56_000 (else allocator never recycles → swa_full_idx_divergence stays 0).
# Pick FULL=81920 (72_384 + headroom), SWA=16384 (2× in-flight floor, well under 56K).
SWA_POOL_SERVER_ARGS: Final[tuple[str, ...]] = (
"--max-total-tokens",
"81920",
"--swa-full-tokens-ratio",
"0.2",
)
DSV4_DEEPEP_CONFIG: Final[str] = (
'{"normal_dispatch":{"num_sms":96},"normal_combine":{"num_sms":96}}'
)
DSV4_POOL_SERVER_ARGS: Final[tuple[str, ...]] = (
"--trust-remote-code",
"--tp",
"4",
"--dp",
"4",
"--enable-dp-attention",
"--moe-a2a-backend",
"deepep",
"--cuda-graph-max-bs",
"128",
"--max-running-requests",
"256",
"--deepep-config",
DSV4_DEEPEP_CONFIG,
"--mem-fraction-static",
"0.7",
"--speculative-algorithm",
"EAGLE",
"--speculative-num-steps",
"3",
"--speculative-eagle-topk",
"1",
"--speculative-num-draft-tokens",
"4",
)
DSV4_POOL_SERVER_ENV: Final[dict[str, str]] = {
"SGLANG_DSV4_FP4_EXPERTS": "0",
"SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "1024",
}
python/sglang/test/kv_canary/e2e_base.py
+13 -3
View File
@@ -18,7 +18,9 @@ from sglang.test.test_utils import (
)
# Long prompt body shared by all canary e2e tests. The repetition count is chosen
# so the tokenised prompt is comfortably long; token count is roughly 7k after BPE.
# so the tokenised prompt comfortably exceeds the SWA sliding window of swa-mode
# fixtures (gemma-4-E2B); short prompts would never exercise the SWA-windowed
# verify path. Token count is roughly 7k after BPE.
_LONG_PROMPT_BODY = ("The quick brown fox jumps over the lazy dog. " * 700).strip()
_UNIQUE_PROMPT_FIRST_CHARS = string.ascii_letters + string.digits
@@ -60,12 +62,14 @@ class CapturedServerE2EBase(CanaryViolationAssertMixin, CustomTestCase):
class CanaryE2EBase(CapturedServerE2EBase):
model_mode: ClassVar[Literal["mha"]]
model_mode: ClassVar[Literal["mha", "swa", "dsv4"]]
kv_canary_mode: ClassVar[CanaryMode]
extra_env: ClassVar[dict[str, str]] = {}
extra_server_args: ClassVar[tuple[str, ...]] = ()
use_unique_prompts: ClassVar[bool] = False
# Number of sequential request batches each test method sends. Default 1 keeps tests fast.
# SWA divergence assertions need slot recycling across batches; setting > 1 makes the
# test methods send N sequential batches so the SWA allocator's full→swa index mapping
# diverges from identity. Default 1 keeps MHA tests fast.
workload_n_batches: ClassVar[int] = 1
_cfg: ClassVar[Optional[_ModeConfig]] = None
@@ -75,6 +79,12 @@ class CanaryE2EBase(CapturedServerE2EBase):
cls._cfg = _MODE_CONFIGS[cls.model_mode]
server_env = os.environ.copy()
server_env.update(cls.extra_env)
if cls.model_mode == "swa":
# SWA mode uses google/gemma-4-E2B-it, whose forward does a
# ``positions += 1`` in-place. canary's WRITE/VERIFY require
# forward_batch.positions to stay 0-indexed, so flip the gemma
# path to out-of-place shift for these tests.
server_env.setdefault("SGLANG_GEMMA_OUT_OF_PLACE_POSITION_MUTATION", "1")
cls._stdout_buf = io.StringIO()
cls._stderr_buf = io.StringIO()
python/sglang/test/kv_canary/fixtures.py
+79
View File
@@ -10,6 +10,7 @@ from sglang.jit_kernel.kv_canary.verify import CANARY_SLOT_BYTES
from sglang.srt.kv_canary.buffer_group import CanaryBufferGroup, PoolKind
from sglang.srt.kv_canary.config import CanaryConfig, CanaryMode
from sglang.srt.kv_canary.pool_patcher.adapters.mha import attach_mha
from sglang.srt.kv_canary.pool_patcher.adapters.swa import attach_swa
from sglang.srt.kv_canary.pool_patcher.api import register_pool_attacher
from sglang.srt.mem_cache.radix_cache import RadixCache, TreeNode
from sglang.srt.model_executor.forward_batch_info import ForwardMode
@@ -35,6 +36,48 @@ class FakeMHAPool:
return ptrs, lens, item_lens
@dataclass
class FakeSwaSubPool:
k_buffer: List[torch.Tensor]
v_buffer: List[torch.Tensor]
@dataclass
class FakeSWAPool:
full_kv_pool: object
swa_kv_pool: object
full_to_swa_index_mapping: torch.Tensor
page_size: int = 1
def get_contiguous_buf_infos(self):
return _kv_buf_infos(
k_buffer=self.full_kv_pool.k_buffer,
v_buffer=self.full_kv_pool.v_buffer,
page_size=self.page_size,
)
def get_state_buf_infos(self):
return _kv_buf_infos(
k_buffer=self.swa_kv_pool.k_buffer,
v_buffer=self.swa_kv_pool.v_buffer,
page_size=self.page_size,
)
def _kv_buf_infos(
*,
k_buffer: List[torch.Tensor],
v_buffer: List[torch.Tensor],
page_size: int,
) -> tuple:
ptrs = [b.data_ptr() for b in k_buffer] + [b.data_ptr() for b in v_buffer]
lens = [b.nbytes for b in k_buffer] + [b.nbytes for b in v_buffer]
item_lens = [b[0].nbytes * page_size for b in k_buffer] + [
b[0].nbytes * page_size for b in v_buffer
]
return ptrs, lens, item_lens
def make_mha_pool(
device: torch.device = DEFAULT_DEVICE,
*,
@@ -53,6 +96,41 @@ def make_mha_pool(
return FakeMHAPool(layer_num=layer_num, k_buffer=k_layers, v_buffer=v_layers)
def make_swa_pool(
device: torch.device = DEFAULT_DEVICE,
*,
full_slots: int = 16,
swa_slots: int = 8,
dim: int = 8,
layer_num: int = 1,
) -> FakeSWAPool:
full = FakeSwaSubPool(
k_buffer=[
torch.zeros(full_slots, dim, dtype=torch.float16, device=device)
for _ in range(layer_num)
],
v_buffer=[
torch.zeros(full_slots, dim, dtype=torch.float16, device=device)
for _ in range(layer_num)
],
)
swa = FakeSwaSubPool(
k_buffer=[
torch.zeros(swa_slots, dim, dtype=torch.float16, device=device)
for _ in range(layer_num)
],
v_buffer=[
torch.zeros(swa_slots, dim, dtype=torch.float16, device=device)
for _ in range(layer_num)
],
)
lut = torch.full((full_slots + 1,), -1, dtype=torch.int64, device=device)
lut[:swa_slots] = torch.arange(swa_slots, dtype=torch.int64, device=device)
return FakeSWAPool(
full_kv_pool=full, swa_kv_pool=swa, full_to_swa_index_mapping=lut
)
def make_base_config() -> CanaryConfig:
return CanaryConfig(
mode=CanaryMode.RAISE,
@@ -189,3 +267,4 @@ def make_radix_cache(
register_pool_attacher(FakeMHAPool, attach_mha)
register_pool_attacher(FakeSWAPool, attach_swa)
python/sglang/test/kv_canary/mode_config.py
+6
View File
@@ -14,4 +14,10 @@ _MODE_CONFIGS: dict[str, _ModeConfig] = {
"mha": _ModeConfig(
model_path="Qwen/Qwen3-0.6B",
),
"swa": _ModeConfig(
model_path="google/gemma-4-E2B-it",
),
"dsv4": _ModeConfig(
model_path="deepseek-ai/DeepSeek-V4-Flash",
),
}