chore: import upstream snapshot with attribution

2026-07-13 13:23:58 +08:00
commit 770d92cb1f
694 changed files with 114634 additions and 0 deletions
@@ -0,0 +1,141 @@
+from mlc_llm.protocol.debug_protocol import DebugConfig
+from mlc_llm.protocol.generation_config import GenerationConfig
+from mlc_llm.serve.sync_engine import EngineConfig, SyncMLCEngine
+from mlc_llm.testing import require_test_model
+
+prompts = [
+    "The meaning of life is",
+    "According to the history of Pittsburgh,",
+    "I have a three-day Seattle travel plan. On the first day,",
+    "Undoubtedly, Alaska is one of the most beautiful places on Earth,",
+    "Explain difference between Lambda calculus and Turing machine is",
+    "To assemble a desktop computer, we need the necessary components of",
+    "Vitamin D is important to human beings, because",
+    "Refer to history, the milk tea is originated from",
+    "In the southernmost place in United States,",
+    "AlphaGo has the capabilities of",
+]
+
+
+def test_engine_system_prompt(engine):
+    system_prompt = "This is a system prompt"
+    system_prompt_tokens = len(engine.tokenizer.encode(system_prompt))
+    max_tokens = 8
+    _, _ = engine.generate(
+        system_prompt,
+        GenerationConfig(
+            temperature=0,
+            max_tokens=max_tokens,
+            debug_config=DebugConfig(pinned_system_prompt=True),
+        ),
+    )
+    metrics = engine.metrics()
+    assert metrics["prefill_tokens_sum"] == system_prompt_tokens
+    sum_prefill_tokens = system_prompt_tokens
+
+    input_token_lens = [len(engine.tokenizer.encode(prompt)) for prompt in prompts]
+
+    generation_config = GenerationConfig(temperature=0, max_tokens=max_tokens)
+    _, _ = engine.generate(prompts, generation_config)
+    metrics = engine.metrics()
+    assert metrics["prefill_tokens_sum"] == sum_prefill_tokens + sum(input_token_lens)
+    sum_prefill_tokens = metrics["prefill_tokens_sum"]
+
+    _, _ = engine.generate(system_prompt + " and why ?", generation_config)
+    metrics = engine.metrics()
+    # system prompt is reused entirely
+    assert metrics["prefill_tokens_sum"] == sum_prefill_tokens + 3
+    sum_prefill_tokens = metrics["prefill_tokens_sum"]
+
+    _, _ = engine.generate(prompts[:4], generation_config)
+    metrics = engine.metrics()
+    # first 4 prompts are removed and need to prefill again
+    assert metrics["prefill_tokens_sum"] == sum_prefill_tokens + sum(input_token_lens[:4])
+
+
+def test_engine_multi_round(engine):
+    num_requests = 10
+    max_tokens = 8
+    generation_config = GenerationConfig(temperature=0, max_tokens=max_tokens)
+    input_token_lens = [len(engine.tokenizer.encode(prompt)) for prompt in prompts[:num_requests]]
+
+    output_texts, _ = engine.generate(prompts[:num_requests], generation_config)
+    metrics = engine.metrics()
+    assert metrics["prefill_tokens_sum"] == sum(input_token_lens)
+    sum_prefill_tokens = metrics["prefill_tokens_sum"]
+    concat_prompt = []
+    for i, output in enumerate(output_texts):
+        concat_prompt.append(prompts[i] + " " + output[0] + " ?")
+    output_texts, _ = engine.generate(concat_prompt[:num_requests], generation_config)
+    metrics = engine.metrics()
+    assert metrics["prefill_tokens_sum"] == sum_prefill_tokens + 2 * num_requests
+
+
+@require_test_model("Llama-2-7b-chat-hf-q0f16-MLC")
+def test_basic_engine_system_prompt(model: str):
+    # Create engine
+    engine = SyncMLCEngine(
+        model=model,
+        mode="local",
+        engine_config=EngineConfig(
+            max_total_sequence_length=4096,
+            prefix_cache_max_num_recycling_seqs=5,
+        ),
+    )
+    test_engine_system_prompt(engine)
+
+
+@require_test_model("Llama-2-7b-chat-hf-q0f16-MLC")
+def test_basic_engine_multi_round(model: str):
+    # Create engine
+    engine = SyncMLCEngine(
+        model=model,
+        mode="server",
+        engine_config=EngineConfig(max_total_sequence_length=4096),
+    )
+    test_engine_multi_round(engine)
+
+
+@require_test_model(
+    "Llama-2-7b-chat-hf-q0f16-MLC",
+    "Llama-2-7b-chat-hf-q4f16_1-MLC",
+)
+def test_engine_spec_multi_round(model: str, small_model: str):
+    # Create engine
+    engine = SyncMLCEngine(
+        model=model,
+        mode="server",
+        engine_config=EngineConfig(
+            max_total_sequence_length=4096,
+            additional_models=[small_model],
+            speculative_mode="small_draft",
+        ),
+    )
+
+    test_engine_multi_round(engine)
+
+
+@require_test_model("Llama-2-7b-chat-hf-q0f16-MLC")
+def test_engine_eagle_multi_round(model: str):
+    # Create engine
+    small_model = "dist/Eagle-llama2-7b-chat-q0f16-MLC"
+    small_model_lib = "dist/Eagle-llama2-7b-chat-q0f16-MLC/Eagle-llama2-7b-chat-q0f16-MLC-cuda.so"
+    engine = SyncMLCEngine(
+        model=model,
+        mode="server",
+        engine_config=EngineConfig(
+            max_total_sequence_length=4096,
+            additional_models=[(small_model, small_model_lib)],
+            speculative_mode="eagle",
+            max_num_sequence=80,
+        ),
+    )
+
+    test_engine_multi_round(engine)
+
+
+if __name__ == "__main__":
+    test_basic_engine_system_prompt()
+    test_basic_engine_multi_round()
+    test_engine_spec_multi_round()
+    test_engine_eagle_multi_round()