[KV Offload] Fix num_tokens_after_batch for different termination types (#49285)
Signed-off-by: Alex <jihuihuang@example.com> Signed-off-by: Alex <jihui.huang@daocloud.io> Signed-off-by: Alex <alex.tech.lab@outlook.com> Signed-off-by: Alex <jihuihuang@users.noreply.github.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Or Ozeri <oro@il.ibm.com>
AlexHuang committed
Jul 26, 2026 at 21:22 UTC
3f1d40960fb79e6f1314755abf2d43d142e33363
1 file changed
+4
-2
vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py
+4
-2
@@ -50,7 +50,7 @@ from vllm.v1.kv_offload.base import (
50
make_offload_key,
51
)
52
from vllm.v1.outputs import KVConnectorOutput
53
-from vllm.v1.request import Request
53
+from vllm.v1.request import Request, RequestStatus
54
55
logger = init_logger(__name__)
56
@@ -976,7 +976,9 @@ class OffloadingConnectorScheduler:
976
continue
977
req = req_status.req
978
979
- if req.is_finished():
979
+ if req.status is RequestStatus.FINISHED_ABORTED:
980
+ num_tokens_after_batch = req.num_computed_tokens
981
+ elif req.is_finished():
982
num_tokens_after_batch = req.num_tokens
983
else:
984
num_scheduled_tokens = scheduler_output.num_scheduled_tokens[req_id]