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 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]