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

[Mamba] Fix long-prefill accuracy drop in radix prefix-cache state restore (#29368)

Browse Source
This commit is contained in:
Aditya Kamat
2026-07-06 13:57:44 -07:00
committed by GitHub
parent ca73c77055
commit 093256aa4b
1 changed files with 0 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/mamba/mamba2_metadata.py
-5
View File
@@ -241,11 +241,6 @@ class Mamba2Metadata(ForwardMetadata):
context_lens_tensor = forward_batch.extend_prefix_lens context_lens_tensor = forward_batch.extend_prefix_lens
assert context_lens_tensor is not None assert context_lens_tensor is not None
has_initial_states = context_lens_tensor > 0 has_initial_states = context_lens_tensor > 0
mamba_track_mask = getattr(forward_batch, "mamba_track_mask", None)
if mamba_track_mask is not None:
has_initial_states = (
has_initial_states & mamba_track_mask[: has_initial_states.shape[0]]
)
prep_initial_states = torch.any(has_initial_states[:num_prefills]).item() prep_initial_states = torch.any(has_initial_states[:num_prefills]).item()
query_start_loc = forward_metadata.query_start_loc[: num_prefills + 1] query_start_loc = forward_metadata.query_start_loc[: num_prefills + 1]