Skip to content

Optimize Gemma 4 MoE router input buffering and scaling:#965

Open
copybara-service[bot] wants to merge 1 commit into
devfrom
test_954067115
Open

Optimize Gemma 4 MoE router input buffering and scaling:#965
copybara-service[bot] wants to merge 1 commit into
devfrom
test_954067115

Conversation

@copybara-service

@copybara-service copybara-service Bot commented Jul 26, 2026

Copy link
Copy Markdown

Optimize Gemma 4 MoE router input buffering and scaling:

  • Changed router_in storage type from BF16 to float in Activations to avoid BF16->float->BF16 round-tripping.
  • Copies pre-FFW activations directly to the float router_in without type conversion.
  • Performs fused inline scaling using Highway SIMD (Decompress1AndCompressInplace) avoiding per-token type roundtrips and temporary vector allocations.

@google-cla

google-cla Bot commented Jul 26, 2026

Copy link
Copy Markdown

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

- Changed `router_in` storage type from `BF16` to `float` in `Activations` to avoid `BF16`->`float`->`BF16` round-tripping.
- Copies pre-FFW activations directly to the `float` `router_in` without type conversion.
- Performs fused inline scaling using Highway SIMD (`Decompress1AndCompressInplace`) avoiding per-token type roundtrips and temporary vector allocations.

PiperOrigin-RevId: 954067115
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.

0 participants