Optimize sequence reordering on TPU by eliminating jnp.take gather#4469
Conversation
|
🤖 Hi @JHCuc3m, I've received your request, and I'm working on it now! You can track my progress in the logs for more details. |
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
There was a problem hiding this comment.
This Pull Request introduces a brilliant performance optimization to the TPU sequence reordering logic by replacing the general jnp.take (gather) operation with strided slicing and native JAX layout transformations. By avoiding the generalized gather operation, XLA can now compile the sequence reordering logic into highly efficient, contiguous memory/layout transformations instead of sequential TensorCore loops, yielding an impressive ~4% overall step-time speedup. The code quality, comments, and test coverage are outstandingly thorough.
🔍 General Feedback
- Exceptional Mathematical Rigor: The new strided slicing implementation for restoring contiguous order (
to_contiguous=True) is mathematically elegant and perfectly preserves functional equivalence while dramatically improving hardware-level efficiency. - Robust Test Coverage: The added equivalence sweep and roundtrip tests are extremely thorough, ensuring correctness across different TPU configurations and tensor shapes (including 1D, standard B-S-H-D, and S-B-H-D formats).
- Clean and Well-Documented Code: The code structure and inline documentation are exceptionally clear and follow JAX development guidelines perfectly.
|
|
||
| # Sweep configurations | ||
| cp_sizes = [2, 4, 8] | ||
| to_contiguous_options = [False, True] |
There was a problem hiding this comment.
| to_contiguous_options = [False, True] | |
| key = random.key(42) |
| shape = (2, 128, 4, 8) | ||
| seq_dim = 1 | ||
| cp_size = 4 | ||
|
|
There was a problem hiding this comment.
| key = random.key(123) |
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
cc43e34 to
ee34456
Compare
|
🤖 Hi @JHCuc3m, I've received your request, and I'm working on it now! You can track my progress in the logs for more details. |
There was a problem hiding this comment.
This Pull Request implements a highly effective performance optimization for sequence reordering on TPUs by replacing generalized jnp.take (gather) operations with structured strided slicing and native JAX layout transformations. This elegantly resolves a key performance bottleneck identified during DeepSeek-V3.2 long-context training sweeps (128K) by allowing XLA to specialize layout transformations rather than emitting expensive, sequential TensorCore while loops. The code is exceptionally clean, well-optimized, and includes high-quality unit tests checking bit-for-bit mathematical equivalence.
🔍 General Feedback
- Sizable Performance Gains: The optimization delivers an impressive and measurable step-time reduction (from 77s down to 74s), demonstrating excellent optimization work via structural JAX layout manipulations.
- Strong Test Suite: The testing strategy is highly rigorous, utilizing dynamic equivalence sweeps and roundtrip verification to guarantee lossless data transformations.
- Robustness Suggestions: The provided comments and code suggestions add robust support for standard Python negative-index dimension targeting (e.g.,
seq_dim = -1), making this utility fully general-purpose.
| if cp_size % 2 != 0: | ||
| raise ValueError(f"{cp_size=} must be a multiple of 2.") | ||
|
|
||
| if seq_len % (cp_size * 2) != 0: | ||
| raise ValueError(f"{tensor.shape=} is not a multiple of {cp_size*2=}") | ||
|
|
||
| ori_tensor_shape = tensor.shape | ||
| reshaped = tensor.reshape( | ||
| *ori_tensor_shape[:seq_dim], | ||
| 2 * cp_size, | ||
| group_size, | ||
| *ori_tensor_shape[seq_dim + 1 :], | ||
| ) | ||
|
|
||
| if not to_contiguous: | ||
| first_half = jnp.arange(cp_size) |
There was a problem hiding this comment.
| if cp_size % 2 != 0: | |
| raise ValueError(f"{cp_size=} must be a multiple of 2.") | |
| if seq_len % (cp_size * 2) != 0: | |
| raise ValueError(f"{tensor.shape=} is not a multiple of {cp_size*2=}") | |
| ori_tensor_shape = tensor.shape | |
| reshaped = tensor.reshape( | |
| *ori_tensor_shape[:seq_dim], | |
| 2 * cp_size, | |
| group_size, | |
| *ori_tensor_shape[seq_dim + 1 :], | |
| ) | |
| if not to_contiguous: | |
| first_half = jnp.arange(cp_size) | |
| ori_tensor_shape = tensor.shape | |
| if seq_dim < 0: | |
| seq_dim = seq_dim % len(ori_tensor_shape) | |
| seq_len = ori_tensor_shape[seq_dim] | |
| group_size = seq_len // (2 * cp_size) | |
| if cp_size % 2 != 0: | |
| raise ValueError(f"{cp_size=} must be a multiple of 2.") | |
| if seq_len % (cp_size * 2) != 0: | |
| raise ValueError(f"{tensor.shape=} is not a multiple of {cp_size*2=}") | |
| reshaped = tensor.reshape( | |
| *ori_tensor_shape[:seq_dim], | |
| 2 * cp_size, | |
| group_size, | |
| *ori_tensor_shape[seq_dim + 1 :], | |
| ) |
| # Optimized TPU path: Strided slicing avoids dimension of size 2 and SC offloading | ||
| first_half = swapped[0::2, ...] | ||
| second_half_reversed = swapped[1::2, ...] | ||
| second_half = second_half_reversed[::-1, ...] |
There was a problem hiding this comment.
To maintain consistency and make the code easier to follow, rename second_half_reversed to second_half_unreversed and second_half to second_half_reversed.
| # Optimized TPU path: Strided slicing avoids dimension of size 2 and SC offloading | |
| first_half = swapped[0::2, ...] | |
| second_half_reversed = swapped[1::2, ...] | |
| second_half = second_half_reversed[::-1, ...] | |
| first_half = swapped[0::2, ...] | |
| second_half_unreversed = swapped[1::2, ...] | |
| second_half_reversed = second_half_unreversed[::-1, ...] | |
| permuted = jnp.concatenate([first_half, second_half_reversed], axis=0) |
| # [S, B, H, D]: [2*cp_size, S/2*cp_size, B, H, D] -> [2, S/2*cp_size, B, H, D] | ||
| ori_tensor_shape = tensor.shape | ||
|
|
||
| reshaped = tensor.reshape( |
There was a problem hiding this comment.
To support negative dimensions robustly, we should normalize seq_dim relative to the tensor rank before any slicing operations.
| reshaped = tensor.reshape( | |
| if seq_dim < 0: | |
| seq_dim = seq_dim % len(ori_tensor_shape) | |
There was a problem hiding this comment.
Could you consider this comment and add a test? i.e. it's possible to use Axis -1 refers to the very last dimension
| ((256, 2, 4, 8), 0), # 4D (S, B, H, D) | ||
| ] | ||
|
|
||
| for cp_size in cp_sizes: | ||
| for to_contiguous in to_contiguous_options: |
There was a problem hiding this comment.
| ((256, 2, 4, 8), 0), # 4D (S, B, H, D) | |
| ] | |
| for cp_size in cp_sizes: | |
| for to_contiguous in to_contiguous_options: | |
| test_cases = [ | |
| ((64,), 0), # 1D | |
| ((64,), -1), # 1D negative index | |
| ((2, 128, 4, 8), 1), # 4D (standard B, S, H, D) | |
| ((2, 128, 4, 8), -3), # 4D standard negative index | |
| ((256, 2, 4, 8), 0), # 4D (S, B, H, D) | |
| ((256, 2, 4, 8), -4), # 4D S negative index | |
| ] |
ee34456 to
edaf223
Compare
|
🤖 Hi @JHCuc3m, I've received your request, and I'm working on it now! You can track my progress in the logs for more details. |
There was a problem hiding this comment.
This pull request presents an exceptional optimization of the reorder_sequence function, replacing expensive generalized gather operations (jnp.take) with structured strided slicing and JAX layout transformations. The code quality is outstanding, offering a significant reduction in step time (from 77s to 74s, a ~4% speedup) for DeepSeek-V3.2 training sweeps at long contexts on TPU.
🔍 General Feedback
- Exceptional Mathematical Rigor: The new implementation achieves bit-for-bit mathematical equivalence with the original gather-based implementation while avoiding the sequential TensorCore
whileloops generated by XLA during gather lowering. - Outstanding Test Coverage: The inclusion of
test_reorder_equivalence_sweepandtest_reorder_roundtripensures complete, robust verification across different Tensor shapes, sequence dimensions, and compute parallelism options. - Idiomatic JAX Patterns: Leveraging JAX layout transformations like
jnp.split,jnp.stack, andjnp.concatenateinstead of index-based gathers is highly idiomatic and enables efficient zero-copy / DMA optimizations in XLA compiler lowering on TPU.
| # Stack and reshape to interleave | ||
| src_indices = jnp.stack([first_half, second_half], axis=1).reshape(-1) | ||
| # Swap seq_dim with axis 0 to perform slicing/concat easily | ||
| swapped = jnp.swapaxes(reshaped, 0, seq_dim) |
There was a problem hiding this comment.
add a comment for shape - e.g,.is it b, t, d?
There was a problem hiding this comment.
+1, let's add a comment on shape of this swapped variable for readability.
edaf223 to
701c2e2
Compare
huytransformer
left a comment
huytransformer
left a comment
There was a problem hiding this comment.
Excellent work!!!
RissyRan
left a comment
RissyRan
left a comment
There was a problem hiding this comment.
Thanks for the change! A few minor comments.
| # Stack and reshape to interleave | ||
| src_indices = jnp.stack([first_half, second_half], axis=1).reshape(-1) | ||
| # Swap seq_dim with axis 0 to perform slicing/concat easily | ||
| swapped = jnp.swapaxes(reshaped, 0, seq_dim) |
There was a problem hiding this comment.
+1, let's add a comment on shape of this swapped variable for readability.
| # [S, B, H, D]: [2*cp_size, S/2*cp_size, B, H, D] -> [2, S/2*cp_size, B, H, D] | ||
| ori_tensor_shape = tensor.shape | ||
|
|
||
| reshaped = tensor.reshape( |
There was a problem hiding this comment.
Could you consider this comment and add a test? i.e. it's possible to use Axis -1 refers to the very last dimension
701c2e2 to
e5e46f6
Compare
|
@RissyRan Applied the changes. Lmk if it looks good to you! |
Replaced the general jnp.take (gather) operation in reorder_sequence with structured strided slicing and JAX layout operations. This optimization completely eliminates the sequential TensorCore 'while' loops on TPU. Tested: 1. Added test_reorder_equivalence_sweep in tests/unit/max_utils_test.py verifying bit-for-bit mathematical equivalence against the original gather-based implementation. 2. Verified against PyTorch deepseek32 indexer reference. TAG=agy CONV=f2dcd166-f4f6-47ab-8984-0f5cdfca2ce1
e5e46f6 to
6ef78e9
Compare
4 participants
Description
Replaces the general
jnp.take(gather) operation inreorder_sequencewith structured strided slicing and JAX layout operations.Context & Problem:
During training sweeps of DeepSeek-V3.2 at long contexts (128K) with CP=128, the sequence reordering logic was identified as a performance bottleneck. XLA was lowering the original
jnp.takegeneralized gather into expensive sequential TensorCorewhileloops.Solution:
Because the permutation indices are perfectly deterministic based on
cp_size, we can natively recreate the permutation mapping using contiguous layout transformations (for load balancing folding) and stride-2 vector slicing (for restoring original order).Tests
Testing performed locally via
pytest:test_reorder_equivalence_sweepintests/unit/max_utils_test.py. It sweepscp_size, direction (to_contiguous=True/False), tensor ranks (1D, 4D), and negative dimension indexing (seq_dim = -1,seq_dim = -3), asserting bit-for-bit mathematical equivalence against the original gather-based implementation.test_reorder_roundtripintests/unit/max_utils_test.pyto ensure that load-balancing a tensor and subsequently restoring it yields the exact original input tensor.Profiles
Baseline: Step time 77s
Improved xprof: Step time 74s
Checklist
Before submitting this PR, please make sure (put X in square brackets):
gemini-reviewlabel.finishes!)
TAG=agy
CONV=f2dcd166-f4f6-47ab-8984-0f5cdfca2ce1