Refactor: Decouple Core Transformer Blocks

[parambole]

TL;DR

• What: This PR refactors the core DecoderLayer and its related components out of layers/models.py and into a new, foundational file: layers/blocks.py.

• Why: To improve the overall code architecture and break potential circular dependencies. This is a necessary prerequisite for adding new, complex modules that also need access to these core building blocks.

• How: By creating layers/blocks.py to house the fundamental Decoder and DecoderLayer classes. Higher-level files like models.py and other future modules now import these components from a single location.

Detailed Description

This pull request introduces a structural refactoring to improve modularity and maintainability.

The primary change is the creation of MaxText/layers/blocks.py, which now serves as the source for fundamental building blocks of the Transformer architecture, such as:

• DecoderLayer

• Decoder

Previously, these classes were located in layers/models.py, which created tight coupling. As we add more features this tight coupling would lead to circular import dependencies.

By decoupling these core components, we establish a clearer hierarchy in the codebase, where high-level modules can depend on these "building blocks" without depending on each other.

Checklist

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

• I have performed a self-review of my code.
• 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.