[Stacked PR 1/5] Fix upstream pyink formatting and GMM v2 compatibility for Qwen3.5 - #5151
[Stacked PR 1/5] Fix upstream pyink formatting and GMM v2 compatibility for Qwen3.5#5151Rohan-Bierneni wants to merge 1 commit into
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request adjusts sublane tiling logic in Megablox TPU kernels and applies formatting cleanups to the Qwen3.5 model definition and tests. In the TGMM kernel, a change was made to compute a common sublane size using min(). Feedback indicates that using min can violate hardware tiling constraints for operands with larger requirements, and suggests using max() instead to ensure compatibility.
| common_sublane = min(size_lhs_sublane, size_rhs_sublane) | ||
| size_lhs_sublane = common_sublane | ||
| size_rhs_sublane = common_sublane |
There was a problem hiding this comment.
Using min to find the common_sublane can violate the hardware sublane tiling constraints of the operand with the larger sublane tiling requirement. For example, if size_lhs_sublane is 128 (e.g., for bf16) and size_rhs_sublane is 8, taking the minimum results in a common_sublane of 8. Setting size_lhs_sublane to 8 violates the 128-alignment requirement for bf16 on TPU, leading to compilation or runtime failures.
Since TPU sublane tilings are always powers of two, their least common multiple is simply their maximum. Using max instead of min ensures that the common sublane size is a multiple of both requirements, satisfying all hardware constraints.
| common_sublane = min(size_lhs_sublane, size_rhs_sublane) | |
| size_lhs_sublane = common_sublane | |
| size_rhs_sublane = common_sublane | |
| common_sublane = max(size_lhs_sublane, size_rhs_sublane) | |
| size_lhs_sublane = common_sublane | |
| size_rhs_sublane = common_sublane |
Rohan-Bierneni
changed the title
…iling - Reformat qwen3.py, qwen3_5.py, pyconfig_test.py, train_tests.py, and train_compile_test.py with pyink to conform with upstream CI style requirements. - Fix Megablox TPU v2 sublane tiling in pallas_mosaic_tpu_v2_gmm_kernel.py and pallas_mosaic_tpu_v2_tgmm_kernel.py to ensure correct tile dimension handling on TPU sublanes.
ee278f5 to
a8aafbd
Compare
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
1 participant
Stacked PR Chain
rbierneni-gdn-1-ci-hygienemainrbierneni-gdn-2-fwd-kernelrbierneni-gdn-1-ci-hygienerbierneni-gdn-3-bwd-kernelrbierneni-gdn-2-fwd-kernelrbierneni-gdn-4-model-integrationrbierneni-gdn-3-bwd-kernelrbierneni-gdnv3-bwdrbierneni-gdn-4-model-integrationDescription
This is PR 1 of 5 in the stacked series enabling the Pallas Gated Delta Net (GDN) backward pass kernel in MaxText.
This PR addresses upstream CI hygiene and kernel compatibility prerequisites:
src/maxtext/models/qwen3_5.py,tests/integration/train_tests.py, andtests/unit/train_compile_test.pywithpyinkfollowing recentmainmerges to conform with repository formatting requirements.src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.pyandsrc/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_tgmm_kernel.pyto ensure proper tile dimension boundaries on TPU sublanes (enforcing minimum sublane size of 16 and matching LHS/RHS sublanes in TGMM) to enable compatibility with the Qwen3.5 MoE model.Files Changed
src/maxtext/models/qwen3_5.py: Pyink formatting conformance.tests/integration/train_tests.py: Pyink formatting conformance.tests/unit/train_compile_test.py: Pyink formatting conformance.src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py: Sublane tiling calculation fix for Qwen3.5 compatibility.src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_tgmm_kernel.py: Sublane tiling calculation fix for Qwen3.5 compatibility.Tests
pre-commit run --files ...(all hooks:codespell,pylint,pyink,yamllintpassed).train_compile_test.py.Checklist