Skip to content

Conversation

@Rohan-Bierneni
Copy link
Collaborator

This pr adds caching support to the Gated Delta Net implementation for Qwen3 Next. This is the first non-standard attention caching support in MaxText and needed changes to both MaxEngine code as well as qwen3 next implementation code.

If the change fixes a bug or a Github issue, please include a link, e.g.,:
FIXES: b/448407748

Tests

Ran the MaxText.decode command and got meaningful output: https://paste.googleplex.com/5431522368815104

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.

Added support for GDN to maxengine but NNX linen incompatible

Merged code from other branch Qwen3-next

Modified to accept dynamic model mode and work with maxengine changes

Fix GDN init with model_mode

Do same cache update during packed prefill as normal prefill

Convert batch to int in init for state
@Rohan-Bierneni Rohan-Bierneni force-pushed the rbierneni-qwen3-next-caching branch from 15f871b to 7ec8a31 Compare January 20, 2026 18:39
@codecov
Copy link

codecov bot commented Jan 20, 2026

Codecov Report

❌ Patch coverage is 38.09524% with 39 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
src/MaxText/layers/decoders.py 9.52% 17 Missing and 2 partials ⚠️
src/MaxText/layers/qwen3.py 60.00% 9 Missing and 5 partials ⚠️
src/MaxText/maxengine.py 14.28% 6 Missing ⚠️

📢 Thoughts on this report? Let us know!

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.

1 participant