Skip to content

GitLab

"vscode:/vscode.git/clone" did not exist on "7a1e1c1cafe4d3f3f935dc7192f9e66d4b2185b3"
Unverified Commit 6222e1c2 authored by Yineng Zhang's avatar Yineng Zhang Committed by GitHub
Browse files

add disable cuda graph unit test for eagle 2 (#3412)

parent fad315cb
Show whitespace changes
Inline Side-by-side
Showing with 62 additions and 44 deletions
test/srt/test_eagle_infer.py
View file @ 6222e1c2
...@@ -30,16 +30,34 @@ class TestEAGLEEngine(unittest.TestCase): ...@@ -30,16 +30,34 @@ class TestEAGLEEngine(unittest.TestCase):
ref_output = ref_engine.generate(prompt, sampling_params)["text"] ref_output = ref_engine.generate(prompt, sampling_params)["text"]
ref_engine.shutdown() ref_engine.shutdown()
# Test cases with different configurations
configs = [
# Original config
{
"model_path": DEFAULT_EAGLE_TARGET_MODEL_FOR_TEST,
"speculative_draft_model_path": DEFAULT_EAGLE_DRAFT_MODEL_FOR_TEST,
"speculative_algorithm": "EAGLE",
"speculative_num_steps": 5,
"speculative_eagle_topk": 8,
"speculative_num_draft_tokens": 64,
"mem_fraction_static": 0.7,
},
# Config with CUDA graph disabled
{
"model_path": DEFAULT_EAGLE_TARGET_MODEL_FOR_TEST,
"speculative_draft_model_path": DEFAULT_EAGLE_DRAFT_MODEL_FOR_TEST,
"speculative_algorithm": "EAGLE",
"speculative_num_steps": 5,
"speculative_eagle_topk": 8,
"speculative_num_draft_tokens": 64,
"mem_fraction_static": 0.7,
"disable_cuda_graph": True,
},
]
for config in configs:
# Launch EAGLE engine # Launch EAGLE engine
engine = sgl.Engine( engine = sgl.Engine(**config)
model_path=DEFAULT_EAGLE_TARGET_MODEL_FOR_TEST,
speculative_draft_model_path=DEFAULT_EAGLE_DRAFT_MODEL_FOR_TEST,
speculative_algorithm="EAGLE",
speculative_num_steps=5,
speculative_eagle_topk=8,
speculative_num_draft_tokens=64,
mem_fraction_static=0.7,
)
# Case 1: Test the output of EAGLE engine is the same as normal engine # Case 1: Test the output of EAGLE engine is the same as normal engine
out1 = engine.generate(prompt, sampling_params)["text"] out1 = engine.generate(prompt, sampling_params)["text"]
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment