From a5d16d59caeb8bc379d7c200188f38b7bdd6c3b0 Mon Sep 17 00:00:00 2001 From: j00628475 Date: Fri, 28 Aug 2026 14:46:35 +0800 Subject: [PATCH 1/2] [feat]support kimi k3 d2rh mooncake transfer offload --- .../test_mooncake_hybrid_connector.py | 19 + .../kv_offload/test_mooncake_swap_staging.py | 644 ++++++++++++++++++ .../kv_p2p/mooncake_hybrid_connector.py | 140 +++- .../utils/mooncake_transfer_engine.py | 43 +- .../kv_offload/mooncake_swap_memory.py | 71 ++ vllm_ascend/worker/model_runner_v1.py | 87 ++- 6 files changed, 964 insertions(+), 40 deletions(-) create mode 100644 tests/ut/kv_offload/test_mooncake_swap_staging.py create mode 100644 vllm_ascend/kv_offload/mooncake_swap_memory.py diff --git a/tests/ut/kv_offload/test_mooncake_hybrid_connector.py b/tests/ut/kv_offload/test_mooncake_hybrid_connector.py index bf2d587d3d11..463f6185b497 100644 --- a/tests/ut/kv_offload/test_mooncake_hybrid_connector.py +++ b/tests/ut/kv_offload/test_mooncake_hybrid_connector.py @@ -21,6 +21,12 @@ KVCacheRecvingThread, MooncakeConnectorScheduler, ) +from vllm_ascend.kv_offload.mooncake_swap_memory import ( # noqa: E402 + clear_swapped_tensors_for_testing, + get_swapped_tensor, + is_swapped_range, + register_swapped_tensor, +) class MockRequest: @@ -138,6 +144,19 @@ def handle_request(req_meta: dict[str, Any]): self.assertEqual(events[0], ("set_device", expected_device.index)) self.assertEqual(events[1][0], "handle") + def test_swap_memory_registry_resolves_transfer_subrange(self): + clear_swapped_tensors_for_testing() + try: + tensor = torch.empty((128,), dtype=torch.int8) + register_swapped_tensor(tensor) + self.assertTrue(is_swapped_range(tensor.data_ptr() + 16, 32)) + resolved = get_swapped_tensor(tensor.data_ptr() + 16, 32) + self.assertIsNotNone(resolved) + self.assertIs(resolved[0], tensor) + self.assertEqual(resolved[1], 16) + finally: + clear_swapped_tensors_for_testing() + def test_submit_request_serializes_same_peer_fifo(self): thread = self._make_thread() release_first_request = threading.Event() diff --git a/tests/ut/kv_offload/test_mooncake_swap_staging.py b/tests/ut/kv_offload/test_mooncake_swap_staging.py new file mode 100644 index 000000000000..aa83bfdeda9d --- /dev/null +++ b/tests/ut/kv_offload/test_mooncake_swap_staging.py @@ -0,0 +1,644 @@ +# SPDX-License-Identifier: Apache-2.0 +"""UT for the Kimi K3 Mooncake swap-memory receive path. + +The tests are split so that they can run on a host without an Ascend NPU: + +* ``TestSwapMemoryRegistry`` and ``TestSwapStagingPlanner`` use CPU tensors and + a fake transfer engine. They cover address planning, the 2 MiB split, batch + windowing and error propagation. +* ``TestSwapStagingOnNPU`` is skipped unless a real NPU with + ``torch_npu.empty_with_swapped_memory`` is present. It covers the parts that + can only be judged against real swap memory. +""" + +import unittest +from unittest import mock + +import torch + +from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake_hybrid_connector import ( + KVCacheRecvingThread, +) +from vllm_ascend.kv_offload.mooncake_swap_memory import ( + clear_swapped_tensors_for_testing, + get_swapped_tensor, + is_swapped_range, + register_swapped_tensor, +) + +ALIGN = 2 * 1024 * 1024 + + +def _npu_swap_available() -> bool: + try: + import torch_npu + except ImportError: + return False + if not hasattr(torch_npu, "empty_with_swapped_memory"): + return False + return bool(getattr(torch, "npu", None)) and torch.npu.is_available() + + +class FakeEngine: + """Records every batch_transfer_sync_read call and replays scripted rets.""" + + def __init__(self, rets=None): + self.calls: list[tuple[str, list[int], list[int], list[int]]] = [] + self._rets = list(rets or []) + + def batch_transfer_sync_read(self, session_id, local_dst, remote_src, lengths): + self.calls.append((session_id, list(local_dst), list(remote_src), list(lengths))) + if self._rets: + return self._rets.pop(0) + return 0 + + +class StagingThreadStub: + """Binds the real staging methods onto a minimal, CPU-friendly object. + + Only the attributes the staging helpers touch are provided, so the tests + exercise the production code without constructing a full connector or + starting a thread. + """ + + _swap_staging_chunk_bytes = KVCacheRecvingThread._swap_staging_chunk_bytes + _ensure_swap_staging = KVCacheRecvingThread._ensure_swap_staging + _copy_staging_to_swapped = KVCacheRecvingThread._copy_staging_to_swapped + _batch_transfer_sync_read_with_swap_staging = ( + KVCacheRecvingThread._batch_transfer_sync_read_with_swap_staging + ) + + def __init__(self, block_len_per_addr, engine, staging_numel): + self.block_len_per_addr = list(block_len_per_addr) + self.engine = engine + # A CPU tensor stands in for the registered NPU staging buffer; the + # planner only needs data_ptr arithmetic and a byte-addressable view. + self._staging_backing = torch.zeros(staging_numel, dtype=torch.int8) + self._swap_staging_storage = self._staging_backing + self._swap_staging_tensor = self._staging_backing + self._swap_staging_bytes = staging_numel + self.kv_caches = {"layer0": self._staging_backing} + self.copies: list[tuple[int, int, int]] = [] + + def record_copy(self, dst, length, staging, staging_offset): + self.copies.append((dst, length, staging_offset)) + + +class TestSwapMemoryRegistry(unittest.TestCase): + def setUp(self): + clear_swapped_tensors_for_testing() + self.addCleanup(clear_swapped_tensors_for_testing) + + def test_resolves_exact_and_sub_ranges(self): + tensor = torch.zeros(4096, dtype=torch.int8) + register_swapped_tensor(tensor) + base = tensor.data_ptr() + + self.assertTrue(is_swapped_range(base, 4096)) + owner, offset = get_swapped_tensor(base + 100, 200) + self.assertIs(owner, tensor) + self.assertEqual(offset, 100) + + def test_rejects_ranges_outside_registered_allocation(self): + tensor = torch.zeros(4096, dtype=torch.int8) + register_swapped_tensor(tensor) + base = tensor.data_ptr() + + # One byte past the end must not resolve, otherwise the connector would + # copy outside the allocation. + self.assertFalse(is_swapped_range(base, 4097)) + self.assertFalse(is_swapped_range(base + 4096, 1)) + self.assertIsNone(get_swapped_tensor(base - 1, 8)) + + def test_unregistered_pointer_is_not_swapped(self): + other = torch.zeros(64, dtype=torch.int8) + self.assertFalse(is_swapped_range(other.data_ptr(), 64)) + + def test_registration_is_idempotent(self): + tensor = torch.zeros(256, dtype=torch.int8) + register_swapped_tensor(tensor) + register_swapped_tensor(tensor) + owner, offset = get_swapped_tensor(tensor.data_ptr(), 256) + self.assertIs(owner, tensor) + self.assertEqual(offset, 0) + + def test_rejects_empty_allocation(self): + with self.assertRaises(ValueError): + register_swapped_tensor(torch.zeros(0, dtype=torch.int8)) + + +class TestSwapStagingPlanner(unittest.TestCase): + """Address planning for _batch_transfer_sync_read_with_swap_staging.""" + + def setUp(self): + clear_swapped_tensors_for_testing() + self.addCleanup(clear_swapped_tensors_for_testing) + self.block_len = 4096 + + def _make(self, engine, staging_slots=4): + slot_bytes = ALIGN + ALIGN # chunk_bytes rounds up to ALIGN, plus window + return StagingThreadStub([self.block_len], engine, slot_bytes * staging_slots) + + def _register_swap_dst(self, numel=ALIGN): + swap = torch.zeros(numel, dtype=torch.int8) + register_swapped_tensor(swap) + return swap + + def test_ordinary_destination_uses_direct_read_only(self): + engine = FakeEngine() + stub = self._make(engine) + plain = torch.zeros(self.block_len, dtype=torch.int8) + + ret = stub._batch_transfer_sync_read_with_swap_staging( + "s1", [plain.data_ptr()], [0x7000], [self.block_len] + ) + + self.assertEqual(ret, 0) + self.assertEqual(len(engine.calls), 1) + _, local, remote, lengths = engine.calls[0] + self.assertEqual(local, [plain.data_ptr()]) + self.assertEqual(remote, [0x7000]) + self.assertEqual(lengths, [self.block_len]) + + def test_swap_destination_reads_into_staging_not_destination(self): + engine = FakeEngine() + stub = self._make(engine) + swap = self._register_swap_dst() + staging_ptr = stub._swap_staging_tensor.data_ptr() + + with mock.patch.object( + StagingThreadStub, "_copy_staging_to_swapped", StagingThreadStub.record_copy + ): + ret = stub._batch_transfer_sync_read_with_swap_staging( + "s1", [swap.data_ptr()], [ALIGN * 3], [self.block_len] + ) + + self.assertEqual(ret, 0) + self.assertEqual(len(engine.calls), 1) + _, local, remote, lengths = engine.calls[0] + # The destination handed to Mooncake must be the staging buffer. + self.assertNotEqual(local, [swap.data_ptr()]) + self.assertEqual(local, [staging_ptr]) + self.assertEqual(remote, [ALIGN * 3]) + self.assertEqual(lengths, [self.block_len]) + # And the data must then be copied into the swap destination. + self.assertEqual(stub.copies, [(swap.data_ptr(), self.block_len, 0)]) + + def test_mixed_batch_splits_direct_and_staged(self): + engine = FakeEngine() + stub = self._make(engine) + swap = self._register_swap_dst() + plain = torch.zeros(self.block_len, dtype=torch.int8) + + with mock.patch.object( + StagingThreadStub, "_copy_staging_to_swapped", StagingThreadStub.record_copy + ): + ret = stub._batch_transfer_sync_read_with_swap_staging( + "s1", + [plain.data_ptr(), swap.data_ptr()], + [ALIGN, ALIGN * 2], + [self.block_len, self.block_len], + ) + + self.assertEqual(ret, 0) + self.assertEqual(len(engine.calls), 2) + self.assertEqual(engine.calls[0][1], [plain.data_ptr()]) + self.assertEqual(engine.calls[1][1], [stub._swap_staging_tensor.data_ptr()]) + self.assertEqual(len(stub.copies), 1) + + def test_misaligned_remote_source_split_preserves_total_length(self): + engine = FakeEngine() + stub = self._make(engine) + # A destination large enough that one block spans an alignment boundary. + swap = self._register_swap_dst(numel=ALIGN * 4) + length = ALIGN + 8192 + remote_base = ALIGN * 5 + 4096 # deliberately not 2 MiB aligned + + with mock.patch.object( + StagingThreadStub, "_copy_staging_to_swapped", StagingThreadStub.record_copy + ): + ret = stub._batch_transfer_sync_read_with_swap_staging( + "s1", [swap.data_ptr()], [remote_base], [length] + ) + + self.assertEqual(ret, 0) + staged_lengths = [n for call in engine.calls for n in call[3]] + self.assertGreater(len(staged_lengths), 1, "expected the span to be split") + self.assertEqual(sum(staged_lengths), length) + self.assertTrue(all(n > 0 for n in staged_lengths)) + + # Remote offsets must stay contiguous and cover the range exactly once. + remotes = [r for call in engine.calls for r in call[2]] + expected = remote_base + for remote, piece in zip(remotes, staged_lengths): + self.assertEqual(remote, expected) + expected += piece + self.assertEqual(expected, remote_base + length) + + # Copy destinations must mirror the same contiguous layout. + self.assertEqual(sum(c[1] for c in stub.copies), length) + expected_dst = swap.data_ptr() + for dst, piece, _ in stub.copies: + self.assertEqual(dst, expected_dst) + expected_dst += piece + + def test_staged_piece_never_exceeds_slot_capacity(self): + engine = FakeEngine() + stub = self._make(engine) + swap = self._register_swap_dst(numel=ALIGN * 4) + chunk = stub._swap_staging_chunk_bytes() + + with mock.patch.object( + StagingThreadStub, "_copy_staging_to_swapped", StagingThreadStub.record_copy + ): + stub._batch_transfer_sync_read_with_swap_staging( + "s1", [swap.data_ptr()], [ALIGN * 7 + 1024], [ALIGN * 2 + 512] + ) + + for call in engine.calls: + for local, piece in zip(call[1], call[3]): + slot_offset = local - stub._swap_staging_tensor.data_ptr() + self.assertLessEqual( + slot_offset % (chunk + ALIGN) + piece, + chunk + ALIGN, + "staged piece must not run past its slot", + ) + + def test_negative_engine_return_skips_copy(self): + engine = FakeEngine(rets=[-1]) + stub = self._make(engine) + swap = self._register_swap_dst() + + with mock.patch.object( + StagingThreadStub, "_copy_staging_to_swapped", StagingThreadStub.record_copy + ): + ret = stub._batch_transfer_sync_read_with_swap_staging( + "s1", [swap.data_ptr()], [ALIGN], [self.block_len] + ) + + self.assertEqual(ret, -1) + self.assertEqual(stub.copies, [], "no copy may run after a failed transfer") + + def test_direct_failure_short_circuits_before_staging(self): + engine = FakeEngine(rets=[-2]) + stub = self._make(engine) + swap = self._register_swap_dst() + plain = torch.zeros(self.block_len, dtype=torch.int8) + + with mock.patch.object( + StagingThreadStub, "_copy_staging_to_swapped", StagingThreadStub.record_copy + ): + ret = stub._batch_transfer_sync_read_with_swap_staging( + "s1", + [plain.data_ptr(), swap.data_ptr()], + [ALIGN, ALIGN * 2], + [self.block_len, self.block_len], + ) + + self.assertEqual(ret, -2) + self.assertEqual(len(engine.calls), 1, "staging must not run after a direct failure") + self.assertEqual(stub.copies, []) + + def test_batches_are_copied_before_window_is_reused(self): + """Slot reuse must not overwrite data that has not been copied yet. + + The staging window is capped at 16 slots, so more than 16 staged + transfers are needed to force a second batch through the same slots. + """ + engine = FakeEngine() + stub = self._make(engine, staging_slots=16) + swap = self._register_swap_dst(numel=ALIGN * 64) + n = 20 + dsts = [swap.data_ptr() + i * self.block_len for i in range(n)] + remotes = [ALIGN * (i + 4) for i in range(n)] + lengths = [self.block_len] * n + + order: list[tuple[str, int]] = [] + + def record_read(session_id, local, remote, length): + order.append(("read", len(local))) + return FakeEngine.batch_transfer_sync_read(engine, session_id, local, remote, length) + + def record_copy(self_, dst, length, staging, staging_offset): + order.append(("copy", dst)) + + engine_wrapper = mock.Mock(side_effect=record_read) + stub.engine = mock.Mock(batch_transfer_sync_read=engine_wrapper) + + with mock.patch.object(StagingThreadStub, "_copy_staging_to_swapped", record_copy): + ret = stub._batch_transfer_sync_read_with_swap_staging("s1", dsts, remotes, lengths) + + self.assertEqual(ret, 0) + self.assertEqual(sum(1 for kind, _ in order if kind == "copy"), n) + reads = sum(1 for kind, _ in order if kind == "read") + self.assertGreater(reads, 1, "expected the window to be reused across batches") + + # Every read must be followed by its copies before the next read, so a + # reused slot can never hold uncopied data. + pending = 0 + for kind, payload in order: + if kind == "read": + self.assertEqual(pending, 0, "a new read started with copies outstanding") + pending = payload + else: + pending -= 1 + self.assertEqual(pending, 0) + + +class TestGlobalTERegistration(unittest.TestCase): + """Incremental (ptr, size, location) registration in GlobalTE.""" + + def _make_te(self, with_location_api=False, ret=0): + from vllm_ascend.distributed.kv_transfer.utils.mooncake_transfer_engine import GlobalTE + + te = GlobalTE() + engine_attrs = {"register_memory.return_value": ret} + engine = mock.Mock(**engine_attrs) + if with_location_api: + engine.register_memory_with_location = mock.Mock(return_value=ret) + else: + # Mock auto-creates attributes, so remove it explicitly. + del engine.register_memory_with_location + te.transfer_engine = engine + return te, engine + + def test_first_registration_calls_register_memory(self): + te, engine = self._make_te() + te.register_buffer([0x1000, 0x2000], [64, 128]) + + engine.register_memory.assert_has_calls( + [mock.call(0x1000, 64), mock.call(0x2000, 128)], any_order=False + ) + self.assertEqual(engine.register_memory.call_count, 2) + + def test_duplicate_registration_is_skipped(self): + te, engine = self._make_te() + te.register_buffer([0x1000], [64]) + te.register_buffer([0x1000], [64]) + + self.assertEqual(engine.register_memory.call_count, 1) + + def test_staging_buffer_registers_after_kv_caches(self): + """The old is_register_buffer flag short-circuited later buffers.""" + te, engine = self._make_te() + te.register_buffer([0x1000], [64]) + self.assertTrue(te.is_register_buffer) + + # A staging buffer allocated later must still be registered. + te.register_buffer([0x9000], [ALIGN]) + + engine.register_memory.assert_called_with(0x9000, ALIGN) + self.assertEqual(engine.register_memory.call_count, 2) + + def test_same_pointer_different_size_registers_again(self): + te, engine = self._make_te() + te.register_buffer([0x1000], [64]) + te.register_buffer([0x1000], [128]) + + self.assertEqual(engine.register_memory.call_count, 2) + + def test_mismatched_size_count_raises(self): + te, _ = self._make_te() + with self.assertRaises(ValueError): + te.register_buffer([0x1000, 0x2000], [64]) + + def test_mismatched_location_count_raises(self): + te, _ = self._make_te() + with self.assertRaises(ValueError): + te.register_buffer([0x1000, 0x2000], [64, 128], ["cpu"]) + + def test_location_registration_uses_location_api(self): + te, engine = self._make_te(with_location_api=True) + te.register_buffer([0x1000], [64], ["cpu:0"]) + + engine.register_memory_with_location.assert_called_once_with(0x1000, 64, "cpu:0") + engine.register_memory.assert_not_called() + + def test_missing_location_api_raises_runtime_error(self): + te, _ = self._make_te(with_location_api=False) + with self.assertRaises(RuntimeError): + te.register_buffer([0x1000], [64], ["cpu:0"]) + + def test_location_and_default_are_tracked_separately(self): + te, engine = self._make_te(with_location_api=True) + te.register_buffer([0x1000], [64]) + te.register_buffer([0x1000], [64], ["cpu:0"]) + + engine.register_memory.assert_called_once_with(0x1000, 64) + engine.register_memory_with_location.assert_called_once_with(0x1000, 64, "cpu:0") + + def test_failed_registration_raises_and_is_not_cached(self): + te, engine = self._make_te(ret=-1) + with self.assertRaises(RuntimeError): + te.register_buffer([0x1000], [64]) + + # A failed range must not be remembered as registered. + engine.register_memory.return_value = 0 + te.register_buffer([0x1000], [64]) + self.assertEqual(engine.register_memory.call_count, 2) + + +def _load_gate_function(): + """Compile _use_kimi_k3_pd_swap_memory straight out of model_runner_v1.py. + + ``model_runner_v1`` pulls in the whole worker stack, which cannot be + imported in a bare UT environment. The gate is self-contained, so the + function is extracted from the real source file and compiled on its own. + That keeps the test bound to the shipped code rather than a copy. + """ + import ast + import pathlib + + import vllm_ascend + + source_path = pathlib.Path(vllm_ascend.__file__).parent / "worker" / "model_runner_v1.py" + tree = ast.parse(source_path.read_text(encoding="utf-8")) + target = "_use_kimi_k3_pd_swap_memory" + for node in ast.walk(tree): + if isinstance(node, ast.FunctionDef) and node.name == target: + module = ast.Module(body=[node], type_ignores=[]) + ast.fix_missing_locations(module) + namespace: dict = {} + exec(compile(module, str(source_path), "exec"), namespace) # noqa: S102 + return namespace[target] + raise AssertionError(f"{target} not found in {source_path}") + + +class _GateStub: + def __init__(self, kv_transfer_config, model_type): + self.vllm_config = mock.Mock(kv_transfer_config=kv_transfer_config) + self.model_config = mock.Mock(hf_config=mock.Mock(model_type=model_type)) + + +def _kv_cfg(role="kv_consumer", connector="MooncakeHybridConnector", extra=None): + return mock.Mock( + kv_role=role, + kv_connector=connector, + kv_connector_extra_config=extra if extra is not None else {}, + ) + + +class TestKimiK3SwapGate(unittest.TestCase): + """Risk 7.4: the gate must not widen to other models or to Prefill.""" + + @classmethod + def setUpClass(cls): + cls.gate = staticmethod(_load_gate_function()) + + def _call(self, kv_transfer_config, model_type): + return type(self).gate(_GateStub(kv_transfer_config, model_type)) + + def test_enabled_for_k3_consumer_with_hybrid_connector(self): + self.assertTrue(self._call(_kv_cfg(), "kimi_k3")) + + def test_disabled_without_kv_transfer_config(self): + self.assertFalse(self._call(None, "kimi_k3")) + + def test_disabled_for_producer_role(self): + self.assertFalse(self._call(_kv_cfg(role="kv_producer"), "kimi_k3")) + + def test_disabled_for_other_connectors(self): + self.assertFalse(self._call(_kv_cfg(connector="MooncakeConnector"), "kimi_k3")) + + def test_disabled_for_kimi_linear_text_config(self): + """The K3 text sub-config keeps model_type kimi_linear.""" + self.assertFalse(self._call(_kv_cfg(), "kimi_linear")) + + def test_disabled_for_unrelated_model(self): + self.assertFalse(self._call(_kv_cfg(), "deepseek_v3")) + + def test_enabled_via_nested_connector_list(self): + extra = {"connectors": [{"kv_connector": "MooncakeHybridConnector"}]} + cfg = _kv_cfg(connector="MultiConnector", extra=extra) + self.assertTrue(self._call(cfg, "kimi_k3")) + + def test_nested_connector_list_without_hybrid_is_disabled(self): + extra = {"connectors": [{"kv_connector": "SomethingElse"}]} + cfg = _kv_cfg(connector="MultiConnector", extra=extra) + self.assertFalse(self._call(cfg, "kimi_k3")) + + def test_connector_name_match_is_case_insensitive(self): + self.assertTrue(self._call(_kv_cfg(connector="mooncakehybridconnector"), "kimi_k3")) + + def test_missing_hf_config_is_disabled(self): + stub = _GateStub(_kv_cfg(), "kimi_k3") + stub.model_config = mock.Mock(spec=[]) + self.assertFalse(type(self).gate(stub)) + + +def _readout(swap_bytes, offset, length): + """Read a swap-memory span without touching it from the host. + + Direct host access to swap memory (``.to("cpu")``, ``tensor[i].item()``, + ``cpu_tensor.copy_(swap)``) faults, so verification must bounce through an + ordinary device tensor first. + """ + out = torch.empty(length, dtype=torch.int8, device=swap_bytes.device) + out.copy_(swap_bytes.narrow(0, offset, length)) + torch.npu.synchronize() + return out.to("cpu") + + +@unittest.skipUnless(_npu_swap_available(), "requires an Ascend NPU with swap-memory support") +class TestSwapStagingOnNPU(unittest.TestCase): + """Checks that can only be judged against real swap memory.""" + + def setUp(self): + clear_swapped_tensors_for_testing() + self.addCleanup(clear_swapped_tensors_for_testing) + torch.npu.set_device(0) + self.device = torch.device("npu:0") + + def _alloc_aligned(self, numel): + """Mirrors NPUModelRunner._allocate_kimi_k3_swap_tensor.""" + from vllm_ascend.kv_offload.mooncake_swap_memory import empty_swapped_memory + + storage = empty_swapped_memory((numel + ALIGN,), dtype=torch.int8) + offset = (-storage.data_ptr()) % ALIGN + view = storage[offset : offset + numel] + self.assertEqual(view.data_ptr() % ALIGN, 0) + register_swapped_tensor(view) + # Keep the backing allocation alive for the duration of the test. + self._storage = storage + return view + + def test_allocation_is_aligned_and_registered(self): + view = self._alloc_aligned(4 * ALIGN) + self.assertEqual(view.device.type, "npu") + self.assertEqual(view.numel(), 4 * ALIGN) + self.assertTrue(is_swapped_range(view.data_ptr(), 4 * ALIGN)) + self.assertTrue(is_swapped_range(view.data_ptr() + ALIGN, 4096)) + self.assertFalse(is_swapped_range(view.data_ptr() + 4 * ALIGN, 1)) + + def test_fresh_allocation_reads_as_zero(self): + """The replaced torch.zeros path guaranteed zeroed cache memory. + + Observed on torch_npu 2.10.0 / Ascend 910: a fresh allocation already + reads back as zero. This is not promised by the op documentation, so + the check is here to catch a regression rather than to rely on it. + """ + view = self._alloc_aligned(2 * ALIGN) + self.assertEqual(int(_readout(view, 0, 64 * 1024).abs().sum()), 0) + + def test_zero_and_fill_are_supported(self): + view = self._alloc_aligned(2 * ALIGN) + view.fill_(5) + torch.npu.synchronize() + self.assertTrue(torch.equal(_readout(view, 0, 4096), torch.full((4096,), 5, dtype=torch.int8))) + view.zero_() + torch.npu.synchronize() + self.assertEqual(int(_readout(view, 0, 4096).abs().sum()), 0) + + def test_copy_staging_to_swapped_writes_exact_bytes(self): + view = self._alloc_aligned(4 * ALIGN) + view.zero_() + torch.npu.synchronize() + + length = 64 * 1024 + pattern = torch.arange(length, dtype=torch.int32).remainder(251).sub(125).to(torch.int8) + staging = pattern.to(self.device) + # A deliberately unaligned destination offset inside the allocation. + dst_offset = ALIGN + 4096 + + stub = object.__new__(KVCacheRecvingThread) + KVCacheRecvingThread._copy_staging_to_swapped( + stub, view.data_ptr() + dst_offset, length, staging, 0 + ) + + got = _readout(view, dst_offset, length) + self.assertTrue(torch.equal(got, pattern)) + + def test_copy_respects_staging_offset_and_leaves_neighbours_intact(self): + view = self._alloc_aligned(4 * ALIGN) + view.zero_() + torch.npu.synchronize() + + length = 32 * 1024 + staging_offset = 8192 + staging = torch.zeros(staging_offset + length, dtype=torch.int8, device=self.device) + pattern = torch.full((length,), 42, dtype=torch.int8) + staging.narrow(0, staging_offset, length).copy_(pattern.to(self.device)) + torch.npu.synchronize() + + dst_offset = 2 * ALIGN + 1024 + stub = object.__new__(KVCacheRecvingThread) + KVCacheRecvingThread._copy_staging_to_swapped( + stub, view.data_ptr() + dst_offset, length, staging, staging_offset + ) + + self.assertTrue(torch.equal(_readout(view, dst_offset, length), pattern)) + # Bytes on both sides of the written span must be untouched. + self.assertEqual(int(_readout(view, dst_offset - 4096, 4096).abs().sum()), 0) + self.assertEqual(int(_readout(view, dst_offset + length, 4096).abs().sum()), 0) + + def test_copy_to_unregistered_destination_raises(self): + self._alloc_aligned(2 * ALIGN) + staging = torch.zeros(4096, dtype=torch.int8, device=self.device) + stub = object.__new__(KVCacheRecvingThread) + + with self.assertRaises(RuntimeError): + KVCacheRecvingThread._copy_staging_to_swapped(stub, 0xDEAD0000, 4096, staging, 0) + + +if __name__ == "__main__": + unittest.main() diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake_hybrid_connector.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake_hybrid_connector.py index 5d07ff499774..1fe240f5fd93 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake_hybrid_connector.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake_hybrid_connector.py @@ -54,6 +54,11 @@ from vllm_ascend.ascend_config import get_ascend_config, init_ascend_config from vllm_ascend.distributed.kv_transfer.utils.mooncake_transfer_engine import global_te from vllm_ascend.distributed.kv_transfer.utils.utils import get_transfer_timeout_value +from vllm_ascend.kv_offload.mooncake_swap_memory import ( + get_swapped_tensor, + is_swapped_range, + iter_tensors, +) from vllm_ascend.utils import enable_custom_op, is_vl_model # isort: off @@ -406,6 +411,9 @@ def __init__( self.mamba_ssm_size = mamba_ssm_size self.remote_te_port: dict[str, dict[int, int]] = SizedDict() self.remote_metadata_lock = threading.Lock() + self._swap_staging_storage: torch.Tensor | None = None + self._swap_staging_tensor: torch.Tensor | None = None + self._swap_staging_bytes = 0 self.request_queue: queue.Queue[Any] = queue.Queue() first_kv_cache = next(iter(self.kv_caches.values())) @@ -463,6 +471,110 @@ def __init__( self.proc_not_transfer_request: dict[str, bool] = {} self.proc_not_transfer_request_lock = threading.Lock() + def _swap_staging_chunk_bytes(self) -> int: + """Use a full aligned block-sized staging slot for each transfer.""" + block_bytes = [int(length) for length in self.block_len_per_addr if int(length) > 0] + alignment = 2 * 1024 * 1024 + chunk_bytes = max(max(block_bytes, default=alignment), alignment) + return (chunk_bytes + alignment - 1) // alignment * alignment + + def _ensure_swap_staging(self, transfer_count: int) -> torch.Tensor: + chunk_bytes = self._swap_staging_chunk_bytes() + alignment = 2 * 1024 * 1024 + # Each slot may start at an arbitrary remote byte offset. Reserve a + # full alignment window in addition to the chunk so that + # ``slot * slot_bytes + remote_offset`` is always in bounds. + slot_bytes = chunk_bytes + alignment + required_bytes = slot_bytes * max(1, transfer_count) + if self._swap_staging_tensor is not None and self._swap_staging_bytes >= required_bytes: + return self._swap_staging_tensor + + first_cache = next(iter(self.kv_caches.values())) + first_tensor = next(iter_tensors(first_cache), None) + if first_tensor is None or first_tensor.device.type == "cpu": + raise RuntimeError("Mooncake swapped-memory receive requires an NPU staging tensor.") + storage = torch.empty(required_bytes + alignment, dtype=torch.int8, device=first_tensor.device) + offset = (-int(storage.data_ptr())) % alignment + staging = storage[offset : offset + required_bytes] + if int(staging.data_ptr()) % alignment != 0: + raise RuntimeError("Mooncake NPU staging tensor is not 2MB aligned.") + global_te.register_buffer([int(staging.data_ptr())], [required_bytes]) + self._swap_staging_storage = storage + self._swap_staging_tensor = staging + self._swap_staging_bytes = required_bytes + return staging + + def _copy_staging_to_swapped(self, dst: int, length: int, staging: torch.Tensor, staging_offset: int) -> None: + resolved = get_swapped_tensor(dst, length) + if resolved is None: + raise RuntimeError(f"Missing swapped tensor for Mooncake destination ptr={dst}, length={length}.") + target, target_offset = resolved + target_bytes = target.view(torch.int8).reshape(-1) + target_bytes.narrow(0, target_offset, length).copy_( + staging.narrow(0, staging_offset, length), non_blocking=False + ) + torch.npu.synchronize() + + def _batch_transfer_sync_read_with_swap_staging( + self, + session_id: str, + local_dst_list: list[int], + remote_src_list: list[int], + length_list: list[int], + ) -> int: + """Read directly for NPU destinations and stage reads for swapped ones.""" + direct_local: list[int] = [] + direct_remote: list[int] = [] + direct_lengths: list[int] = [] + staged: list[tuple[int, int, int]] = [] + chunk_bytes = self._swap_staging_chunk_bytes() + alignment = 2 * 1024 * 1024 + + for local_dst, remote_src, length in zip(local_dst_list, remote_src_list, length_list): + if not is_swapped_range(local_dst, length): + direct_local.append(local_dst) + direct_remote.append(remote_src) + direct_lengths.append(length) + continue + remaining = int(length) + offset = 0 + while remaining: + remote_offset = (int(remote_src) + offset) % alignment + piece_len = min(remaining, chunk_bytes - remote_offset) + staged.append((int(local_dst) + offset, int(remote_src) + offset, piece_len)) + offset += piece_len + remaining -= piece_len + + if direct_local: + ret = self.engine.batch_transfer_sync_read(session_id, direct_local, direct_remote, direct_lengths) + if ret < 0: + return ret + if not staged: + return 0 + + staging = self._ensure_swap_staging(min(len(staged), 16)) + staging_ptr = int(staging.data_ptr()) + slot_bytes = chunk_bytes + alignment + window = max(1, min(16, len(staging) // slot_bytes)) + for batch_start in range(0, len(staged), window): + batch = staged[batch_start : batch_start + window] + local_staging: list[int] = [] + remote_src: list[int] = [] + lengths: list[int] = [] + copies: list[tuple[int, int, int]] = [] + for slot, (local_dst, remote, length) in enumerate(batch): + staging_offset = slot * slot_bytes + remote % alignment + local_staging.append(staging_ptr + staging_offset) + remote_src.append(remote) + lengths.append(length) + copies.append((local_dst, length, staging_offset)) + ret = self.engine.batch_transfer_sync_read(session_id, local_staging, remote_src, lengths) + if ret < 0: + return ret + for local_dst, length, staging_offset in copies: + self._copy_staging_to_swapped(local_dst, length, staging, staging_offset) + return 0 + def add_request( self, request_id: str, @@ -696,7 +808,7 @@ def _transfer_kv_cache_all_groups(self, req_meta: dict[str, Any]): dst_list.append(dst) length_list.append(length) - ret = self.engine.batch_transfer_sync_read(session_id, src_list, dst_list, length_list) + ret = self._batch_transfer_sync_read_with_swap_staging(session_id, src_list, dst_list, length_list) if ret < 0: logger.error( "Mooncake transfer failed for request. remote_request_id=%s, ret=%d. ", @@ -796,7 +908,7 @@ def _transfer_kv_cache(self, req_meta: dict[str, Any]): dst_list.append(dst) length_list.append(length) - ret = self.engine.batch_transfer_sync_read(session_id, src_list, dst_list, length_list) + ret = self._batch_transfer_sync_read_with_swap_staging(session_id, src_list, dst_list, length_list) if ret < 0: logger.error( "Mooncake transfer failed for request. remote_request_id=%s, ret=%d. ", @@ -1618,6 +1730,18 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): self.addr_group_idx: list[int] = [] ptrs = [] lengths = [] + + def add_registration(ptr: int, length: int) -> None: + # Swap-backed tensors are the final Decode destination. AscendDirect + # must not register/write those ranges; the receive thread stages + # into a registered NPU buffer and copies into swap memory instead. + # A hybrid cache view can begin in the middle of the raw allocation + # (for example after a padding prefix), so test the start address + # independently of the requested registration span. + if get_swapped_tensor(ptr, 1) is None: + ptrs.append(ptr) + lengths.append(length) + if not self.use_hybrid: for layer_name, kv_cache_tuple in kv_caches.items(): if isinstance(kv_cache_tuple, (list, tuple)) is False: @@ -1630,8 +1754,10 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): single_kv_cache.element_size() * math.prod(block_shape) * block_size_scale ) self.kv_caches_base_addr.append(single_kv_cache.data_ptr()) - ptrs.append(single_kv_cache.data_ptr()) - lengths.append(single_kv_cache.element_size() * math.prod(single_kv_cache.shape)) + add_registration( + single_kv_cache.data_ptr(), + single_kv_cache.element_size() * math.prod(single_kv_cache.shape), + ) elif self.use_mamba: for kv_cache_tensor in self.kv_cache_config.kv_cache_tensors: share_tensor_addr = [] @@ -1651,8 +1777,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): self.kv_caches_base_addr.append(single_kv_cache.data_ptr()) share_tensor_addr.append(single_kv_cache.data_ptr()) if share_tensor_addr: - ptrs.append(min(share_tensor_addr)) - lengths.append(kv_cache_tensor.size) + add_registration(min(share_tensor_addr), kv_cache_tensor.size) self.block_stride_per_addr.extend(self.block_len_per_addr) elif self.use_compress: layer_group_idx = dict[str, int]() @@ -1681,8 +1806,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): self.addr_group_idx.append(cur_tensor_group_idx) # type: ignore[arg-type] self.block_stride_per_addr.append(share_tensor_stride[0]) self.block_len_per_addr.append(share_tensor_stride[0]) - ptrs.append(min(share_tensor_addr)) - lengths.append(kv_cache_tensor.size) + add_registration(min(share_tensor_addr), kv_cache_tensor.size) else: raise TypeError("Mooncake connector does not support this type kv_cache now.") diff --git a/vllm_ascend/distributed/kv_transfer/utils/mooncake_transfer_engine.py b/vllm_ascend/distributed/kv_transfer/utils/mooncake_transfer_engine.py index 7cf309b0ee51..06830c09d8d3 100644 --- a/vllm_ascend/distributed/kv_transfer/utils/mooncake_transfer_engine.py +++ b/vllm_ascend/distributed/kv_transfer/utils/mooncake_transfer_engine.py @@ -7,6 +7,7 @@ def __init__(self): self.is_register_buffer: bool = False self.transfer_engine_lock = threading.Lock() self.register_buffer_lock = threading.Lock() + self.registered_buffers: set[tuple[int, int, str | None]] = set() def get_transfer_engine(self, hostname: str, device_name: str | None): if self.transfer_engine is None: @@ -28,15 +29,45 @@ def get_transfer_engine(self, hostname: str, device_name: str | None): raise RuntimeError(f"TransferEngine initialization failed with ret_value: {ret_value}") return self.transfer_engine - def register_buffer(self, ptrs: list[int], sizes: list[int]): + def register_buffer( + self, + ptrs: list[int], + sizes: list[int], + locations: list[str | None] | None = None, + ): with self.register_buffer_lock: assert self.transfer_engine is not None, "Transfer engine must be initialized" - if self.is_register_buffer: - return - for ptr, size in zip(ptrs, sizes): - ret_value = self.transfer_engine.register_memory(ptr, size) + if len(ptrs) != len(sizes): + raise ValueError("Mooncake register pointer/size counts differ.") + + if locations is None: + # Registration is incremental: the receive path can allocate + # an NPU staging buffer after the KV cache buffers have been + # registered. Keep the legacy call shape while consulting + # the per-buffer registry below instead of short-circuiting + # on the historical boolean flag. + locations = [None] * len(ptrs) + elif len(locations) != len(ptrs): + raise ValueError("Mooncake register locations must match ptr count.") + + register_with_location = getattr(self.transfer_engine, "register_memory_with_location", None) + for ptr, size, location in zip(ptrs, sizes, locations): + key = (int(ptr), int(size), location) + if key in self.registered_buffers: + continue + if location is not None: + if register_with_location is None: + raise RuntimeError( + "Mooncake TransferEngine does not support location-aware memory registration." + ) + ret_value = register_with_location(ptr, size, location) + else: + ret_value = self.transfer_engine.register_memory(ptr, size) if ret_value != 0: - raise RuntimeError("Mooncake memory registration failed.") + raise RuntimeError( + f"Mooncake memory registration failed. ptr={ptr} size={size} location={location}" + ) + self.registered_buffers.add(key) self.is_register_buffer = True diff --git a/vllm_ascend/kv_offload/mooncake_swap_memory.py b/vllm_ascend/kv_offload/mooncake_swap_memory.py new file mode 100644 index 000000000000..568d0c403d0d --- /dev/null +++ b/vllm_ascend/kv_offload/mooncake_swap_memory.py @@ -0,0 +1,71 @@ +"""Swap-memory helpers for Mooncake PD receive targets. + +``torch_npu.empty_with_swapped_memory`` returns an NPU tensor backed by +host-side swap memory. Mooncake cannot reliably write that allocation +directly with AscendDirect, so the connector uses this registry to identify +staged receive destinations and copies the data from an NPU staging buffer +into the swapped allocation after each transfer. +""" + +from __future__ import annotations + +from typing import Any, Iterator + +import torch + + +_SWAPPED_TENSORS: list[tuple[int, int, torch.Tensor]] = [] + + +def iter_tensors(value: Any) -> Iterator[torch.Tensor]: + """Yield tensors from a tensor/list/tuple cache structure.""" + if isinstance(value, torch.Tensor): + yield value + elif isinstance(value, (list, tuple)): + for item in value: + yield from iter_tensors(item) + + +def register_swapped_tensor(tensor: torch.Tensor) -> None: + """Keep the allocation alive and resolve transfer sub-ranges to it.""" + ptr = int(tensor.data_ptr()) + size = int(tensor.numel() * tensor.element_size()) + if ptr <= 0 or size <= 0: + raise ValueError(f"Invalid swapped tensor range: ptr={ptr}, size={size}") + if not any(start == ptr and length == size for start, length, _ in _SWAPPED_TENSORS): + _SWAPPED_TENSORS.append((ptr, size, tensor)) + + +def get_swapped_tensor(ptr: int, size: int) -> tuple[torch.Tensor, int] | None: + """Return the owning swapped tensor and byte offset for a sub-range.""" + ptr = int(ptr) + size = int(size) + end = ptr + size + for start, length, tensor in _SWAPPED_TENSORS: + if ptr >= start and end <= start + length: + return tensor, ptr - start + return None + + +def is_swapped_range(ptr: int, size: int) -> bool: + return get_swapped_tensor(ptr, size) is not None + + +def clear_swapped_tensors_for_testing() -> None: + _SWAPPED_TENSORS.clear() + + +def empty_swapped_memory(shape: tuple[int, ...], *, dtype: torch.dtype) -> torch.Tensor: + """Allocate an NPU tensor whose storage is host-side swap memory.""" + try: + import torch_npu + except ImportError as exc: + raise RuntimeError( + "Mooncake swap-memory receive requires torch_npu.empty_with_swapped_memory." + ) from exc + + allocator = getattr(torch_npu, "empty_with_swapped_memory", None) + if allocator is None: + raise RuntimeError("Mooncake swap-memory receive requires torch_npu.empty_with_swapped_memory.") + return allocator(shape, dtype=dtype, device="npu") + diff --git a/vllm_ascend/worker/model_runner_v1.py b/vllm_ascend/worker/model_runner_v1.py index f528043f9a96..8072d48168a8 100644 --- a/vllm_ascend/worker/model_runner_v1.py +++ b/vllm_ascend/worker/model_runner_v1.py @@ -146,6 +146,10 @@ from vllm_ascend.eplb.core.eplb_device_transfer_loader import D2DExpertWeightLoader from vllm_ascend.eplb.core.eplb_worker import EplbProcess from vllm_ascend.eplb.eplb_updator import EplbUpdator +from vllm_ascend.kv_offload.mooncake_swap_memory import ( + empty_swapped_memory, + register_swapped_tensor, +) from vllm_ascend.model_executor.offloader import create_offloader from vllm_ascend.ops.rotary_embedding import set_cos_and_sin, update_cos_sin from vllm_ascend.ops.triton.spec_decode.ngram import triton_ngram_spec_decode @@ -3939,6 +3943,50 @@ def _allocate_int8_cache_tensor( ) return self._align_memory(raw_tensor, alignment)[:numel] + def _use_kimi_k3_pd_swap_memory(self) -> bool: + """Whether K3 decode must receive Mooncake KV into swap memory. + + ``KimiK3TextConfig`` keeps the ``kimi_linear`` model type, so the + outer multimodal config is the reliable K3 discriminator. Restrict + this allocation to the Decode side of the hybrid connector; Prefill + and other connectors retain the normal device allocation path. + """ + kv_transfer_config = self.vllm_config.kv_transfer_config + if kv_transfer_config is None: + return False + if str(getattr(kv_transfer_config, "kv_role", "")) != "kv_consumer": + return False + connector_names = [str(getattr(kv_transfer_config, "kv_connector", ""))] + extra_config = getattr(kv_transfer_config, "kv_connector_extra_config", None) or {} + nested_connectors = extra_config.get("connectors", []) + if isinstance(nested_connectors, dict): + nested_connectors = nested_connectors.values() + for nested in nested_connectors: + if isinstance(nested, dict): + nested_name = nested.get("kv_connector", "") + else: + nested_name = getattr(nested, "kv_connector", "") + connector_names.append(str(nested_name)) + if not any(name.lower() == "mooncakehybridconnector" for name in connector_names): + return False + hf_config = getattr(self.model_config, "hf_config", None) + return str(getattr(hf_config, "model_type", "")) == "kimi_k3" + + def _allocate_kimi_k3_swap_tensor(self, numel: int, alignment: int) -> torch.Tensor: + """Allocate an aligned int8 raw tensor backed by swap memory.""" + if numel <= 0: + raise ValueError(f"Invalid cache tensor size: {numel}") + # Allocate extra bytes so the logical view can satisfy Mooncake's + # 2-MiB address requirement, just like the regular device path. + storage = empty_swapped_memory((numel + alignment,), dtype=torch.int8) + tensor = self._align_memory(storage, alignment)[:numel] + if int(tensor.data_ptr()) % alignment != 0: + raise RuntimeError( + "Kimi K3 Mooncake swap-memory tensor is not aligned to 2 MiB." + ) + register_swapped_tensor(tensor) + return tensor + def _allocate_sparse_c8_indexer_tensors( self, dsa_k_tensor_size: int, @@ -4014,6 +4062,13 @@ def _allocate_kv_cache_tensors(self, kv_cache_config: KVCacheConfig) -> dict[str kv_cache_raw_tensors: dict[str, torch.Tensor | tuple[torch.Tensor, ...]] = {} # prefill disaggregation need the addr of cache tensor be aligned with 2M alignment = 2 * 1024 * 1024 + use_kimi_k3_swap_memory = self._use_kimi_k3_pd_swap_memory() + + def allocate_raw_tensor(numel: int) -> torch.Tensor: + if use_kimi_k3_swap_memory: + return self._allocate_kimi_k3_swap_tensor(numel, alignment) + return self._allocate_int8_cache_tensor(numel, alignment) + layer_kv_cache_spec = self._get_layer_kv_cache_specs(kv_cache_config) # If some tensors are shared by linear layers and attention layers, # the same tensor format must be maintained even if some layers @@ -4048,22 +4103,12 @@ def _allocate_kv_cache_tensors(self, kv_cache_config: KVCacheConfig) -> dict[str is_hidden_state_cache_spec(layer_kv_cache_spec.get(ln)) for ln in kv_cache_tensor.shared_by ) - if self.vllm_config.kv_transfer_config is None: - tensor = torch.zeros(kv_cache_tensor.size, dtype=torch.int8, device=self.device) - else: - cache_size_aligned = kv_cache_tensor.size + alignment - tensor = torch.zeros(cache_size_aligned, dtype=torch.int8, device=self.device) - tensor = self._align_memory(tensor, alignment)[: kv_cache_tensor.size] + tensor = allocate_raw_tensor(kv_cache_tensor.size) if has_mamba and has_hidden: # Allocate separate tensor for HiddenStateCacheSpec layers # so ssm_state writes don't corrupt hidden-states data - if self.vllm_config.kv_transfer_config is None: - tensor_hs = torch.zeros(kv_cache_tensor.size, dtype=torch.int8, device=self.device) - else: - cache_size_aligned = kv_cache_tensor.size + alignment - tensor_hs = torch.zeros(cache_size_aligned, dtype=torch.int8, device=self.device) - tensor_hs = self._align_memory(tensor_hs, alignment)[: kv_cache_tensor.size] + tensor_hs = allocate_raw_tensor(kv_cache_tensor.size) for layer_name_inner in kv_cache_tensor.shared_by: if is_hidden_state_cache_spec(layer_kv_cache_spec.get(layer_name_inner)): kv_cache_raw_tensors[layer_name_inner] = tensor_hs @@ -4074,14 +4119,7 @@ def _allocate_kv_cache_tensors(self, kv_cache_config: KVCacheConfig) -> dict[str kv_cache_raw_tensors[layer_name_inner] = tensor elif "attn" in layer_name and self.use_compress and layer_name not in kv_cache_raw_tensors: - if self.vllm_config.kv_transfer_config is None: - tensor = torch.zeros(kv_cache_tensor.size, - dtype=torch.int8, - device=self.device) - else: - cache_size_aligned = kv_cache_tensor.size + alignment - tensor = torch.zeros(cache_size_aligned, dtype=torch.int8, device=self.device) - tensor = self._align_memory(tensor, alignment)[: kv_cache_tensor.size] + tensor = allocate_raw_tensor(kv_cache_tensor.size) for layer_name_inner in kv_cache_tensor.shared_by: # shared the kvcache between the self_attn specs in the same group kv_cache_raw_tensors[layer_name_inner] = tensor @@ -4117,9 +4155,8 @@ def _allocate_kv_cache_tensors(self, kv_cache_config: KVCacheConfig) -> dict[str ) raw_cache = (k_tensor, scale_tensor) else: - k_tensor = self._allocate_int8_cache_tensor( + k_tensor = allocate_raw_tensor( k_tensor_size, - alignment, ) raw_cache = (k_tensor,) @@ -4174,14 +4211,12 @@ def _allocate_kv_cache_tensors(self, kv_cache_config: KVCacheConfig) -> dict[str # are allocated as int8 raw bytes first and then viewed as # the target dtype in _reshape_kv_cache_tensors. v_tensor = None - k_tensor = self._allocate_int8_cache_tensor( + k_tensor = allocate_raw_tensor( k_tensor_size, - alignment, ) if v_tensor_size is not None: - v_tensor = self._allocate_int8_cache_tensor( + v_tensor = allocate_raw_tensor( v_tensor_size, - alignment, ) for layer_name_inner in kv_cache_tensor.shared_by: From 0b388a42bae728e2dc71e424599af84073bd55aa Mon Sep 17 00:00:00 2001 From: j00628475 Date: Fri, 28 Aug 2026 17:49:01 +0800 Subject: [PATCH 2/2] [feat]support kimi k3 d2rh mooncake transfer offload --- .../kv_offload/test_mooncake_swap_staging.py | 299 +++++++++++++++--- .../kv_p2p/mooncake_hybrid_connector.py | 128 +++++--- .../kv_offload/mooncake_swap_memory.py | 17 +- 3 files changed, 352 insertions(+), 92 deletions(-) diff --git a/tests/ut/kv_offload/test_mooncake_swap_staging.py b/tests/ut/kv_offload/test_mooncake_swap_staging.py index aa83bfdeda9d..40ec6bb57529 100644 --- a/tests/ut/kv_offload/test_mooncake_swap_staging.py +++ b/tests/ut/kv_offload/test_mooncake_swap_staging.py @@ -11,22 +11,35 @@ can only be judged against real swap memory. """ +import sys +import threading +import types import unittest +from concurrent.futures import ThreadPoolExecutor from unittest import mock import torch +from vllm_ascend.distributed.kv_transfer.kv_p2p import ( + mooncake_hybrid_connector as connector_module, +) from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake_hybrid_connector import ( + SWAP_STAGING_ALIGNMENT, + SWAP_STAGING_WINDOW, KVCacheRecvingThread, + KVCacheTaskTracker, + MooncakeConnector, + MooncakeConnectorWorker, ) from vllm_ascend.kv_offload.mooncake_swap_memory import ( clear_swapped_tensors_for_testing, + empty_swapped_memory, get_swapped_tensor, is_swapped_range, register_swapped_tensor, ) -ALIGN = 2 * 1024 * 1024 +ALIGN = SWAP_STAGING_ALIGNMENT def _npu_swap_available() -> bool: @@ -64,22 +77,27 @@ class StagingThreadStub: _swap_staging_chunk_bytes = KVCacheRecvingThread._swap_staging_chunk_bytes _ensure_swap_staging = KVCacheRecvingThread._ensure_swap_staging _copy_staging_to_swapped = KVCacheRecvingThread._copy_staging_to_swapped - _batch_transfer_sync_read_with_swap_staging = ( - KVCacheRecvingThread._batch_transfer_sync_read_with_swap_staging - ) + _batch_transfer_sync_read_with_swap_staging = KVCacheRecvingThread._batch_transfer_sync_read_with_swap_staging - def __init__(self, block_len_per_addr, engine, staging_numel): + def __init__(self, block_len_per_addr, engine, staging_numel=None): self.block_len_per_addr = list(block_len_per_addr) self.engine = engine + self._swap_staging_pools: dict[threading.Thread, tuple[torch.Tensor, torch.Tensor]] = {} + self._swap_staging_pool_lock = threading.Lock() # A CPU tensor stands in for the registered NPU staging buffer; the # planner only needs data_ptr arithmetic and a byte-addressable view. - self._staging_backing = torch.zeros(staging_numel, dtype=torch.int8) - self._swap_staging_storage = self._staging_backing - self._swap_staging_tensor = self._staging_backing - self._swap_staging_bytes = staging_numel + self._staging_backing = torch.zeros(staging_numel or 1, dtype=torch.int8) + if staging_numel is not None: + self._swap_staging_pools[threading.current_thread()] = ( + self._staging_backing, + self._staging_backing, + ) self.kv_caches = {"layer0": self._staging_backing} self.copies: list[tuple[int, int, int]] = [] + def current_staging(self): + return self._swap_staging_pools[threading.current_thread()][1] + def record_copy(self, dst, length, staging, staging_offset): self.copies.append((dst, length, staging_offset)) @@ -149,9 +167,7 @@ def test_ordinary_destination_uses_direct_read_only(self): stub = self._make(engine) plain = torch.zeros(self.block_len, dtype=torch.int8) - ret = stub._batch_transfer_sync_read_with_swap_staging( - "s1", [plain.data_ptr()], [0x7000], [self.block_len] - ) + ret = stub._batch_transfer_sync_read_with_swap_staging("s1", [plain.data_ptr()], [0x7000], [self.block_len]) self.assertEqual(ret, 0) self.assertEqual(len(engine.calls), 1) @@ -164,11 +180,9 @@ def test_swap_destination_reads_into_staging_not_destination(self): engine = FakeEngine() stub = self._make(engine) swap = self._register_swap_dst() - staging_ptr = stub._swap_staging_tensor.data_ptr() + staging_ptr = stub.current_staging().data_ptr() - with mock.patch.object( - StagingThreadStub, "_copy_staging_to_swapped", StagingThreadStub.record_copy - ): + with mock.patch.object(StagingThreadStub, "_copy_staging_to_swapped", StagingThreadStub.record_copy): ret = stub._batch_transfer_sync_read_with_swap_staging( "s1", [swap.data_ptr()], [ALIGN * 3], [self.block_len] ) @@ -190,9 +204,7 @@ def test_mixed_batch_splits_direct_and_staged(self): swap = self._register_swap_dst() plain = torch.zeros(self.block_len, dtype=torch.int8) - with mock.patch.object( - StagingThreadStub, "_copy_staging_to_swapped", StagingThreadStub.record_copy - ): + with mock.patch.object(StagingThreadStub, "_copy_staging_to_swapped", StagingThreadStub.record_copy): ret = stub._batch_transfer_sync_read_with_swap_staging( "s1", [plain.data_ptr(), swap.data_ptr()], @@ -203,7 +215,7 @@ def test_mixed_batch_splits_direct_and_staged(self): self.assertEqual(ret, 0) self.assertEqual(len(engine.calls), 2) self.assertEqual(engine.calls[0][1], [plain.data_ptr()]) - self.assertEqual(engine.calls[1][1], [stub._swap_staging_tensor.data_ptr()]) + self.assertEqual(engine.calls[1][1], [stub.current_staging().data_ptr()]) self.assertEqual(len(stub.copies), 1) def test_misaligned_remote_source_split_preserves_total_length(self): @@ -214,12 +226,8 @@ def test_misaligned_remote_source_split_preserves_total_length(self): length = ALIGN + 8192 remote_base = ALIGN * 5 + 4096 # deliberately not 2 MiB aligned - with mock.patch.object( - StagingThreadStub, "_copy_staging_to_swapped", StagingThreadStub.record_copy - ): - ret = stub._batch_transfer_sync_read_with_swap_staging( - "s1", [swap.data_ptr()], [remote_base], [length] - ) + with mock.patch.object(StagingThreadStub, "_copy_staging_to_swapped", StagingThreadStub.record_copy): + ret = stub._batch_transfer_sync_read_with_swap_staging("s1", [swap.data_ptr()], [remote_base], [length]) self.assertEqual(ret, 0) staged_lengths = [n for call in engine.calls for n in call[3]] @@ -248,16 +256,14 @@ def test_staged_piece_never_exceeds_slot_capacity(self): swap = self._register_swap_dst(numel=ALIGN * 4) chunk = stub._swap_staging_chunk_bytes() - with mock.patch.object( - StagingThreadStub, "_copy_staging_to_swapped", StagingThreadStub.record_copy - ): + with mock.patch.object(StagingThreadStub, "_copy_staging_to_swapped", StagingThreadStub.record_copy): stub._batch_transfer_sync_read_with_swap_staging( "s1", [swap.data_ptr()], [ALIGN * 7 + 1024], [ALIGN * 2 + 512] ) for call in engine.calls: for local, piece in zip(call[1], call[3]): - slot_offset = local - stub._swap_staging_tensor.data_ptr() + slot_offset = local - stub.current_staging().data_ptr() self.assertLessEqual( slot_offset % (chunk + ALIGN) + piece, chunk + ALIGN, @@ -269,12 +275,8 @@ def test_negative_engine_return_skips_copy(self): stub = self._make(engine) swap = self._register_swap_dst() - with mock.patch.object( - StagingThreadStub, "_copy_staging_to_swapped", StagingThreadStub.record_copy - ): - ret = stub._batch_transfer_sync_read_with_swap_staging( - "s1", [swap.data_ptr()], [ALIGN], [self.block_len] - ) + with mock.patch.object(StagingThreadStub, "_copy_staging_to_swapped", StagingThreadStub.record_copy): + ret = stub._batch_transfer_sync_read_with_swap_staging("s1", [swap.data_ptr()], [ALIGN], [self.block_len]) self.assertEqual(ret, -1) self.assertEqual(stub.copies, [], "no copy may run after a failed transfer") @@ -285,9 +287,7 @@ def test_direct_failure_short_circuits_before_staging(self): swap = self._register_swap_dst() plain = torch.zeros(self.block_len, dtype=torch.int8) - with mock.patch.object( - StagingThreadStub, "_copy_staging_to_swapped", StagingThreadStub.record_copy - ): + with mock.patch.object(StagingThreadStub, "_copy_staging_to_swapped", StagingThreadStub.record_copy): ret = stub._batch_transfer_sync_read_with_swap_staging( "s1", [plain.data_ptr(), swap.data_ptr()], @@ -345,6 +345,219 @@ def record_copy(self_, dst, length, staging, staging_offset): self.assertEqual(pending, 0) +class TestSwapStagingPools(unittest.TestCase): + def setUp(self): + clear_swapped_tensors_for_testing() + self.addCleanup(clear_swapped_tensors_for_testing) + + @staticmethod + def _fake_allocations(): + real_empty = torch.empty + allocations: list[int] = [] + allocations_lock = threading.Lock() + + def allocate_on_cpu(numel, *, dtype, device): + del device + with allocations_lock: + allocations.append(numel) + return real_empty(numel, dtype=dtype, device="cpu") + + fake_npu_tensor = mock.Mock(device=types.SimpleNamespace(type="npu")) + return allocations, allocate_on_cpu, fake_npu_tensor + + def test_same_worker_allocates_and_registers_full_pool_once(self): + stub = StagingThreadStub([4096], FakeEngine()) + allocations, allocate_on_cpu, fake_npu_tensor = self._fake_allocations() + expected_bytes = (4096 + 1) * SWAP_STAGING_WINDOW + + with ( + mock.patch.object(connector_module, "SWAP_STAGING_ALIGNMENT", 1), + mock.patch.object( + connector_module, + "iter_tensors", + side_effect=lambda _: iter([fake_npu_tensor]), + ), + mock.patch.object( + connector_module.torch, + "empty", + side_effect=allocate_on_cpu, + ), + mock.patch.object(connector_module.global_te, "register_buffer") as register, + ): + first = stub._ensure_swap_staging() + second = stub._ensure_swap_staging() + + self.assertIs(first, second) + self.assertEqual(first.numel(), expected_bytes) + self.assertEqual(allocations, [expected_bytes + 1]) + register.assert_called_once_with([first.data_ptr()], [expected_bytes]) + self.assertEqual(len(stub._swap_staging_pools), 1) + + def test_concurrent_workers_use_distinct_pools_without_overwrite(self): + class ConcurrentEngine: + def __init__(self): + self.owner = None + self.barrier = threading.Barrier(2) + self.staging_ptrs: dict[str, int] = {} + self.lock = threading.Lock() + + def batch_transfer_sync_read(self, session_id, local_dst, remote_src, lengths): + del remote_src + with self.lock: + self.staging_ptrs[session_id] = local_dst[0] + # Both workers must own a pool before either staged copy runs. + self.barrier.wait(timeout=5) + assert self.owner is not None + with self.owner._swap_staging_pool_lock: + pools = list(self.owner._swap_staging_pools.values()) + for _, staging in pools: + start = staging.data_ptr() + offset = local_dst[0] - start + if offset >= 0 and offset + lengths[0] <= staging.numel(): + staging.narrow(0, offset, lengths[0]).fill_({"worker-1": 11, "worker-2": 22}[session_id]) + break + else: + raise AssertionError("staging pointer did not resolve to a worker pool") + self.barrier.wait(timeout=5) + return 0 + + engine = ConcurrentEngine() + stub = StagingThreadStub([4096], engine) + engine.owner = stub + first_dst = torch.zeros(4096, dtype=torch.int8) + second_dst = torch.zeros(4096, dtype=torch.int8) + register_swapped_tensor(first_dst) + register_swapped_tensor(second_dst) + allocations, allocate_on_cpu, fake_npu_tensor = self._fake_allocations() + + def copy_without_npu_sync(self_, dst, length, staging, staging_offset): + resolved = get_swapped_tensor(dst, length) + assert resolved is not None + target, target_offset = resolved + target.view(torch.int8).reshape(-1).narrow(0, target_offset, length).copy_( + staging.narrow(0, staging_offset, length) + ) + + def transfer(session_id, destination): + return stub._batch_transfer_sync_read_with_swap_staging( + session_id, + [destination.data_ptr()], + [ALIGN * 3], + [destination.numel()], + ) + + with ( + mock.patch.object(connector_module, "SWAP_STAGING_ALIGNMENT", 1), + mock.patch.object( + connector_module, + "iter_tensors", + side_effect=lambda _: iter([fake_npu_tensor]), + ), + mock.patch.object( + connector_module.torch, + "empty", + side_effect=allocate_on_cpu, + ), + mock.patch.object(connector_module.global_te, "register_buffer") as register, + mock.patch.object( + StagingThreadStub, + "_copy_staging_to_swapped", + copy_without_npu_sync, + ), + ThreadPoolExecutor(max_workers=2) as executor, + ): + futures = [ + executor.submit(transfer, "worker-1", first_dst), + executor.submit(transfer, "worker-2", second_dst), + ] + self.assertEqual([future.result(timeout=10) for future in futures], [0, 0]) + + self.assertEqual(len(stub._swap_staging_pools), 2) + self.assertNotEqual(engine.staging_ptrs["worker-1"], engine.staging_ptrs["worker-2"]) + self.assertTrue(torch.equal(first_dst, torch.full_like(first_dst, 11))) + self.assertTrue(torch.equal(second_dst, torch.full_like(second_dst, 22))) + expected_allocation = (4096 + 1) * SWAP_STAGING_WINDOW + 1 + self.assertEqual(allocations, [expected_allocation, expected_allocation]) + self.assertEqual(register.call_count, 2) + + +class FailureThreadStub: + _handle_request = KVCacheRecvingThread._handle_request + get_and_clear_invalid_block_ids = KVCacheRecvingThread.get_and_clear_invalid_block_ids + _is_failed_recv_request = KVCacheRecvingThread._is_failed_recv_request + _mark_failed_recv_request = KVCacheRecvingThread._mark_failed_recv_request + _clear_failed_recv_request = KVCacheRecvingThread._clear_failed_recv_request + + def __init__(self, task_done_results): + self.failed_recv_requests: set[str] = set() + self.invalid_block_ids: set[int] = set() + self.failed_recv_requests_lock = threading.Lock() + self.use_hybrid = False + self._transfer_kv_cache = mock.Mock(side_effect=RuntimeError("transfer failed")) + self._transfer_kv_cache_all_groups = mock.Mock() + self._mark_request_task_done = mock.Mock(side_effect=task_done_results) + self._send_done_signal_to_free_remote_port = mock.Mock() + self._send_done_recv_signal = mock.Mock() + self.task_tracker = KVCacheTaskTracker() + self.proc_not_transfer_request: dict[str, bool] = {} + self.proc_not_transfer_request_lock = threading.Lock() + self.request_queue = mock.Mock() + + +def _recv_task(local_block_ids, *, all_task_done): + return { + "request_id": "request-1", + "remote_request_id": "remote-1", + "remote_host": "127.0.0.1", + "remote_handshake_port": 1234, + "remote_port_send_num": {}, + "local_block_ids": local_block_ids, + "all_task_done": all_task_done, + } + + +class TestMooncakeLoadFailureReporting(unittest.TestCase): + def test_failed_blocks_are_returned_and_cleared_through_connector_api(self): + recv_thread = FailureThreadStub([True]) + recv_thread.task_tracker.add_req_to_process("request-1") + recv_thread._handle_request(_recv_task(((3, 7), (11,)), all_task_done=True)) + + worker = MooncakeConnectorWorker.__new__(MooncakeConnectorWorker) + worker.kv_role = "kv_consumer" + worker.kv_recv_thread = recv_thread + connector = MooncakeConnector.__new__(MooncakeConnector) + connector.connector_worker = worker + + self.assertEqual(connector.get_block_ids_with_load_errors(), {3, 7, 11}) + self.assertEqual(connector.get_block_ids_with_load_errors(), set()) + self.assertEqual(recv_thread.get_and_clear_finished_requests(), {"request-1"}) + + def test_later_task_skips_transfer_after_request_failure(self): + recv_thread = FailureThreadStub([False, True]) + recv_thread.task_tracker.add_req_to_process("request-1") + + recv_thread._handle_request(_recv_task(((1, 2), (3,)), all_task_done=False)) + recv_thread._handle_request(_recv_task(((4,), (5, 6)), all_task_done=True)) + + recv_thread._transfer_kv_cache.assert_called_once() + self.assertEqual(recv_thread.get_and_clear_invalid_block_ids(), {1, 2, 3, 4, 5, 6}) + self.assertEqual(recv_thread.get_and_clear_finished_requests(), {"request-1"}) + + +class TestSwappedMemoryAllocation(unittest.TestCase): + def test_allocator_result_is_explicitly_zeroed(self): + allocated = torch.full((32,), 9, dtype=torch.int8) + allocator = mock.Mock(return_value=allocated) + fake_torch_npu = types.SimpleNamespace(empty_with_swapped_memory=allocator) + + with mock.patch.dict(sys.modules, {"torch_npu": fake_torch_npu}): + result = empty_swapped_memory((32,), dtype=torch.int8) + + self.assertIs(result, allocated) + self.assertTrue(torch.equal(result, torch.zeros_like(result))) + allocator.assert_called_once_with((32,), dtype=torch.int8, device="npu") + + class TestGlobalTERegistration(unittest.TestCase): """Incremental (ptr, size, location) registration in GlobalTE.""" @@ -366,9 +579,7 @@ def test_first_registration_calls_register_memory(self): te, engine = self._make_te() te.register_buffer([0x1000, 0x2000], [64, 128]) - engine.register_memory.assert_has_calls( - [mock.call(0x1000, 64), mock.call(0x2000, 128)], any_order=False - ) + engine.register_memory.assert_has_calls([mock.call(0x1000, 64), mock.call(0x2000, 128)], any_order=False) self.assertEqual(engine.register_memory.call_count, 2) def test_duplicate_registration_is_skipped(self): @@ -601,9 +812,7 @@ def test_copy_staging_to_swapped_writes_exact_bytes(self): dst_offset = ALIGN + 4096 stub = object.__new__(KVCacheRecvingThread) - KVCacheRecvingThread._copy_staging_to_swapped( - stub, view.data_ptr() + dst_offset, length, staging, 0 - ) + KVCacheRecvingThread._copy_staging_to_swapped(stub, view.data_ptr() + dst_offset, length, staging, 0) got = _readout(view, dst_offset, length) self.assertTrue(torch.equal(got, pattern)) diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake_hybrid_connector.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake_hybrid_connector.py index 1fe240f5fd93..6109de3247b4 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake_hybrid_connector.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake_hybrid_connector.py @@ -77,6 +77,8 @@ # number of peers is larger than max_workers. Yield after a small FIFO batch so # other peers already waiting in the global executor queue can make progress. MAX_REQUESTS_PER_PEER_HANDLER = 5 +SWAP_STAGING_ALIGNMENT = 2 * 1024 * 1024 +SWAP_STAGING_WINDOW = 16 class RemotePortInfo(TypedDict): @@ -411,9 +413,11 @@ def __init__( self.mamba_ssm_size = mamba_ssm_size self.remote_te_port: dict[str, dict[int, int]] = SizedDict() self.remote_metadata_lock = threading.Lock() - self._swap_staging_storage: torch.Tensor | None = None - self._swap_staging_tensor: torch.Tensor | None = None - self._swap_staging_bytes = 0 + # Executor workers can transfer from different peers concurrently. + # Keep one fixed-size staging pool per worker so ADXL writes from one + # request cannot overwrite another worker's pending copy. + self._swap_staging_pools: dict[threading.Thread, tuple[torch.Tensor, torch.Tensor]] = {} + self._swap_staging_pool_lock = threading.Lock() self.request_queue: queue.Queue[Any] = queue.Queue() first_kv_cache = next(iter(self.kv_caches.values())) @@ -470,39 +474,67 @@ def __init__( self.num_kv_heads = max(self.model_config.hf_text_config.num_key_value_heads // self.tp_size, 1) self.proc_not_transfer_request: dict[str, bool] = {} self.proc_not_transfer_request_lock = threading.Lock() + self.failed_recv_requests: set[str] = set() + self.invalid_block_ids: set[int] = set() + self.failed_recv_requests_lock = threading.Lock() def _swap_staging_chunk_bytes(self) -> int: """Use a full aligned block-sized staging slot for each transfer.""" block_bytes = [int(length) for length in self.block_len_per_addr if int(length) > 0] - alignment = 2 * 1024 * 1024 + alignment = SWAP_STAGING_ALIGNMENT chunk_bytes = max(max(block_bytes, default=alignment), alignment) return (chunk_bytes + alignment - 1) // alignment * alignment - def _ensure_swap_staging(self, transfer_count: int) -> torch.Tensor: - chunk_bytes = self._swap_staging_chunk_bytes() - alignment = 2 * 1024 * 1024 - # Each slot may start at an arbitrary remote byte offset. Reserve a - # full alignment window in addition to the chunk so that - # ``slot * slot_bytes + remote_offset`` is always in bounds. - slot_bytes = chunk_bytes + alignment - required_bytes = slot_bytes * max(1, transfer_count) - if self._swap_staging_tensor is not None and self._swap_staging_bytes >= required_bytes: - return self._swap_staging_tensor - - first_cache = next(iter(self.kv_caches.values())) - first_tensor = next(iter_tensors(first_cache), None) - if first_tensor is None or first_tensor.device.type == "cpu": - raise RuntimeError("Mooncake swapped-memory receive requires an NPU staging tensor.") - storage = torch.empty(required_bytes + alignment, dtype=torch.int8, device=first_tensor.device) - offset = (-int(storage.data_ptr())) % alignment - staging = storage[offset : offset + required_bytes] - if int(staging.data_ptr()) % alignment != 0: - raise RuntimeError("Mooncake NPU staging tensor is not 2MB aligned.") - global_te.register_buffer([int(staging.data_ptr())], [required_bytes]) - self._swap_staging_storage = storage - self._swap_staging_tensor = staging - self._swap_staging_bytes = required_bytes - return staging + def _ensure_swap_staging(self) -> torch.Tensor: + """Return the current executor worker's fixed-size staging pool.""" + worker = threading.current_thread() + with self._swap_staging_pool_lock: + existing = self._swap_staging_pools.get(worker) + if existing is not None: + return existing[1] + + chunk_bytes = self._swap_staging_chunk_bytes() + alignment = SWAP_STAGING_ALIGNMENT + # Each slot may start at an arbitrary remote byte offset. Reserve + # one alignment window after the chunk so the whole staged piece + # remains inside its slot. + slot_bytes = chunk_bytes + alignment + required_bytes = slot_bytes * SWAP_STAGING_WINDOW + + first_cache = next(iter(self.kv_caches.values())) + first_tensor = next(iter_tensors(first_cache), None) + if first_tensor is None or first_tensor.device.type == "cpu": + raise RuntimeError("Mooncake swapped-memory receive requires an NPU staging tensor.") + storage = torch.empty(required_bytes + alignment, dtype=torch.int8, device=first_tensor.device) + offset = (-int(storage.data_ptr())) % alignment + staging = storage[offset : offset + required_bytes] + if int(staging.data_ptr()) % alignment != 0: + raise RuntimeError("Mooncake NPU staging tensor is not 2MB aligned.") + global_te.register_buffer([int(staging.data_ptr())], [required_bytes]) + # The receiver owns every pool for its full lifetime. Registered + # memory is never replaced or freed while TE can still use it. + self._swap_staging_pools[worker] = (storage, staging) + return staging + + def get_and_clear_invalid_block_ids(self) -> set[int]: + """Return block IDs whose Mooncake load failed.""" + with self.failed_recv_requests_lock: + invalid_block_ids = self.invalid_block_ids + self.invalid_block_ids = set() + return invalid_block_ids + + def _is_failed_recv_request(self, request_id: str) -> bool: + with self.failed_recv_requests_lock: + return request_id in self.failed_recv_requests + + def _mark_failed_recv_request(self, request_id: str, local_block_ids: BlockIds) -> None: + with self.failed_recv_requests_lock: + self.failed_recv_requests.add(request_id) + self.invalid_block_ids.update(block_id for group in local_block_ids for block_id in group) + + def _clear_failed_recv_request(self, request_id: str) -> None: + with self.failed_recv_requests_lock: + self.failed_recv_requests.discard(request_id) def _copy_staging_to_swapped(self, dst: int, length: int, staging: torch.Tensor, staging_offset: int) -> None: resolved = get_swapped_tensor(dst, length) @@ -528,7 +560,7 @@ def _batch_transfer_sync_read_with_swap_staging( direct_lengths: list[int] = [] staged: list[tuple[int, int, int]] = [] chunk_bytes = self._swap_staging_chunk_bytes() - alignment = 2 * 1024 * 1024 + alignment = SWAP_STAGING_ALIGNMENT for local_dst, remote_src, length in zip(local_dst_list, remote_src_list, length_list): if not is_swapped_range(local_dst, length): @@ -552,10 +584,10 @@ def _batch_transfer_sync_read_with_swap_staging( if not staged: return 0 - staging = self._ensure_swap_staging(min(len(staged), 16)) + staging = self._ensure_swap_staging() staging_ptr = int(staging.data_ptr()) slot_bytes = chunk_bytes + alignment - window = max(1, min(16, len(staging) // slot_bytes)) + window = max(1, min(SWAP_STAGING_WINDOW, len(staging) // slot_bytes)) for batch_start in range(0, len(staged), window): batch = staged[batch_start : batch_start + window] local_staging: list[int] = [] @@ -707,16 +739,23 @@ def _handle_request(self, req_meta: dict[str, Any]): remote_handshake_port = req_meta["remote_handshake_port"] remote_port_send_num = req_meta["remote_port_send_num"] all_task_done = req_meta["all_task_done"] + transfer_failed = self._is_failed_recv_request(request_id) try: - logger.debug("Starting to transfer KV cache for request %s.", remote_request_id) - if not self.use_hybrid: - self._transfer_kv_cache(req_meta) + if transfer_failed: + self._mark_failed_recv_request(request_id, req_meta["local_block_ids"]) + logger.warning("Skipping KV cache transfer for failed request %s.", remote_request_id) else: - self._transfer_kv_cache_all_groups(req_meta) - logger.debug("Finished transferring KV cache for request %s.", remote_request_id) - except Exception: - logger.exception("Failed to transfer KV cache for request %s.", remote_request_id) + logger.debug("Starting to transfer KV cache for request %s.", remote_request_id) + if not self.use_hybrid: + self._transfer_kv_cache(req_meta) + else: + self._transfer_kv_cache_all_groups(req_meta) + logger.debug("Finished transferring KV cache for request %s.", remote_request_id) + except Exception as exc: + transfer_failed = True + self._mark_failed_recv_request(request_id, req_meta["local_block_ids"]) + logger.exception("Failed to transfer KV cache for request %s: %s", remote_request_id, exc) finally: self._send_done_signal_to_free_remote_port(remote_request_id, remote_host, remote_port_send_num) if self._mark_request_task_done(request_id, all_task_done): @@ -724,6 +763,7 @@ def _handle_request(self, req_meta: dict[str, Any]): self.task_tracker.update_done_task_count(request_id) with self.proc_not_transfer_request_lock: self.proc_not_transfer_request.pop(remote_request_id, None) + self._clear_failed_recv_request(request_id) self.request_queue.task_done() # Always send the done signal to the remote host to ensure proper # resource cleanup. Failing to do so may cause a memory leak on the @@ -1245,6 +1285,11 @@ def get_finished(self, finished_req_ids: set[str]) -> tuple[set[str], set[str]]: assert self.connector_worker is not None return self.connector_worker.get_finished() + def get_block_ids_with_load_errors(self) -> set[int]: + """Get block IDs whose Mooncake load failed.""" + assert self.connector_worker is not None + return self.connector_worker.get_block_ids_with_load_errors() + def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None: assert self.connector_worker is not None assert isinstance(self._connector_metadata, MooncakeConnectorMetadata) @@ -1894,6 +1939,11 @@ def get_finished(self) -> tuple[set[str], set[str]]: ) return done_sending, done_recving + def get_block_ids_with_load_errors(self) -> set[int]: + if self.kv_role == "kv_consumer" and self.kv_recv_thread is not None: + return self.kv_recv_thread.get_and_clear_invalid_block_ids() + return set() + def start_load_kv(self, metadata: MooncakeConnectorMetadata): """Start loading KV blocks from remote engine.""" for req_id, meta in metadata.requests.items(): diff --git a/vllm_ascend/kv_offload/mooncake_swap_memory.py b/vllm_ascend/kv_offload/mooncake_swap_memory.py index 568d0c403d0d..877690a68daf 100644 --- a/vllm_ascend/kv_offload/mooncake_swap_memory.py +++ b/vllm_ascend/kv_offload/mooncake_swap_memory.py @@ -9,11 +9,11 @@ from __future__ import annotations -from typing import Any, Iterator +from collections.abc import Iterator +from typing import Any import torch - _SWAPPED_TENSORS: list[tuple[int, int, torch.Tensor]] = [] @@ -56,16 +56,17 @@ def clear_swapped_tensors_for_testing() -> None: def empty_swapped_memory(shape: tuple[int, ...], *, dtype: torch.dtype) -> torch.Tensor: - """Allocate an NPU tensor whose storage is host-side swap memory.""" + """Allocate zero-initialized NPU tensor storage in host-side swap memory.""" try: import torch_npu except ImportError as exc: - raise RuntimeError( - "Mooncake swap-memory receive requires torch_npu.empty_with_swapped_memory." - ) from exc + raise RuntimeError("Mooncake swap-memory receive requires torch_npu.empty_with_swapped_memory.") from exc allocator = getattr(torch_npu, "empty_with_swapped_memory", None) if allocator is None: raise RuntimeError("Mooncake swap-memory receive requires torch_npu.empty_with_swapped_memory.") - return allocator(shape, dtype=dtype, device="npu") - + tensor = allocator(shape, dtype=dtype, device="npu") + # The replaced KV-cache path used torch.zeros. Do not depend on the + # version-specific observation that swapped memory starts zeroed. + tensor.zero_() + return tensor