Add ragged sort kernel fallback mechanism and version guard

[NuojCheng]

Description

This PR

• Introduces two flags, ragged_gather_fallback and ragged_gather_reduce_fallback. When they are true, a fully JAX, non-ragged version instead of the kernel version gonna be used
• Add version guard protecting ragged gather kernel.

Tests

Based on xprofs, both flags work effectively.

Fall back ragged gather reduce only

smoke_train model_name=deepseek2-16b ici_expert_parallelism=4 per_device_batch_size=1 max_target_length=4096 use_random_routing=true use_ring_of_experts=true use_ragged_sort=true ragged_gather_reduce_fallback=true debug_sharding=false profiler=xplane ragged_gather_fallback=false enable_tpu_profiling_options=true

xprof

Fall back ragged gather only

smoke_train model_name=deepseek2-16b ici_expert_parallelism=4 per_device_batch_size=1 max_target_length=4096 use_random_routing=true use_ring_of_experts=true use_ragged_sort=true ragged_gather_reduce_fallback=false debug_sharding=false profiler=xplane ragged_gather_fallback=true enable_tpu_profiling_options=true

xprof

Fall back both kernels

smoke_train model_name=deepseek2-16b ici_expert_parallelism=4 per_device_batch_size=1 max_target_length=4096 use_random_routing=true use_ring_of_experts=true use_ragged_sort=true ragged_gather_reduce_fallback=true debug_sharding=false profiler=xplane ragged_gather_fallback=true enable_tpu_profiling_options=true

xprof

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.

[codecov]

Codecov Report

❌ Patch coverage is 93.33333% with 1 line in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/kernels/ragged/ragged_gather_reduce.py	66.66%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!

[gobbleturk]

this is a bit surprising to see a non ragged fallback inside of a function called ragged_gather - should the decision to use ragged_gather vs a fallback happen at a higher level function (e.g. somewhere higher in the call stack?)

[NuojCheng]

The new ragged_gather_fallback and ragged_gather_reduce_fallback flags control kernel-level fallback logic. We avoid using these flags directly in moe.py because we may need the native JAX implementation for the forward pass and the custom kernel for the backward pass. While we could theoretically use them in ragged_sort.py, a fallback mechanism already exists in the primary kernel wrapper. We leverage that existing structure instead of introducing redundant if-else conditions.

[gobbleturk]

why skip on tpu7x?

[NuojCheng]

I think there were some issues using ragged sort on bloom. @darisoy do we have a buganizer tracking this?

[gobbleturk]

is this a ragged implementation? Doesn't this grow with the full buffer size or no?

[NuojCheng]

There are only two options:

1. Jax sort, no raggedness
2. sparse core kernels, ragged
   Update the comments to better reflect JAX is non-ragged.

[Shuwen-Fang]

could this just be:
ragged_gather: true --> use ragged kernel
ragged_gather:false --> use non ragged fallback

[NuojCheng]

There are two operations:
1: ragged permute (fwd: ragged gather; bwd: ragged gather reduce)
2. ragged unpermute (fwd: ragged gather reduce, bwd: ragged gather)

Technically we can introduce 4 flags, controlling whether we want kernel version/JAX version respectively. However, it is probably not necessary. Instead, we have one flag use_ragged_sort indicating whether we want to use kernel version for all, and ragged_gather_fallback and ragged_gather_reduce_fallback controlling kernels respectively.

I think you are suggesting re-naming these two flags, with opposite meanings? I don't have strong opinion on this, but I can update if you think they are much better...