FazBrowse GitHub Viewer | Trending |
URL:
| Home
Tools: [Download Repo ZIP]   [Original HTTPS Page]

Improve performance of gemma4 MoE inference. by NicoGrande · Pull Request #3875 · AI-Hypercomputer/maxtext · GitHub

Improve performance of gemma4 MoE inference. - #3875

Merged
copybara-service[bot] merged 1 commit into
mainfrom
nicogrande/improve-gemma4-vllm-perf
May 12, 2026
Merged

Improve performance of gemma4 MoE inference.#3875
copybara-service[bot] merged 1 commit into
mainfrom
nicogrande/improve-gemma4-vllm-perf

Conversation

Copy link
Copy Markdown
Collaborator

Description

This PR introduces a small, inference specific optimization to reduce overall step time. Specifically, this PR moves the per-expert scale application to the model initialization instead of during the call method, given that these parameters are static during inference.

Tests

vllm_decode.py before and after.

Step time without this PR: 13.59 ms
Step time with this PR: 7.975 ms

Checklist

Before submitting this PR, please make sure (put X in square brackets):

  • I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
  • I have necessary comments in my code, particularly in hard-to-understand areas.
  • I have run end-to-end tests tests and provided workload links above if applicable.
  • I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

Comment thread src/maxtext/layers/moe.py Outdated
NicoGrande force-pushed the nicogrande/improve-gemma4-vllm-perf branch from 1bea1ad to e3bc236 Compare May 11, 2026 22:28

codecov Bot commented May 11, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 0% with 3 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
src/maxtext/layers/moe.py 0.00% 1 Missing and 2 partials ⚠️

📢 Thoughts on this report? Let us know!

NicoGrande force-pushed the nicogrande/improve-gemma4-vllm-perf branch from 2f91e1b to 5b0ae2f Compare May 11, 2026 23:58
copybara-service Bot merged commit 4d9f390 into main May 12, 2026
37 of 38 checks passed
copybara-service Bot deleted the nicogrande/improve-gemma4-vllm-perf branch May 12, 2026 18:56
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants


Back | FazBrowse Home | New Git URL