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

[PD] Add PD support for hybrid model (Qwen3-Next, DeepSeek V3.2 Exp) (#10912)

Browse Source 
Signed-off-by: Shangming Cai <csmthu@gmail.com>
Co-authored-by: hzh0425 <hzh0425@apache.org>
Co-authored-by: ZeldaHuang <hzm414167@alibaba-inc.com>
This commit is contained in:
Shangming Cai
2025-10-16 09:59:14 +08:00
committed by GitHub
parent 97d857c096
commit 868403f642
13 changed files with 730 additions and 189 deletions
Download Patch File Download Diff File Expand all files Collapse all files

391
python/sglang/srt/mem_cache/memory_pool.py
View File

@@ -142,72 +142,93 @@ class MambaPool:
ssm_dtype = cache_params.dtype.temporal
num_mamba_layers = len(cache_params.layers)
# assume conv_state = (dim, state_len)
assert conv_state_shape[0] > conv_state_shape[1]
conv_state = torch.zeros(
size=(num_mamba_layers, size + 1) + conv_state_shape,
dtype=conv_dtype,
device=device,
# for disagg with nvlink
self.enable_custom_mem_pool = get_bool_env_var(
"SGLANG_MOONCAKE_CUSTOM_MEM_POOL", "false"
)
temporal_state = torch.zeros(
size=(num_mamba_layers, size + 1) + temporal_state_shape,
dtype=ssm_dtype,
device=device,
)
if speculative_num_draft_tokens is not None:
# Cache intermediate SSM states per draft token during target verify
# Shape: [num_layers, size + 1, speculative_num_draft_tokens, HV, K, V]
intermediate_ssm_state_cache = torch.zeros(
size=(
num_mamba_layers,
size + 1,
speculative_num_draft_tokens,
temporal_state_shape[0],
temporal_state_shape[1],
temporal_state_shape[2],
),
dtype=ssm_dtype,
device="cuda",
)
# Cache intermediate conv windows (last K-1 inputs) per draft token during target verify
# Shape: [num_layers, size + 1, speculative_num_draft_tokens, dim, K-1]
intermediate_conv_window_cache = torch.zeros(
size=(
num_mamba_layers,
size + 1,
speculative_num_draft_tokens,
conv_state_shape[0],
conv_state_shape[1],
),
dtype=conv_dtype,
device="cuda",
)
self.mamba_cache = self.SpeculativeState(
conv=conv_state,
temporal=temporal_state,
intermediate_ssm=intermediate_ssm_state_cache,
intermediate_conv_window=intermediate_conv_window_cache,
)
logger.info(
f"Mamba Cache is allocated. "
f"max_mamba_cache_size: {size}, "
f"conv_state size: {get_tensor_size_bytes(conv_state) / GB:.2f}GB, "
f"ssm_state size: {get_tensor_size_bytes(temporal_state) / GB:.2f}GB "
f"intermediate_ssm_state_cache size: {get_tensor_size_bytes(intermediate_ssm_state_cache) / GB:.2f}GB "
f"intermediate_conv_window_cache size: {get_tensor_size_bytes(intermediate_conv_window_cache) / GB:.2f}GB "
)
if self.enable_custom_mem_pool:
# TODO(shangming): abstract custom allocator class for more backends
from mooncake.allocator import NVLinkAllocator
allocator = NVLinkAllocator.get_allocator(self.device)
self.custom_mem_pool = torch.cuda.MemPool(allocator.allocator())
else:
self.mamba_cache = self.State(conv=conv_state, temporal=temporal_state)
logger.info(
f"Mamba Cache is allocated. "
f"max_mamba_cache_size: {size}, "
f"conv_state size: {get_tensor_size_bytes(conv_state) / GB:.2f}GB, "
f"ssm_state size: {get_tensor_size_bytes(temporal_state) / GB:.2f}GB "
self.custom_mem_pool = None
with (
torch.cuda.use_mem_pool(self.custom_mem_pool)
if self.enable_custom_mem_pool
else nullcontext()
):
# assume conv_state = (dim, state_len)
assert conv_state_shape[0] > conv_state_shape[1]
conv_state = torch.zeros(
size=(num_mamba_layers, size + 1) + conv_state_shape,
dtype=conv_dtype,
device=device,
)
self.size = size
self.device = device
self.free_slots = torch.arange(self.size, dtype=torch.int64, device=self.device)
self.mem_usage = self.mamba_cache.mem_usage_bytes() / GB
temporal_state = torch.zeros(
size=(num_mamba_layers, size + 1) + temporal_state_shape,
dtype=ssm_dtype,
device=device,
)
if speculative_num_draft_tokens is not None:
# Cache intermediate SSM states per draft token during target verify
# Shape: [num_layers, size + 1, speculative_num_draft_tokens, HV, K, V]
intermediate_ssm_state_cache = torch.zeros(
size=(
num_mamba_layers,
size + 1,
speculative_num_draft_tokens,
temporal_state_shape[0],
temporal_state_shape[1],
temporal_state_shape[2],
),
dtype=ssm_dtype,
device="cuda",
)
# Cache intermediate conv windows (last K-1 inputs) per draft token during target verify
# Shape: [num_layers, size + 1, speculative_num_draft_tokens, dim, K-1]
intermediate_conv_window_cache = torch.zeros(
size=(
num_mamba_layers,
size + 1,
speculative_num_draft_tokens,
conv_state_shape[0],
conv_state_shape[1],
),
dtype=conv_dtype,
device="cuda",
)
self.mamba_cache = self.SpeculativeState(
conv=conv_state,
temporal=temporal_state,
intermediate_ssm=intermediate_ssm_state_cache,
intermediate_conv_window=intermediate_conv_window_cache,
)
logger.info(
f"Mamba Cache is allocated. "
f"max_mamba_cache_size: {size}, "
f"conv_state size: {get_tensor_size_bytes(conv_state) / GB:.2f}GB, "
f"ssm_state size: {get_tensor_size_bytes(temporal_state) / GB:.2f}GB "
f"intermediate_ssm_state_cache size: {get_tensor_size_bytes(intermediate_ssm_state_cache) / GB:.2f}GB "
f"intermediate_conv_window_cache size: {get_tensor_size_bytes(intermediate_conv_window_cache) / GB:.2f}GB "
)
else:
self.mamba_cache = self.State(conv=conv_state, temporal=temporal_state)
logger.info(
f"Mamba Cache is allocated. "
f"max_mamba_cache_size: {size}, "
f"conv_state size: {get_tensor_size_bytes(conv_state) / GB:.2f}GB, "
f"ssm_state size: {get_tensor_size_bytes(temporal_state) / GB:.2f}GB "
)
self.size = size
self.device = device
self.free_slots = torch.arange(
self.size, dtype=torch.int64, device=self.device
)
self.mem_usage = self.mamba_cache.mem_usage_bytes() / GB
self.num_mamba_layers = num_mamba_layers
def get_speculative_mamba2_params_all_layers(self) -> SpeculativeState:
assert isinstance(self.mamba_cache, self.SpeculativeState)
@@ -253,6 +274,22 @@ class MambaPool:
self.copy_from(src_index, dst_index)
return dst_index
def get_contiguous_buf_infos(self):
state_tensors = [
getattr(self.mamba_cache, field) for field in vars(self.mamba_cache)
]
data_ptrs, data_lens, item_lens = [], [], []
for _, state_tensor in enumerate(state_tensors):
data_ptrs += [
state_tensor[i].data_ptr() for i in range(self.num_mamba_layers)
]
data_lens += [state_tensor[i].nbytes for i in range(self.num_mamba_layers)]
item_lens += [
state_tensor[i][0].nbytes for i in range(self.num_mamba_layers)
]
return data_ptrs, data_lens, item_lens
class HybridReqToTokenPool(ReqToTokenPool):
"""A memory pool that maps a request to its token locations."""
@@ -274,13 +311,26 @@ class HybridReqToTokenPool(ReqToTokenPool):
device=device,
enable_memory_saver=enable_memory_saver,
)
self.mamba_pool = MambaPool(
self._init_mamba_pool(
size=mamba_size,
cache_params=cache_params,
device=device,
speculative_num_draft_tokens=speculative_num_draft_tokens,
)
def _init_mamba_pool(
self,
size: int,
cache_params: "Mamba2CacheParams",
device: str,
speculative_num_draft_tokens: int = None,
):
self.mamba_pool = MambaPool(
size=size,
cache_params=cache_params,
device=device,
speculative_num_draft_tokens=speculative_num_draft_tokens,
)
self.mamba_map = {layer_id: i for i, layer_id in enumerate(cache_params.layers)}
self.device = device
@@ -375,6 +425,19 @@ class KVCache(abc.ABC):
# default state for optional layer-wise transfer control
self.layer_transfer_counter = None
# for disagg with nvlink
self.enable_custom_mem_pool = get_bool_env_var(
"SGLANG_MOONCAKE_CUSTOM_MEM_POOL", "false"
)
if self.enable_custom_mem_pool:
# TODO(shangming): abstract custom allocator class for more backends
from mooncake.allocator import NVLinkAllocator
allocator = NVLinkAllocator.get_allocator(self.device)
self.custom_mem_pool = torch.cuda.MemPool(allocator.allocator())
else:
self.custom_mem_pool = None
def _finalize_allocation_log(self, num_tokens: int):
"""Common logging and mem_usage computation for KV cache allocation.
Supports both tuple (K, V) size returns and single KV size returns.
@@ -426,6 +489,9 @@ class KVCache(abc.ABC):
def load_cpu_copy(self, kv_cache_cpu, indices):
raise NotImplementedError()
def maybe_get_custom_mem_pool(self):
return self.custom_mem_pool
class MHATokenToKVPool(KVCache):
@@ -456,19 +522,6 @@ class MHATokenToKVPool(KVCache):
self.head_num = head_num
self.head_dim = head_dim
# for disagg with nvlink
self.enable_custom_mem_pool = get_bool_env_var(
"SGLANG_MOONCAKE_CUSTOM_MEM_POOL", "false"
)
if self.enable_custom_mem_pool:
# TODO(shangming): abstract custom allocator class for more backends
from mooncake.allocator import NVLinkAllocator
allocator = NVLinkAllocator.get_allocator(self.device)
self.custom_mem_pool = torch.cuda.MemPool(allocator.allocator())
else:
self.custom_mem_pool = None
self._create_buffers()
self.device_module = torch.get_device_module(self.device)
@@ -611,9 +664,6 @@ class MHATokenToKVPool(KVCache):
]
return kv_data_ptrs, kv_data_lens, kv_item_lens
def maybe_get_custom_mem_pool(self):
return self.custom_mem_pool
def get_cpu_copy(self, indices):
torch.cuda.synchronize()
kv_cache_cpu = []
@@ -756,12 +806,18 @@ class HybridLinearKVPool(KVCache):
full_attention_layer_ids: List[int],
enable_kvcache_transpose: bool,
device: str,
mamba_pool: MambaPool,
):
self.size = size
self.dtype = dtype
self.device = device
self.full_layer_nums = len(full_attention_layer_ids)
self.page_size = page_size
# TODO support pp?
self.start_layer = 0
self.head_num = head_num
self.head_dim = head_dim
self.mamba_pool = mamba_pool
# TODO MHATransposedTokenToKVPool if enable_kvcache_transpose is True
assert not enable_kvcache_transpose
if _is_npu:
@@ -790,6 +846,15 @@ class HybridLinearKVPool(KVCache):
def get_contiguous_buf_infos(self):
return self.full_kv_pool.get_contiguous_buf_infos()
def get_state_buf_infos(self):
mamba_data_ptrs, mamba_data_lens, mamba_item_lens = (
self.mamba_pool.get_contiguous_buf_infos()
)
return mamba_data_ptrs, mamba_data_lens, mamba_item_lens
def maybe_get_custom_mem_pool(self):
return self.full_kv_pool.maybe_get_custom_mem_pool()
def _transfer_full_attention_id(self, layer_id: int):
if layer_id not in self.full_attention_layer_id_mapping:
raise ValueError(
@@ -841,22 +906,47 @@ class SWAKVPool(KVCache):
size: int,
size_swa: int,
dtype: torch.dtype,
head_num: int,
head_dim: int,
swa_attention_layer_ids: List[int],
full_attention_layer_ids: List[int],
enable_kvcache_transpose: bool,
device: str,
token_to_kv_pool_class: KVCache = MHATokenToKVPool,
**kwargs,
):
self.size = size
self.size_swa = size_swa
self.dtype = dtype
self.head_num = head_num
self.head_dim = head_dim
self.device = device
self.swa_layer_nums = len(swa_attention_layer_ids)
self.full_layer_nums = len(full_attention_layer_ids)
self.start_layer = 0
self.page_size = 1
kwargs["page_size"] = 1
kwargs["enable_memory_saver"] = False
kwargs["head_num"] = head_num
kwargs["head_dim"] = head_dim
kwargs["device"] = device
# TODO MHATransposedTokenToKVPool if enable_kvcache_transpose is True
assert not enable_kvcache_transpose
# for disagg with nvlink
self.enable_custom_mem_pool = get_bool_env_var(
"SGLANG_MOONCAKE_CUSTOM_MEM_POOL", "false"
)
if self.enable_custom_mem_pool:
# TODO(shangming): abstract custom allocator class for more backends
from mooncake.allocator import NVLinkAllocator
allocator = NVLinkAllocator.get_allocator(self.device)
self.custom_mem_pool = torch.cuda.MemPool(allocator.allocator())
else:
self.custom_mem_pool = None
self.swa_kv_pool = token_to_kv_pool_class(
size=size_swa,
dtype=dtype,
@@ -878,6 +968,9 @@ class SWAKVPool(KVCache):
k_size, v_size = self.get_kv_size_bytes()
self.mem_usage = (k_size + v_size) / GB
logger.info(
f"SWAKVPool mem usage: {self.mem_usage} GB, swa size: {self.size_swa}, full size: {self.size}"
)
def get_kv_size_bytes(self):
k_size, v_size = self.full_kv_pool.get_kv_size_bytes()
@@ -888,15 +981,19 @@ class SWAKVPool(KVCache):
full_kv_data_ptrs, full_kv_data_lens, full_kv_item_lens = (
self.full_kv_pool.get_contiguous_buf_infos()
)
kv_data_ptrs = full_kv_data_ptrs
kv_data_lens = full_kv_data_lens
kv_item_lens = full_kv_item_lens
return kv_data_ptrs, kv_data_lens, kv_item_lens
def get_state_buf_infos(self):
swa_kv_data_ptrs, swa_kv_data_lens, swa_kv_item_lens = (
self.swa_kv_pool.get_contiguous_buf_infos()
)
kv_data_ptrs = full_kv_data_ptrs + swa_kv_data_ptrs
kv_data_lens = full_kv_data_lens + swa_kv_data_lens
kv_item_lens = full_kv_item_lens + swa_kv_item_lens
return kv_data_ptrs, kv_data_lens, kv_item_lens
return swa_kv_data_ptrs, swa_kv_data_lens, swa_kv_item_lens
def get_key_buffer(self, layer_id: int):
layer_id_pool, is_swa = self.layers_mapping[layer_id]
@@ -1152,19 +1249,6 @@ class MLATokenToKVPool(KVCache):
else (kv_lora_rank + qk_rope_head_dim)
)
# for disagg with nvlink
self.enable_custom_mem_pool = get_bool_env_var(
"SGLANG_MOONCAKE_CUSTOM_MEM_POOL", "false"
)
if self.enable_custom_mem_pool:
# TODO(shangming): abstract custom allocator class for more backends
from mooncake.allocator import NVLinkAllocator
allocator = NVLinkAllocator.get_allocator(self.device)
self.custom_mem_pool = torch.cuda.MemPool(allocator.allocator())
else:
self.custom_mem_pool = None
with self.memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE):
with (
torch.cuda.use_mem_pool(self.custom_mem_pool)
@@ -1207,9 +1291,6 @@ class MLATokenToKVPool(KVCache):
]
return kv_data_ptrs, kv_data_lens, kv_item_lens
def maybe_get_custom_mem_pool(self):
return self.custom_mem_pool
def get_key_buffer(self, layer_id: int):
if self.layer_transfer_counter is not None:
self.layer_transfer_counter.wait_until(layer_id - self.start_layer)
@@ -1346,24 +1427,31 @@ class NSATokenToKVPool(MLATokenToKVPool):
assert index_head_dim == 128
assert self.page_size == 64
self.index_k_with_scale_buffer = [
torch.zeros(
# Layout:
# ref: test_attention.py :: kv_cache_cast_to_fp8
# shape: (num_pages, page_size 64 * head_dim 128 + page_size 64 * fp32_nbytes 4)
# data: for page i,
# * buf[i, :page_size * head_dim] for fp8 data
# * buf[i, page_size * head_dim:].view(float32) for scale
(
(size + page_size + 1) // self.page_size,
self.page_size
* (index_head_dim + index_head_dim // self.quant_block_size * 4),
),
dtype=self.index_k_with_scale_buffer_dtype,
device=device,
)
for _ in range(layer_num)
]
with (
torch.cuda.use_mem_pool(self.custom_mem_pool)
if self.custom_mem_pool
else nullcontext()
):
self.index_k_with_scale_buffer = [
torch.zeros(
# Layout:
# ref: test_attention.py :: kv_cache_cast_to_fp8
# shape: (num_pages, page_size 64 * head_dim 128 + page_size 64 * fp32_nbytes 4)
# data: for page i,
# * buf[i, :page_size * head_dim] for fp8 data
# * buf[i, page_size * head_dim:].view(float32) for scale
(
(size + page_size + 1) // self.page_size,
self.page_size
* (
index_head_dim + index_head_dim // self.quant_block_size * 4
),
),
dtype=self.index_k_with_scale_buffer_dtype,
device=device,
)
for _ in range(layer_num)
]
self._finalize_allocation_log(size)
def get_index_k_with_scale_buffer(self, layer_id: int) -> torch.Tensor:
@@ -1406,6 +1494,18 @@ class NSATokenToKVPool(MLATokenToKVPool):
pool=self, buf=buf, loc=loc, index_k=index_k, index_k_scale=index_k_scale
)
def get_state_buf_infos(self):
data_ptrs = [
self.index_k_with_scale_buffer[i].data_ptr() for i in range(self.layer_num)
]
data_lens = [
self.index_k_with_scale_buffer[i].nbytes for i in range(self.layer_num)
]
item_lens = [
self.index_k_with_scale_buffer[i][0].nbytes for i in range(self.layer_num)
]
return data_ptrs, data_lens, item_lens
def get_kv_size_bytes(self):
kv_size_bytes = super().get_kv_size_bytes()
for index_k_cache in self.index_k_with_scale_buffer:
@@ -1636,27 +1736,38 @@ class DoubleSparseTokenToKVPool(KVCache):
)
with self.memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE):
# [size, head_num, head_dim] for each layer
self.k_buffer = [
torch.zeros(
(size + page_size, head_num, head_dim), dtype=dtype, device=device
)
for _ in range(layer_num)
]
self.v_buffer = [
torch.zeros(
(size + page_size, head_num, head_dim), dtype=dtype, device=device
)
for _ in range(layer_num)
]
with (
torch.cuda.use_mem_pool(self.custom_mem_pool)
if self.enable_custom_mem_pool
else nullcontext()
):
# [size, head_num, head_dim] for each layer
self.k_buffer = [
torch.zeros(
(size + page_size, head_num, head_dim),
dtype=dtype,
device=device,
)
for _ in range(layer_num)
]
self.v_buffer = [
torch.zeros(
(size + page_size, head_num, head_dim),
dtype=dtype,
device=device,
)
for _ in range(layer_num)
]
# [size, head_num, heavy_channel_num] for each layer
self.label_buffer = [
torch.zeros(
(size + 1, head_num, heavy_channel_num), dtype=dtype, device=device
)
for _ in range(layer_num)
]
# [size, head_num, heavy_channel_num] for each layer
self.label_buffer = [
torch.zeros(
(size + 1, head_num, heavy_channel_num),
dtype=dtype,
device=device,
)
for _ in range(layer_num)
]
def get_key_buffer(self, layer_id: int):
return self.k_buffer[layer_id - self.start_layer]