Skip to content

Pull requests: AI-Hypercomputer/maxtext

New pull request New
223 Open 2,371 Closed
223 Open 2,371 Closed
Author
Filter by author
Loading
Label
Filter by label
Loading
Use alt + click/return to exclude labels
or ⇧ + click/return for logical OR
Projects
Filter by project
Loading
Milestones
Filter by milestone
Loading
Reviews
Filter by reviews
No reviews Review required Approved review Changes requested
Assignee
Filter by who’s assigned
Assigned to nobody Loading
Sort
Sort by
Newest Oldest Most commented Least commented Recently updated Least recently updated Best match
Most reactions
👍 👎 😄 🎉 😕 ❤️ 🚀 👀

Pull requests list

More UXR fixes
#2794 opened Dec 6, 2025 by hengtaoguo • Draft
4 tasks done
This change introduces a pure JAX implementation of flash attention to Maxtext, designed as a drop-in replacement for the existing Pallas kernel. In this cl we set up the stage by integrating it with maxtext in fsdp mode. We have plans for further optimizations to close the gap with pallas using different techniques such as:
#2793 opened Dec 5, 2025 by copybara-service bot Loading…
Migrate multi_token_prediction to NNX
#2788 opened Dec 4, 2025 by hsuan-lun-chiang • Draft
4 tasks
Add concat_then_split packing in Grain pipeline
#2786 opened Dec 4, 2025 by aireenmei Loading…
4 tasks done
4
fix bug in grain pipeline and refactor
#2785 opened Dec 4, 2025 by aireenmei Loading…
4 tasks done
1
4
Onboard explicit sharding to deepseek [split version] gemini-review
#2783 opened Dec 3, 2025 by NuojCheng Loading…
4 tasks done
10
Fixes Issue #2781 for gcloud command in getting number of workers
#2782 opened Dec 3, 2025 by MrGeislinger Loading…
4 tasks done
1
add markdown linting as a pre-commit hook
#2780 opened Dec 3, 2025 by ev-br Loading…
4 tasks done
Enable to autotune maxtext extra_flags and maxtext_flags.
#2779 opened Dec 2, 2025 by copybara-service bot Loading…
Support Custom MaxText model (with vLLM engine) in RL rollouts.
#2778 opened Dec 2, 2025 by NicoGrande Loading…
4 tasks done
[src/MaxText/configs/types.py] Error loudly on invalid keys ; Add missing base_config meta field ; [src/MaxText/pyconfig.py] Enable pydantic to be constructed in public API ; [tests/configs_value_test.py] Test ValueError is raised
#2775 opened Dec 2, 2025 by SamuelMarks Loading…
4 tasks done
4
Use MaxText max_segments_per_seq config variable to control Grain batch packing
#2774 opened Dec 2, 2025 by gabeweisz Loading…
4 tasks done
2
Support checkpointing TE quantizations with new remat policies
#2773 opened Dec 2, 2025 by jberchtold-nvidia Loading…
4 tasks done
Internal change
#2771 opened Dec 2, 2025 by copybara-service bot Loading…
@SurbhiJainUSC
Migrate Pipeline to use NNX module.
#2770 opened Dec 2, 2025 by ecnal-cienet • Draft
4 tasks done
Docs: Improve TPU Runtime & Colab setup guide
#2768 opened Dec 1, 2025 by RexBearIU Loading…
4 tasks done
@RexBearIU
35
Support Custom MaxText model (with vLLM engine) in RL rollouts.
#2767 opened Dec 1, 2025 by gagika • Draft
4 tasks done
Add vLLM Support for GPT OSS and its mapping generator for tunix.
#2766 opened Dec 1, 2025 by abhinavclemson Loading…
4 tasks done
2
feat: ported nnx model converter
#2765 opened Dec 1, 2025 by mesakhcienet • Draft
4 tasks
Update documentation links and remove an old requirements file.
#2763 opened Nov 28, 2025 by copybara-service bot Loading…
@shralex
[pyproject.toml] Use already-in-use hatch plugin over custom hook
#2762 opened Nov 27, 2025 by SamuelMarks Loading…
4 tasks done
Fix tuple unpacking regression in Decoder Layers
#2753 opened Nov 25, 2025 by gagika Loading…
4 tasks done
@shralex @khatwanimohit @NuojCheng
3
Implement option to perform yarn embedding using a rotation matrix in real vector space.
#2752 opened Nov 25, 2025 by copybara-service bot Loading…
Adding option for TE Dot with BF16 + unittest
#2750 opened Nov 25, 2025 by phu0ngng Loading…
4 tasks done
Fix the grpo on llama3.1-8b-instruction and its notebook example. gemini-review
#2746 opened Nov 24, 2025 by mathczh Loading…
4 tasks done
@hengtaoguo
27
ProTip! Type g i on any issue or pull request to go back to the issue listing page.