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

support page size large than 64 for mamba radix cache (#16657)

Browse Source 
Co-authored-by: Hanming Lu <hanming@x.ai>
This commit is contained in:
Yi Zhang
2026-01-07 22:52:24 +08:00
committed by GitHub
parent 8729ad5e6c
commit 7fc12e0bfa
3 changed files with 18 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
python/sglang/srt/managers/schedule_batch.py
View File

@@ -1621,7 +1621,10 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
mamba_track_indices_cpu: List[int],
mamba_track_seqlens_cpu: List[int],
):
mask = (req.extend_input_len // FLA_CHUNK_SIZE) * FLA_CHUNK_SIZE > 0
MAMBA_CACHE_V2_CHUNK_SIZE = max(
FLA_CHUNK_SIZE, self.token_to_kv_pool_allocator.page_size
)
mask = req.extend_input_len >= MAMBA_CACHE_V2_CHUNK_SIZE
mamba_track_mask_cpu.append(mask)
mamba_track_indices_cpu.append(
req.mamba_ping_pong_track_buffer[req.mamba_next_track_idx].item()
@@ -1640,7 +1643,8 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
# mamba radix cache to track which seqlen this mamba state should store at.
mamba_track_seqlen_aligned = (
len(req.prefix_indices)
+ (req.extend_input_len // FLA_CHUNK_SIZE) * FLA_CHUNK_SIZE
+ (req.extend_input_len // MAMBA_CACHE_V2_CHUNK_SIZE)
* MAMBA_CACHE_V2_CHUNK_SIZE
)
req.mamba_next_track_idx = (
self.req_to_token_pool.get_mamba_ping_pong_other_idx(
@@ -1652,7 +1656,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
# is within the current extend batch.
branching_seqlen_aligned_mask = (
req.mamba_branching_seqlen - len(req.prefix_indices)
) % FLA_CHUNK_SIZE == 0
) % MAMBA_CACHE_V2_CHUNK_SIZE == 0
if (
req.mamba_branching_seqlen > len(req.prefix_indices)
and req.mamba_branching_seqlen < mamba_track_seqlen

11
python/sglang/srt/mem_cache/mamba_radix_cache.py
View File

@@ -955,11 +955,14 @@ class MambaRadixCache(BasePrefixCache):
# Calculate the branching point. It is defined as the last aligned position that
# does not have a mamba value.
if len(value) > best_value_len:
fla_chunk_aligned_seqlen = (
sum(len(v) for v in value) // FLA_CHUNK_SIZE
) * FLA_CHUNK_SIZE
MAMBA_CACHE_V2_CHUNK_SIZE = max(FLA_CHUNK_SIZE, self.page_size)
mamba_cache_v2_chunk_aligned_seqlen = (
sum(len(v) for v in value) // MAMBA_CACHE_V2_CHUNK_SIZE
) * MAMBA_CACHE_V2_CHUNK_SIZE
mamba_branching_seqlen = (
fla_chunk_aligned_seqlen if fla_chunk_aligned_seqlen > 0 else None
mamba_cache_v2_chunk_aligned_seqlen
if mamba_cache_v2_chunk_aligned_seqlen > 0
else None
)
else:
mamba_branching_seqlen = None

6
python/sglang/srt/server_args.py
View File

@@ -1471,8 +1471,10 @@ class ServerArgs:
self.mamba_track_interval % self.page_size == 0
), f"mamba_track_interval {self.mamba_track_interval} must be divisible by page_size {self.page_size}"
assert (
FLA_CHUNK_SIZE % self.page_size == 0
), f"Page size for hybrid GDN model must be divisible by {FLA_CHUNK_SIZE}, got {self.page_size}"
max(FLA_CHUNK_SIZE, self.page_size)
% min(FLA_CHUNK_SIZE, self.page_size)
== 0
), f"max(FLA_CHUNK_SIZE, self.page_size) % min(FLA_CHUNK_SIZE, self.page_size) should be 0, got {FLA_CHUNK_SIZE=}, {self.page_size=}"
elif not self.disable_radix_cache:
logger.warning(