Skip to content

feat: Add Gemma-4 Multi-token-prediction Guide for TPUs - #520

Draft
syeda-anjum wants to merge 13 commits into
mainfrom
sanjum-dspark-mtp
Draft

feat: Add Gemma-4 Multi-token-prediction Guide for TPUs#520
syeda-anjum wants to merge 13 commits into
mainfrom
sanjum-dspark-mtp

Conversation

@syeda-anjum

Copy link
Copy Markdown
Collaborator

Overview:
This PR introduces a complete, end-to-end reference implementation for running Multi-Token Prediction (MTP) speculative decoding using Gemma-4 on Trillium TPUs. It provides the necessary vLLM manifests, inference-perf benchmarking configurations, and comprehensive architectural guidance.

Key Changes:

  • Inference Overlays: Added the v6e-gemma-4-31b-it-mtp Kustomize overlay to support MTP on TPUs, including updating the base vllm-tpu image to v0.25.0 and ensuring the deployment is distinctly separated from the base Gemma-4 workload.
  • Benchmarking Suite: Created dedicated inference-perf benchmark manifests tailored for sd-mtp to accurately measure MTP latency improvements.
  • Comprehensive Documentation: Added the vllm-gemma-4-mtp.md guide, which includes:
    • Step-by-step instructions for tfvars configuration and downloading the assistant/drafter model.
    • A comparative analysis of MTP versus Standard Speculative Decoding, Draft-Head architectures (EAGLE), dSpark (SGLang), and N-gram prompt lookup.
    • Guidance on selecting optimal benchmarking datasets (e.g., CNN Daily Mail, HumanEval) to maximize speculative decoding acceptance rates.

kubectl --namespace=${ira_online_tpu_kubernetes_namespace_name} logs -l app=inference-perf --tail 10"
```

- Delete the benchmark workload.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

How the result compares with standard inference (without MTP)?

Can you make an analysis here and share this with customers? Is it 30% better for this use case, 50% better? how customer analyse the results?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants