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

support page size large than 64 for mamba radix cache with fix (#16768)

Browse Source
This commit is contained in:
Hanming Lu
2026-01-08 21:28:37 -08:00
committed by GitHub
parent ccd0fb3291
commit 41609b52fe
5 changed files with 61 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py
View File

@@ -34,6 +34,7 @@ from sglang.srt.layers.radix_attention import RadixAttention
from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool, MambaPool
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
from sglang.srt.model_executor.model_runner import ModelRunner
from sglang.srt.server_args import get_global_server_args
from sglang.srt.speculative.eagle_info import EagleDraftInput, EagleVerifyInput
from sglang.srt.speculative.spec_info import SpecInput
from sglang.srt.utils import is_cuda, is_npu
@@ -284,7 +285,8 @@ class MambaAttnBackendBase(AttentionBackend):
lens_to_track = (
forward_batch.mamba_track_seqlens - forward_batch.extend_prefix_lens
)
aligned_len = (lens_to_track // FLA_CHUNK_SIZE) * FLA_CHUNK_SIZE
mamba_cache_chunk_size = get_global_server_args().mamba_cache_chunk_size
aligned_len = (lens_to_track // mamba_cache_chunk_size) * mamba_cache_chunk_size
start_indices = query_start_loc[:-1] + aligned_len - conv_state_len
start_indices = start_indices[forward_batch.mamba_track_mask]

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

@@ -1618,10 +1618,19 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
mamba_track_indices_cpu: List[int],
mamba_track_seqlens_cpu: List[int],
):
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
def _force_track_h(i: int) -> int:
assert i % FLA_CHUNK_SIZE == 0
# There are 3 cases for mamba_track_seqlen passed to mamba_track_seqlens_cpu:
# 1) aligned with FLA_CHUNK_SIZE-> retrieve from last_recurrent_state
# a) is the last position -> retrieve from last_recurrent_state
# b) is NOT the last position -> retrieve from h
# 2) unaligned with FLA_CHUNK_SIZE -> retrieve from h
# Currently, the math calculation only supports case 1a and 2. So for 1b, we need to add 1
# to force the math calculation to retrieve the correct mamba state from h.
return i + 1
mamba_cache_chunk_size = get_global_server_args().mamba_cache_chunk_size
mask = req.extend_input_len >= mamba_cache_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()
@@ -1636,13 +1645,28 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
# We need to pass the non-aligned seqlen to the calculation. Even though
# we pass in mamba_track_seqlen, the actual tracked seqlen is mamba_last_track_seqlen.
mamba_track_seqlen = len(req.prefix_indices) + req.extend_input_len
# mamba_last_track_seqlen is actual tracked seqlen. Used to pass to
# mamba_track_seqlen_aligned/mamba_last_track_seqlen is actual tracked seqlen. Used to pass to
# 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 // MAMBA_CACHE_V2_CHUNK_SIZE)
* MAMBA_CACHE_V2_CHUNK_SIZE
+ (req.extend_input_len // mamba_cache_chunk_size)
* mamba_cache_chunk_size
)
# mamba_track_fla_chunk_aligned is the aligned seqlen based on FLA_CHUNK_SIZE
# If mamba_track_fla_chunk_aligned != mamba_track_seqlen_aligned, which can be true when
# page_size > FLA_CHUNK_SIZE, we need to force the math calculation to retrieve the correct mamba state from h
# by _force_track_h()
mamba_track_fla_chunk_aligned = (
len(req.prefix_indices)
+ (req.extend_input_len // FLA_CHUNK_SIZE) * FLA_CHUNK_SIZE
)
if mamba_track_fla_chunk_aligned != mamba_track_seqlen_aligned:
# We want to track mamba_track_seqlen_aligned, and it's not the last position,
# so we need to add 1 to the seqlen to retrieve the correct mamba state from h.
mamba_track_seqlen = _force_track_h(mamba_track_seqlen_aligned)
req.mamba_next_track_idx = (
self.req_to_token_pool.get_mamba_ping_pong_other_idx(
req.mamba_next_track_idx
@@ -1653,18 +1677,16 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
# is within the current extend batch.
branching_seqlen_aligned_mask = (
req.mamba_branching_seqlen - len(req.prefix_indices)
) % MAMBA_CACHE_V2_CHUNK_SIZE == 0
) % mamba_cache_chunk_size == 0
if (
req.mamba_branching_seqlen > len(req.prefix_indices)
and req.mamba_branching_seqlen < mamba_track_seqlen
and branching_seqlen_aligned_mask
):
# NOTE: See the comment above for mamba_track_seqlen, the +1 is necessary
# because the branching point is not the last aligned position, so we need
# to retrieve its state from h. Adding 1 will give us the correct index in h,
# otherwise the calculation will retrieve the state from the last_recurrent_state,
# which is not correct.
mamba_track_seqlen = req.mamba_branching_seqlen + 1
# We want to track mamba_track_seqlen_aligned, and it's not the last position,
# so we need to add 1 to the seqlen to retrieve the correct mamba state from h.
# See _force_track_h() for more details.
mamba_track_seqlen = _force_track_h(req.mamba_branching_seqlen)
mamba_track_seqlen_aligned = req.mamba_branching_seqlen
req.mamba_last_track_seqlen = mamba_track_seqlen_aligned
mamba_track_seqlens_cpu.append(mamba_track_seqlen)

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

@@ -41,6 +41,7 @@ from sglang.srt.mem_cache.radix_cache import (
_key_match_paged,
get_child_key,
)
from sglang.srt.server_args import get_global_server_args
if TYPE_CHECKING:
from sglang.srt.managers.schedule_batch import Req
@@ -955,13 +956,13 @@ 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:
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_cache_chunk_size = get_global_server_args().mamba_cache_chunk_size
mamba_cache_chunk_aligned_seqlen = (
sum(len(v) for v in value) // mamba_cache_chunk_size
) * mamba_cache_chunk_size
mamba_branching_seqlen = (
mamba_cache_v2_chunk_aligned_seqlen
if mamba_cache_v2_chunk_aligned_seqlen > 0
mamba_cache_chunk_aligned_seqlen
if mamba_cache_chunk_aligned_seqlen > 0
else None
)
else:

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

@@ -1475,7 +1475,7 @@ class ServerArgs:
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=}"
), f"For SSM models with extra buffer, either FLA_CHUNK_SIZE or page_size must be divisible by the other, got {FLA_CHUNK_SIZE=}, {self.page_size=}"
elif not self.disable_radix_cache:
logger.warning(
@@ -4630,6 +4630,12 @@ class ServerArgs:
def enable_mamba_extra_buffer(self) -> bool:
return self.mamba_scheduler_strategy == "extra_buffer"
@property
def mamba_cache_chunk_size(self) -> int:
# For mamba cache with extra buffer, the chunk size is the max of FLA_CHUNK_SIZE and page_size.
# It is used to determine the caching point in a sequence during prefill.
return max(FLA_CHUNK_SIZE, self.page_size)
def check_server_args(self):
# Check parallel size constraints
assert (