feat(spec_decode): run the Kimi-K3 DSpark draft on a dedicated remote GPU

tests/v1/spec_decode/test_acceptance_length_controller.py

@@ -260,6 +260,22 @@ def test_runner_v2_limits_drafts_to_adaptive_depth():
     )
 
 
+def test_runner_v2_allows_scheduler_to_disable_speculation():
+    draft_tokens = torch.tensor([[1, 2, 3], [4, 5, 6]])
+
+    limited = limit_draft_tokens(
+        draft_tokens,
+        num_speculative_tokens=0,
+        max_num_speculative_tokens=3,
+    )
+
+    assert limited.shape == (2, 0)
+    assert (
+        limited.untyped_storage().data_ptr()
+        == draft_tokens.untyped_storage().data_ptr()
+    )
+
+
 def test_synthetic_scheduler_output_uses_default_speculative_depth():
     output = SchedulerOutput.make_empty()
 

tests/v1/spec_decode/test_k3_dspark_remote_speculator.py

@@ -0,0 +1,152 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+import torch
+
+from vllm.v1.worker.gpu.spec_decode.dspark.remote_speculator import (
+    RemoteK3DSparkSpeculator,
+    _anchor_positions_from_context,
+    _build_valid_context_plan,
+    _contiguous_draft_output,
+    _RetainedRequestPrefix,
+)
+
+
+def test_build_valid_context_plan_drops_rejected_tail_rows():
+    batch = SimpleNamespace(
+        num_reqs=2,
+        num_scheduled_tokens=np.array([4, 3], dtype=np.int32),
+        num_computed_tokens_np=np.array([10, 20], dtype=np.int32),
+    )
+
+    indices, counts = _build_valid_context_plan(batch, [2, 0])
+
+    assert indices == [0, 1, 4, 5, 6]
+    assert counts == [2, 3]
+
+
+def test_anchor_positions_follow_actual_valid_context_rows():
+    positions = torch.tensor([24, 25, 26, 80, 81], dtype=torch.int64)
+
+    anchors = _anchor_positions_from_context([3, 2], positions)
+
+    assert anchors == [27, 82]
+
+
+def test_remote_tokens_copy_supports_adaptive_depth():
+    proxy = RemoteK3DSparkSpeculator.__new__(RemoteK3DSparkSpeculator)
+    proxy.device = torch.device("cpu")
+    proxy.draft_tokens = torch.full((3, 8), -1, dtype=torch.int64)
+
+    proxy._copy_tokens_from_response(
+        {"tokens": [[11, 12], [21, 22]]},
+        active_indices=[0, 2],
+        num_speculative_tokens=2,
+    )
+
+    assert proxy.draft_tokens.tolist() == [
+        [11, 12, -1, -1, -1, -1, -1, -1],
+        [-1, -1, -1, -1, -1, -1, -1, -1],
+        [21, 22, -1, -1, -1, -1, -1, -1],
+    ]
+
+
+def test_remote_speculator_accepts_scheduler_selected_zero_depth():
+    proxy = RemoteK3DSparkSpeculator.__new__(RemoteK3DSparkSpeculator)
+    proxy.num_speculative_steps = 3
+    proxy.draft_tokens = torch.full((4, 3), -1, dtype=torch.int64)
+    batch = SimpleNamespace(num_reqs=2)
+    empty = torch.empty(0)
+
+    output = proxy.propose(
+        batch,
+        {},
+        {},
+        empty,
+        None,
+        empty,
+        empty,
+        empty,
+        empty,
+        empty,
+        empty,
+        num_speculative_tokens=0,
+    )
+
+    assert output.shape == (2, 0)
+    assert output.is_contiguous()
+
+
+def test_adaptive_depth_output_is_contiguous_for_tp_broadcast():
+    draft_tokens = torch.arange(24, dtype=torch.int64).view(3, 8)
+
+    output = _contiguous_draft_output(draft_tokens, 2, 3)
+
+    assert output.is_contiguous()
+    assert output.tolist() == [[0, 1, 2], [8, 9, 10]]
+
+
+@pytest.mark.parametrize("rejected", [[5, 0], [-1, 0]])
+def test_build_valid_context_plan_rejects_invalid_counts(rejected):
+    batch = SimpleNamespace(
+        num_reqs=2,
+        num_scheduled_tokens=np.array([4, 3], dtype=np.int32),
+        num_computed_tokens_np=np.array([0, 0], dtype=np.int32),
+    )
+
+    with pytest.raises(ValueError, match="Invalid valid-context length"):
+        _build_valid_context_plan(batch, rejected)
+
+
+def _make_prefix_matcher() -> RemoteK3DSparkSpeculator:
+    proxy = RemoteK3DSparkSpeculator.__new__(RemoteK3DSparkSpeculator)
+    proxy._known_requests = {"old"}
+    proxy._remote_block_size = 16
+    proxy._remote_window_size = 32
+    proxy._remote_prefix_cache_tokens = 128
+    proxy._retained_prefixes = {
+        "old": _RetainedRequestPrefix(
+            token_ids=torch.arange(96, dtype=torch.int32),
+            committed_end=96,
+            context_start=0,
+            serial=1,
+        )
+    }
+    return proxy
+
+
+def test_remote_prefix_match_requires_exact_token_identity():
+    proxy = _make_prefix_matcher()
+    matching = torch.arange(80, dtype=torch.int32)
+
+    assert proxy._find_reconnect_source(matching, 80, {"new"}) == "old"
+
+    mismatched = matching.clone()
+    mismatched[40] = -1
+    assert proxy._find_reconnect_source(mismatched, 80, {"new"}) is None
+
+
+def test_remote_prefix_match_rejects_range_evicted_from_projected_cache():
+    proxy = _make_prefix_matcher()
+    proxy._remote_prefix_cache_tokens = 48
+    matching = torch.arange(40, dtype=torch.int32)
+
+    assert proxy._find_reconnect_source(matching, 40, {"new"}) is None
+
+
+def test_remote_prefix_match_rejects_history_before_cold_bootstrap():
+    proxy = _make_prefix_matcher()
+    proxy._retained_prefixes["old"].context_start = 64
+
+    assert (
+        proxy._find_reconnect_source(torch.arange(80, dtype=torch.int32), 80, {"new"})
+        is None
+    )
+    assert (
+        proxy._find_reconnect_source(torch.arange(96, dtype=torch.int32), 96, {"new"})
+        == "old"
+    )

tests/v1/spec_decode/test_k3_dspark_standalone.py

@@ -0,0 +1,172 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import json
+
+import pytest
+import torch
+
+from vllm.entrypoints.k3_dspark_rpc import (
+    DraftKVSlotAllocator,
+    ProjectedContextCache,
+)
+from vllm.entrypoints.k3_dspark_standalone import (
+    EMBED_TENSOR,
+    LM_HEAD_TENSOR,
+    resolve_shared_weight_files,
+)
+
+
+def _write_index(root, weight_map):
+    (root / "model.safetensors.index.json").write_text(
+        json.dumps({"weight_map": weight_map})
+    )
+
+
+def test_resolve_shared_weight_files_requires_both_target_tensors(tmp_path):
+    shard = tmp_path / "shared.safetensors"
+    shard.touch()
+    _write_index(tmp_path, {EMBED_TENSOR: shard.name})
+
+    with pytest.raises(KeyError, match=LM_HEAD_TENSOR):
+        resolve_shared_weight_files(tmp_path)
+
+
+def test_resolve_shared_weight_files_resolves_checkpoint_shards(tmp_path):
+    shard = tmp_path / "shared.safetensors"
+    shard.touch()
+    _write_index(
+        tmp_path,
+        {
+            EMBED_TENSOR: shard.name,
+            LM_HEAD_TENSOR: shard.name,
+        },
+    )
+
+    resolved = resolve_shared_weight_files(tmp_path)
+
+    assert resolved == {
+        EMBED_TENSOR: shard.resolve(),
+        LM_HEAD_TENSOR: shard.resolve(),
+    }
+
+
+def test_resolve_shared_weight_files_rejects_path_escape(tmp_path):
+    outside = tmp_path.parent / "outside.safetensors"
+    outside.touch()
+    _write_index(
+        tmp_path,
+        {
+            EMBED_TENSOR: f"../{outside.name}",
+            LM_HEAD_TENSOR: f"../{outside.name}",
+        },
+    )
+
+    with pytest.raises(ValueError, match="escapes"):
+        resolve_shared_weight_files(tmp_path)
+
+
+def test_draft_kv_slot_allocator_keeps_rolling_blocks_unique():
+    allocator = DraftKVSlotAllocator(
+        num_cache_blocks=11,
+        block_size=4,
+        window_size=16,
+        max_requests=2,
+    )
+    first, created = allocator.get_or_allocate("first")
+    second, _ = allocator.get_or_allocate("second")
+
+    assert created
+    assert allocator.physical_block_range(first) == slice(1, 6)
+    assert allocator.physical_block_range(second) == slice(6, 11)
+    blocks, local_len = allocator.block_table(first, 21)
+    assert blocks == [2, 3, 4, 5, 1]
+    assert len(blocks) == len(set(blocks))
+    assert local_len == 17
+
+
+def test_draft_kv_slot_allocator_reuses_freed_request_slot():
+    allocator = DraftKVSlotAllocator(
+        num_cache_blocks=6,
+        block_size=4,
+        window_size=16,
+        max_requests=1,
+    )
+    original, _ = allocator.get_or_allocate("original")
+    with pytest.raises(RuntimeError, match="capacity exhausted"):
+        allocator.get_or_allocate("other")
+
+    assert allocator.free("original") is original
+    replacement, created = allocator.get_or_allocate("replacement")
+    assert created
+    assert replacement.slot == original.slot
+
+
+def test_draft_kv_slot_allocator_rebinds_without_changing_slot():
+    allocator = DraftKVSlotAllocator(
+        num_cache_blocks=6,
+        block_size=4,
+        window_size=16,
+        max_requests=1,
+    )
+    original, _ = allocator.get_or_allocate("original")
+
+    rebound = allocator.rebind("original", "replacement")
+
+    assert rebound is original
+    assert rebound.request_id == "replacement"
+    assert allocator.get("original") is None
+    assert allocator.get("replacement") is rebound
+
+
+def test_projected_context_cache_rewinds_and_overwrites_exact_prefix():
+    cache = ProjectedContextCache(hidden_size=2, max_tokens=8, chunk_size=4)
+    initial = torch.arange(16, dtype=torch.bfloat16).view(8, 2)
+    cache.append(0, initial)
+
+    replacement = torch.tensor([[100, 101], [102, 103]], dtype=torch.bfloat16)
+    cache.append(5, replacement)
+
+    assert cache.start_position == 0
+    assert cache.end_position == 7
+    assert torch.equal(cache.read(0, 5), initial[:5])
+    assert torch.equal(cache.read(5, 7), replacement)
+
+
+def test_projected_context_cache_evicts_old_rows_without_regaining_them():
+    cache = ProjectedContextCache(hidden_size=1, max_tokens=6, chunk_size=4)
+    cache.append(0, torch.arange(8, dtype=torch.bfloat16).view(8, 1))
+
+    assert cache.start_position == 2
+    assert cache.has_range(2, 8)
+    assert not cache.has_range(1, 8)
+
+    cache.append(5, torch.tensor([[50], [60]], dtype=torch.bfloat16))
+    assert cache.start_position == 2
+    assert cache.end_position == 7
+    with pytest.raises(ValueError, match="unavailable"):
+        cache.read(1, 7)
+
+
+def test_projected_context_cache_tracks_configured_device():
+    cache = ProjectedContextCache(
+        hidden_size=2,
+        max_tokens=4,
+        chunk_size=2,
+        device=torch.device("cpu"),
+    )
+    states = torch.arange(8, dtype=torch.bfloat16).view(4, 2)
+
+    cache.append(0, states)
+
+    assert cache.device == torch.device("cpu")
+    assert cache.read(0, 4).device == cache.device
+    assert cache.allocated_bytes == states.numel() * states.element_size()
+
+
+def test_projected_context_cache_rejects_device_mismatch():
+    cache = ProjectedContextCache(hidden_size=2, max_tokens=4)
+    states = torch.empty((1, 2), dtype=torch.bfloat16, device="meta")
+
+    with pytest.raises(ValueError, match="device mismatch"):
+        cache.append(0, states)