include llama.cpp since there were some edits

llama.cpp/tools/server/tests/unit/test_speculative.py

@@ -0,0 +1,267 @@
+import pytest
+from utils import *
+
+# We use a F16 MOE gguf as main model, and q4_0 as draft model
+
+server = ServerPreset.stories15m_moe()
+
+MODEL_DRAFT_FILE_URL = "https://huggingface.co/ggml-org/tiny-llamas/resolve/main/stories15M-q4_0.gguf"
+
+def create_server():
+    global server
+    server = ServerPreset.stories15m_moe()
+    # set default values
+    server.model_draft = download_file(MODEL_DRAFT_FILE_URL)
+    server.spec_type = "draft-simple"
+    server.spec_draft_n_min = 4
+    server.spec_draft_n_max = 8
+    server.fa = "off"
+
+
+@pytest.fixture(autouse=True)
+def fixture_create_server():
+    return create_server()
+
+
+def test_with_and_without_draft():
+    global server
+    request = {
+        "prompt": "I believe the meaning of life is",
+        "temperature": 0.2,
+        "top_k": 5,
+        "seed": 4242,
+        "n_predict": 16,
+        "return_tokens": True,
+    }
+
+    server.model_draft = None  # disable draft model
+    server.spec_type = None
+    server.start()
+    res = server.make_request("POST", "/completion", data=request)
+    assert res.status_code == 200
+    tokens_no_draft = res.body["tokens"]
+    server.stop()
+
+    # create new server with draft model
+    create_server()
+    server.start()
+    res = server.make_request("POST", "/completion", data=request)
+    assert res.status_code == 200
+    assert res.body["timings"]["draft_n"] > 0
+    tokens_draft = res.body["tokens"]
+
+    assert tokens_no_draft == tokens_draft
+
+    server.stop()
+    create_server()
+    assert server.spec_draft_n_max is not None
+    server.spec_synth_rates = [0.0] * server.spec_draft_n_max
+    server.start()
+    res = server.make_request("POST", "/completion", data=request)
+
+    assert res.status_code == 200
+    assert res.body["timings"]["draft_n"] > 0
+    assert res.body["timings"]["draft_n_accepted"] == 0
+    assert res.body["tokens"] == tokens_no_draft
+
+
+def test_different_draft_min_draft_max():
+    global server
+    test_values = [
+        (1, 2),
+        (1, 4),
+        (4, 8),
+        (4, 12),
+        (8, 16),
+    ]
+    last_content = None
+    for draft_min, draft_max in test_values:
+        server.stop()
+        server.spec_draft_n_min = draft_min
+        server.spec_draft_n_max = draft_max
+        server.start()
+        res = server.make_request("POST", "/completion", data={
+            "prompt": "I believe the meaning of life is",
+            "temperature": 0.0,
+            "top_k": 1,
+            "n_predict": 16,
+        })
+        assert res.status_code == 200
+        if last_content is not None:
+            assert last_content == res.body["content"]
+        last_content = res.body["content"]
+
+
+def test_synth_is_deterministic():
+    global server
+    assert server.spec_draft_n_max is not None
+    server.spec_synth_rates = [0.75 ** (i + 1) for i in range(server.spec_draft_n_max)]
+    server.start()
+
+    request = {
+        "prompt": "I believe the meaning of life is",
+        "temperature": 0.2,
+        "top_k": 5,
+        "seed": 4242,
+        "n_predict": 32,
+    }
+    responses = [server.make_request("POST", "/completion", data=request) for _ in range(2)]
+
+    for res in responses:
+        assert res.status_code == 200
+        assert res.body["timings"]["draft_n"] > 0
+    assert responses[0].body["timings"]["draft_n"] == responses[1].body["timings"]["draft_n"]
+    assert responses[0].body["timings"]["draft_n_accepted"] == responses[1].body["timings"]["draft_n_accepted"]
+
+
+def test_synth_ignores_target_tokens():
+    global server
+    assert server.spec_draft_n_max is not None
+    server.spec_synth_rates = [1.0] * server.spec_draft_n_max
+    server.start()
+
+    res = server.make_request("POST", "/completion", data={
+        "prompt": "I believe the meaning of life is",
+        "temperature": 0.0,
+        "seed": 4242,
+        "n_predict": 32,
+    })
+
+    assert res.status_code == 200
+    assert res.body["timings"]["draft_n"] > 0
+    assert res.body["timings"]["draft_n_accepted"] == res.body["timings"]["draft_n"]
+
+    res = server.make_request("POST", "/completion", data={
+        "prompt": "I believe the meaning of life is",
+        "temperature": 0.0,
+        "seed": 4242,
+        "n_predict": 6,
+        "grammar": 'root ::= "a"{5,5}',
+    })
+    assert res.status_code == 200, res.body
+
+    res = server.make_request("POST", "/completion", data={
+        "prompt": "Respond with only: OK",
+        "temperature": 0.0,
+        "seed": 4242,
+        "n_predict": 64,
+        "ignore_eos": True,
+    })
+    assert res.status_code == 200, res.body
+    assert res.body["tokens_predicted"] == 64
+    assert res.body["stop_type"] == "limit"
+
+
+def test_draft_n_ubatch_smaller_than_target():
+    global server
+    request = {
+        "prompt": "I believe the meaning of life is",
+        "temperature": 0.0,
+        "top_k": 1,
+        "n_predict": 16,
+    }
+
+    # draft with the default ubatch size of the draft context
+    server.start()
+    res_default = server.make_request("POST", "/completion", data=request)
+    server.stop()
+    assert res_default.status_code == 200
+    assert res_default.body["timings"]["draft_n"] > 0
+
+    # a smaller draft ubatch splits the draft batches into more ubatches, but must not change the result
+    server.spec_draft_n_ubatch = 8
+    server.start()
+    res_small = server.make_request("POST", "/completion", data=request)
+    server.stop()
+    assert res_small.status_code == 200
+    assert res_small.body["timings"]["draft_n"] > 0
+    assert res_small.body["tokens"] == res_default.body["tokens"]
+
+
+def test_slot_ctx_not_exceeded():
+    global server
+    server.n_ctx = 256
+    server.start()
+    res = server.make_request("POST", "/completion", data={
+        "prompt": "Hello " * 248,
+        "temperature": 0.0,
+        "top_k": 1,
+        "speculative.p_min": 0.0,
+    })
+    assert res.status_code == 200
+    assert len(res.body["content"]) > 0
+
+
+def test_with_ctx_shift():
+    global server
+    server.n_ctx = 256
+    server.enable_ctx_shift = True
+    server.start()
+    res = server.make_request("POST", "/completion", data={
+        "prompt": "Hello " * 248,
+        "temperature": 0.0,
+        "top_k": 1,
+        "n_predict": 256,
+        "speculative.p_min": 0.0,
+    })
+    assert res.status_code == 200
+    assert len(res.body["content"]) > 0
+    assert res.body["tokens_predicted"] == 256
+    assert res.body["truncated"] == True
+
+
+@pytest.mark.parametrize("n_slots,n_requests", [
+    (1, 2),
+    (2, 2),
+])
+def test_multi_requests_parallel(n_slots: int, n_requests: int):
+    global server
+    server.n_slots = n_slots
+    server.start()
+    tasks = []
+    for _ in range(n_requests):
+        tasks.append((server.make_request, ("POST", "/completion", {
+            "prompt": "I believe the meaning of life is",
+            "temperature": 0.0,
+            "top_k": 1,
+        })))
+    results = parallel_function_calls(tasks)
+    for res in results:
+        assert res.status_code == 200
+        assert match_regex("(wise|kind|owl|answer)+", res.body["content"])
+
+
+def test_draft_token_probs():
+    # tokens accepted from the draft must carry the target model's probabilities,
+    # same as without speculative decoding (#24271)
+    global server
+    request = {
+        "prompt": "I believe the meaning of life is",
+        "temperature": 0.0,
+        "top_k": 1,
+        "n_predict": 16,
+        "n_probs": 4,
+    }
+
+    server.model_draft = None  # disable draft model
+    server.spec_type = None
+    server.start()
+    res = server.make_request("POST", "/completion", data=request)
+    assert res.status_code == 200
+    probs_no_draft = res.body["completion_probabilities"]
+    server.stop()
+
+    create_server()
+    server.start()
+    res = server.make_request("POST", "/completion", data=request)
+    assert res.status_code == 200
+    assert res.body["timings"]["draft_n"] > 0
+    probs_draft = res.body["completion_probabilities"]
+
+    assert len(probs_draft) == len(probs_no_draft)
+    for tok_draft, tok_no_draft in zip(probs_draft, probs_no_draft):
+        assert tok_draft["id"] == tok_no_draft["id"]
+        assert len(tok_draft["top_logprobs"]) == len(tok_no_draft["top_logprobs"])
+        assert abs(tok_draft["logprob"] - tok_no_draft["logprob"]) < 0.25
+    # real logprobs are not all exactly 0.0 (prob 1.0)
+    assert any(tok["logprob"] < -1e-4 for tok in probs_draft)