Fix hybrid swa chunked prefill oom (#23174)

2026-04-21 10:46:45 +08:00
parent ab3ce02de9
commit 50fc2c9e23
2 changed files with 64 additions and 1 deletions
@@ -631,10 +631,16 @@ class PrefillAdder:
        else:
            _rem_tokens = min(self.rem_chunk_tokens, int(self.rem_total_tokens))
            if self.is_hybrid_swa:
-                _rem_tokens = min(_rem_tokens, int(self.rem_swa_tokens))
+                # alloc_extend needs extend_num_tokens + page_size per request,
+                # so reserve one page here to avoid OOM
+                _rem_tokens = min(
+                    _rem_tokens, int(self.rem_swa_tokens) - self.page_size
+                )
            # The chunked_req must be added to the list; otherwise, it will cause a memory leak.
            # Therefore, in certain cases where _rem_tokens <= 0, it should be replaced with rem_chunk_tokens.
            if _rem_tokens <= 0:
+                if self.is_hybrid_swa:
+                    return req
                _rem_tokens = self.rem_chunk_tokens

        truncated = req.extend_input_len > _rem_tokens