patch_vllm_gemma4_mtp_embedding_share.py 1.4 KB

1234567891011121314151617181920212223242526272829303132
  1. #!/usr/bin/env python3
  2. """Keep the EAGLE embedding-width guard from breaking Gemma4 MTP."""
  3. import sys
  4. from pathlib import Path
  5. source_root = Path(sys.argv[1]) if len(sys.argv) > 1 else Path.cwd()
  6. target = source_root / "vllm/v1/spec_decode/llm_base_proposer.py"
  7. old = """ if share_embeddings:
  8. draft_embed = self.model.model.embed_tokens
  9. # Only share when both models use the same embedding width.
  10. # Guard with isinstance so non-Tensor weights (e.g. in tests)
  11. """
  12. new = (
  13. ' if share_embeddings and hasattr(self.model, "has_own_embed_tokens"):\n'
  14. " draft_embed = self.model.model.embed_tokens\n"
  15. " # Only share when both models use the same embedding width.\n"
  16. " # Guard with isinstance so non-Tensor weights (e.g. in tests)\n"
  17. )
  18. if not target.exists():
  19. print(f"{target} not found; skipping Gemma4 MTP embedding-share workaround")
  20. else:
  21. text = target.read_text()
  22. if 'if share_embeddings and hasattr(self.model, "has_own_embed_tokens"):' in text:
  23. print("Gemma4 MTP embedding-share workaround already present; skipping")
  24. elif old in text:
  25. target.write_text(text.replace(old, new, 1))
  26. print("Applied Gemma4 MTP embedding-share workaround")
  27. else:
  28. print("Known Gemma4 MTP embedding-share pattern not found; skipping")