A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 03:45 UTC
[KV Connector] Opt DecodeBenchConnector into SupportsHMA (#41770)
Signed-off-by: Zijing Liu <liuzijing2014@gmail.com> Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Zijing Liu committed
May 7, 2026 at 16:10 UTC
09a7cc5ba94c5c1b1c908a29f28ed41c793f1e9f
1 file changed
+16
-2
vllm/distributed/kv_transfer/kv_connector/v1/decode_bench_connector.py
+16
-2
@@ -40,7 +40,10 @@ from vllm.distributed.kv_transfer.kv_connector.v1 import (
40
KVConnectorBase_V1,
41
KVConnectorRole,
42
)
43
-from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorMetadata
43
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (
44
+ KVConnectorMetadata,
45
+ SupportsHMA,
46
+)
47
from vllm.logger import init_logger
48
from vllm.utils.math_utils import cdiv
49
from vllm.v1.attention.backend import AttentionMetadata
@@ -71,7 +74,7 @@ class DecodeBenchConnectorMetadata(KVConnectorMetadata):
74
reqs_to_fill: dict[str, tuple[tuple[list[int], ...], int]]
75
76
74
-class DecodeBenchConnector(KVConnectorBase_V1):
77
+class DecodeBenchConnector(KVConnectorBase_V1, SupportsHMA):
78
"""
79
A KV Connector for decode instance performance testing.
80
@@ -164,6 +167,17 @@ class DecodeBenchConnector(KVConnectorBase_V1):
167
self.connector_scheduler.request_finished(request)
168
return False, None
169
170
+ def request_finished_all_groups(
171
+ self,
172
+ request: "Request",
173
+ block_ids: tuple[list[int], ...],
174
+ ) -> tuple[bool, dict[str, Any] | None]:
175
+ # HMA-enabled path: same cleanup as the single-group variant since
176
+ # this connector owns no external state per block.
177
+ assert self.connector_scheduler is not None
178
+ self.connector_scheduler.request_finished(request)
179
+ return False, None
180
+
181
182
class DecodeBenchConnectorScheduler:
183
"""Scheduler-side implementation for DecodeBenchConnector."""