[Bugfix] Release Mamba cache after PP dynamic chunk profiling (#31321)

2026-07-16 02:32:07 +08:00
parent 2d00e20a52
commit 7e7129acd7
1 changed files with 3 additions and 1 deletions
@@ -714,12 +714,14 @@ class SchedulerPPMixin:
                seq_lens.append(len(input_ids))
                latencies.append(latency_ms)

-                # Release KV cache
+                # Release KV and Mamba cache
                if req.req_pool_idx is not None:
                    kv_indices = self.req_to_token_pool.req_to_token[
                        req.req_pool_idx, : req.extend_range.end
                    ]
                    self.token_to_kv_pool_allocator.free(kv_indices)
+                    if req.mamba_pool_idx is not None:
+                        self.req_to_token_pool.free_mamba_cache(req)
                    self.req_to_token_pool.free(req)
                    req.kv = None