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 5, 11:00 UTC
[KVConnector] Support 3FS KVConnector (#37636)
Signed-off-by: wuchenxin <wuchenxin.wcx@alibaba-inc.com> Signed-off-by: ibifrost <47308427+ibifrost@users.noreply.github.com> Co-authored-by: Simon Mo <simon.mo@hey.com>
ibifrost committed
Apr 7, 2026 at 23:46 UTC
96b5004b71105b7b191a55de710736eea315f871
14 files changed
+3353
-2
setup.py
+1
@@ -1013,6 +1013,7 @@ package_data = {
1013
"model_executor/layers/quantization/utils/configs/*.json",
1014
"entrypoints/serve/instrumentator/static/*.js",
1015
"entrypoints/serve/instrumentator/static/*.css",
1016
+ "distributed/kv_transfer/kv_connector/v1/hf3fs/utils/*.cpp",
1017
]
1018
}
1019
tests/v1/kv_connector/unit/test_hf3fs_client.py
new
+284
@@ -0,0 +1,284 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+Tests for resource management in hf3fs_client.py: constructor failure cleanup
5
+and idempotent close(). Tests use mock to replace real I/O operations
6
+(hf3fs_fuse.io, SharedMemory, os, CUDA).
7
+Requires hf3fs_fuse.io to be installed; skipped otherwise.
8
+"""
9
+
10
+from typing import Any
11
+from unittest.mock import MagicMock, patch
12
+
13
+import pytest
14
+
15
+HF3FS_AVAILABLE = True
16
+try:
17
+ from hf3fs_fuse.io import ( # noqa: F401
18
+ deregister_fd,
19
+ extract_mount_point,
20
+ make_ioring,
21
+ make_iovec,
22
+ register_fd,
23
+ )
24
+
25
+ from vllm.distributed.kv_transfer.kv_connector.v1.hf3fs.hf3fs_client import (
26
+ Hf3fsClient,
27
+ )
28
+except Exception:
29
+ HF3FS_AVAILABLE = False
30
+
31
+requires_hf3fs = pytest.mark.skipif(
32
+ not HF3FS_AVAILABLE,
33
+ reason="hf3fs_fuse.io is not available on this machine",
34
+)
35
+
36
+
37
+# ---------------------------------------------------------------------------
38
+# Helpers
39
+# ---------------------------------------------------------------------------
40
+
41
+
42
+class _FakeShm:
43
+ """Shared-memory stub matching the multiprocessing.shared_memory.SharedMemory
44
+ interface used by Hf3fsClient:
45
+
46
+ Attributes accessed by the constructor:
47
+ .buf – memoryview / buffer-protocol object consumed by torch.frombuffer
48
+ Methods called during normal lifetime:
49
+ .unlink() – called right after the iovec is set up
50
+ .close() – called in _release_resources()
51
+ """
52
+
53
+ def __init__(self, size: int = 1024):
54
+ self._data = bytearray(size)
55
+ self.buf = memoryview(self._data)
56
+ self.closed = False
57
+ self.close_call_count = 0
58
+ self.unlink_call_count = 0
59
+
60
+ def close(self):
61
+ self.closed = True
62
+ self.close_call_count += 1
63
+
64
+ def unlink(self):
65
+ self.unlink_call_count += 1
66
+
67
+
68
+# ===========================================================================
69
+# TestHf3fsClientResourceManagement
70
+# ===========================================================================
71
+
72
+
73
+@requires_hf3fs
74
+class TestHf3fsClientResourceManagement:
75
+ """Tests for constructor failure cleanup and idempotent close()."""
76
+
77
+ _MOD = "vllm.distributed.kv_transfer.kv_connector.v1.hf3fs.hf3fs_client"
78
+
79
+ # ------------------------------------------------------------------
80
+ # Helper: build a minimal Hf3fsClient bypassing all real I/O so that
81
+ # we can fully control its internal state.
82
+ # ------------------------------------------------------------------
83
+
84
+ def _make_client(self, tmp_path):
85
+ """Return a fully-mocked Hf3fsClient with controllable internals."""
86
+ fake_shm_r = _FakeShm()
87
+ fake_shm_w = _FakeShm()
88
+
89
+ patcher_list: list[Any] = [
90
+ patch(f"{self._MOD}.HF3FS_AVAILABLE", True),
91
+ patch(f"{self._MOD}.register_fd"),
92
+ patch(f"{self._MOD}.deregister_fd"),
93
+ patch(f"{self._MOD}.extract_mount_point", return_value="/mnt/hf3fs"),
94
+ patch(f"{self._MOD}.make_ioring", return_value=MagicMock()),
95
+ patch(f"{self._MOD}.make_iovec", return_value=MagicMock()),
96
+ patch(
97
+ "multiprocessing.shared_memory.SharedMemory",
98
+ side_effect=[fake_shm_r, fake_shm_w],
99
+ ),
100
+ patch("os.open", return_value=99),
101
+ patch("os.ftruncate"),
102
+ patch("os.close"),
103
+ patch("os.fsync"),
104
+ patch("torch.cuda.Stream", return_value=MagicMock()),
105
+ patch("torch.frombuffer", return_value=MagicMock()),
106
+ patch("torch.empty", return_value=MagicMock()),
107
+ ]
108
+ for p in patcher_list:
109
+ p.start()
110
+
111
+ try:
112
+ client = Hf3fsClient(
113
+ path=str(tmp_path / "test.bin"),
114
+ size=1024,
115
+ bytes_per_page=256,
116
+ entries=4,
117
+ )
118
+ finally:
119
+ for p in patcher_list:
120
+ p.stop()
121
+
122
+ # Manually point internal handles to our controllable fakes so that
123
+ # assertions after close() can inspect them directly.
124
+ client.shm_r = fake_shm_r
125
+ client.shm_w = fake_shm_w
126
+ client.file = 99
127
+ return client, fake_shm_r, fake_shm_w
128
+
129
+ # ------------------------------------------------------------------
130
+ # close() idempotency
131
+ # ------------------------------------------------------------------
132
+
133
+ def test_close_idempotent_and_handles_cleared(self, tmp_path):
134
+ """Multiple close() calls must not raise; deregister_fd called exactly
135
+ once, all handles set to None, shm.close() invoked."""
136
+ client, shm_r, shm_w = self._make_client(tmp_path)
137
+
138
+ with (
139
+ patch(f"{self._MOD}.deregister_fd") as mock_dereg,
140
+ patch("os.close"),
141
+ ):
142
+ client.close() # first close
143
+ client.close() # second close — must be no-op
144
+ client.close() # third close — must be no-op
145
+
146
+ assert client._closed is True
147
+ assert mock_dereg.call_count == 1, (
148
+ f"deregister_fd called {mock_dereg.call_count} times; expected 1"
149
+ )
150
+ for attr in ("iov_r", "iov_w", "ior_r", "ior_w", "shm_r", "shm_w", "file"):
151
+ assert getattr(client, attr) is None, f"{attr} should be None after close()"
152
+ assert shm_r.closed is True
153
+ assert shm_w.closed is True
154
+
155
+ def test_flush_after_close_is_noop(self, tmp_path):
156
+ """flush() after close() must silently do nothing (no fsync call)."""
157
+ client, _, _ = self._make_client(tmp_path)
158
+
159
+ with (
160
+ patch(f"{self._MOD}.deregister_fd"),
161
+ patch("os.close"),
162
+ patch("os.fsync") as mock_fsync,
163
+ ):
164
+ client.close()
165
+ client.flush()
166
+
167
+ mock_fsync.assert_not_called()
168
+
169
+ # ------------------------------------------------------------------
170
+ # Constructor failure leaves no leaked resources
171
+ # ------------------------------------------------------------------
172
+
173
+ def test_constructor_failure_after_file_open_cleans_file(self, tmp_path):
174
+ """If the constructor raises after os.open(), the fd must be closed."""
175
+ with (
176
+ patch(f"{self._MOD}.HF3FS_AVAILABLE", True),
177
+ patch(f"{self._MOD}.register_fd"),
178
+ patch(f"{self._MOD}.deregister_fd"),
179
+ patch(
180
+ f"{self._MOD}.extract_mount_point",
181
+ side_effect=RuntimeError("mount point not found"),
182
+ ),
183
+ patch("os.open", return_value=55),
184
+ patch("os.ftruncate"),
185
+ patch("os.close") as mock_os_close,
186
+ patch("torch.cuda.Stream", return_value=MagicMock()),
187
+ pytest.raises(RuntimeError, match="mount point not found"),
188
+ ):
189
+ Hf3fsClient(
190
+ path=str(tmp_path / "fail.bin"),
191
+ size=1024,
192
+ bytes_per_page=256,
193
+ entries=4,
194
+ )
195
+
196
+ mock_os_close.assert_called_once_with(55)
197
+
198
+ def test_constructor_failure_after_shm_alloc_closes_shm(self, tmp_path):
199
+ """Constructor raises after SharedMemory creation → both shm objects closed."""
200
+ fake_shm_r = _FakeShm()
201
+ fake_shm_w = _FakeShm()
202
+
203
+ with (
204
+ patch(f"{self._MOD}.HF3FS_AVAILABLE", True),
205
+ patch(f"{self._MOD}.register_fd"),
206
+ patch(f"{self._MOD}.deregister_fd"),
207
+ patch(f"{self._MOD}.extract_mount_point", return_value="/mnt/hf3fs"),
208
+ patch(
209
+ "multiprocessing.shared_memory.SharedMemory",
210
+ side_effect=[fake_shm_r, fake_shm_w],
211
+ ),
212
+ patch("os.open", return_value=66),
213
+ patch("os.ftruncate"),
214
+ patch("os.close"),
215
+ patch("torch.frombuffer", return_value=MagicMock()),
216
+ patch("torch.empty", return_value=MagicMock()),
217
+ patch(
218
+ f"{self._MOD}.make_ioring",
219
+ side_effect=RuntimeError("ioring init failed"),
220
+ ),
221
+ patch(f"{self._MOD}.make_iovec", return_value=MagicMock()),
222
+ patch("torch.cuda.Stream", return_value=MagicMock()),
223
+ pytest.raises(RuntimeError, match="ioring init failed"),
224
+ ):
225
+ Hf3fsClient(
226
+ path=str(tmp_path / "fail2.bin"),
227
+ size=1024,
228
+ bytes_per_page=256,
229
+ entries=4,
230
+ )
231
+
232
+ assert fake_shm_r.closed is True, (
233
+ "shm_r was not closed after constructor failure"
234
+ )
235
+ assert fake_shm_w.closed is True, (
236
+ "shm_w was not closed after constructor failure"
237
+ )
238
+
239
+ def test_constructor_failure_does_not_close_unallocated_shm(self, tmp_path):
240
+ """Failure before SharedMemory is created must not raise AttributeError
241
+ or TypeError from cleanup."""
242
+ with (
243
+ patch(f"{self._MOD}.HF3FS_AVAILABLE", True),
244
+ patch(f"{self._MOD}.register_fd"),
245
+ patch(f"{self._MOD}.deregister_fd"),
246
+ patch(
247
+ f"{self._MOD}.extract_mount_point",
248
+ side_effect=RuntimeError("early failure"),
249
+ ),
250
+ patch("os.open", return_value=77),
251
+ patch("os.ftruncate"),
252
+ patch("os.close"),
253
+ patch("torch.cuda.Stream", return_value=MagicMock()),
254
+ pytest.raises(RuntimeError, match="early failure"),
255
+ ):
256
+ Hf3fsClient(
257
+ path=str(tmp_path / "early_fail.bin"),
258
+ size=1024,
259
+ bytes_per_page=256,
260
+ entries=4,
261
+ )
262
+
263
+ # ------------------------------------------------------------------
264
+ # _release_resources on already-cleared state must be a no-op
265
+ # ------------------------------------------------------------------
266
+
267
+ def test_release_resources_on_empty_state_is_safe(self, tmp_path):
268
+ """_release_resources() on a fully-cleared client must not raise."""
269
+ client, _, _ = self._make_client(tmp_path)
270
+
271
+ with (
272
+ patch(f"{self._MOD}.deregister_fd"),
273
+ patch("os.close"),
274
+ ):
275
+ client.close() # clears all handles
276
+
277
+ with (
278
+ patch(f"{self._MOD}.deregister_fd") as mock_dereg2,
279
+ patch("os.close") as mock_os_close2,
280
+ ):
281
+ client._release_resources() # must not raise
282
+
283
+ mock_dereg2.assert_not_called()
284
+ mock_os_close2.assert_not_called()
tests/v1/kv_connector/unit/test_hf3fs_connector.py
new
+230
@@ -0,0 +1,230 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+Tests for HF3FS KV Connector high-level components:
5
+ - TestHf3fsMockClient : file-backed mock client I/O correctness
6
+ - TestHF3FSKVConnectorStats: metric collection, aggregation, serialisation
7
+"""
8
+
9
+import os
10
+from unittest.mock import MagicMock
11
+
12
+import pytest
13
+import torch
14
+
15
+from vllm.distributed.kv_transfer.kv_connector.v1.hf3fs.hf3fs_connector import (
16
+ HF3FSKVConnectorStats,
17
+)
18
+from vllm.distributed.kv_transfer.kv_connector.v1.hf3fs.utils.hf3fs_mock_client import (
19
+ Hf3fsClient as MockHf3fsClient,
20
+)
21
+
22
+# ---------------------------------------------------------------------------
23
+# Helpers
24
+# ---------------------------------------------------------------------------
25
+
26
+
27
+@pytest.fixture
28
+def hf3fs_stats():
29
+ """Fresh HF3FSKVConnectorStats instance."""
30
+ return HF3FSKVConnectorStats()
31
+
32
+
33
+def _make_cuda_event():
34
+ """Return a real CUDA event when available, otherwise a MagicMock."""
35
+ if torch.cuda.is_available():
36
+ return torch.cuda.Event()
37
+ return MagicMock()
38
+
39
+
40
+# ===========================================================================
41
+# TestHf3fsMockClient
42
+# ===========================================================================
43
+
44
+
45
+class TestHf3fsMockClient:
46
+ """Tests for hf3fs_mock_client.Hf3fsClient (file-backend mock)."""
47
+
48
+ def test_init_creates_file(self, tmp_path):
49
+ """Initializing the client should create the backing file."""
50
+ path = str(tmp_path / "test_file")
51
+ client = MockHf3fsClient(path=path, size=4096, bytes_per_page=512, entries=4)
52
+ assert os.path.exists(path), "Backing file should be created on init"
53
+ assert os.path.getsize(path) == 4096
54
+ client.close()
55
+
56
+ @pytest.mark.parametrize(
57
+ "dtype, bytes_per_page",
58
+ [
59
+ (torch.float32, 512),
60
+ (torch.float16, 256),
61
+ (torch.bfloat16, 256),
62
+ ],
63
+ ids=["float32", "float16", "bfloat16"],
64
+ )
65
+ def test_batch_write_and_read_dtype(self, tmp_path, dtype, bytes_per_page):
66
+ """Write a tensor of the given dtype and verify round-trip correctness."""
67
+ path = str(tmp_path / f"rw_{dtype}")
68
+ client = MockHf3fsClient(
69
+ path=path, size=bytes_per_page * 8, bytes_per_page=bytes_per_page, entries=4
70
+ )
71
+ elem_size = torch.tensor([], dtype=dtype).element_size()
72
+ numel = bytes_per_page // elem_size
73
+ tensor_write = torch.arange(numel, dtype=dtype)
74
+ event = _make_cuda_event()
75
+
76
+ results = client.batch_write([0], [tensor_write], event)
77
+ assert results == [bytes_per_page], f"Write should succeed, got {results}"
78
+
79
+ tensor_read = torch.zeros(numel, dtype=dtype)
80
+ results = client.batch_read([0], [tensor_read])
81
+ assert results == [bytes_per_page], f"Read should succeed, got {results}"
82
+ assert torch.equal(tensor_write, tensor_read), (
83
+ "Read tensor should match written tensor"
84
+ )
85
+ client.close()
86
+
87
+ def test_batch_read_empty_file_returns_error(self, tmp_path):
88
+ """Reading out-of-bounds offset should return -1."""
89
+ bytes_per_page = 128
90
+ size = bytes_per_page * 4
91
+ path = str(tmp_path / "empty_read")
92
+ client = MockHf3fsClient(
93
+ path=path, size=size, bytes_per_page=bytes_per_page, entries=4
94
+ )
95
+ numel = bytes_per_page // 4
96
+ tensor_read = torch.zeros(numel, dtype=torch.float32)
97
+ results = client.batch_read([size], [tensor_read]) # offset == size => OOB
98
+ assert results[0] == -1, "Out-of-bounds read should return -1"
99
+ client.close()
100
+
101
+ def test_batch_write_out_of_bounds_returns_error(self, tmp_path):
102
+ """Writing at an offset beyond file size should return -1."""
103
+ bytes_per_page = 128
104
+ size = bytes_per_page * 4
105
+ path = str(tmp_path / "oob_write")
106
+ client = MockHf3fsClient(
107
+ path=path, size=size, bytes_per_page=bytes_per_page, entries=4
108
+ )
109
+ numel = bytes_per_page // 4
110
+ tensor = torch.ones(numel, dtype=torch.float32)
111
+ event = _make_cuda_event()
112
+ results = client.batch_write([size], [tensor], event) # OOB offset
113
+ assert results[0] == -1, "Out-of-bounds write should return -1"
114
+ client.close()
115
+
116
+ def test_multiple_tensors_rw(self, tmp_path):
117
+ """Write multiple tensors at different offsets, then read all back."""
118
+ bytes_per_page = 128
119
+ n = 4
120
+ path = str(tmp_path / "multi_rw")
121
+ client = MockHf3fsClient(
122
+ path=path,
123
+ size=bytes_per_page * n * 2,
124
+ bytes_per_page=bytes_per_page,
125
+ entries=8,
126
+ )
127
+ tensors_write = [
128
+ torch.full((bytes_per_page // 4,), float(i), dtype=torch.float32)
129
+ for i in range(n)
130
+ ]
131
+ offsets = [i * bytes_per_page for i in range(n)]
132
+ event = _make_cuda_event()
133
+
134
+ results = client.batch_write(offsets, tensors_write, event)
135
+ assert all(r == bytes_per_page for r in results)
136
+
137
+ tensors_read = [
138
+ torch.zeros(bytes_per_page // 4, dtype=torch.float32) for _ in range(n)
139
+ ]
140
+ results = client.batch_read(offsets, tensors_read)
141
+ assert all(r == bytes_per_page for r in results)
142
+
143
+ for i, (tw, tr) in enumerate(zip(tensors_write, tensors_read)):
144
+ assert torch.allclose(tw, tr), f"Tensor {i} mismatch after round-trip"
145
+ client.close()
146
+
147
+ def test_flush_and_close_no_error(self, tmp_path):
148
+ """flush() and close() should not raise exceptions."""
149
+ path = str(tmp_path / "flush_close")
150
+ client = MockHf3fsClient(path=path, size=1024, bytes_per_page=128, entries=4)
151
+ client.flush()
152
+ client.close()
153
+
154
+
155
+# ===========================================================================
156
+# TestHF3FSKVConnectorStats
157
+# ===========================================================================
158
+
159
+
160
+class TestHF3FSKVConnectorStats:
161
+ """Tests for HF3FSKVConnectorStats metric collection and aggregation."""
162
+
163
+ def test_initial_is_empty(self, hf3fs_stats):
164
+ """Fresh stats object should report is_empty() == True."""
165
+ assert hf3fs_stats.is_empty() is True
166
+
167
+ @pytest.mark.parametrize(
168
+ "task_type, duration_key",
169
+ [
170
+ ("Saved", "save_duration"),
171
+ ("Loaded", "load_duration"),
172
+ ],
173
+ ids=["save", "load"],
174
+ )
175
+ def test_record_success_duration(self, hf3fs_stats, task_type, duration_key):
176
+ """Recording a successful task should update duration list and total count."""
177
+ hf3fs_stats.record_success_task_duration(task_type, 0.5)
178
+ assert not hf3fs_stats.is_empty()
179
+ assert len(hf3fs_stats.data[duration_key]) == 1
180
+ assert hf3fs_stats.data[duration_key][0] == pytest.approx(0.5)
181
+ assert hf3fs_stats.data["num_transfer_task"] == 1
182
+
183
+ @pytest.mark.parametrize(
184
+ "task_type, failed_key",
185
+ [
186
+ ("Saved", "num_failed_save"),
187
+ ("Loaded", "num_failed_load"),
188
+ ],
189
+ ids=["save", "load"],
190
+ )
191
+ def test_record_failed_task(self, hf3fs_stats, task_type, failed_key):
192
+ """Recording a failed task should increment the corresponding counter."""
193
+ hf3fs_stats.record_failed_task_count(task_type)
194
+ assert hf3fs_stats.data[failed_key] == 1
195
+ assert hf3fs_stats.data["num_transfer_task"] == 1
196
+
197
+ def test_aggregate_two_stats(self):
198
+ """aggregate() should merge save/load duration lists and sum counters."""
199
+ stats1 = HF3FSKVConnectorStats()
200
+ stats1.record_success_task_duration("Saved", 0.1)
201
+ stats1.record_success_task_duration("Loaded", 0.2)
202
+
203
+ stats2 = HF3FSKVConnectorStats()
204
+ stats2.record_success_task_duration("Saved", 0.3)
205
+ stats2.record_failed_task_count("Loaded")
206
+
207
+ stats1.aggregate(stats2)
208
+ assert stats1.data["save_duration"] == pytest.approx([0.1, 0.3])
209
+ assert stats1.data["load_duration"] == pytest.approx([0.2])
210
+ assert stats1.data["num_failed_load"] == 1
211
+ assert stats1.data["num_transfer_task"] == 4
212
+
213
+ def test_reduce_with_data(self):
214
+ """reduce() computes correct averages when data is present."""
215
+ stats = HF3FSKVConnectorStats()
216
+ stats.record_success_task_duration("Saved", 1.0)
217
+ stats.record_success_task_duration("Saved", 3.0)
218
+ result = stats.reduce()
219
+ assert result["Num save task success"] == pytest.approx(2.0, rel=0.01)
220
+ assert result["Num save task failed"] == pytest.approx(0.0, rel=0.01)
221
+ assert result["Avg save duration (ms)"] == pytest.approx(2000.0, rel=0.01)
222
+
223
+ def test_clone_and_reset(self, hf3fs_stats):
224
+ """clone_and_reset() returns a copy with data and resets the original."""
225
+ hf3fs_stats.record_success_task_duration("Saved", 0.7)
226
+ hf3fs_stats.record_success_task_duration("Loaded", 0.4)
227
+
228
+ clone = hf3fs_stats.clone_and_reset()
229
+ assert clone.data["num_transfer_task"] == 2
230
+ assert hf3fs_stats.is_empty()
tests/v1/kv_connector/unit/test_hf3fs_metadata_server.py
new
+193
@@ -0,0 +1,193 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+Tests for HF3FS metadata server data structures and allocation logic:
5
+ - RankFileMetadata : page allocation / release primitives
6
+ - KeyMetadata : per-key rank-page tracking and completion detection
7
+ - GlobalMetadataState : coordinated allocation with cache-hit semantics
8
+"""
9
+
10
+import pytest
11
+
12
+from vllm.distributed.kv_transfer.kv_connector.v1.hf3fs.hf3fs_metadata_server import (
13
+ GlobalMetadataState,
14
+ KeyMetadata,
15
+ RankFileMetadata,
16
+)
17
+
18
+# ===========================================================================
19
+# TestRankFileMetadata
20
+# ===========================================================================
21
+
22
+
23
+class TestRankFileMetadata:
24
+ """Unit tests for RankFileMetadata page allocation primitives."""
25
+
26
+ @pytest.mark.parametrize(
27
+ "alloc_count, expected_pages",
28
+ [(3, 3), (5, 0)],
29
+ ids=["alloc_partial", "alloc_exceeds"],
30
+ )
31
+ def test_allocate_pages(self, alloc_count, expected_pages):
32
+ """allocate_pages returns correct pages or empty list when insufficient."""
33
+ rank_meta = RankFileMetadata(rank_id=0, num_pages=3, free_pages=list(range(3)))
34
+ pages = rank_meta.allocate_pages(alloc_count)
35
+ assert len(pages) == expected_pages
36
+ if expected_pages > 0:
37
+ rank_meta.release_pages(pages)
38
+ assert rank_meta.get_free_page_count() == 3
39
+
40
+ def test_release_pages_restores_count(self):
41
+ """Releasing allocated pages returns them to the free pool."""
42
+ rank_meta = RankFileMetadata(rank_id=0, num_pages=4, free_pages=list(range(4)))
43
+ pages = rank_meta.allocate_pages(2)
44
+ assert rank_meta.get_free_page_count() == 2
45
+ rank_meta.release_pages(pages)
46
+ assert rank_meta.get_free_page_count() == 4
47
+
48
+ def test_release_pages_no_duplicates(self):
49
+ """Releasing the same page twice must not create duplicates."""
50
+ rank_meta = RankFileMetadata(rank_id=0, num_pages=3, free_pages=list(range(3)))
51
+ rank_meta.allocate_pages(1) # takes page 0
52
+ rank_meta.release_pages([0])
53
+ rank_meta.release_pages([0]) # second release of the same page
54
+ assert rank_meta.get_free_page_count() == 3
55
+
56
+
57
+# ===========================================================================
58
+# TestKeyMetadata
59
+# ===========================================================================
60
+
61
+
62
+class TestKeyMetadata:
63
+ """Unit tests for KeyMetadata completion tracking."""
64
+
65
+ def test_is_complete_false_until_all_ranks(self):
66
+ """is_complete() returns True only when all ranks confirmed."""
67
+ key_meta = KeyMetadata(key="k", rank_to_page={}, tp_world_size=2)
68
+ assert key_meta.is_complete() is False
69
+ key_meta.add_rank_page(0, 5)
70
+ assert key_meta.is_complete() is False
71
+ key_meta.add_rank_page(1, 10)
72
+ assert key_meta.is_complete() is True
73
+
74
+ def test_get_rank_page_returns_none_for_missing_rank(self):
75
+ """get_rank_page() returns None when the rank has no entry."""
76
+ key_meta = KeyMetadata(key="k", rank_to_page={0: 3}, tp_world_size=2)
77
+ assert key_meta.get_rank_page(0) == 3
78
+ assert key_meta.get_rank_page(1) is None
79
+
80
+ def test_get_all_pages(self):
81
+ """get_all_pages() returns all (rank, page) pairs."""
82
+ key_meta = KeyMetadata(key="k", rank_to_page={0: 1, 1: 2}, tp_world_size=2)
83
+ pairs = key_meta.get_all_pages()
84
+ assert set(pairs) == {(0, 1), (1, 2)}
85
+
86
+
87
+# ===========================================================================
88
+# TestGlobalMetadataStateAllocation
89
+# ===========================================================================
90
+
91
+
92
+class TestGlobalMetadataStateAllocation:
93
+ """Tests for GlobalMetadataState allocation and cache-hit semantics."""
94
+
95
+ def test_uninitialized_rank_raises_on_allocate(self):
96
+ """allocate_pages_for_keys raises ValueError for unknown rank."""
97
+ state = GlobalMetadataState()
98
+ with pytest.raises((ValueError, Exception)):
99
+ state.allocate_pages_for_keys(99, [("key", "")])
100
+
101
+ def test_uninitialized_rank_raises_on_get_locations(self):
102
+ """get_key_locations raises ValueError for unknown rank."""
103
+ state = GlobalMetadataState()
104
+ with pytest.raises((ValueError, Exception)):
105
+ state.get_key_locations(99, ["any_key"])
106
+
107
+ def test_basic_allocation_and_confirm(self):
108
+ """Allocating a page and confirming it marks the key as complete."""
109
+ state = GlobalMetadataState()
110
+ state.initialize_rank(0, 4)
111
+
112
+ results = state.allocate_pages_for_keys(0, [("K", "")])
113
+ assert results["K"] >= 0
114
+
115
+ state.confirm_write_for_keys(0, [("K", results["K"])])
116
+ assert state.batch_key_exists(["K"]) == [True]
117
+ locations = state.get_key_locations(0, ["K"])
118
+ assert locations == [results["K"]]
119
+
120
+ def test_allocate_pages_cache_hit_does_not_leak_pages(self):
121
+ """Cache-hit key must not consume a page from the free pool;
122
+ the pre-allocated slot must be returned before reusing the existing page.
123
+ """
124
+ state = GlobalMetadataState()
125
+ state.initialize_rank(0, 5) # 5 free pages: [0,1,2,3,4]
126
+
127
+ # Simulate a key that has already been fully written and confirmed.
128
+ state.key_metadata["K_cached"] = KeyMetadata(
129
+ key="K_cached", rank_to_page={0: 2}, tp_world_size=1
130
+ )
131
+
132
+ free_before = state.rank_metadata[0].get_free_page_count() # 5
133
+
134
+ results = state.allocate_pages_for_keys(0, [("K_cached", ""), ("K_new", "")])
135
+
136
+ free_after = state.rank_metadata[0].get_free_page_count()
137
+
138
+ # Cache-hit key must reuse its existing page.
139
+ assert results["K_cached"] == 2, (
140
+ f"Cache-hit key should reuse page 2, got {results['K_cached']}"
141
+ )
142
+ # New key must receive a valid page.
143
+ assert results["K_new"] >= 0, (
144
+ f"New key should get a valid page, got {results['K_new']}"
145
+ )
146
+ # Exactly one page consumed from the free pool.
147
+ assert free_before - free_after == 1, (
148
+ f"Expected 1 page consumed, got delta={free_before - free_after}"
149
+ )
150
+
151
+ def test_allocate_pages_all_cache_hits_frees_all_slots(self):
152
+ """When every key in the batch is a cache hit, no pages are consumed."""
153
+ state = GlobalMetadataState()
154
+ state.initialize_rank(0, 5)
155
+
156
+ for key, page in (("K1", 0), ("K2", 1)):
157
+ state.key_metadata[key] = KeyMetadata(
158
+ key=key, rank_to_page={0: page}, tp_world_size=1
159
+ )
160
+
161
+ free_before = state.rank_metadata[0].get_free_page_count()
162
+ results = state.allocate_pages_for_keys(0, [("K1", ""), ("K2", "")])
163
+ free_after = state.rank_metadata[0].get_free_page_count()
164
+
165
+ assert results["K1"] == 0
166
+ assert results["K2"] == 1
167
+ assert free_after == free_before, (
168
+ f"All-cache-hit batch must not consume free pages; "
169
+ f"before={free_before}, after={free_after}"
170
+ )
171
+
172
+ def test_allocate_returns_minus_one_when_pool_exhausted(self):
173
+ """If the free pool is exhausted, all new keys receive -1."""
174
+ state = GlobalMetadataState()
175
+ state.initialize_rank(0, 1) # only 1 free page
176
+
177
+ results = state.allocate_pages_for_keys(0, [("K1", ""), ("K2", "")])
178
+ # allocate_pages uses all-or-nothing: 2 needed but only 1 available → []
179
+ assert all(v == -1 for v in results.values()), f"Expected all -1, got {results}"
180
+
181
+ def test_confirm_write_releases_pages(self):
182
+ """confirm_write_for_keys with pages_to_release returns them to pool."""
183
+ state = GlobalMetadataState()
184
+ state.initialize_rank(0, 3)
185
+
186
+ results = state.allocate_pages_for_keys(0, [("K", "")])
187
+ page = results["K"]
188
+ free_after_alloc = state.rank_metadata[0].get_free_page_count()
189
+
190
+ state.confirm_write_for_keys(0, [("K", page)], pages_to_release=[page])
191
+ free_after_release = state.rank_metadata[0].get_free_page_count()
192
+
193
+ assert free_after_release == free_after_alloc + 1
vllm/distributed/kv_transfer/kv_connector/factory.py
+5
-2
@@ -211,15 +211,18 @@ KVConnectorFactory.register_connector(
211
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector",
212
"MooncakeConnector",
213
)
214
-
214
KVConnectorFactory.register_connector(
215
"FlexKVConnectorV1",
216
"vllm.distributed.kv_transfer.kv_connector.v1.flexkv_connector",
217
"FlexKVConnectorV1",
218
)
220
-
219
KVConnectorFactory.register_connector(
220
"SimpleCPUOffloadConnector",
221
"vllm.distributed.kv_transfer.kv_connector.v1.simple_cpu_offload_connector",
222
"SimpleCPUOffloadConnector",
223
)
224
+KVConnectorFactory.register_connector(
225
+ "HF3FSKVConnector",
226
+ "vllm.distributed.kv_transfer.kv_connector.v1.hf3fs.hf3fs_connector",
227
+ "HF3FSKVConnector",
228
+)
vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/__init__.py
vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/hf3fs_client.py
new
+298
@@ -0,0 +1,298 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import logging
4
+import multiprocessing
5
+import os
6
+import threading
7
+from functools import wraps
8
+from pathlib import Path
9
+
10
+import torch
11
+import torch.utils.cpp_extension
12
+from torch.utils.cpp_extension import load
13
+
14
+root = Path(__file__).parent.resolve()
15
+cuda_include_path = os.path.join(torch.utils.cpp_extension.CUDA_HOME, "include")
16
+hf3fs_utils = load(
17
+ name="hf3fs_utils",
18
+ sources=[f"{root}/utils/hf3fs_utils.cpp"],
19
+ extra_include_paths=[cuda_include_path],
20
+)
21
+
22
+logger = logging.getLogger(__name__)
23
+
24
+HF3FS_AVAILABLE = True
25
+try:
26
+ from hf3fs_fuse.io import (
27
+ deregister_fd,
28
+ extract_mount_point,
29
+ make_ioring,
30
+ make_iovec,
31
+ register_fd,
32
+ )
33
+except ImportError:
34
+ HF3FS_AVAILABLE = False
35
+
36
+
37
+def rsynchronized():
38
+ def _decorator(func):
39
+ @wraps(func)
40
+ def wrapper(self, *args, **kwargs):
41
+ with self.rlock:
42
+ return func(self, *args, **kwargs)
43
+
44
+ return wrapper
45
+
46
+ return _decorator
47
+
48
+
49
+def wsynchronized():
50
+ def _decorator(func):
51
+ @wraps(func)
52
+ def wrapper(self, *args, **kwargs):
53
+ with self.wlock:
54
+ return func(self, *args, **kwargs)
55
+
56
+ return wrapper
57
+
58
+ return _decorator
59
+
60
+
61
+class Hf3fsClient:
62
+ def __init__(self, path: str, size: int, bytes_per_page: int, entries: int):
63
+ """Initialize the HF3FS client with hf3fs_fuse.
64
+
65
+ Args:
66
+ path: Path to the file used for storage
67
+ size: Total size of the storage file in bytes
68
+ bytes_per_page: Size of each page in bytes
69
+ entries: Maximum number of concurrent operations
70
+ """
71
+ if not HF3FS_AVAILABLE:
72
+ raise ImportError(
73
+ "hf3fs_fuse.io is not available. Please install the hf3fs_fuse package."
74
+ )
75
+
76
+ self.path = path
77
+ self.size = size
78
+ self.bytes_per_page = bytes_per_page
79
+ self.entries = entries
80
+
81
+ self._closed = False
82
+
83
+ self.file = None
84
+ self.shm_r = None
85
+ self.shm_w = None
86
+ self.ior_r = None
87
+ self.ior_w = None
88
+ self.iov_r = None
89
+ self.iov_w = None
90
+ try:
91
+ # Create the file if it doesn't exist and set its size
92
+ self.file = os.open(self.path, os.O_RDWR | os.O_CREAT)
93
+ os.ftruncate(self.file, size)
94
+ register_fd(self.file)
95
+
96
+ self.hf3fs_mount_point = extract_mount_point(path)
97
+ self.bs = self.bytes_per_page
98
+ self.shm_r = multiprocessing.shared_memory.SharedMemory(
99
+ size=self.bs * self.entries, create=True
100
+ )
101
+ self.shm_w = multiprocessing.shared_memory.SharedMemory(
102
+ size=self.bs * self.entries, create=True
103
+ )
104
+
105
+ self.shm_r_tensor = torch.frombuffer(self.shm_r.buf, dtype=torch.uint8)
106
+ self.shm_w_tensor = torch.frombuffer(self.shm_w.buf, dtype=torch.uint8)
107
+
108
+ numel = self.bs * self.entries
109
+ self.r_pinned = torch.empty(
110
+ numel,
111
+ dtype=torch.uint8,
112
+ device="cpu",
113
+ pin_memory=True,
114
+ )
115
+ self.w_pinned = torch.empty(
116
+ numel,
117
+ dtype=torch.uint8,
118
+ device="cpu",
119
+ pin_memory=True,
120
+ )
121
+
122
+ self.numa = -1
123
+ self.ior_r = make_ioring(
124
+ self.hf3fs_mount_point,
125
+ self.entries,
126
+ for_read=True,
127
+ timeout=1,
128
+ numa=self.numa,
129
+ )
130
+ self.ior_w = make_ioring(
131
+ self.hf3fs_mount_point,
132
+ self.entries,
133
+ for_read=False,
134
+ timeout=1,
135
+ numa=self.numa,
136
+ )
137
+ self.iov_r = make_iovec(self.shm_r, self.hf3fs_mount_point)
138
+ self.iov_w = make_iovec(self.shm_w, self.hf3fs_mount_point)
139
+ self.shm_r.unlink()
140
+ self.shm_w.unlink()
141
+
142
+ self.rlock = threading.RLock()
143
+ self.wlock = threading.RLock()
144
+
145
+ self.stream = torch.cuda.Stream()
146
+ self.stream_ptr_int = self.stream.cuda_stream
147
+
148
+ except Exception:
149
+ self._release_resources()
150
+ raise
151
+
152
+ logger.debug(
153
+ "Initialized HF3FS client with file: %s, size: %s bytes", path, size
154
+ )
155
+
156
+ def _release_resources(self) -> None:
157
+ """Release all acquired resources safely"""
158
+ # iov must be released before ioring and shm
159
+ for attr in ("iov_r", "iov_w", "ior_r", "ior_w"):
160
+ obj = getattr(self, attr, None)
161
+ if obj is not None:
162
+ del obj
163
+ setattr(self, attr, None)
164
+
165
+ for attr in ("shm_r", "shm_w"):
166
+ shm = getattr(self, attr, None)
167
+ if shm is not None:
168
+ try:
169
+ shm.close()
170
+ except Exception as e:
171
+ logger.warning("Failed to close %s: %s", attr, e)
172
+ setattr(self, attr, None)
173
+
174
+ if self.file is not None:
175
+ try:
176
+ deregister_fd(self.file)
177
+ except Exception as e:
178
+ logger.warning("deregister_fd failed: %s", e)
179
+ try:
180
+ os.close(self.file)
181
+ except OSError as e:
182
+ logger.warning("os.close failed: %s", e)
183
+ self.file = None
184
+
185
+ @rsynchronized()
186
+ def batch_read(self, offsets: list[int], tensors: list[torch.Tensor]) -> list[int]:
187
+ """Read data from the file at specified offsets into tensors.
188
+
189
+ Args:
190
+ offsets: List of byte offsets to read from
191
+ tensors: List of tensors to read data into
192
+
193
+ Returns:
194
+ List of operation results (0 for success, non-zero for error)
195
+ """
196
+ self.check(offsets, tensors)
197
+ assert self.ior_r is not None
198
+ assert self.iov_r is not None
199
+
200
+ # prepare
201
+ current = 0
202
+ for offset, tensor in zip(offsets, tensors):
203
+ size = tensor.numel() * tensor.itemsize
204
+ self.ior_r.prepare(
205
+ self.iov_r[current : current + size], True, self.file, offset
206
+ )
207
+ current += size
208
+
209
+ # submit
210
+ ionum = len(offsets)
211
+ resv = self.ior_r.submit().wait(min_results=ionum)
212
+
213
+ # results
214
+ with torch.cuda.stream(self.stream):
215
+ hf3fs_utils.read_shm(
216
+ self.shm_r_tensor, self.r_pinned, tensors, self.stream_ptr_int
217
+ )
218
+ results = [res.result for res in resv]
219
+
220
+ return results
221
+
222
+ @wsynchronized()
223
+ def batch_write(
224
+ self, offsets: list[int], tensors: list[torch.Tensor], event: torch.cuda.Event
225
+ ) -> list[int]:
226
+ """Write data from tensors to the file at specified offsets.
227
+
228
+ Args:
229
+ offsets: List of byte offsets to write to
230
+ tensors: List of tensors containing data to write
231
+
232
+ Returns:
233
+ List of operation results (0 for success, non-zero for error)
234
+ """
235
+
236
+ self.check(offsets, tensors)
237
+ assert self.ior_w is not None
238
+ assert self.iov_w is not None
239
+
240
+ # prepare
241
+ with torch.cuda.stream(self.stream):
242
+ self.stream.wait_event(event)
243
+ hf3fs_utils.write_shm(
244
+ tensors, self.shm_w_tensor, self.w_pinned, self.stream_ptr_int
245
+ )
246
+
247
+ current = 0
248
+ for offset, tensor in zip(offsets, tensors):
249
+ size = tensor.numel() * tensor.itemsize
250
+ self.ior_w.prepare(
251
+ self.iov_w[current : current + size], False, self.file, offset
252
+ )
253
+ current += size
254
+
255
+ # submit
256
+ ionum = len(offsets)
257
+ resv = self.ior_w.submit().wait(min_results=ionum)
258
+
259
+ # results
260
+ results = [res.result for res in resv]
261
+
262
+ return results
263
+
264
+ def check(self, offsets: list[int], tensors: list[torch.Tensor]) -> None:
265
+ sizes = [t.numel() * t.itemsize for t in tensors]
266
+ if any(
267
+ [
268
+ len(offsets) > self.entries,
269
+ len(offsets) != len(sizes),
270
+ any(
271
+ offset < 0 or offset + size > self.size
272
+ for offset, size in zip(offsets, sizes)
273
+ ),
274
+ any(size > self.bytes_per_page for size in sizes),
275
+ ]
276
+ ):
277
+ self.close()
278
+ raise ValueError("Hf3fsClient.check Failed")
279
+
280
+ def get_size(self) -> int:
281
+ """Get the total size of the storage file.
282
+
283
+ Returns:
284
+ Size of the file in bytes
285
+ """
286
+ return self.size
287
+
288
+ def close(self) -> None:
289
+ """Close the client and clean up resources."""
290
+ if self._closed:
291
+ return
292
+ self._closed = True
293
+ self._release_resources()
294
+
295
+ def flush(self) -> None:
296
+ """Flush any pending writes to disk."""
297
+ if not self._closed and self.file is not None:
298
+ os.fsync(self.file)
vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/hf3fs_connector.py
new
+1195
@@ -0,0 +1,1195 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+HF3FS KV Connector Implementation for vLLM.
5
+
6
+This module implements a KV connector that uses
7
+the 3FS for storing and retrieving KV cache data.
8
+
9
+Key components:
10
+1. HF3FSConnector: Main connector implementation
11
+ 2.1 AsyncOperationManager: Manages async save/load operations with background threads
12
+ 2.2 HF3FSConnectorMetadata: Container for connector metadata
13
+3. HF3FSMetadataServer: Mini Metadata server for HF3FS connector
14
+4. HF3FSClient: 3FS Client Implementation
15
+"""
16
+
17
+import atexit
18
+import concurrent
19
+import copy
20
+import hashlib
21
+import os
22
+import queue
23
+import signal
24
+import threading
25
+import time
26
+from concurrent.futures import Future
27
+from dataclasses import dataclass
28
+from queue import Empty
29
+from typing import Any, Optional
30
+
31
+import numpy as np
32
+import torch
33
+
34
+from vllm.config import VllmConfig
35
+from vllm.distributed.kv_transfer.kv_connector.v1.base import (
36
+ KVConnectorBase_V1,
37
+ KVConnectorMetadata,
38
+ KVConnectorRole,
39
+)
40
+from vllm.distributed.kv_transfer.kv_connector.v1.hf3fs.hf3fs_metadata_server import (
41
+ Hf3fsGlobalMetadataClient as Hf3fsMetadataClient,
42
+)
43
+from vllm.distributed.kv_transfer.kv_connector.v1.hf3fs.utils import (
44
+ gather_scatter_helper,
45
+)
46
+from vllm.distributed.kv_transfer.kv_connector.v1.hf3fs.utils.common import (
47
+ AtomicCounter,
48
+ HF3FSConnectorMetadata,
49
+ HF3FSRequestMetadata,
50
+ LoadBlockInfo,
51
+ RequestSchedulingState,
52
+)
53
+from vllm.distributed.kv_transfer.kv_connector.v1.hf3fs.utils.gather_scatter_helper import ( # noqa: E501
54
+ CopyBufferAllocator,
55
+)
56
+from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
57
+ KVConnectorPromMetrics,
58
+ KVConnectorStats,
59
+ PromMetric,
60
+ PromMetricT,
61
+)
62
+from vllm.distributed.parallel_state import get_tensor_model_parallel_rank
63
+from vllm.forward_context import ForwardContext
64
+from vllm.logger import init_logger
65
+from vllm.v1.attention.backend import AttentionMetadata
66
+from vllm.v1.core.kv_cache_manager import KVCacheBlocks
67
+from vllm.v1.core.sched.output import SchedulerOutput
68
+from vllm.v1.kv_cache_interface import KVCacheConfig
69
+from vllm.v1.metrics.utils import create_metric_per_engine
70
+from vllm.v1.request import Request
71
+
72
+HF3FS_AVAILABLE = True
73
+Hf3fsClient = None
74
+try:
75
+ from hf3fs_fuse.io import deregister_fd # noqa: F401
76
+
77
+ from vllm.distributed.kv_transfer.kv_connector.v1.hf3fs.hf3fs_client import (
78
+ Hf3fsClient as _RealClient,
79
+ )
80
+
81
+ Hf3fsClient = _RealClient
82
+except Exception:
83
+ HF3FS_AVAILABLE = False
84
+ from vllm.distributed.kv_transfer.kv_connector.v1.hf3fs.utils.hf3fs_mock_client import ( # noqa: E501
85
+ Hf3fsClient as _MockClient,
86
+ )
87
+
88
+ Hf3fsClient = _MockClient # type: ignore
89
+
90
+# Constants
91
+DEFAULT_MAX_IO_ENTRIES = 8
92
+
93
+logger = init_logger(__name__)
94
+
95
+
96
+# ============================================================================
97
+# Async Operation Management
98
+# ============================================================================
99
+
100
+
101
+class AsyncOperationManager:
102
+ """
103
+ Manages async save/load operations with background threads.
104
+ """
105
+
106
+ def __init__(self, connector: "HF3FSKVConnector"):
107
+ # Store connector reference and extract commonly used attributes
108
+ self._connector = connector
109
+ self._device = connector._device
110
+ self._dtype = connector._dtype
111
+ self._shape_per_page = connector._shape_per_page
112
+ self._bytes_per_page = connector._bytes_per_page
113
+ self._rank = connector._rank
114
+ self._numjobs = connector._numjobs
115
+ self._max_device_buffer_count = connector._max_device_buffer_count
116
+
117
+ # Operation tracking
118
+ self._save_futures: dict[str, list[Future]] = {}
119
+ self._load_futures: dict[str, Future] = {}
120
+ self._pending_finished_requests: set[str] = set()
121
+
122
+ # Initialize resources
123
+ self._init_cuda_resources()
124
+ self._init_worker_threads()
125
+
126
+ # Metrics
127
+ self.hf3fs_stats = HF3FSKVConnectorStats()
128
+
129
+ logger.info("AsyncOperationManager initialized for rank %d", self._rank)
130
+
131
+ def _init_cuda_resources(self) -> None:
132
+ """Initialize CUDA streams, events and buffer allocators."""
133
+ # CUDA streams for async operations
134
+ self._save_stream = torch.cuda.Stream()
135
+ self._load_stream = torch.cuda.Stream()
136
+ self._save_event = torch.cuda.Event()
137
+
138
+ # Buffer allocators for data copying
139
+ self._save_buffer_allocator = CopyBufferAllocator(
140
+ self._device,
141
+ self._dtype,
142
+ self._shape_per_page,
143
+ self._max_device_buffer_count,
144
+ )
145
+ self._load_buffer_allocator = CopyBufferAllocator(
146
+ self._device,
147
+ self._dtype,
148
+ self._shape_per_page,
149
+ self._max_device_buffer_count,
150
+ )
151
+
152
+ def _init_worker_threads(self) -> None:
153
+ """Initialize worker threads and I/O executor."""
154
+ # Thread synchronization
155
+ self._stop_event = threading.Event()
156
+ self._save_queue: queue.Queue[Any] = queue.Queue()
157
+ self._load_queue: queue.Queue[Any] = queue.Queue()
158
+
159
+ # I/O thread pool
160
+ self._io_executor = concurrent.futures.ThreadPoolExecutor(
161
+ max_workers=self._numjobs,
162
+ thread_name_prefix=f"HF3FS-Rank{self._rank}",
163
+ )
164
+
165
+ # Background worker threads
166
+ self._save_thread = threading.Thread(target=self._save_worker, daemon=True)
167
+ self._load_thread = threading.Thread(target=self._load_worker, daemon=True)
168
+ self._save_thread.start()
169
+ self._load_thread.start()
170
+
171
+ def submit_save_operation(self, request_id: str, block_ids, block_hashes) -> Future:
172
+ """Submit a save operation for async execution."""
173
+ future: Future[Any] = Future()
174
+ main_stream_event = torch.cuda.Event()
175
+ main_stream_event.record()
176
+ task = (request_id, block_ids, block_hashes, future, main_stream_event)
177
+ self._save_queue.put(task)
178
+
179
+ if request_id not in self._save_futures:
180
+ self._save_futures[request_id] = []
181
+ self._save_futures[request_id].append(future)
182
+ return future
183
+
184
+ def submit_load_operation(self, request_id: str, block_ids, block_hashes) -> Future:
185
+ """Submit a load operation for async execution."""
186
+ future: Future[Any] = Future()
187
+ task = (request_id, block_ids, block_hashes, future)
188
+ self._load_queue.put(task)
189
+ self._load_futures[request_id] = future
190
+ return future
191
+
192
+ def get_finished_operations(
193
+ self, finished_req_ids: set[str]
194
+ ) -> tuple[set[str], set[str]]:
195
+ completed_saves = self._check_completed_saves(finished_req_ids)
196
+ completed_loads = self._check_completed_loads()
197
+
198
+ if completed_saves or completed_loads:
199
+ logger.info(
200
+ "HF3FS Connector Completed: %d saves, %d loads operations",
201
+ len(completed_saves),
202
+ len(completed_loads),
203
+ )
204
+
205
+ return completed_saves, completed_loads
206
+
207
+ def _check_completed_saves(self, finished_req_ids: set[str]) -> set[str]:
208
+ """Check for completed save operations."""
209
+ completed = set()
210
+
211
+ # Check pending finished requests first
212
+ for request_id in list(self._pending_finished_requests):
213
+ if request_id in self._save_futures and self._all_saves_done(request_id):
214
+ completed.add(request_id)
215
+ self._save_futures.pop(request_id)
216
+ self._pending_finished_requests.remove(request_id)
217
+
218
+ # Process newly finished requests
219
+ for request_id in finished_req_ids:
220
+ if request_id in self._save_futures:
221
+ if self._all_saves_done(request_id):
222
+ completed.add(request_id)
223
+ self._save_futures.pop(request_id)
224
+ else:
225
+ self._pending_finished_requests.add(request_id)
226
+ else:
227
+ completed.add(request_id)
228
+
229
+ return completed
230
+
231
+ def _check_completed_loads(self) -> set[str]:
232
+ """Check for completed load operations."""
233
+ completed = set()
234
+ for request_id in list(self._load_futures):
235
+ if self._load_futures[request_id].done():
236
+ completed.add(request_id)
237
+ self._load_futures.pop(request_id)
238
+ return completed
239
+
240
+ def _all_saves_done(self, request_id: str) -> bool:
241
+ """Check if all save operations for a request are completed."""
242
+ return all(future.done() for future in self._save_futures[request_id])
243
+
244
+ def _save_worker(self) -> None:
245
+ """Background worker for handling save operations."""
246
+ torch.accelerator.set_device_index(self._device)
247
+ while not self._stop_event.is_set():
248
+ try:
249
+ task = self._save_queue.get(block=True, timeout=1)
250
+ self._handle_save_task(task)
251
+ except Empty:
252
+ continue
253
+ except Exception as e:
254
+ logger.error("Save worker error: %s", e)
255
+
256
+ def _load_worker(self) -> None:
257
+ """Background worker for handling load operations."""
258
+ torch.accelerator.set_device_index(self._device)
259
+ while not self._stop_event.is_set():
260
+ try:
261
+ task = self._load_queue.get(block=True, timeout=1)
262
+ self._handle_load_task(task)
263
+ except Empty:
264
+ continue
265
+ except Exception as e:
266
+ logger.error("Load worker error: %s", e)
267
+
268
+ def _handle_save_task(self, task) -> None:
269
+ """Handle individual save task with proper stream synchronization."""
270
+ request_id, block_ids, block_hashes, future, main_stream_event = task
271
+ start_time = time.perf_counter()
272
+ buffers = None
273
+ try:
274
+ # Step1: Allocate storage pages
275
+ key_pairs = [(hash_val, "") for hash_val in block_hashes]
276
+ allocation_results = (
277
+ self._connector._metadata_client.allocate_pages_for_keys(
278
+ self._rank, key_pairs
279
+ )
280
+ )
281
+
282
+ if any(result[1] < 0 for result in allocation_results):
283
+ return self._fail_task(
284
+ "Saved", "Page allocation failed", request_id, future
285
+ )
286
+
287
+ page_indices = [result[1] for result in allocation_results]
288
+ offsets = [idx * self._bytes_per_page for idx in page_indices]
289
+
290
+ # Step2: Allocate buffers and gather KV cache data
291
+ buffers = self._save_buffer_allocator.alloc_buffer(len(block_ids))
292
+ if buffers is None:
293
+ return self._fail_task(
294
+ "Saved",
295
+ f"Buffer allocation failed for {len(block_ids)} blocks",
296
+ request_id,
297
+ future,
298
+ )
299
+
300
+ # Synchronize streams and gather data
301
+ with torch.cuda.stream(self._save_stream):
302
+ self._save_stream.wait_event(main_stream_event) # Wait for main stream
303
+ self._connector._gather_or_scatter_kv_caches(
304
+ block_ids, buffers, "gather"
305
+ )
306
+
307
+ save_stream_event = torch.cuda.Event()
308
+ save_stream_event.record(self._save_stream) # Record gather completion
309
+
310
+ # Step3: Write data in batches
311
+ write_futures = []
312
+ for i in range(0, len(offsets), DEFAULT_MAX_IO_ENTRIES):
313
+ batch_offsets = offsets[i : i + DEFAULT_MAX_IO_ENTRIES]
314
+ batch_buffers = buffers[i : i + DEFAULT_MAX_IO_ENTRIES]
315
+ client = self._connector._clients[self._connector._ac.next()]
316
+ write_future = self._io_executor.submit(
317
+ client.batch_write, batch_offsets, batch_buffers, save_stream_event
318
+ )
319
+ write_futures.append(write_future)
320
+
321
+ # Check write results
322
+ write_success = all(
323
+ result == self._bytes_per_page
324
+ for write_future in write_futures
325
+ for result in write_future.result()
326
+ )
327
+
328
+ # Step4: Confirm writes to metadata server
329
+ if write_success:
330
+ written_keys = list(zip(block_hashes, page_indices))
331
+ self._connector._metadata_client.confirm_write_for_keys(
332
+ self._rank, written_keys, []
333
+ )
334
+ self._save_buffer_allocator.free_buffer(buffers)
335
+ return self._succeed_task(
336
+ "Saved", start_time, request_id, len(block_ids), future
337
+ )
338
+ else:
339
+ self._connector._metadata_client.confirm_write_for_keys(
340
+ self._rank, [], page_indices
341
+ )
342
+ self._save_buffer_allocator.free_buffer(buffers)
343
+ return self._fail_task(
344
+ "Saved", "Write operation failed", request_id, future
345
+ )
346
+
347
+ except Exception as e:
348
+ if buffers is not None:
349
+ self._save_buffer_allocator.free_buffer(buffers)
350
+ return self._fail_task(
351
+ "Saved", f"Task execution error: {e}", request_id, future
352
+ )
353
+
354
+ def _handle_load_task(self, task) -> None:
355
+ """Handle individual load task."""
356
+ request_id, block_ids, block_hashes, future = task
357
+ start_time = time.perf_counter()
358
+ buffers = None
359
+ try:
360
+ # Step1: Get block locations from metadata server
361
+ page_indices = self._connector._metadata_client.get_key_locations(
362
+ self._rank, block_hashes
363
+ )
364
+
365
+ if any(idx is None for idx in page_indices):
366
+ return self._fail_task("Loaded", "Blocks not found", request_id, future)
367
+
368
+ # Allocate read buffer
369
+ buffers = self._load_buffer_allocator.alloc_buffer(len(block_ids))
370
+ if buffers is None:
371
+ return self._fail_task(
372
+ "Loaded",
373
+ f"Buffer allocation failed for {len(block_ids)} blocks",
374
+ request_id,
375
+ future,
376
+ )
377
+
378
+ # Step2: Read data in batches
379
+ offsets = [idx * self._bytes_per_page for idx in page_indices]
380
+ read_futures = []
381
+ for i in range(0, len(offsets), DEFAULT_MAX_IO_ENTRIES):
382
+ batch_offsets = offsets[i : i + DEFAULT_MAX_IO_ENTRIES]
383
+ batch_buffers = buffers[i : i + DEFAULT_MAX_IO_ENTRIES]
384
+ client = self._connector._clients[self._connector._ac.next()]
385
+ read_future = self._io_executor.submit(
386
+ client.batch_read, batch_offsets, batch_buffers
387
+ )
388
+ read_futures.append(read_future)
389
+
390
+ # Check read results
391
+ read_success = all(
392
+ result == self._bytes_per_page
393
+ for read_future in read_futures
394
+ for result in read_future.result()
395
+ )
396
+
397
+ if not read_success:
398
+ self._load_buffer_allocator.free_buffer(buffers)
399
+ return self._fail_task(
400
+ "Loaded", "Read operation failed", request_id, future
401
+ )
402
+
403
+ # Step3: Scatter data back to KV cache
404
+ with torch.cuda.stream(self._load_stream):
405
+ self._connector._gather_or_scatter_kv_caches(
406
+ block_ids, buffers, "scatter"
407
+ )
408
+
409
+ self._load_stream.synchronize()
410
+ self._load_buffer_allocator.free_buffer(buffers)
411
+ return self._succeed_task(
412
+ "Loaded", start_time, request_id, len(block_ids), future
413
+ )
414
+
415
+ except Exception as e:
416
+ if buffers is not None:
417
+ self._load_buffer_allocator.free_buffer(buffers)
418
+ return self._fail_task(
419
+ "Loaded", f"Task execution error: {e}", request_id, future
420
+ )
421
+
422
+ def _fail_task(
423
+ self, operation: str, error_msg: str, request_id: str, future: Future
424
+ ) -> None:
425
+ """Helper to fail task with error logging."""
426
+ logger.error(
427
+ "%s for %s request %s",
428
+ error_msg,
429
+ operation,
430
+ request_id,
431
+ )
432
+ self.hf3fs_stats.record_failed_task_count(operation)
433
+ future.set_result(False)
434
+
435
+ def _succeed_task(
436
+ self,
437
+ operation: str,
438
+ start_time: float,
439
+ request_id: str,
440
+ block_count: int,
441
+ future: Future,
442
+ ) -> None:
443
+ """Helper to succeed task with logging."""
444
+ duration = time.perf_counter() - start_time
445
+ logger.info(
446
+ "%s %s: %d blocks in %.2fs",
447
+ operation,
448
+ request_id,
449
+ block_count,
450
+ duration,
451
+ )
452
+ self.hf3fs_stats.record_success_task_duration(operation, duration)
453
+ future.set_result(True)
454
+
455
+ def shutdown(self) -> None:
456
+ """Clean shutdown of all background threads and resources."""
457
+ self._stop_event.set()
458
+ self._save_thread.join()
459
+ self._load_thread.join()
460
+ self._io_executor.shutdown(wait=True)
461
+ logger.info("AsyncOperationManager shutdown completed")
462
+
463
+
464
+# ============================================================================
465
+# HF3FS Connector
466
+# ============================================================================
467
+
468
+
469
+class HF3FSKVConnector(KVConnectorBase_V1):
470
+ """HF3FS KV Connector implementation."""
471
+
472
+ def __init__(
473
+ self,
474
+ vllm_config: "VllmConfig",
475
+ role: KVConnectorRole,
476
+ kv_cache_config: "KVCacheConfig",
477
+ ):
478
+ super().__init__(
479
+ vllm_config=vllm_config, role=role, kv_cache_config=kv_cache_config
480
+ )
481
+
482
+ # Core configuration
483
+ self._vllm_config = vllm_config
484
+ self._role = role
485
+ self._block_size = vllm_config.cache_config.block_size
486
+ self._use_mla = vllm_config.model_config.use_mla
487
+ self._model_config = vllm_config.model_config
488
+
489
+ logger.info("Using MLA: %s", self._use_mla)
490
+
491
+ # HF3FS configuration
492
+ kv_config = vllm_config.kv_transfer_config
493
+ assert kv_config is not None
494
+
495
+ self._storage_path = kv_config.get_from_extra_config(
496
+ "hf3fs_storage_path", "/vllm-workspace/mnt/hf3fs"
497
+ )
498
+ self._metadata_server_url = kv_config.get_from_extra_config(
499
+ "hf3fs_metadata_server_url", "http://localhost:18000"
500
+ )
501
+ self._file_size = kv_config.get_from_extra_config(
502
+ "hf3fs_file_size", 1024 * 1024 * 1024
503
+ )
504
+ self._numjobs = kv_config.get_from_extra_config("hf3fs_client_numjobs", 16)
505
+ self._max_device_buffer_count = kv_config.get_from_extra_config(
506
+ "hf3fs_max_device_buffer_count", 128
507
+ )
508
+ self._max_device_buffer_count = max(
509
+ self._max_device_buffer_count, self._numjobs * DEFAULT_MAX_IO_ENTRIES
510
+ )
511
+
512
+ if self._role == KVConnectorRole.SCHEDULER:
513
+ self._scheduling_states: dict[str, RequestSchedulingState] = {}
514
+ self._metadata_client = Hf3fsMetadataClient()
515
+ self._metadata_client.initialize(0, role="scheduler")
516
+
517
+ atexit.register(self.close)
518
+ signal.signal(signal.SIGINT, lambda sig, frame: self.close())
519
+ signal.signal(signal.SIGTERM, lambda sig, frame: self.close())
520
+ signal.signal(signal.SIGQUIT, lambda sig, frame: self.close())
521
+
522
+ logger.info(
523
+ "HF3FSKVConnector initialized: path=%s, role=%s",
524
+ self._storage_path,
525
+ self._role.name,
526
+ )
527
+
528
+ ############################################################
529
+ # Worker Side Methods
530
+ ############################################################
531
+
532
+ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]) -> None:
533
+ self._kv_caches = kv_caches
534
+ self._setup_kv_cache_config()
535
+ self._setup_storage_clients()
536
+ self._async_manager = AsyncOperationManager(self)
537
+
538
+ def _setup_kv_cache_config(self):
539
+ first_cache = next(iter(self._kv_caches.values()))
540
+ self._device = first_cache.device
541
+ self._dtype = first_cache.dtype
542
+ element_size = first_cache.element_size()
543
+
544
+ if self._use_mla:
545
+ assert len(first_cache.shape) == 3, "MLA format should have 3 dimensions"
546
+ # MLA format: [num_blocks, block_size, head_size]
547
+ num_blocks, block_size, head_size = first_cache.shape
548
+ num_heads = 1
549
+ else:
550
+ # MHA format: [2, num_blocks, block_size, num_heads, head_size]
551
+ _, num_blocks, block_size, num_heads, head_size = first_cache.shape
552
+
553
+ self._local_total_tokens = num_blocks * block_size
554
+ self._local_block_size = block_size
555
+
556
+ if self._use_mla:
557
+ layer_block_size = block_size * head_size * element_size
558
+ self._bytes_per_page = layer_block_size * len(self._kv_caches)
559
+ self._shape_per_page = [
560
+ len(self._kv_caches),
561
+ block_size,
562
+ head_size,
563
+ ]
564
+ else:
565
+ layer_block_size = 2 * block_size * num_heads * head_size * element_size
566
+ self._bytes_per_page = layer_block_size * len(self._kv_caches)
567
+ self._shape_per_page = [
568
+ len(self._kv_caches),
569
+ 2,
570
+ block_size,
571
+ num_heads * head_size,
572
+ ]
573
+
574
+ self._kvcache_ptrs = torch.tensor(
575
+ [cache.data_ptr() for cache in self._kv_caches.values()],
576
+ dtype=torch.int64,
577
+ device=self._device,
578
+ )
579
+
580
+ def _setup_storage_clients(self):
581
+ os.makedirs(self._storage_path, exist_ok=True)
582
+
583
+ self._rank = get_tensor_model_parallel_rank()
584
+ file_path = os.path.join(
585
+ self._storage_path, f"hf3fs_vllm_data_file_{self._rank}"
586
+ )
587
+
588
+ try:
589
+ # Initialize HF3FS clients
590
+ self._ac = AtomicCounter(self._numjobs)
591
+ assert Hf3fsClient is not None
592
+ self._clients = [
593
+ Hf3fsClient(
594
+ path=file_path,
595
+ size=self._file_size,
596
+ bytes_per_page=self._bytes_per_page,
597
+ entries=DEFAULT_MAX_IO_ENTRIES,
598
+ )
599
+ for _ in range(self._numjobs)
600
+ ]
601
+
602
+ # Initialize metadata client
603
+ num_pages = self._file_size // self._bytes_per_page
604
+ self._metadata_client = Hf3fsMetadataClient()
605
+ self._metadata_client.initialize(self._rank, num_pages, role="worker")
606
+ except Exception as e:
607
+ logger.error("HF3FS client initialization failed: %s", e)
608
+ raise
609
+
610
+ def save_kv_layer(
611
+ self,
612
+ layer_name: str,
613
+ kv_layer: torch.Tensor,
614
+ attn_metadata: "AttentionMetadata",
615
+ **kwargs,
616
+ ) -> None:
617
+ """HF3FSConnector does not do layerwise saving."""
618
+ pass
619
+
620
+ def wait_for_save(self) -> None:
621
+ metadata = self._get_connector_metadata()
622
+ if not isinstance(metadata, HF3FSConnectorMetadata):
623
+ logger.error("Invalid metadata type: %s", type(metadata))
624
+ return
625
+
626
+ for request in metadata.requests:
627
+ if request.save_block_op is None:
628
+ continue
629
+
630
+ skip_blocks = request.save_block_op.skip_leading_blocks
631
+ block_hashes = self._generate_block_hashes(request.token_ids, skip_blocks)
632
+ block_ids = request.block_ids[skip_blocks : skip_blocks + len(block_hashes)]
633
+
634
+ for i in range(0, len(block_ids), self._max_device_buffer_count):
635
+ batch_block_ids = block_ids[i : i + self._max_device_buffer_count]
636
+ batch_block_hashes = block_hashes[i : i + self._max_device_buffer_count]
637
+ self._async_manager.submit_save_operation(
638
+ request.request_id, batch_block_ids, batch_block_hashes
639
+ )
640
+
641
+ def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None:
642
+ metadata = self._get_connector_metadata()
643
+ if not isinstance(metadata, HF3FSConnectorMetadata):
644
+ logger.error("Invalid metadata type for loading")
645
+ return
646
+
647
+ for request in metadata.requests:
648
+ if request.load_block_op is None:
649
+ continue
650
+
651
+ load_op = request.load_block_op
652
+ block_ids = request.block_ids[: load_op.num_blocks_to_load]
653
+ block_hashes = self._generate_block_hashes(
654
+ request.token_ids, load_op.num_computed_blocks, len(block_ids)
655
+ )
656
+
657
+ for i in range(0, len(block_ids), self._max_device_buffer_count):
658
+ batch_block_ids = block_ids[i : i + self._max_device_buffer_count]
659
+ batch_block_hashes = block_hashes[i : i + self._max_device_buffer_count]
660
+ self._async_manager.submit_load_operation(
661
+ request.request_id, batch_block_ids, batch_block_hashes
662
+ )
663
+
664
+ def wait_for_layer_load(self, layer_name: str) -> None:
665
+ pass
666
+
667
+ def get_finished(
668
+ self, finished_req_ids: set[str]
669
+ ) -> tuple[set[str] | None, set[str] | None]:
670
+ return self._async_manager.get_finished_operations(finished_req_ids)
671
+
672
+ def get_kv_connector_stats(self) -> Optional["KVConnectorStats"]:
673
+ """
674
+ Get the KV connector stats collected during the last interval.
675
+ """
676
+ # Clear stats for next iteration
677
+ if (
678
+ hasattr(self, "_async_manager")
679
+ and not self._async_manager.hf3fs_stats.is_empty()
680
+ ):
681
+ return self._async_manager.hf3fs_stats.clone_and_reset()
682
+ return None
683
+
684
+ ############################################################
685
+ # Scheduler Side Methods
686
+ ############################################################
687
+
688
+ def request_finished(
689
+ self,
690
+ request: "Request",
691
+ block_ids: list[int],
692
+ ) -> tuple[bool, dict[str, Any] | None]:
693
+ return True, None
694
+
695
+ def get_num_new_matched_tokens(
696
+ self, request: "Request", num_computed_tokens: int
697
+ ) -> tuple[int, bool]:
698
+ """Get number of new tokens that can be loaded from external cache."""
699
+ try:
700
+ state = self._get_or_create_scheduling_state(request.request_id)
701
+ state.request = request
702
+ assert request.prompt_token_ids is not None
703
+
704
+ num_tokens_to_check = self._align_to_block_size(
705
+ len(request.prompt_token_ids) - 1
706
+ )
707
+
708
+ if num_tokens_to_check <= num_computed_tokens:
709
+ state.load_op = LoadBlockInfo(
710
+ num_computed_blocks=num_computed_tokens // self._block_size,
711
+ num_blocks_to_load=0,
712
+ need_fetch_block_ids=[],
713
+ )
714
+ return 0, False
715
+
716
+ token_ids_to_check = request.prompt_token_ids[:num_tokens_to_check]
717
+ block_hashes = self._generate_block_hashes(token_ids_to_check, 0)
718
+
719
+ # Check existence
720
+ exists_results = self._metadata_client.batch_key_exists(block_hashes)
721
+
722
+ # Count consecutive matches
723
+ matched_blocks = next(
724
+ (i for i, exists in enumerate(exists_results) if not exists),
725
+ len(exists_results),
726
+ )
727
+ matched_tokens = matched_blocks * self._block_size
728
+ new_hit_tokens = max(0, matched_tokens - num_computed_tokens)
729
+
730
+ # Store load operation
731
+ state.load_op = LoadBlockInfo(
732
+ num_computed_blocks=num_computed_tokens // self._block_size,
733
+ num_blocks_to_load=new_hit_tokens // self._block_size,
734
+ need_fetch_block_ids=[],
735
+ )
736
+
737
+ logger.info(
738
+ (
739
+ "Token matching for %s: "
740
+ "%d matched (%d blocks), "
741
+ "%d new hits, "
742
+ "prompt len %d"
743
+ ),
744
+ request.request_id,
745
+ matched_tokens,
746
+ matched_blocks,
747
+ new_hit_tokens,
748
+ len(request.prompt_token_ids),
749
+ )
750
+ return new_hit_tokens, new_hit_tokens > 0
751
+
752
+ except Exception as e:
753
+ logger.error(
754
+ "Error calculating matches for request %s: %s", request.request_id, e
755
+ )
756
+ return 0, False
757
+
758
+ def update_state_after_alloc(
759
+ self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int
760
+ ) -> None:
761
+ """Update state after block allocation."""
762
+ state = self._get_or_create_scheduling_state(request.request_id)
763
+ state.request = request
764
+
765
+ if num_external_tokens <= 0 or not state.needs_loading():
766
+ return
767
+
768
+ # Validate block allocation
769
+ assert state.load_op is not None
770
+ expected_blocks = state.load_op.num_blocks_to_load
771
+ actual_blocks = num_external_tokens // self._block_size
772
+ assert actual_blocks == expected_blocks, (
773
+ f"Block count mismatch for {request.request_id}: "
774
+ f"expected {expected_blocks}, got {actual_blocks}"
775
+ )
776
+
777
+ # Update load operation with allocated block IDs
778
+ if actual_blocks > 0:
779
+ local_block_ids = blocks.get_unhashed_block_ids()
780
+ state.load_op.need_fetch_block_ids.extend(local_block_ids)
781
+ state.phase = "WAITING_TO_LOAD"
782
+
783
+ def build_connector_meta(
784
+ self, scheduler_output: SchedulerOutput
785
+ ) -> KVConnectorMetadata:
786
+ """Build connector metadata for scheduling step."""
787
+ metadata = HF3FSConnectorMetadata()
788
+
789
+ for request_id in scheduler_output.finished_req_ids:
790
+ self._scheduling_states.pop(request_id, None)
791
+
792
+ # Process requests by phase
793
+ self._process_waiting_to_load_requests(metadata)
794
+ self._process_new_requests(scheduler_output, metadata)
795
+ self._process_cached_requests(scheduler_output, metadata)
796
+
797
+ return metadata
798
+
799
+ def _process_waiting_to_load_requests(
800
+ self, metadata: HF3FSConnectorMetadata
801
+ ) -> None:
802
+ """Process requests waiting to load."""
803
+ for state in list(self._scheduling_states.values()):
804
+ if not state.is_ready_to_load():
805
+ continue
806
+ assert state.load_op is not None
807
+ assert (
808
+ state.request is not None and state.request.prompt_token_ids is not None
809
+ )
810
+ # Create load request metadata
811
+ num_cached_blocks = (
812
+ state.load_op.num_computed_blocks + state.load_op.num_blocks_to_load
813
+ )
814
+ num_tokens_to_compute = num_cached_blocks * self._block_size
815
+
816
+ # Initialize token_ids and allocated_block_ids for loading
817
+ state.token_ids = state.request.prompt_token_ids[
818
+ :num_tokens_to_compute
819
+ ].copy()
820
+ state.allocated_block_ids = state.load_op.need_fetch_block_ids.copy()
821
+
822
+ request_metadata = HF3FSRequestMetadata.from_scheduling_state(
823
+ state, self._block_size, state.load_op, num_cached_blocks
824
+ )
825
+
826
+ if request_metadata:
827
+ metadata.add_request(request_metadata)
828
+ state.phase = "ACTIVE"
829
+
830
+ def _process_new_requests(
831
+ self, scheduler_output: SchedulerOutput, metadata: HF3FSConnectorMetadata
832
+ ) -> None:
833
+ """Process new requests."""
834
+ for request in scheduler_output.scheduled_new_reqs:
835
+ state = self._get_or_create_scheduling_state(request.req_id)
836
+
837
+ # Calculate tokens to compute
838
+ num_tokens_to_compute = (
839
+ request.num_computed_tokens
840
+ + scheduler_output.num_scheduled_tokens[request.req_id]
841
+ )
842
+ self._initialize_state_from_new_request(
843
+ state, request, num_tokens_to_compute
844
+ )
845
+
846
+ # Create save metadata (skip cached blocks if any)
847
+ num_cached_blocks = None
848
+ if state.load_op:
849
+ num_cached_blocks = (
850
+ state.load_op.num_computed_blocks + state.load_op.num_blocks_to_load
851
+ )
852
+
853
+ request_metadata = HF3FSRequestMetadata.from_scheduling_state(
854
+ state, self._block_size, None, num_cached_blocks
855
+ )
856
+
857
+ if request_metadata:
858
+ metadata.add_request(request_metadata)
859
+ state.phase = "ACTIVE"
860
+
861
+ def _process_cached_requests(
862
+ self, scheduler_output: SchedulerOutput, metadata: HF3FSConnectorMetadata
863
+ ) -> None:
864
+ """Process cached requests."""
865
+ cached_reqs = scheduler_output.scheduled_cached_reqs
866
+ for i, request_id in enumerate(cached_reqs.req_ids):
867
+ state = self._get_or_create_scheduling_state(request_id)
868
+ assert state.request is not None
869
+
870
+ # Update with new tokens and blocks
871
+ num_new_tokens = scheduler_output.num_scheduled_tokens[request_id]
872
+ num_current_tokens = len(state.token_ids)
873
+ new_token_ids = state.request.all_token_ids[
874
+ num_current_tokens : num_current_tokens + num_new_tokens
875
+ ]
876
+ new_block_ids = cached_reqs.new_block_ids[i]
877
+
878
+ state.update_tokens_and_blocks(new_token_ids, new_block_ids)
879
+
880
+ # Create save metadata
881
+ request_metadata = HF3FSRequestMetadata.from_scheduling_state(
882
+ state, self._block_size, None
883
+ )
884
+
885
+ if request_metadata:
886
+ metadata.add_request(request_metadata)
887
+
888
+ @classmethod
889
+ def build_kv_connector_stats(
890
+ cls, data: dict[str, Any] | None = None
891
+ ) -> Optional["KVConnectorStats"]:
892
+ """
893
+ KVConnectorStats resolution method. This method allows dynamically
894
+ registered connectors to return their own KVConnectorStats object,
895
+ which can implement custom aggregation logic on the data dict.
896
+ """
897
+ return (
898
+ HF3FSKVConnectorStats(data=data)
899
+ if data is not None
900
+ else HF3FSKVConnectorStats()
901
+ )
902
+
903
+ @classmethod
904
+ def build_prom_metrics(
905
+ cls,
906
+ vllm_config: VllmConfig,
907
+ metric_types: dict[type[PromMetric], type[PromMetricT]],
908
+ labelnames: list[str],
909
+ per_engine_labelvalues: dict[int, list[object]],
910
+ ) -> KVConnectorPromMetrics:
911
+ return HF3FSPromMetrics(
912
+ vllm_config, metric_types, labelnames, per_engine_labelvalues
913
+ )
914
+
915
+ def close(self) -> None:
916
+ try:
917
+ if hasattr(self, "_async_manager"):
918
+ self._async_manager.shutdown()
919
+
920
+ if hasattr(self, "_clients"):
921
+ for client in self._clients:
922
+ client.close()
923
+ logger.info("HF3FS clients closed")
924
+ except Exception as e:
925
+ logger.error("Connector shutdown error: %s", e)
926
+
927
+ ############################################################
928
+ # Utility Methods
929
+ ############################################################
930
+
931
+ def _get_or_create_scheduling_state(
932
+ self, request_id: str
933
+ ) -> RequestSchedulingState:
934
+ """Get existing or create new scheduling state."""
935
+ if request_id not in self._scheduling_states:
936
+ self._scheduling_states[request_id] = RequestSchedulingState(
937
+ request_id=request_id
938
+ )
939
+ return self._scheduling_states[request_id]
940
+
941
+ def _initialize_state_from_new_request(
942
+ self, state: RequestSchedulingState, request, num_tokens_to_compute: int
943
+ ) -> None:
944
+ """Initialize state from new request data."""
945
+ # Handle different block_ids formats in vLLM 0.9.0+
946
+ if isinstance(request.block_ids[0], list):
947
+ unfolded_block_ids = request.block_ids[0].copy()
948
+ else:
949
+ unfolded_block_ids = request.block_ids.copy()
950
+
951
+ state.token_ids = request.prompt_token_ids[:num_tokens_to_compute].copy()
952
+ state.allocated_block_ids = unfolded_block_ids
953
+ state.num_saved_blocks = 0
954
+
955
+ def _generate_block_hashes(
956
+ self,
957
+ token_ids: list[int],
958
+ start_block_id: int,
959
+ max_blocks_count: int | None = None,
960
+ ) -> list[str]:
961
+ """Generate block hashes for token sequence."""
962
+ block_hashes = []
963
+ previous_hash = ""
964
+
965
+ for start_idx in range(0, len(token_ids), self._block_size):
966
+ if start_idx + self._block_size > len(token_ids):
967
+ break
968
+
969
+ end_idx = start_idx + self._block_size
970
+ block_hash = self._compute_prefix_hash(
971
+ token_ids[start_idx:end_idx], previous_hash
972
+ )
973
+
974
+ block_index = start_idx // self._block_size
975
+ if block_index >= start_block_id:
976
+ block_hashes.append(block_hash)
977
+
978
+ if max_blocks_count and len(block_hashes) >= max_blocks_count:
979
+ break
980
+ previous_hash = block_hash
981
+
982
+ return block_hashes
983
+
984
+ def _gather_or_scatter_kv_caches(
985
+ self, block_ids: list[int], block_buffers, operation: str
986
+ ):
987
+ for buffer_tensor, block_id in zip(block_buffers, block_ids):
988
+ start_idx = block_id * self._local_block_size
989
+ token_indices = list(range(start_idx, start_idx + self._local_block_size))
990
+ if operation == "gather":
991
+ gather_scatter_helper.gather_kv_caches(
992
+ self._kvcache_ptrs,
993
+ self._local_total_tokens,
994
+ buffer_tensor,
995
+ token_indices,
996
+ is_mla=self._use_mla,
997
+ )
998
+ else:
999
+ gather_scatter_helper.scatter_kv_caches(
1000
+ self._kvcache_ptrs,
1001
+ self._local_total_tokens,
1002
+ buffer_tensor,
1003
+ token_indices,
1004
+ is_mla=self._use_mla,
1005
+ )
1006
+
1007
+ def _compute_prefix_hash(
1008
+ self, token_ids: list[int], previous_hash: str = ""
1009
+ ) -> str:
1010
+ """Compute prefix hash for token block."""
1011
+ combined_string = f"{previous_hash}_{token_ids}"
1012
+ return hashlib.md5(combined_string.encode()).hexdigest()
1013
+
1014
+ def _align_to_block_size(self, num_tokens: int) -> int:
1015
+ """Align token count to block size."""
1016
+ return (num_tokens // self._block_size) * self._block_size
1017
+
1018
+
1019
+@dataclass
1020
+class HF3FSKVConnectorStats(KVConnectorStats):
1021
+ """Container for transfer performance metrics"""
1022
+
1023
+ def __post_init__(self):
1024
+ if not self.data:
1025
+ # Empty container init, no data is passed in.
1026
+ self.reset()
1027
+
1028
+ def reset(self):
1029
+ # Must be serializable
1030
+ self.data: dict[str, Any] = {
1031
+ "save_duration": [],
1032
+ "load_duration": [],
1033
+ "num_failed_save": 0,
1034
+ "num_failed_load": 0,
1035
+ "num_transfer_task": 0,
1036
+ }
1037
+
1038
+ def aggregate(self, other: "KVConnectorStats") -> "KVConnectorStats":
1039
+ if not other.is_empty():
1040
+ for k, v in other.data.items():
1041
+ accumulator = self.data[k]
1042
+ if isinstance(accumulator, list):
1043
+ accumulator.extend(v)
1044
+ else: # int
1045
+ self.data[k] += v
1046
+ return self
1047
+
1048
+ def reduce(self) -> dict[str, int | float]:
1049
+ # Compute compact representative stats suitable for CLI logging
1050
+ if self.is_empty():
1051
+ return {
1052
+ "Num transfers task": 0,
1053
+ "Num save task success": 0,
1054
+ "Num save task failed": 0,
1055
+ "Num load task success": 0,
1056
+ "Num load task failed": 0,
1057
+ "Avg save duration (ms)": 0,
1058
+ "P90 save duration (ms)": 0,
1059
+ "Avg load duration (ms)": 0,
1060
+ "P90 load duration (ms)": 0,
1061
+ }
1062
+ num_success_save = len(self.data["save_duration"] or [])
1063
+ num_success_load = len(self.data["load_duration"] or [])
1064
+ num_failed_save = self.data["num_failed_save"]
1065
+ num_failed_load = self.data["num_failed_load"]
1066
+ if num_success_save == 0:
1067
+ save_duration = np.zeros(1)
1068
+ else:
1069
+ save_duration = np.asarray(self.data["save_duration"])
1070
+ if num_success_load == 0:
1071
+ load_duration = np.zeros(1)
1072
+ else:
1073
+ load_duration = np.asarray(self.data["load_duration"])
1074
+
1075
+ return {
1076
+ "Num transfers task": self.data["num_transfer_task"],
1077
+ "Num save task success": num_success_save,
1078
+ "Num save task failed": num_failed_save,
1079
+ "Num load task success": num_success_load,
1080
+ "Num load task failed": num_failed_load,
1081
+ "Avg save duration (ms)": round(save_duration.mean() * 1e3, 3),
1082
+ "P90 save duration (ms)": round(np.percentile(save_duration, 90) * 1e3, 3),
1083
+ "Avg load duration (ms)": round(load_duration.mean() * 1e3, 3),
1084
+ "P90 load duration (ms)": round(np.percentile(load_duration, 90) * 1e3, 3),
1085
+ }
1086
+
1087
+ def is_empty(self) -> bool:
1088
+ return self.data["num_transfer_task"] == 0
1089
+
1090
+ def record_success_task_duration(self, operation, duration):
1091
+ if operation == "Saved":
1092
+ self.data["save_duration"].append(duration)
1093
+ elif operation == "Loaded":
1094
+ self.data["load_duration"].append(duration)
1095
+ self.data["num_transfer_task"] += 1
1096
+
1097
+ def record_failed_task_count(self, operation):
1098
+ if operation == "Saved":
1099
+ self.data["num_failed_save"] += 1
1100
+ elif operation == "Loaded":
1101
+ self.data["num_failed_load"] += 1
1102
+ self.data["num_transfer_task"] += 1
1103
+
1104
+ def clone_and_reset(self):
1105
+ old = copy.copy(self)
1106
+ self.reset()
1107
+ return old
1108
+
1109
+
1110
+class HF3FSPromMetrics(KVConnectorPromMetrics):
1111
+ def __init__(
1112
+ self,
1113
+ vllm_config: VllmConfig,
1114
+ metric_types: dict[type[PromMetric], type[PromMetricT]],
1115
+ labelnames: list[str],
1116
+ per_engine_labelvalues: dict[int, list[object]],
1117
+ ):
1118
+ super().__init__(vllm_config, metric_types, labelnames, per_engine_labelvalues)
1119
+ buckets = [
1120
+ 0.001,
1121
+ 0.005,
1122
+ 0.01,
1123
+ 0.025,
1124
+ 0.05,
1125
+ 0.075,
1126
+ 0.1,
1127
+ 0.2,
1128
+ 0.3,
1129
+ 0.5,
1130
+ 0.75,
1131
+ 1.0,
1132
+ 5.0,
1133
+ ]
1134
+ hf3fs_save_duration = self._histogram_cls(
1135
+ name="vllm:hf3fs_save_duration_seconds",
1136
+ documentation="Histogram of save duration for HF3FSKVConnector.",
1137
+ buckets=buckets,
1138
+ labelnames=labelnames,
1139
+ )
1140
+ self.hf3fs_save_duration = create_metric_per_engine(
1141
+ hf3fs_save_duration, self.per_engine_labelvalues
1142
+ )
1143
+
1144
+ hf3fs_load_duration = self._histogram_cls(
1145
+ name="vllm:hf3fs_load_duration_seconds",
1146
+ documentation="Histogram of load duration for HF3FSKVConnector.",
1147
+ buckets=buckets,
1148
+ labelnames=labelnames,
1149
+ )
1150
+ self.hf3fs_load_duration = create_metric_per_engine(
1151
+ hf3fs_load_duration, self.per_engine_labelvalues
1152
+ )
1153
+
1154
+ hf3fs_num_failed_save = self._counter_cls(
1155
+ name="vllm:hf3fs_num_failed_save",
1156
+ documentation="Number of failed HF3FS KV save.",
1157
+ labelnames=labelnames,
1158
+ )
1159
+ self.hf3fs_num_failed_save = create_metric_per_engine(
1160
+ hf3fs_num_failed_save, self.per_engine_labelvalues
1161
+ )
1162
+
1163
+ hf3fs_num_failed_load = self._counter_cls(
1164
+ name="vllm:hf3fs_num_failed_load",
1165
+ documentation="Number of failed HF3FS KV load.",
1166
+ labelnames=labelnames,
1167
+ )
1168
+ self.hf3fs_num_failed_load = create_metric_per_engine(
1169
+ hf3fs_num_failed_load, self.per_engine_labelvalues
1170
+ )
1171
+
1172
+ def observe(self, transfer_stats_data: dict[str, Any], engine_idx: int = 0):
1173
+ for prom_obj, list_item_key in zip(
1174
+ [
1175
+ self.hf3fs_save_duration,
1176
+ self.hf3fs_load_duration,
1177
+ ],
1178
+ [
1179
+ "save_duration",
1180
+ "load_duration",
1181
+ ],
1182
+ ):
1183
+ for list_item in transfer_stats_data[list_item_key]:
1184
+ prom_obj[engine_idx].observe(list_item)
1185
+ for counter_obj, counter_item_key in zip(
1186
+ [
1187
+ self.hf3fs_num_failed_save,
1188
+ self.hf3fs_num_failed_load,
1189
+ ],
1190
+ [
1191
+ "num_failed_save",
1192
+ "num_failed_load",
1193
+ ],
1194
+ ):
1195
+ counter_obj[engine_idx].inc(transfer_stats_data[counter_item_key])
vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/hf3fs_metadata_server.py
new
+530
@@ -0,0 +1,530 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+HF3FS Metadata Server with key-based organization.
5
+"""
6
+
7
+import argparse
8
+import logging
9
+import threading
10
+from abc import ABC, abstractmethod
11
+from dataclasses import dataclass
12
+
13
+try:
14
+ import orjson
15
+
16
+ HAS_ORJSON = True
17
+except ImportError:
18
+ import json as orjson # type: ignore
19
+
20
+ HAS_ORJSON = False
21
+
22
+import requests
23
+from fastapi import FastAPI, HTTPException, Request, Response
24
+from fastapi.responses import ORJSONResponse
25
+from requests.adapters import HTTPAdapter
26
+from urllib3.util.retry import Retry
27
+
28
+logging.basicConfig(
29
+ level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
30
+)
31
+logger = logging.getLogger(__name__)
32
+
33
+
34
+@dataclass
35
+class RankFileMetadata:
36
+ """Manages file page allocation for a single rank."""
37
+
38
+ rank_id: int
39
+ num_pages: int
40
+ free_pages: list[int]
41
+
42
+ def allocate_pages(self, num_pages: int) -> list[int]:
43
+ """Allocate specified number of free pages."""
44
+ if len(self.free_pages) < num_pages:
45
+ return []
46
+
47
+ allocated = self.free_pages[:num_pages]
48
+ self.free_pages = self.free_pages[num_pages:]
49
+ return allocated
50
+
51
+ def release_pages(self, page_indices: list[int]) -> None:
52
+ """Release pages back to free pool."""
53
+ for page_idx in page_indices:
54
+ if page_idx not in self.free_pages:
55
+ self.free_pages.append(page_idx)
56
+
57
+ def get_free_page_count(self) -> int:
58
+ """Get current number of free pages."""
59
+ return len(self.free_pages)
60
+
61
+
62
+@dataclass
63
+class KeyMetadata:
64
+ """Manages metadata for a single key across multiple ranks."""
65
+
66
+ key: str
67
+ rank_to_page: dict[int, int] # rank -> allocated page index
68
+ tp_world_size: int
69
+
70
+ def add_rank_page(self, rank: int, page_index: int) -> None:
71
+ """Add page allocation for a specific rank."""
72
+ self.rank_to_page[rank] = page_index
73
+
74
+ def get_all_pages(self) -> list[tuple[int, int]]:
75
+ """Get all (rank, page) pairs for this key."""
76
+ return [(rank, page) for rank, page in self.rank_to_page.items()]
77
+
78
+ def get_rank_page(self, rank: int) -> int | None:
79
+ """Get page index for a specific rank."""
80
+ return self.rank_to_page.get(rank)
81
+
82
+ def is_complete(self) -> bool:
83
+ """Check if all ranks in the TP world have allocated pages."""
84
+ return len(self.rank_to_page) == self.tp_world_size
85
+
86
+
87
+class GlobalMetadataState:
88
+ """Manages global metadata state across all ranks and keys."""
89
+
90
+ def __init__(self):
91
+ self.global_lock = threading.RLock()
92
+ self.rank_metadata: dict[int, RankFileMetadata] = {}
93
+ self.key_metadata: dict[str, KeyMetadata] = {}
94
+
95
+ def clear(self) -> None:
96
+ """Clear all metadata state."""
97
+ with self.global_lock:
98
+ self.rank_metadata.clear()
99
+ self.key_metadata.clear()
100
+ logger.info("Cleared all metadata state")
101
+
102
+ def initialize_rank(self, rank: int, num_pages: int) -> None:
103
+ """Initialize a new rank with specified number of pages."""
104
+ with self.global_lock:
105
+ if rank not in self.rank_metadata:
106
+ self.rank_metadata[rank] = RankFileMetadata(
107
+ rank, num_pages, list(range(num_pages))
108
+ )
109
+ logger.info("Initialized rank %s with %s pages", rank, num_pages)
110
+
111
+ def allocate_pages_for_keys(
112
+ self, rank: int, keys: list[tuple[str, str]]
113
+ ) -> dict[str, int]:
114
+ """Allocate one page for each key on the specified rank.
115
+
116
+ Args:
117
+ rank: Rank ID to allocate pages on
118
+ keys: List of keys to allocate pages for
119
+
120
+ Returns:
121
+ Dictionary mapping key -> allocated page index
122
+ """
123
+ with self.global_lock:
124
+ if rank not in self.rank_metadata:
125
+ raise ValueError(f"Rank {rank} not initialized")
126
+
127
+ # Batch allocate pages for all keys
128
+ num_pages_needed = len(keys)
129
+ allocated_pages = self.rank_metadata[rank].allocate_pages(num_pages_needed)
130
+
131
+ if len(allocated_pages) < num_pages_needed:
132
+ logger.warning(
133
+ "Rank %s only allocated %s pages for %s keys",
134
+ rank,
135
+ len(allocated_pages),
136
+ num_pages_needed,
137
+ )
138
+
139
+ allocation_results = {}
140
+ for i, (key, prefix_key) in enumerate(keys):
141
+ if key in self.key_metadata:
142
+ key_meta = self.key_metadata[key]
143
+ if key_meta.is_complete() and rank in key_meta.rank_to_page:
144
+ # key is already fully written, reuse the existing page
145
+ # and release the allocated pages back to the free pool.
146
+ if i < len(allocated_pages):
147
+ self.rank_metadata[rank].release_pages([allocated_pages[i]])
148
+ allocation_results[key] = key_meta.rank_to_page[rank]
149
+ continue
150
+
151
+ if i < len(allocated_pages):
152
+ allocation_results[key] = allocated_pages[i]
153
+ else:
154
+ allocation_results[key] = -1 # No pages available
155
+
156
+ return allocation_results
157
+
158
+ def confirm_write_for_keys(
159
+ self,
160
+ rank: int,
161
+ key_confirmations: list[tuple[str, int]],
162
+ pages_to_release: list[int] | None = None,
163
+ ) -> None:
164
+ """Confirm write operations for keys and update metadata.
165
+
166
+ Args:
167
+ rank: Rank ID that confirmed the writes
168
+ key_confirmations: List of (key, page_index) tuples
169
+ pages_to_release: List of page indices to release back to free pool
170
+ """
171
+ with self.global_lock:
172
+ # Confirm successful writes
173
+ for key, page_index in key_confirmations:
174
+ if key not in self.key_metadata:
175
+ # Need to determine tp_world_size from rank_metadata
176
+ tp_world_size = len(self.rank_metadata)
177
+ self.key_metadata[key] = KeyMetadata(key, {}, tp_world_size)
178
+
179
+ # Add confirmed page to key metadata
180
+ self.key_metadata[key].add_rank_page(rank, page_index)
181
+
182
+ # Release specified pages back to free pool
183
+ if pages_to_release:
184
+ self.rank_metadata[rank].release_pages(pages_to_release)
185
+ logger.debug(
186
+ "Released %s pages on rank %s: %s",
187
+ len(pages_to_release),
188
+ rank,
189
+ pages_to_release,
190
+ )
191
+
192
+ def batch_key_exists(self, keys: list[str]) -> list[bool]:
193
+ """Check if keys exist in metadata and all ranks have confirmed writes.
194
+
195
+ Args:
196
+ keys: List of keys to check
197
+
198
+ Returns:
199
+ List of boolean values indicating key existence and completion
200
+ """
201
+ with self.global_lock:
202
+ results = []
203
+ for key in keys:
204
+ if key not in self.key_metadata:
205
+ results.append(False)
206
+ else:
207
+ # Check if all ranks in the TP world have confirmed writes
208
+ key_meta = self.key_metadata[key]
209
+ results.append(key_meta.is_complete())
210
+ return results
211
+
212
+ def get_key_locations(self, rank: int, keys: list[str]) -> list[int | None]:
213
+ """Get page indices for keys on a specific rank.
214
+
215
+ Args:
216
+ rank: Rank ID to query
217
+ keys: List of keys to look up
218
+
219
+ Returns:
220
+ List of page indices in the same order as input keys (None if key not found)
221
+ """
222
+ with self.global_lock:
223
+ if rank not in self.rank_metadata:
224
+ raise ValueError(f"Rank {rank} not initialized")
225
+
226
+ results = []
227
+ for key in keys:
228
+ if key in self.key_metadata:
229
+ key_meta = self.key_metadata[key]
230
+ if key_meta.is_complete():
231
+ page_index = key_meta.get_rank_page(rank)
232
+ else:
233
+ page_index = None
234
+
235
+ results.append(page_index)
236
+ else:
237
+ results.append(None)
238
+
239
+ return results
240
+
241
+
242
+class Hf3fsMetadataServer:
243
+ """HF3FS Metadata Server with improved key-based organization."""
244
+
245
+ def __init__(self, persistence_path: str | None = None, save_interval: int = 60):
246
+ self.state = GlobalMetadataState()
247
+ if HAS_ORJSON:
248
+ self.app = FastAPI(default_response_class=ORJSONResponse)
249
+ else:
250
+ self.app = FastAPI()
251
+ self._setup_routes()
252
+
253
+ async def _read_json(self, request: Request) -> dict:
254
+ """Parse request JSON using orjson if available."""
255
+ body = await request.body()
256
+ return orjson.loads(body)
257
+
258
+ def _json_response(self, content: dict):
259
+ """Return ORJSONResponse when available to bypass jsonable_encoder."""
260
+ if HAS_ORJSON:
261
+ return ORJSONResponse(content)
262
+ else:
263
+ return content
264
+
265
+ def _setup_routes(self):
266
+ """Setup FastAPI routes for new API design."""
267
+ self.app.post("/rank/{rank}/initialize")(self.initialize_rank)
268
+ self.app.post("/keys/batch_allocate")(self.batch_allocate_pages_for_keys)
269
+ self.app.post("/keys/confirm_write")(self.confirm_write_for_keys)
270
+ self.app.post("/keys/batch_exists")(self.batch_key_exists)
271
+ self.app.post("/keys/get_locations")(self.get_key_locations)
272
+ self.app.post("/clear")(self.clear)
273
+
274
+ async def initialize_rank(self, rank: int, request: Request):
275
+ """Initialize a rank with specified number of pages."""
276
+ data = await self._read_json(request)
277
+ role = data.get("role", "worker")
278
+ num_pages = data.get("num_pages", 0)
279
+
280
+ if role == "scheduler":
281
+ return self._json_response(
282
+ {"message": "Scheduler role does not require initialization"}
283
+ )
284
+
285
+ if role == "worker" and num_pages > 0:
286
+ self.state.initialize_rank(rank, num_pages)
287
+ return self._json_response(
288
+ {"message": f"Rank {rank} initialized with {num_pages} pages"}
289
+ )
290
+ else:
291
+ raise HTTPException(
292
+ status_code=400, detail="Invalid initialization parameters"
293
+ )
294
+
295
+ async def batch_allocate_pages_for_keys(self, request: Request):
296
+ """Allocate one page for each key on a specific rank."""
297
+ data = await self._read_json(request)
298
+ rank = data.get("rank")
299
+ keys = data.get("keys", [])
300
+
301
+ # Validate input format
302
+ if rank is None or not isinstance(keys, list):
303
+ raise HTTPException(
304
+ status_code=400, detail="Invalid request format: need 'rank' and 'keys'"
305
+ )
306
+
307
+ try:
308
+ # Perform allocation
309
+ results = self.state.allocate_pages_for_keys(rank, keys)
310
+
311
+ # Convert results to response format
312
+ response = {"rank": rank, "results": list(results.items())}
313
+ return self._json_response(response)
314
+ except Exception as e:
315
+ raise HTTPException(
316
+ status_code=500, detail=f"Allocation failed: {str(e)}"
317
+ ) from e
318
+
319
+ async def confirm_write_for_keys(self, request: Request):
320
+ """Confirm write operations for keys."""
321
+ data = await self._read_json(request)
322
+ rank = data.get("rank")
323
+ confirmations = data.get("confirmations", [])
324
+ pages_to_release = data.get("pages_to_release", [])
325
+
326
+ # Validate input format
327
+ if rank is None or not isinstance(confirmations, list):
328
+ raise HTTPException(
329
+ status_code=400,
330
+ detail="Invalid request format: need 'rank' and 'confirmations'",
331
+ )
332
+
333
+ try:
334
+ self.state.confirm_write_for_keys(rank, confirmations, pages_to_release)
335
+
336
+ return Response(status_code=204)
337
+
338
+ except Exception as e:
339
+ logger.error("Confirm write for keys failed: %s", e)
340
+ raise HTTPException(
341
+ status_code=500, detail=f"Confirmation failed: {str(e)}"
342
+ ) from e
343
+
344
+ async def batch_key_exists(self, request: Request):
345
+ """Check if multiple keys exist in metadata."""
346
+ data = await self._read_json(request)
347
+ keys = data.get("keys", [])
348
+
349
+ if not isinstance(keys, list):
350
+ raise HTTPException(status_code=400, detail="Invalid keys format")
351
+
352
+ try:
353
+ exists_results = self.state.batch_key_exists(keys)
354
+ return self._json_response({"exists": exists_results})
355
+ except Exception as e:
356
+ raise HTTPException(
357
+ status_code=500, detail=f"Key existence check failed: {str(e)}"
358
+ ) from e
359
+
360
+ async def get_key_locations(self, request: Request):
361
+ """Get page indices for keys on a specific rank."""
362
+ data = await self._read_json(request)
363
+ rank = data.get("rank")
364
+ keys = data.get("keys", [])
365
+
366
+ # Validate input format
367
+ if rank is None or not isinstance(keys, list):
368
+ raise HTTPException(
369
+ status_code=400, detail="Invalid request format: need 'rank' and 'keys'"
370
+ )
371
+
372
+ try:
373
+ # Get key locations
374
+ locations = self.state.get_key_locations(rank, keys)
375
+ return self._json_response({"locations": locations})
376
+ except Exception as e:
377
+ raise HTTPException(
378
+ status_code=500, detail=f"Failed to get key locations: {str(e)}"
379
+ ) from e
380
+
381
+ async def clear(self, request: Request):
382
+ """Clear the metadata server."""
383
+ self.state.clear()
384
+ return Response(status_code=204)
385
+
386
+ def run(self, host: str = "0.0.0.0", port: int = 18000):
387
+ """Run the metadata server."""
388
+ import uvicorn
389
+
390
+ logger.info("Starting improved metadata server on http://%s:%s", host, port)
391
+ uvicorn.run(self.app, host=host, port=port)
392
+
393
+
394
+# --- Client implementation ---
395
+class Hf3fsMetadataInterface(ABC):
396
+ """Interface for HF3FS metadata operations."""
397
+
398
+ @abstractmethod
399
+ def initialize(self, rank: int, num_pages: int = 0, role: str = "worker") -> None:
400
+ """Initialize the metadata service with specified number of pages."""
401
+ pass
402
+
403
+ @abstractmethod
404
+ def allocate_pages_for_keys(
405
+ self, rank: int, keys: list[tuple[str, str]]
406
+ ) -> list[tuple[str, int]]:
407
+ """Allocate one page for each key on the specified rank."""
408
+ pass
409
+
410
+ @abstractmethod
411
+ def confirm_write_for_keys(
412
+ self,
413
+ rank: int,
414
+ key_confirmations: list[tuple[str, int]],
415
+ pages_to_release: list[int] | None = None,
416
+ ) -> None:
417
+ """Confirm write operations for keys and optionally release pages."""
418
+ pass
419
+
420
+ @abstractmethod
421
+ def batch_key_exists(self, keys: list[str]) -> list[bool]:
422
+ """Check if keys exist and are complete across all ranks."""
423
+ pass
424
+
425
+ @abstractmethod
426
+ def get_key_locations(self, rank: int, keys: list[str]) -> list[int]:
427
+ """Get page indices for keys on a specific rank."""
428
+ pass
429
+
430
+
431
+class Hf3fsGlobalMetadataClient(Hf3fsMetadataInterface):
432
+ """Global HTTP metadata client for HF3FS."""
433
+
434
+ def __init__(self, base_url: str = "http://localhost:18000", max_retries: int = 3):
435
+ self.base_url = base_url.rstrip("/")
436
+ self._session = requests.Session()
437
+
438
+ retry_strategy = Retry(
439
+ total=max_retries,
440
+ backoff_factor=0.3,
441
+ status_forcelist=[500, 502, 503, 504],
442
+ allowed_methods=["GET", "POST"],
443
+ )
444
+ adapter = HTTPAdapter(max_retries=retry_strategy)
445
+ self._session.mount("http://", adapter)
446
+
447
+ def _post(self, endpoint: str, json_data: dict) -> dict:
448
+ """Make POST request to metadata server."""
449
+ try:
450
+ url = f"{self.base_url}/{endpoint}"
451
+ headers = {"Content-Type": "application/json"}
452
+ if HAS_ORJSON:
453
+ payload = orjson.dumps(json_data)
454
+ else:
455
+ import json
456
+
457
+ payload = json.dumps(json_data).encode("utf-8")
458
+ response = self._session.post(url, data=payload, headers=headers)
459
+ response.raise_for_status()
460
+
461
+ if response.status_code == 204 or not response.content:
462
+ return {}
463
+ if HAS_ORJSON:
464
+ return orjson.loads(response.content)
465
+ else:
466
+ return response.json()
467
+ except requests.exceptions.RequestException as e:
468
+ logger.error("Failed to POST to %s after retries: %s", endpoint, e)
469
+ raise RuntimeError(f"Failed to connect to metadata server: {e}") from e
470
+
471
+ def initialize(self, rank: int, num_pages: int = 0, role: str = "worker") -> None:
472
+ """Initialize a rank with specified number of pages."""
473
+ self._post(f"rank/{rank}/initialize", {"num_pages": num_pages, "role": role})
474
+
475
+ def allocate_pages_for_keys(
476
+ self, rank: int, keys: list[tuple[str, str]]
477
+ ) -> list[tuple[str, int]]:
478
+ """Allocate pages for keys on the specified rank."""
479
+ response = self._post("keys/batch_allocate", {"rank": rank, "keys": keys})
480
+
481
+ # Convert response to expected format
482
+ return response.get("results", {})
483
+
484
+ def confirm_write_for_keys(
485
+ self,
486
+ rank: int,
487
+ key_confirmations: list[tuple[str, int]],
488
+ pages_to_release: list[int] | None = None,
489
+ ) -> None:
490
+ """Confirm write operations for keys and optionally release pages."""
491
+ payload = {
492
+ "rank": rank,
493
+ "confirmations": key_confirmations,
494
+ "pages_to_release": pages_to_release or [],
495
+ }
496
+
497
+ self._post("keys/confirm_write", payload)
498
+
499
+ def batch_key_exists(self, keys: list[str]) -> list[bool]:
500
+ """Check if keys exist and are complete across all ranks."""
501
+ response = self._post("keys/batch_exists", {"keys": keys})
502
+ return response.get("exists", [])
503
+
504
+ def get_key_locations(self, rank: int, keys: list[str]) -> list[int]:
505
+ """Get page indices for keys on a specific rank."""
506
+ response = self._post("keys/get_locations", {"rank": rank, "keys": keys})
507
+ return response.get("locations", [])
508
+
509
+
510
+def run_metadata_server(
511
+ host: str = "0.0.0.0",
512
+ port: int = 18000,
513
+):
514
+ """Run the improved HF3FS metadata server."""
515
+ server = Hf3fsMetadataServer()
516
+ server.run(host=host, port=port)
517
+
518
+
519
+# --- Main Execution ---
520
+if __name__ == "__main__":
521
+ parser = argparse.ArgumentParser(description="Improved HF3FS Metadata Server")
522
+ parser.add_argument(
523
+ "--host", type=str, default="0.0.0.0", help="Host to bind the server to."
524
+ )
525
+ parser.add_argument(
526
+ "--port", type=int, default=18000, help="Port to run the server on."
527
+ )
528
+ args = parser.parse_args()
529
+
530
+ run_metadata_server(args.host, args.port)
vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/utils/__init__.py
vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/utils/common.py
new
+139
@@ -0,0 +1,139 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import threading
4
+from dataclasses import dataclass, field
5
+from typing import Optional
6
+
7
+from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorMetadata
8
+from vllm.v1.request import Request
9
+
10
+
11
+class AtomicCounter:
12
+ """Thread-safe atomic counter for round-robin operations."""
13
+
14
+ def __init__(self, n: int):
15
+ assert n > 0, "Counter size must be positive"
16
+ self._n = n
17
+ self._value = 0
18
+ self._lock = threading.Lock()
19
+
20
+ def next(self) -> int:
21
+ """Get next value in round-robin fashion."""
22
+ with self._lock:
23
+ current = self._value
24
+ self._value = (current + 1) % self._n
25
+ return current
26
+
27
+
28
+@dataclass
29
+class LoadBlockInfo:
30
+ """Operation for loading blocks from external storage."""
31
+
32
+ num_computed_blocks: int
33
+ num_blocks_to_load: int
34
+ need_fetch_block_ids: list[int]
35
+
36
+
37
+@dataclass
38
+class SaveBlockInfo:
39
+ """Operation for saving blocks to external storage."""
40
+
41
+ skip_leading_blocks: int
42
+
43
+
44
+@dataclass
45
+class RequestSchedulingState:
46
+ """Unified request scheduling state management."""
47
+
48
+ request_id: str
49
+ request: Request | None = None
50
+
51
+ # Token and block tracking
52
+ token_ids: list[int] = field(default_factory=list)
53
+ allocated_block_ids: list[int] = field(default_factory=list)
54
+ num_saved_blocks: int = 0
55
+
56
+ # Load operation info
57
+ load_op: LoadBlockInfo | None = None
58
+
59
+ # Scheduling phase
60
+ phase: str = "NEW" # NEW -> WAITING_TO_LOAD -> ACTIVE -> FINISHED
61
+
62
+ def needs_loading(self) -> bool:
63
+ """Check if request needs loading."""
64
+ return self.load_op is not None and self.load_op.num_blocks_to_load > 0
65
+
66
+ def is_ready_to_load(self) -> bool:
67
+ """Check if request is ready for loading."""
68
+ return self.phase == "WAITING_TO_LOAD" and self.needs_loading()
69
+
70
+ def update_tokens_and_blocks(self, new_token_ids: list[int], new_block_ids) -> None:
71
+ """Update with new tokens and blocks."""
72
+ if new_token_ids:
73
+ self.token_ids.extend(new_token_ids)
74
+
75
+ if new_block_ids is not None:
76
+ normalized_block_ids = self._normalize_block_ids(new_block_ids)
77
+ self.allocated_block_ids.extend(normalized_block_ids)
78
+
79
+ def _normalize_block_ids(self, block_ids) -> list[int]:
80
+ """Normalize block_ids to list format."""
81
+ if not block_ids:
82
+ return []
83
+ if isinstance(block_ids, tuple):
84
+ return block_ids[0] if block_ids else []
85
+ if isinstance(block_ids, list):
86
+ return block_ids
87
+ return []
88
+
89
+
90
+@dataclass
91
+class HF3FSRequestMetadata:
92
+ """Metadata for a single request in HF3FS connector."""
93
+
94
+ request_id: str
95
+ token_ids: list[int]
96
+ block_ids: list[int]
97
+ load_block_op: LoadBlockInfo | None = None
98
+ save_block_op: SaveBlockInfo | None = None
99
+
100
+ @staticmethod
101
+ def from_scheduling_state(
102
+ state: "RequestSchedulingState",
103
+ block_size: int,
104
+ load_op: LoadBlockInfo | None = None,
105
+ skip_leading_blocks: int | None = None,
106
+ ) -> Optional["HF3FSRequestMetadata"]:
107
+ """Create request metadata from scheduling state."""
108
+ token_count = len(state.token_ids)
109
+ total_blocks = token_count // block_size
110
+
111
+ skip_blocks = (
112
+ state.num_saved_blocks
113
+ if skip_leading_blocks is None
114
+ else skip_leading_blocks
115
+ )
116
+
117
+ new_blocks_to_save = total_blocks - state.num_saved_blocks
118
+ if new_blocks_to_save <= 0 and load_op is None:
119
+ return None
120
+
121
+ state.num_saved_blocks = total_blocks
122
+ return HF3FSRequestMetadata(
123
+ request_id=state.request_id,
124
+ token_ids=state.token_ids,
125
+ block_ids=state.allocated_block_ids,
126
+ load_block_op=load_op,
127
+ save_block_op=SaveBlockInfo(skip_leading_blocks=skip_blocks),
128
+ )
129
+
130
+
131
+class HF3FSConnectorMetadata(KVConnectorMetadata):
132
+ """Container for HF3FS connector metadata."""
133
+
134
+ def __init__(self):
135
+ self.requests: list[HF3FSRequestMetadata] = []
136
+
137
+ def add_request(self, request_metadata: HF3FSRequestMetadata) -> None:
138
+ """Add request to metadata."""
139
+ self.requests.append(request_metadata)
vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/utils/gather_scatter_helper.py
new
+288
@@ -0,0 +1,288 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import torch
4
+
5
+from vllm.logger import init_logger
6
+from vllm.triton_utils import tl, triton
7
+
8
+
9
+@triton.jit
10
+def kv_cache_scatter_kernel(
11
+ kv_cache_ptrs_ptr,
12
+ source_ptr,
13
+ token_indices_ptr,
14
+ num_tokens_in_block,
15
+ hidden_size,
16
+ total_token_in_kvcache,
17
+ num_layers,
18
+ is_mla,
19
+ BLOCK_SIZE: tl.constexpr,
20
+):
21
+ layer_idx = tl.program_id(0)
22
+ token_pos = tl.program_id(1)
23
+
24
+ if layer_idx >= num_layers or token_pos >= num_tokens_in_block:
25
+ return
26
+
27
+ token_idx = tl.load(token_indices_ptr + token_pos)
28
+ kv_cache_ptr = tl.cast(tl.load(kv_cache_ptrs_ptr + layer_idx), source_ptr.dtype)
29
+
30
+ if token_idx >= total_token_in_kvcache:
31
+ return
32
+
33
+ if is_mla:
34
+ # MLA format: source [num_layers, num_tokens_in_block, hidden_size]
35
+ # MLA format: target [total_token_in_kvcache, hidden_size] (per layer)
36
+ source_offset = (layer_idx * num_tokens_in_block + token_pos) * hidden_size
37
+ target_offset = token_idx * hidden_size
38
+
39
+ for i in range(0, hidden_size, BLOCK_SIZE):
40
+ offset = i + tl.arange(0, BLOCK_SIZE)
41
+ mask = offset < hidden_size
42
+ val = tl.load(source_ptr + source_offset + offset, mask=mask)
43
+ tl.store(kv_cache_ptr + target_offset + offset, val, mask=mask)
44
+ else:
45
+ # MHA format: source [num_layers, 2, num_tokens_in_block, hidden_size]
46
+ # MHA format: target [2, total_token_in_kvcache, hidden_size]
47
+ source_offset_k = (
48
+ layer_idx * num_tokens_in_block * 2 + token_pos
49
+ ) * hidden_size
50
+ source_offset_v = (
51
+ layer_idx * num_tokens_in_block * 2 + num_tokens_in_block + token_pos
52
+ ) * hidden_size
53
+
54
+ target_offset_k = token_idx * hidden_size
55
+ target_offset_v = (total_token_in_kvcache + token_idx) * hidden_size
56
+
57
+ for i in range(0, hidden_size, BLOCK_SIZE):
58
+ offset = i + tl.arange(0, BLOCK_SIZE)
59
+ mask = offset < hidden_size
60
+
61
+ val_k = tl.load(source_ptr + source_offset_k + offset, mask=mask)
62
+ val_v = tl.load(source_ptr + source_offset_v + offset, mask=mask)
63
+
64
+ tl.store(kv_cache_ptr + target_offset_k + offset, val_k, mask=mask)
65
+ tl.store(kv_cache_ptr + target_offset_v + offset, val_v, mask=mask)
66
+
67
+
68
+@triton.jit
69
+def kv_cache_gather_kernel(
70
+ kv_cache_ptrs_ptr,
71
+ dst_ptr,
72
+ token_indices_ptr,
73
+ num_tokens_in_block,
74
+ hidden_size,
75
+ total_token_in_kvcache,
76
+ num_layers,
77
+ is_mla,
78
+ BLOCK_SIZE: tl.constexpr,
79
+):
80
+ layer_idx = tl.program_id(0)
81
+ token_pos = tl.program_id(1)
82
+
83
+ if layer_idx >= num_layers or token_pos >= num_tokens_in_block:
84
+ return
85
+
86
+ token_idx = tl.load(token_indices_ptr + token_pos)
87
+ kv_cache_ptr = tl.cast(tl.load(kv_cache_ptrs_ptr + layer_idx), dst_ptr.dtype)
88
+
89
+ if token_idx >= total_token_in_kvcache:
90
+ return
91
+
92
+ if is_mla:
93
+ # MLA format: source [total_token_in_kvcache, hidden_size] (per layer)
94
+ # MLA format: dst [num_layers, num_tokens_in_block, hidden_size]
95
+ kvcache_offset = token_idx * hidden_size
96
+ dst_offset = (layer_idx * num_tokens_in_block + token_pos) * hidden_size
97
+
98
+ for i in range(0, hidden_size, BLOCK_SIZE):
99
+ offset = i + tl.arange(0, BLOCK_SIZE)
100
+ mask = offset < hidden_size
101
+ val = tl.load(kv_cache_ptr + kvcache_offset + offset, mask=mask)
102
+ tl.store(dst_ptr + dst_offset + offset, val, mask=mask)
103
+ else:
104
+ # MHA format: source [2, total_token_in_kvcache, hidden_size]
105
+ # MHA format: dst [num_layers, 2, num_tokens_in_block, hidden_size]
106
+ dst_offset_k = (layer_idx * num_tokens_in_block * 2 + token_pos) * hidden_size
107
+ dst_offset_v = (
108
+ layer_idx * num_tokens_in_block * 2 + num_tokens_in_block + token_pos
109
+ ) * hidden_size
110
+
111
+ kvcache_offset_k = token_idx * hidden_size
112
+ kvcache_offset_v = (total_token_in_kvcache + token_idx) * hidden_size
113
+
114
+ for i in range(0, hidden_size, BLOCK_SIZE):
115
+ offset = i + tl.arange(0, BLOCK_SIZE)
116
+ mask = offset < hidden_size
117
+
118
+ val_k = tl.load(kv_cache_ptr + kvcache_offset_k + offset, mask=mask)
119
+ val_v = tl.load(kv_cache_ptr + kvcache_offset_v + offset, mask=mask)
120
+
121
+ tl.store(dst_ptr + dst_offset_k + offset, val_k, mask=mask)
122
+ tl.store(dst_ptr + dst_offset_v + offset, val_v, mask=mask)
123
+
124
+
125
+def scatter_kv_caches(
126
+ kv_caches_ptrs: torch.Tensor,
127
+ total_token_in_kvcache: int,
128
+ src_tensor: torch.Tensor,
129
+ token_indices: list[int],
130
+ is_mla: bool = False,
131
+) -> None:
132
+ """Scatter KV cache data from source tensor to KV cache storage.
133
+
134
+ Args:
135
+ kv_caches_ptrs: Tensor of KV cache pointers (one per layer)
136
+ total_token_in_kvcache: Total number of tokens in KV cache
137
+ src_tensor: Source tensor containing data to scatter
138
+ - MHA format: [num_layers, 2, num_tokens_in_block, hidden_size]
139
+ - MLA format: [num_layers, num_tokens_in_block, hidden_size]
140
+ token_indices: List of token positions to update
141
+ is_mla: Whether using MLA model format
142
+ """
143
+ num_layers = len(kv_caches_ptrs)
144
+ num_tokens_in_block = len(token_indices)
145
+
146
+ if is_mla:
147
+ # MLA: src_tensor is [num_layers, num_tokens_in_block, hidden_size]
148
+ assert len(src_tensor.shape) == 3, (
149
+ f"MLA src_tensor should be 3D, got {src_tensor.shape}"
150
+ )
151
+ hidden_size = src_tensor.shape[2]
152
+ else:
153
+ # MHA: src_tensor is [num_layers, 2, num_tokens_in_block, hidden_size]
154
+ assert len(src_tensor.shape) == 4, (
155
+ f"MHA src_tensor should be 4D, got {src_tensor.shape}"
156
+ )
157
+ hidden_size = src_tensor.shape[3]
158
+
159
+ device = src_tensor.device
160
+ token_indices_tensor = torch.tensor(
161
+ token_indices, dtype=torch.int32, device="cpu"
162
+ ).to(device, non_blocking=True)
163
+
164
+ grid = (num_layers, num_tokens_in_block)
165
+ BLOCK_SIZE = 128
166
+
167
+ kv_cache_scatter_kernel[grid](
168
+ kv_caches_ptrs,
169
+ src_tensor,
170
+ token_indices_tensor,
171
+ num_tokens_in_block,
172
+ hidden_size,
173
+ total_token_in_kvcache,
174
+ num_layers,
175
+ is_mla,
176
+ BLOCK_SIZE=BLOCK_SIZE,
177
+ )
178
+
179
+
180
+def gather_kv_caches(
181
+ kv_caches_ptrs: torch.Tensor,
182
+ total_token_in_kvcache: int,
183
+ dst_tensor: torch.Tensor,
184
+ token_indices: list[int],
185
+ is_mla: bool = False,
186
+) -> None:
187
+ """Gather KV cache data from KV cache storage to destination tensor.
188
+
189
+ Args:
190
+ kv_caches_ptrs: Tensor of KV cache pointers (one per layer)
191
+ total_token_in_kvcache: Total number of tokens in KV cache
192
+ dst_tensor: Destination tensor to store gathered data
193
+ - MHA format: [num_layers, 2, num_tokens_in_block, hidden_size]
194
+ - MLA format: [num_layers, num_tokens_in_block, hidden_size]
195
+ token_indices: List of token positions to gather
196
+ is_mla: Whether using MLA model format
197
+ """
198
+ num_layers = kv_caches_ptrs.shape[0]
199
+ num_tokens_in_block = len(token_indices)
200
+
201
+ if is_mla:
202
+ # MLA: dst_tensor is [num_layers, num_tokens_in_block, hidden_size]
203
+ assert len(dst_tensor.shape) == 3, (
204
+ f"MLA dst_tensor should be 3D, got {dst_tensor.shape}"
205
+ )
206
+ assert dst_tensor.shape[0] == num_layers, (
207
+ f"Layer count mismatch: {dst_tensor.shape[0]} vs {num_layers}"
208
+ )
209
+ assert dst_tensor.shape[1] == num_tokens_in_block, (
210
+ f"Token count mismatch: {dst_tensor.shape[1]} vs {num_tokens_in_block}"
211
+ )
212
+ hidden_size = dst_tensor.shape[2]
213
+ else:
214
+ # MHA: dst_tensor is [num_layers, 2, num_tokens_in_block, hidden_size]
215
+ assert len(dst_tensor.shape) == 4, (
216
+ f"MHA dst_tensor should be 4D, got {dst_tensor.shape}"
217
+ )
218
+ assert dst_tensor.shape[0] == num_layers, (
219
+ f"Layer count mismatch: {dst_tensor.shape[0]} vs {num_layers}"
220
+ )
221
+ assert dst_tensor.shape[1] == 2, (
222
+ f"MHA should have 2 (K,V) components, got {dst_tensor.shape[1]}"
223
+ )
224
+ assert dst_tensor.shape[2] == num_tokens_in_block, (
225
+ f"Token count mismatch: {dst_tensor.shape[2]} vs {num_tokens_in_block}"
226
+ )
227
+ hidden_size = dst_tensor.shape[3]
228
+
229
+ device = dst_tensor.device
230
+ token_indices_tensor = torch.tensor(
231
+ token_indices, dtype=torch.int32, device="cpu"
232
+ ).to(device, non_blocking=True)
233
+
234
+ grid = (num_layers, num_tokens_in_block)
235
+ BLOCK_SIZE = 128
236
+
237
+ kv_cache_gather_kernel[grid](
238
+ kv_caches_ptrs,
239
+ dst_tensor,
240
+ token_indices_tensor,
241
+ num_tokens_in_block,
242
+ hidden_size,
243
+ total_token_in_kvcache,
244
+ num_layers,
245
+ is_mla,
246
+ BLOCK_SIZE=BLOCK_SIZE,
247
+ )
248
+
249
+
250
+class CopyBufferAllocator:
251
+ """Memory pool for tensor buffers to avoid frequent allocation/deallocation."""
252
+
253
+ def __init__(
254
+ self, device: torch.device, dtype: torch.dtype, shape: list, max_count: int
255
+ ):
256
+ self._shape = shape
257
+ self._max_count = max_count
258
+ self._device = device
259
+ self._free_buffers = [
260
+ torch.empty(shape, dtype=dtype, device=device) for _ in range(max_count)
261
+ ]
262
+ self._inuse_count = 0
263
+
264
+ def alloc_buffer(self, count: int) -> list[torch.Tensor] | None:
265
+ """Allocate buffers from the pool."""
266
+ if count == 0:
267
+ return []
268
+
269
+ if self._inuse_count + count <= self._max_count:
270
+ self._inuse_count += count
271
+ result = self._free_buffers[-count:]
272
+ del self._free_buffers[-count:]
273
+ return result
274
+ return None
275
+
276
+ def free_buffer(self, buffers: list[torch.Tensor]) -> None:
277
+ """Return buffers to the pool."""
278
+ if not buffers:
279
+ return
280
+
281
+ if self._inuse_count >= len(buffers):
282
+ self._inuse_count -= len(buffers)
283
+ self._free_buffers.extend(buffers)
284
+ else:
285
+ raise RuntimeError("Attempted to free more buffers than allocated")
286
+
287
+
288
+logger = init_logger(__name__)
vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/utils/hf3fs_mock_client.py
new
+133
@@ -0,0 +1,133 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import logging
4
+import os
5
+
6
+import torch
7
+
8
+logger = logging.getLogger(__name__)
9
+HF3FS_AVAILABLE = True
10
+
11
+
12
+class Hf3fsClient:
13
+ """Mock HF3FS client using file backend for debugging and testing."""
14
+
15
+ def __init__(self, path: str, size: int, bytes_per_page: int, entries: int):
16
+ self._size = size
17
+ self._bytes_per_page = bytes_per_page
18
+ self._entries = entries
19
+ self._file_path = path
20
+
21
+ self._ensure_file_exists()
22
+ logger.debug("Initialized mock HF3FS client: %s (%d bytes)", path, size)
23
+
24
+ def _ensure_file_exists(self) -> None:
25
+ """Create file if it doesn't exist."""
26
+ if not os.path.exists(self._file_path):
27
+ with open(self._file_path, "w+b") as f:
28
+ f.truncate(self._size)
29
+
30
+ def batch_read(self, offsets: list[int], tensors: list[torch.Tensor]) -> list[int]:
31
+ """Read data from file at specified offsets into tensors."""
32
+ results = []
33
+
34
+ try:
35
+ with open(self._file_path, "rb") as f:
36
+ for offset, tensor in zip(offsets, tensors):
37
+ num_bytes = tensor.numel() * tensor.element_size()
38
+
39
+ if offset < 0 or offset + num_bytes > self._size:
40
+ results.append(-1)
41
+ continue
42
+
43
+ f.seek(offset)
44
+ buffer_data = f.read(num_bytes)
45
+
46
+ if len(buffer_data) == num_bytes == self._bytes_per_page:
47
+ tensor_data = self._convert_buffer_to_tensor(
48
+ buffer_data, tensor.dtype
49
+ )
50
+ tensor.copy_(
51
+ tensor_data.reshape(tensor.shape).to(tensor.device)
52
+ )
53
+ results.append(self._bytes_per_page)
54
+ else:
55
+ logger.error(
56
+ "Read size mismatch: got %d, expected %d",
57
+ len(buffer_data),
58
+ num_bytes,
59
+ )
60
+ results.append(-1)
61
+ except Exception as e:
62
+ logger.error("Batch read error: %s", e)
63
+ results.extend([-1] * (len(offsets) - len(results)))
64
+
65
+ return results
66
+
67
+ def _convert_buffer_to_tensor(
68
+ self, buffer_data: bytes, dtype: torch.dtype
69
+ ) -> torch.Tensor:
70
+ """Convert buffer data to tensor with proper dtype handling."""
71
+ if dtype == torch.bfloat16:
72
+ tensor_data = torch.frombuffer(buffer_data, dtype=torch.uint16)
73
+ return tensor_data.view(dtype=torch.bfloat16)
74
+ else:
75
+ return torch.frombuffer(buffer_data, dtype=dtype)
76
+
77
+ def batch_write(
78
+ self, offsets: list[int], tensors: list[torch.Tensor], event: torch.cuda.Event
79
+ ) -> list[int]:
80
+ """Write data from tensors to file at specified offsets."""
81
+ results = []
82
+
83
+ try:
84
+ torch.cuda.current_stream().wait_event(event)
85
+
86
+ # Convert tensors to bytes
87
+ data_bytes_list = [self._tensor_to_bytes(tensor) for tensor in tensors]
88
+
89
+ # Write to file
90
+ with open(self._file_path, "r+b") as f:
91
+ for offset, data_bytes in zip(offsets, data_bytes_list):
92
+ if offset < 0 or offset + len(data_bytes) > self._size:
93
+ results.append(-1)
94
+ continue
95
+
96
+ f.seek(offset)
97
+ bytes_written = f.write(data_bytes)
98
+
99
+ if bytes_written == len(data_bytes) == self._bytes_per_page:
100
+ results.append(self._bytes_per_page)
101
+ else:
102
+ logger.error(
103
+ "Write size mismatch: wrote %d, expected %d",
104
+ bytes_written,
105
+ self._bytes_per_page,
106
+ )
107
+ results.append(-1)
108
+
109
+ except Exception as e:
110
+ logger.error("Batch write error: %s", e)
111
+ results.extend([-1] * (len(offsets) - len(results)))
112
+
113
+ return results
114
+
115
+ def _tensor_to_bytes(self, tensor: torch.Tensor) -> bytes:
116
+ """Convert tensor to bytes with proper dtype handling."""
117
+ cpu_tensor = tensor.cpu()
118
+ if cpu_tensor.dtype == torch.bfloat16:
119
+ return cpu_tensor.view(dtype=torch.uint16).numpy().tobytes()
120
+ else:
121
+ return cpu_tensor.numpy().tobytes()
122
+
123
+ def get_size(self) -> int:
124
+ """Get the total size of the storage file."""
125
+ return self._size
126
+
127
+ def close(self) -> None:
128
+ """Close the client (no-op for file backend)."""
129
+ pass
130
+
131
+ def flush(self) -> None:
132
+ """Flush any pending writes (no-op for file backend)."""
133
+ pass
vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/utils/hf3fs_utils.cpp
new
+57
@@ -0,0 +1,57 @@
1
+#include <cuda_runtime.h>
2
+#include <torch/extension.h>
3
+#include <cstring>
4
+#include <vector>
5
+
6
+void read_shm(const torch::Tensor& shm, const torch::Tensor& pin,
7
+ std::vector<torch::Tensor> dst, uint64_t stream_ptr) {
8
+ py::gil_scoped_release release;
9
+
10
+ cudaStream_t stream = reinterpret_cast<cudaStream_t>(stream_ptr);
11
+
12
+ // Copy from shared memory to pinned memory
13
+ char* shm_ptr = static_cast<char*>(shm.data_ptr());
14
+ char* src_ptr = static_cast<char*>(pin.data_ptr());
15
+ std::memcpy(src_ptr, shm_ptr, shm.numel() * shm.element_size());
16
+
17
+ // Copy from pinned memory to GPU tensors
18
+ size_t current = 0;
19
+ for (size_t i = 0; i < dst.size(); ++i) {
20
+ auto& t = dst[i];
21
+ size_t t_bytes = t.numel() * t.element_size();
22
+ char* dst_ptr = static_cast<char*>(t.data_ptr());
23
+ cudaMemcpyAsync(dst_ptr, src_ptr + current, t_bytes, cudaMemcpyHostToDevice,
24
+ stream);
25
+ current += t_bytes;
26
+ }
27
+ cudaStreamSynchronize(stream);
28
+}
29
+
30
+void write_shm(const std::vector<torch::Tensor> src, torch::Tensor& shm,
31
+ const torch::Tensor& pin, uint64_t stream_ptr) {
32
+ py::gil_scoped_release release;
33
+
34
+ cudaStream_t stream = reinterpret_cast<cudaStream_t>(stream_ptr);
35
+
36
+ // Copy from GPU tensors to pinned memory
37
+ char* dst_ptr = static_cast<char*>(pin.data_ptr());
38
+ size_t current = 0;
39
+ for (size_t i = 0; i < src.size(); ++i) {
40
+ auto& t = src[i];
41
+ size_t t_bytes = t.numel() * t.element_size();
42
+ char* src_ptr = static_cast<char*>(t.data_ptr());
43
+ cudaMemcpyAsync(dst_ptr + current, src_ptr, t_bytes, cudaMemcpyDeviceToHost,
44
+ stream);
45
+ current += t_bytes;
46
+ }
47
+ cudaStreamSynchronize(stream);
48
+
49
+ // Copy from pinned memory to shared memory
50
+ char* shm_ptr = static_cast<char*>(shm.data_ptr());
51
+ std::memcpy(shm_ptr, dst_ptr, shm.numel() * shm.element_size());
52
+}
53
+
54
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
55
+ m.def("read_shm", &read_shm, "Read tensors from shared memory");
56
+ m.def("write_shm", &write_shm, "Write tensors to shared memory");
57
+}
\ No newline at end of file