[JAX] TE Permutation integration to Maxtext #2672

tdophung · 2026-02-11T00:09:09Z

Description

Changes needed on TE side to make maxtext integration works

Issue # 2585

Type of change

Documentation change (change only to the documentation, either a fix or a new content)
Bug fix (non-breaking change which fixes an issue)
New feature (non-breaking change which adds functionality)
Breaking change (fix or feature that would cause existing functionality to not work as expected)
Infra/Build change
Code refactoring

Changes

Please list the changes introduced in this PR:

Masking out padding tokens on each local EP (will be used in local permute step)
Pass along split_sizes and sorted_indices in residual of sort_chunks_by_index (local_permute) to avoid mismatch in size issue during tracing when EP>1

Checklist:

I have read and followed the contributing guidelines
The functionality is complete
I have commented my code, particularly in hard-to-understand areas
I have made corresponding changes to the documentation
My changes generate no new warnings
I have added tests that prove my fix is effective or that my feature works
New and existing unit tests pass locally with my changes

for more information, see https://pre-commit.ci

Signed-off-by: tdophung <[email protected]>

…ger than num tokens Signed-off-by: tdophung <[email protected]>

Signed-off-by: JAX Toolbox <[email protected]>

for more information, see https://pre-commit.ci

tdophung · 2026-02-11T00:11:06Z

This PR contain changes cherry-picked from #2651 . I can wait until this gets merged and then merge mine, but if my PR is needed more urgently, happy to remove the cherry picked change

greptile-apps · 2026-02-11T00:12:28Z

Greptile Overview

Greptile Summary

This PR adds necessary changes to support MaxText integration with TE permutation operations when Expert Parallelism (EP) > 1.

Key Changes:

Padding token masking in Triton kernel: Modified _make_chunk_sort_map_kernel to compute total_valid_tokens and apply identity mapping for padding tokens (positions beyond the valid token count), preventing corruption of valid output positions
VJP tracing fix: Removed nondiff_argnums=(1, 2) from @jax.custom_vjp decorator and instead passed split_sizes and sorted_indices through residuals with zero gradients, resolving size mismatch issues during JAX tracing with EP>1

Technical Implementation:

The Triton kernel now checks pid < total_valid_tokens before applying the permutation mapping, using tl.where(pid < total_valid_tokens, dst_row, pid) to preserve padding tokens at their original positions
The backward pass now correctly unpacks split_sizes and sorted_indices from residuals and returns zero gradients for these integer arrays, maintaining gradient flow correctness

Confidence Score: 4/5

This PR is safe to merge with minor risk - changes are well-targeted to specific integration issues
The changes are focused and address specific technical issues (padding token handling and JAX tracing) for MaxText integration. The padding logic is sound and the VJP changes follow standard JAX patterns. Confidence reduced from 5 to 4 due to lack of new tests mentioned in the checklist.
No files require special attention - both changes are straightforward bug fixes

Important Files Changed

Filename	Overview
transformer_engine/common/triton/permutation.py	Added padding token masking logic to `_make_chunk_sort_map_kernel` to handle buffers larger than valid token count with identity mapping
transformer_engine/jax/permutation.py	Removed `nondiff_argnums` from `@jax.custom_vjp` and passed `split_sizes`/`sorted_indices` through residuals with zero gradients to fix tracing issues

Sequence Diagram

sequenceDiagram
    participant User as MaxText User
    participant API as sort_chunks_by_index
    participant FWD as _sort_chunks_by_index_fwd_rule
    participant Kernel as _make_chunk_sort_map_kernel
    participant BWD as _sort_chunks_by_index_bwd_rule

    User->>API: sort_chunks_by_index(inp, split_sizes, sorted_indices)
    API->>FWD: Forward pass
    FWD->>Kernel: Generate row_id_map with padding handling
    Note over Kernel: Compute total_valid_tokens<br/>Apply identity mapping for pid >= total_valid_tokens
    Kernel-->>FWD: row_id_map (with padding masked)
    FWD->>FWD: sort_chunks_by_map(inp, row_id_map)
    FWD-->>API: (output, row_id_map), residuals
    Note over FWD: residuals now include split_sizes<br/>and sorted_indices (not nondiff_argnums)
    
    User->>BWD: Backward pass (gradient)
    BWD->>BWD: Extract split_sizes, sorted_indices from residuals
    BWD->>BWD: sort_chunks_by_map(output_grad, row_id_map, is_forward=False)
    BWD-->>User: (inp_grad, zeros_like(split_sizes), zeros_like(sorted_indices))

_{Last reviewed commit: 11a45d3}

greptile-apps

_{6 files reviewed, 3 comments}

_{Edit Code Review Agent Settings | Greptile}

greptile-apps · 2026-02-11T00:12:32Z

transformer_engine/jax/inspect.py

+_inspect.defvjp(_inspect_fwd_rule, _inspect_bwd_rule)
+
+
+def inspect_array(x: jnp.ndarray, name: str) -> jnp.ndarray:


name parameter is unused - not passed to C++ backend or used in filename

Suggested change

def inspect_array(x: jnp.ndarray, name: str) -> jnp.ndarray:

def inspect_array(x: jnp.ndarray) -> jnp.ndarray:

greptile-apps · 2026-02-11T00:12:33Z

transformer_engine/jax/csrc/extensions/amax.cpp

+  std::ofstream file(filename, std::ios::binary);
+  if (file.is_open()) {
+    file.write(reinterpret_cast<const char *>(input_data.data()), input_data.size());
+    file.close();
+  }


No error handling if file fails to open - silently continues without writing data

Suggested change

std::ofstream file(filename, std::ios::binary);

if (file.is_open()) {

file.write(reinterpret_cast<const char *>(input_data.data()), input_data.size());

file.close();

}

std::ofstream file(filename, std::ios::binary);

if (!file.is_open()) {

return ffi::Error(ffi::ErrorCode::kInternal, "Failed to open file for writing");

}

file.write(reinterpret_cast<const char *>(input_data.data()), input_data.size());

file.close();

transformer_engine/jax/permutation.py

tdophung · 2026-02-11T21:07:14Z

/te-ci

tdophung · 2026-02-11T21:07:41Z

/te_ci

transformer_engine/common/triton/permutation.py

transformer_engine/jax/permutation.py

Signed-off-by: JAX Toolbox <[email protected]>

…nsformerEngine into maxtext_integ_2

for more information, see https://pre-commit.ci

greptile-apps

_{2 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

jberchtold-nvidia

LGTM, thanks!

Signed-off-by: JAX Toolbox <[email protected]>

…rmerEngine into maxtext_integ_2

greptile-apps

_{2 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

tdophung · 2026-02-13T05:04:23Z

/te_ci

tdophung · 2026-02-13T05:04:29Z

/te-ci

pre-commit-ci bot and others added 6 commits February 5, 2026 13:55

[pre-commit.ci] auto fixes from pre-commit.com hooks

6d33449

for more information, see https://pre-commit.ci

adding more stuff missing from cherry picky jeremy PR for inspecting

3e121b5

Signed-off-by: tdophung <[email protected]>

fix some tracing issues when intergating to maxtext

a76763a

Signed-off-by: tdophung <[email protected]>

Have sort_chunks_by_index handle situations where input buffer is lar…

1059c2b

…ger than num tokens Signed-off-by: tdophung <[email protected]>

remove unnecessary assert and comments

738cce9

Signed-off-by: JAX Toolbox <[email protected]>

[pre-commit.ci] auto fixes from pre-commit.com hooks

5d8d06c

for more information, see https://pre-commit.ci

greptile-apps bot reviewed Feb 11, 2026

View reviewed changes

tdophung requested a review from jberchtold-nvidia February 11, 2026 21:06

jberchtold-nvidia reviewed Feb 11, 2026

View reviewed changes

transformer_engine/common/triton/permutation.py Show resolved Hide resolved

transformer_engine/jax/permutation.py Show resolved Hide resolved

nvjax and others added 3 commits February 12, 2026 14:30

remove Jeremy's PR for inspect ffi

81461d9

Signed-off-by: JAX Toolbox <[email protected]>

Qq:qMerge branch 'maxtext_integ_2' of https://github.com/tdophung/Tra…

9879032

…nsformerEngine into maxtext_integ_2

[pre-commit.ci] auto fixes from pre-commit.com hooks

d4228a8

for more information, see https://pre-commit.ci

greptile-apps bot reviewed Feb 12, 2026

View reviewed changes

jberchtold-nvidia previously approved these changes Feb 13, 2026

View reviewed changes

nvjax added 2 commits February 12, 2026 17:04

untouch the amax file, also change comment on te

0029e2c

Signed-off-by: JAX Toolbox <[email protected]>

Merge branch 'maxtext_integ_2' of https://github.com/tdophung/Transfo…

11a45d3

…rmerEngine into maxtext_integ_2

tdophung dismissed jberchtold-nvidia’s stale review via 11a45d3 February 13, 2026 01:05

jberchtold-nvidia approved these changes Feb 13, 2026

View reviewed changes

greptile-apps bot reviewed Feb 13, 2026

View reviewed changes

		_inspect.defvjp(_inspect_fwd_rule, _inspect_bwd_rule)


		def inspect_array(x: jnp.ndarray, name: str) -> jnp.ndarray:

	def inspect_array(x: jnp.ndarray, name: str) -> jnp.ndarray:
	def inspect_array(x: jnp.ndarray) -> jnp.ndarray:

[JAX] TE Permutation integration to Maxtext #2672

Are you sure you want to change the base?

[JAX] TE Permutation integration to Maxtext #2672

Conversation

tdophung commented Feb 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Type of change

Changes

Checklist:

Uh oh!

tdophung commented Feb 11, 2026

Uh oh!

greptile-apps bot commented Feb 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Greptile Overview

Greptile Summary

Confidence Score: 4/5

Important Files Changed

Sequence Diagram

Uh oh!

greptile-apps bot left a comment

Choose a reason for hiding this comment

Uh oh!

greptile-apps bot Feb 11, 2026

Choose a reason for hiding this comment

Uh oh!

greptile-apps bot Feb 11, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

tdophung commented Feb 11, 2026

Uh oh!

tdophung commented Feb 11, 2026

Uh oh!

Uh oh!

Uh oh!

greptile-apps bot left a comment

Choose a reason for hiding this comment

Uh oh!

jberchtold-nvidia left a comment

Choose a reason for hiding this comment

Uh oh!

greptile-apps bot left a comment

Choose a reason for hiding this comment

Uh oh!

tdophung commented Feb 13, 2026

Uh oh!

tdophung commented Feb 13, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

tdophung commented Feb 11, 2026 •

edited

Loading

greptile-apps bot commented Feb 11, 2026 •

edited

Loading