Generative AI Google Brain

open-models-benchmarking-eagle-on-vertex-ai

4 min read

Accelerate LLM Inference with EAGLE Speculative Decoding on Vertex AI

Source notebook

Repo path: open-models/benchmarking_eagle_on_vertex_ai.ipynb · Open on GitHub · advanced

Benchmarks EAGLE speculative decoding for Llama 4 Scout on Vertex AI against a baseline endpoint.

Summary

This notebook teaches how to deploy two Llama 4 Scout endpoints on Vertex AI: one standard baseline and one EAGLE-enabled endpoint. It prepares ShareGPT conversations and local Hugging Face tokenizer artifacts, patches vLLM benchmark requests for Vertex AI compatibility, then runs smoke tests and concurrency sweeps. It compares TTFT, TPOT, throughput, and request throughput across concurrency levels to quantify EAGLE speedups.

Key code patterns

Initialize Vertex AI

PROJECT_ID = "[your-project-id]"
LOCATION = "asia-southeast1"
if not PROJECT_ID or PROJECT_ID == "[your-project-id]":
    PROJECT_ID = str(os.environ.get("GOOGLE_CLOUD_PROJECT"))
vertexai.init(project=PROJECT_ID, location=LOCATION)

Sets the project and region before using Model Garden deployment APIs.

Deploy baseline endpoint

baseline_args = [
    f"--model={MODEL_GCS_PATH}",
    "--attention-backend=fa3",
    "--context-length=131072",
    "--tp=8",
    "--enable-multimodal",
]
baseline_model = model_garden.OpenModel(MODEL_NAME)
baseline_endpoint = baseline_model.deploy(
    machine_type="a3-highgpu-8g",
    accelerator_type="NVIDIA_H100_80GB",
    accelerator_count=8,
    serving_container_args=baseline_args,
)

Creates the no-EAGLE comparison endpoint on the same 8x H100 hardware.

Enable EAGLE speculative decoding

eagle_args = baseline_args + [
    "--speculative-algo=EAGLE3",
    "--speculative-draft-model-path=gs://vertex-model-garden-restricted-us/llama4/Llama-4-Scout-17B-16E-Instruct-EAGLE3-20250829/",
    "--speculative-num-steps=3",
    "--speculative-eagle-topk=4",
    "--speculative-num-draft-tokens=8",
]
eagle_endpoint = eagle_model.deploy(
    machine_type="a3-highgpu-8g",
    accelerator_type="NVIDIA_H100_80GB",
    accelerator_count=8,
    serving_container_args=eagle_args,
)

Adds only the EAGLE-specific serving arguments so the comparison remains controlled.

Patch vLLM for Vertex AI

payload = {
    "model": request_func_input.model_name or request_func_input.model,
    "messages": [{"role": "user", "content": content}],
    "temperature": 0.0,
    "max_tokens": request_func_input.output_len,
    "stream": True,
}
headers = {
    "Content-Type": "application/json",
    "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}",
}
endpoint_func.ASYNC_REQUEST_FUNCS["openai-chat"] = patched_async_request_openai_chat_completions

Adapts the OpenAI-style vLLM benchmark client to Vertex AI by using max_tokens and omitting stream_options.

Run concurrency sweep

for concurrency in [1, 2, 4, 6, 8, 10]:
    os.environ["OPENAI_API_KEY"] = get_fresh_token()
    subprocess.run([
        "vllm", "bench", "serve",
        "--backend", "openai-chat",
        "--base-url", baseline_url,
        "--endpoint", "/chat/completions",
        "--tokenizer", model_path,
        "--dataset-name", "sharegpt",
        "--dataset-path", dataset_path,
        "--num-prompts", "1000",
        "--max-concurrency", str(concurrency),
        "--save-result",
    ], check=True)

Measures endpoint behavior under increasing concurrent request load.

Compare benchmark metrics

ttft_improvement = ((baseline["TTFT (ms)"] - eagle["TTFT (ms)"]) / baseline["TTFT (ms)"]) * 100
throughput_improvement = ((eagle["Throughput (tok/s)"] - baseline["Throughput (tok/s)"]) / baseline["Throughput (tok/s)"]) * 100
baseline_df = pd.DataFrame(baseline_metrics)
eagle_df = pd.DataFrame(eagle_metrics)
improvements_df = pd.DataFrame(improvements)

Turns saved JSON benchmark results into latency and throughput comparisons.

Models & APIs used

  • Models: meta/llama4@llama-4-scout-17b-16e-instruct, meta-llama/Llama-4-Scout-17B-16E-Instruct
  • APIs / services: Vertex AI, Vertex AI Model Garden, Cloud Storage
  • SDKs / libraries: google-cloud-aiplatform, vertexai, google.auth, huggingface_hub, vllm, pandas, matplotlib, seaborn, aiohttp, hf-transfer

When to use this

Use this pattern when benchmarking speculative decoding for an open model on Vertex AI under realistic concurrent workloads.

Gotchas & caveats

  • Vertex AI API and billing must be enabled on the Google Cloud project.
  • The selected region needs quota and availability for 8x NVIDIA H100 80GB GPUs.
  • The package installation cell requires a runtime or kernel restart before continuing.
  • Colab users must authenticate, while Vertex AI Workbench users can skip that authentication cell.
  • Hugging Face access to meta-llama/Llama-4-Scout-17B-16E-Instruct requires a token and Meta license acceptance.
  • vLLM’s OpenAI chat benchmark function needs a Vertex AI compatibility patch for max_tokens, stream_options, and longer timeout.
  • Google auth tokens are refreshed before long smoke tests and benchmark runs.
  • The notebook states an estimated 90-120 minutes and about $250 cost, mostly from deployment and GPU hours.

Best practices

  • Pin package versions for reproducibility.
  • Deploy baseline and EAGLE endpoints on identical hardware for a fair comparison.
  • Run 100-prompt smoke tests before the full benchmark.
  • Use ShareGPT conversations as a realistic benchmark workload.
  • Use the model tokenizer locally so prompt and response token lengths are measured accurately.
  • Refresh authentication tokens before each long benchmark run.
  • Run the same concurrency levels for baseline and EAGLE endpoints.
  • Use median latency metrics because medians are robust to outliers.

Graph View

Backlinks