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

[Qwen3-next] support mamba radix cache for overlap scheduler (#14792)

Browse Source
This commit is contained in:
Hanming Lu
2025-12-14 18:54:16 -08:00
committed by GitHub
parent 36e7c8c59f
commit e61dabf5e4
30 changed files with 1414 additions and 204 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
python/sglang/srt/model_executor/cuda_graph_runner.py
View File

@@ -321,6 +321,11 @@ class CudaGraphRunner:
num_tokens_per_bs=self.num_tokens_per_bs,
)
enable_mamba_track = (
self.model_runner.server_args.enable_mamba_extra_buffer()
and self.model_runner.spec_algorithm.is_none()
)
if self.require_gathered_buffer:
assert self.require_mlp_tp_gather or self.require_attn_tp_gather
self.buffers: GraphInputBuffers = GraphInputBuffers.create(
@@ -338,6 +343,7 @@ class CudaGraphRunner:
encoder_len_fill_value=self.encoder_len_fill_value,
num_tokens_per_bs=self.num_tokens_per_bs,
cache_loc_dtype=self._cache_loc_dtype(),
enable_mamba_track=enable_mamba_track,
)
self.tbo_plugin = TboCudaGraphRunnerPlugin()
@@ -537,7 +543,7 @@ class CudaGraphRunner:
def capture_one_batch_size(
self, bs: int, forward: Callable, stream_idx: Optional[int] = None
):
buffers = self.buffers
buffers: GraphInputBuffers = self.buffers
graph = self._create_device_graph()
stream = self.stream
num_tokens = bs * self.num_tokens_per_bs
@@ -611,6 +617,18 @@ class CudaGraphRunner:
else:
lora_ids = None
# mamba state tracking
mamba_track_indices = (
buffers.mamba_track_indices[:bs]
if buffers.mamba_track_indices is not None
else None
)
mamba_track_mask = (
buffers.mamba_track_mask[:bs]
if buffers.mamba_track_mask is not None
else None
)
if stream_idx is None:
attn_backend = self.model_runner.attn_backend
else:
@@ -631,6 +649,9 @@ class CudaGraphRunner:
attn_backend=attn_backend,
out_cache_loc=out_cache_loc,
seq_lens_sum=seq_lens.sum().item(),
mamba_track_indices=mamba_track_indices,
mamba_track_mask=mamba_track_mask,
mamba_track_seqlens=None, # Prefill only
encoder_lens=encoder_lens,
return_logprob=False,
positions=positions,

19
python/sglang/srt/model_executor/forward_batch_info.py
View File

@@ -249,6 +249,12 @@ class ForwardBatch:
# The indices of output tokens in the token_to_kv_pool_swa
# TODO(shiyang, biao): integrate out_cache_loc_swa into multiple attention backends
out_cache_loc_swa: Optional[torch.Tensor] = None
# The indices to track mamba state with
mamba_track_indices: Optional[torch.Tensor] = None # shape: [b], int64
# The mask to track mamba state if needed
mamba_track_mask: Optional[torch.Tensor] = None # shape: [b], bool
# The seqlens to track mamba state if masked, prefill only.
mamba_track_seqlens: Optional[torch.Tensor] = None # shape: [b], int64
# Optional seq_lens on cpu
seq_lens_cpu: Optional[torch.Tensor] = None
@@ -398,6 +404,9 @@ class ForwardBatch:
req_pool_indices=batch.req_pool_indices,
seq_lens=batch.seq_lens,
out_cache_loc=batch.out_cache_loc,
mamba_track_indices=batch.mamba_track_indices,
mamba_track_mask=batch.mamba_track_mask,
mamba_track_seqlens=batch.mamba_track_seqlens,
mm_inputs=batch.multimodal_inputs,
encoder_cached=batch.encoder_cached,
encoder_lens=batch.encoder_lens,
@@ -881,6 +890,16 @@ class ForwardBatch:
if self.encoder_lens is not None:
self.encoder_lens = self._pad_tensor_to_size(self.encoder_lens, bs)
self.positions = self._pad_tensor_to_size(self.positions, num_tokens)
if self.mamba_track_indices is not None:
self.mamba_track_indices = self._pad_tensor_to_size(
self.mamba_track_indices, bs
)
if self.mamba_track_mask is not None:
self.mamba_track_mask = self._pad_tensor_to_size(self.mamba_track_mask, bs)
if self.mamba_track_seqlens is not None:
self.mamba_track_seqlens = self._pad_tensor_to_size(
self.mamba_track_seqlens, bs
)
if self.mrope_positions is not None:
self.mrope_positions = self._pad_tensor_to_size(self.mrope_positions, bs)

28
python/sglang/srt/model_executor/input_buffers.py
View File

@@ -25,6 +25,8 @@ class GraphInputBuffers:
num_token_non_padded: torch.Tensor
custom_mask: torch.Tensor
next_token_logits_buffer: torch.Tensor
mamba_track_indices: Optional[torch.Tensor]
mamba_track_mask: Optional[torch.Tensor]
global_num_tokens_gpu: torch.Tensor
global_num_tokens_for_logprob_gpu: torch.Tensor
encoder_lens: Optional[torch.Tensor]
@@ -48,6 +50,7 @@ class GraphInputBuffers:
encoder_len_fill_value: int,
num_tokens_per_bs: int,
cache_loc_dtype: torch.dtype,
enable_mamba_track: bool,
) -> "GraphInputBuffers":
with torch.device(device):
input_ids = torch.zeros((max_num_token,), dtype=torch.int64)
@@ -66,6 +69,14 @@ class GraphInputBuffers:
(max_num_token, vocab_size),
dtype=torch.float,
)
mamba_track_indices = (
torch.zeros((max_bs,), dtype=torch.int64)
if enable_mamba_track
else None
)
mamba_track_mask = (
torch.zeros((max_bs,), dtype=torch.bool) if enable_mamba_track else None
)
if pp_size > 1:
pp_proxy_tensors = {
@@ -111,6 +122,8 @@ class GraphInputBuffers:
num_token_non_padded=num_token_non_padded,
custom_mask=custom_mask,
next_token_logits_buffer=next_token_logits_buffer,
mamba_track_indices=mamba_track_indices,
mamba_track_mask=mamba_track_mask,
encoder_lens=encoder_lens,
global_num_tokens_gpu=global_num_tokens_gpu,
global_num_tokens_for_logprob_gpu=global_num_tokens_for_logprob_gpu,
@@ -134,6 +147,10 @@ class GraphInputBuffers:
if bs != raw_bs:
self.seq_lens.fill_(seq_len_fill_value)
self.out_cache_loc.zero_()
if self.mamba_track_indices is not None:
self.mamba_track_indices.zero_()
if self.mamba_track_mask is not None:
self.mamba_track_mask.fill_(False)
# Common inputs
self.input_ids[:raw_num_token].copy_(forward_batch.input_ids)
@@ -142,6 +159,17 @@ class GraphInputBuffers:
self.out_cache_loc[:raw_num_token].copy_(forward_batch.out_cache_loc)
self.positions[:raw_num_token].copy_(forward_batch.positions)
if (
self.mamba_track_indices is not None
and forward_batch.mamba_track_indices is not None
):
self.mamba_track_indices[:raw_bs].copy_(forward_batch.mamba_track_indices)
if (
self.mamba_track_mask is not None
and forward_batch.mamba_track_mask is not None
):
self.mamba_track_mask[:raw_bs].copy_(forward_batch.mamba_track_mask)
seq_lens_cpu: Optional[torch.Tensor] = None
if forward_batch.seq_lens_cpu is not None:
if bs != raw_bs:

61
python/sglang/srt/model_executor/model_runner.py
View File

@@ -238,8 +238,10 @@ def add_chunked_prefix_cache_attention_backend(backend_name):
# Detect stragger ranks in model loading
UNBALANCED_MODEL_LOADING_TIMEOUT_S = 480 # leave more time for post data processing
# the ratio of mamba cache pool size to max_running_requests, it will be safe when it is larger than 2 (yizhang2077)
# the ratio of mamba cache pool size to max_running_requests
MAMBA_CACHE_SIZE_MAX_RUNNING_REQUESTS_RATIO = 3
MAMBA_CACHE_V2_ADDITIONAL_RATIO_OVERLAP = 2
MAMBA_CACHE_V2_ADDITIONAL_RATIO_NO_OVERLAP = 1
logger = logging.getLogger(__name__)
@@ -1446,14 +1448,9 @@ class ModelRunner:
server_args = self.server_args
assert config is not None
speculativa_ratio = (
0
if server_args.speculative_num_draft_tokens is None
else server_args.speculative_num_draft_tokens
)
if (
server_args.disable_radix_cache
or config.mamba2_cache_params.mamba_cache_per_req == 0
or server_args.max_mamba_cache_size is not None
):
# with disable radix cache, sets the max_mamba_cache_size based on the max_running_requests
if server_args.max_mamba_cache_size is None:
@@ -1461,7 +1458,25 @@ class ModelRunner:
server_args.max_mamba_cache_size = server_args.max_running_requests
else:
server_args.max_mamba_cache_size = 512
server_args.max_mamba_cache_size = server_args.max_mamba_cache_size // (
server_args.dp_size if server_args.enable_dp_attention else 1
)
else:
assert config.mamba2_cache_params.mamba_cache_per_req > 0
# reserve the memory for the intermediate mamba states used for spec dec
if not self.spec_algorithm.is_none():
assert server_args.speculative_num_draft_tokens is not None
assert server_args.max_running_requests is not None
mamba_state_intermediate_size = (
config.mamba2_cache_params.mamba_cache_per_req
* server_args.max_running_requests
* server_args.speculative_num_draft_tokens
)
total_rest_memory = total_rest_memory - (
mamba_state_intermediate_size / (1 << 30)
)
# allocate the memory based on the ratio between mamba state memory vs. full kv cache memory
# solve the equations:
# 1. mamba_state_memory + full_kv_cache_memory == total_rest_memory
@@ -1475,21 +1490,22 @@ class ModelRunner:
server_args.max_mamba_cache_size = int(
(mamba_state_memory_raw * (1 << 30))
// config.mamba2_cache_params.mamba_cache_per_req
// (1 + speculativa_ratio)
)
if self.hybrid_gdn_config is not None:
server_args.max_mamba_cache_size = server_args.max_mamba_cache_size // (
server_args.dp_size if server_args.enable_dp_attention else 1
)
mamba_state_memory = (
server_args.max_mamba_cache_size
* config.mamba2_cache_params.mamba_cache_per_req
* (1 + speculativa_ratio)
/ (1 << 30)
)
return total_rest_memory - mamba_state_memory
@property
def qwen3_next_config(self):
config = self.model_config.hf_config
if isinstance(config, Qwen3NextConfig):
return config
return None
@property
def hybrid_gdn_config(self):
config = self.model_config.hf_config
@@ -1683,11 +1699,18 @@ class ModelRunner:
)
if self.mambaish_config is not None:
ratio = (
MAMBA_CACHE_SIZE_MAX_RUNNING_REQUESTS_RATIO
if not self.server_args.disable_radix_cache
else 1
)
additional_ratio = 0
if (
self.server_args.enable_mamba_extra_buffer()
and not self.spec_algorithm.is_none()
):
additional_ratio = MAMBA_CACHE_V2_ADDITIONAL_RATIO_NO_OVERLAP
else:
additional_ratio = MAMBA_CACHE_V2_ADDITIONAL_RATIO_OVERLAP
if self.server_args.disable_radix_cache:
ratio = 1
else:
ratio = MAMBA_CACHE_SIZE_MAX_RUNNING_REQUESTS_RATIO + additional_ratio
max_num_reqs = min(
max_num_reqs, self.server_args.max_mamba_cache_size // ratio
)
@@ -1789,11 +1812,13 @@ class ModelRunner:
self.req_to_token_pool = HybridReqToTokenPool(
size=max_num_reqs,
mamba_size=self.server_args.max_mamba_cache_size,
mamba_spec_state_size=max_num_reqs,
max_context_len=self.model_config.context_len
+ extra_max_context_len,
device=self.device,
enable_memory_saver=self.server_args.enable_memory_saver,
cache_params=config.mamba2_cache_params,
enable_mamba_extra_buffer=self.server_args.enable_mamba_extra_buffer(),
speculative_num_draft_tokens=self.server_args.speculative_num_draft_tokens,
)
else: