Skip to content

Move src/MaxText/kernels to src/maxtext/kernels #3032

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Open
dipannita08 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Move src/MaxText/kernels to src/maxtext/kernels #3032

dipannita08 wants to merge 1 commit into from

Conversation

@dipannita08
Copy link
Collaborator

@dipannita08 dipannita08 commented Jan 28, 2026 •
edited
Loading

Description

This PR moves the kernels modules into the restructured location.
Details here

Tests

  • Relevant unit tests
  • Manual tests which exercise moved modules.
  • Presubmits

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.

SurbhiJainUSC
SurbhiJainUSC reviewed Jan 28, 2026
View reviewed changes
SurbhiJainUSC
SurbhiJainUSC reviewed Jan 28, 2026
View reviewed changes
@codecov
Copy link

codecov bot commented Jan 29, 2026 •
edited
Loading

Codecov Report

❌ Patch coverage is 28.00000% with 18 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
src/maxtext/kernels/jax_flash_attention.py 5.55% 17 Missing ⚠️
src/maxtext/kernels/megablox/ops.py 50.00% 1 Missing ⚠️

📢 Thoughts on this report? Let us know!

SurbhiJainUSC
SurbhiJainUSC reviewed Jan 30, 2026
View reviewed changes
src/maxtext/kernels/ragged_attention.py
Copy link
Collaborator

@SurbhiJainUSC SurbhiJainUSC Jan 30, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you move all three attention modules to src/maxtext/kernels/attention?
https://screenshot.googleplex.com/5jZeGVHLhcWJQhA

@SurbhiJainUSC SurbhiJainUSC self-requested a review January 30, 2026 21:25
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@NuojCheng NuojCheng NuojCheng approved these changes

@jacoguzo jacoguzo Awaiting requested review from jacoguzo

@bvandermoon bvandermoon Awaiting requested review from bvandermoon bvandermoon is a code owner

@richjames0 richjames0 Awaiting requested review from richjames0 richjames0 is a code owner

@shralex shralex Awaiting requested review from shralex shralex is a code owner

@gobbleturk gobbleturk Awaiting requested review from gobbleturk gobbleturk is a code owner

@RissyRan RissyRan Awaiting requested review from RissyRan RissyRan is a code owner

@gagika gagika Awaiting requested review from gagika gagika is a code owner

@A9isha A9isha Awaiting requested review from A9isha A9isha is a code owner

@jiangjy1982 jiangjy1982 Awaiting requested review from jiangjy1982 jiangjy1982 is a code owner

@vipannalla vipannalla Awaiting requested review from vipannalla vipannalla is a code owner

@michelle-yooh michelle-yooh Awaiting requested review from michelle-yooh michelle-yooh is a code owner

@suexu1025 suexu1025 Awaiting requested review from suexu1025 suexu1025 is a code owner

@jesselu-google jesselu-google Awaiting requested review from jesselu-google jesselu-google is a code owner

@patemotter patemotter Awaiting requested review from patemotter patemotter is a code owner

@parambole parambole Awaiting requested review from parambole parambole is a code owner

@shuningjin shuningjin Awaiting requested review from shuningjin shuningjin is a code owner

@NicoGrande NicoGrande Awaiting requested review from NicoGrande NicoGrande is a code owner

@khatwanimohit khatwanimohit Awaiting requested review from khatwanimohit khatwanimohit is a code owner

@hengtaoguo hengtaoguo Awaiting requested review from hengtaoguo hengtaoguo is a code owner

@aireenmei aireenmei Awaiting requested review from aireenmei aireenmei is a code owner

@SurbhiJainUSC SurbhiJainUSC Awaiting requested review from SurbhiJainUSC SurbhiJainUSC is a code owner

At least 2 approving reviews are required to merge this pull request.

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@dipannita08 @SurbhiJainUSC @NuojCheng