Skip to content

fix: cancel KVTask immediately when match/put condition not met to prevent memory leak#176

Open
xianweihuihuan wants to merge 2 commits into
taco-project:mainfrom
xianweihuihuan:fix/cancel-kvtask-on-zero-match
Open

fix: cancel KVTask immediately when match/put condition not met to prevent memory leak#176
xianweihuihuan wants to merge 2 commits into
taco-project:mainfrom
xianweihuihuan:fix/cancel-kvtask-on-zero-match

Conversation

@xianweihuihuan
Copy link
Copy Markdown
Contributor

Problem

In _get_match() and the put path, flexkv_manager.get_match() /
get_put() unconditionally creates a KVTask stored in
KVTaskManager.tasks. When the match/put condition is not met
(num_new_matched_tokens == 0 or num_unmatched_tokens == 0), the
task is neither added to req_id_to_task_dict nor cancelled,
causing it to leak in self.tasks until ExpiringDict auto-expires
it after 30 minutes.

Under high concurrency with frequent cache misses, each request
produces a leaked KVTask holding token_ids, slot_mapping, and
token_mask arrays, leading to significant memory pressure.

Fix

Call self.flexkv_manager.cancel(task_ids=[task_id]) immediately
in the early-return branches where the task would otherwise be
abandoned, ensuring prompt cleanup of task resources.

@xianweihuihuan
Copy link
Copy Markdown
Contributor Author

Updated: the get path (_get_match) already handles this case
in the original code. This PR only adds the missing cancel call
in the put path (_put_match) when num_unmatched_tokens == 0.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant