diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 8bb10b257..69eb22f30 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -612,7 +612,12 @@ class SchedulerDisaggregationPrefillMixin: else: self.send_kv_chunk(self.chunked_req) # chunked request keeps its rid but will get a new req_pool_idx - self.req_to_token_pool.free(self.chunked_req.req_pool_idx) + if self.tp_worker.model_runner.mambaish_config is not None: + self.req_to_token_pool.free( + self.chunked_req.req_pool_idx, free_mamba_cache=False + ) + else: + self.req_to_token_pool.free(self.chunked_req.req_pool_idx) self.running_batch.batch_is_full = False def send_kv_chunk(