Skip to content

Expose sa_fuse_reciprocal and sa_use_base2_exp flags in MaxText#4264

Open
dandragona wants to merge 1 commit into
mainfrom
expose-tokamax-splash-flags
Open

Expose sa_fuse_reciprocal and sa_use_base2_exp flags in MaxText#4264
dandragona wants to merge 1 commit into
mainfrom
expose-tokamax-splash-flags

Conversation

@dandragona

@dandragona dandragona commented Jun 25, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Description

This PR exposes two flags that are supported by the underlying tokamax SplashConfig:

  • sa_fuse_reciprocal: Maps to fuse_reciprocal in SplashConfig (defaults to True).
  • sa_use_base2_exp: Maps to use_base2_exp in SplashConfig (defaults to True).

These are passed to tokamax_splash_kernel.SplashConfig.

Tests

  • Config validation tests were run: configs_test.py and pyconfig_test.py passed.

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.

TAG=agy
CONV=9b0f857c-78d2-4359-be96-73df07e21957
CHANGE_STEWARD=true

@dandragona dandragona force-pushed the expose-tokamax-splash-flags branch from e7df328 to 6e56e27 Compare June 25, 2026 02:34
@codecov

codecov Bot commented Jun 25, 2026 •
edited
Loading

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

@dandragona dandragona force-pushed the expose-tokamax-splash-flags branch from 6e56e27 to 464416f Compare June 25, 2026 02:39
@dandragona
Expose sa_fuse_reciprocal and sa_use_base2_exp flags in MaxText
675052d 
This change exposes two flags that are supported by the underlying tokamax SplashConfig:
- sa_fuse_reciprocal: Maps to fuse_reciprocal in SplashConfig (defaults to True).
- sa_use_base2_exp: Maps to use_base2_exp in SplashConfig (defaults to True).

These are passed to tokamax_splash_kernel.SplashConfig.

TAG=agy
CONV=9b0f857c-78d2-4359-be96-73df07e21957
CHANGE_STEWARD=true
@dandragona dandragona force-pushed the expose-tokamax-splash-flags branch from 464416f to 675052d Compare June 25, 2026 02:43
@dandragona dandragona marked this pull request as ready for review June 25, 2026 03:41
@dandragona dandragona assigned dandragona and unassigned dandragona Jun 25, 2026
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
@parambole parambole Awaiting requested review from parambole parambole is a code owner
@shuningjin shuningjin Awaiting requested review from shuningjin shuningjin is a code owner
@RissyRan RissyRan Awaiting requested review from RissyRan RissyRan is a code owner
@suexu1025 suexu1025 Awaiting requested review from suexu1025 suexu1025 is a code owner
@jiangjy1982 jiangjy1982 Awaiting requested review from jiangjy1982 jiangjy1982 is a code owner
@gobbleturk gobbleturk Awaiting requested review from gobbleturk gobbleturk is a code owner
@bvandermoon bvandermoon Awaiting requested review from bvandermoon bvandermoon is a code owner
@gagika gagika Awaiting requested review from gagika gagika is a code owner
@shralex shralex Awaiting requested review from shralex shralex is a code owner
@richjames0 richjames0 Awaiting requested review from richjames0 richjames0 is a code owner
@NicoGrande NicoGrande Awaiting requested review from NicoGrande NicoGrande is a code owner
@abhinavclemson abhinavclemson Awaiting requested review from abhinavclemson abhinavclemson is a code owner
@khatwanimohit khatwanimohit Awaiting requested review from khatwanimohit khatwanimohit is a code owner
@vipannalla vipannalla Awaiting requested review from vipannalla vipannalla is a code owner
@SurbhiJainUSC SurbhiJainUSC Awaiting requested review from SurbhiJainUSC SurbhiJainUSC is a code owner
@hengtaoguo hengtaoguo Awaiting requested review from hengtaoguo hengtaoguo is a code owner
@A9isha A9isha Awaiting requested review from A9isha A9isha is a code owner
@aireenmei aireenmei Awaiting requested review from aireenmei aireenmei is a code owner
@dipannita08 dipannita08 Awaiting requested review from dipannita08 dipannita08 is a code owner
@igorts-git igorts-git Awaiting requested review from igorts-git igorts-git is a code owner
@darisoy darisoy Awaiting requested review from darisoy darisoy 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.

2 participants

@dandragona @NuojCheng