mesh: add mesh_ring_axis to give a mesh axis physical rings on a torus without wraparound

[systalyze-ai]

Description

Adds mesh_ring_axis to arrange an even-sized logical mesh axis into physical rings built from k-by-2 torus blocks. This provides a physical closing link when the axis is mapped to a slice dimension without wraparound.

The change affects device placement only. AOT HLO was byte-identical by md5 with the option disabled and enabled.

Performance

• On v6e-128 with DeepSeek-V4-Flash 284B LoRA at seq 16384, enabling the option changed step time from 4.632 to 4.138 s/step (1.12x).

Setup: ici_expert_parallelism=16, ici_tensor_parallelism=8, ici_fsdp_parallelism=1, max_target_length=16384, per_device_batch_size=0.125 (global batch 16), LIBTPU_INIT_ARGS=--xla_tpu_scoped_vmem_limit_kib=98304, dynamic-splash attention path enabled. mesh_ring_axis=tensor places the 8-way tensor axis on physical rings.

Tests

• The command python3 -m pytest tests/unit/max_utils_test.py -k RingAxisDeviceMeshTest passed 3 tests.

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.

[gemini-code-assist]

Code Review

This pull request introduces the mesh_ring_axis configuration to lay out a logical mesh axis as a physical TPU ring for single-slice setups, implementing the create_ring_axis_device_mesh utility and corresponding unit tests. The review feedback highlights a potential issue in create_ring_axis_device_mesh where the hardcoded transposition check can cause a ValueError on certain valid TPU topologies (such as a 2x16 grid with ring=8), and suggests a more robust, dynamic tiling check.

[gemini-code-assist]

The current implementation of create_ring_axis_device_mesh only transposes the physical grid if physical.shape[0] > physical.shape[1]. However, this hardcoded check is too restrictive and can cause a ValueError on valid TPU topologies (such as a 2x16 grid with ring=8, where ring_height=4 and ring_width=2). In this case, the height of the block (4) is greater than the height of the grid (2), so tiling fails in the default orientation, but would succeed if the grid were transposed to 16x2.

To make the mesh creation more robust, we should dynamically check if transposing the physical grid is required to make it divisible by ring_height and ring_width.

  ring_height, ring_width = ring // 2, 2
  # Determine if transposing is needed/preferred to allow valid tiling.
  can_tile_normal = (physical.shape[0] % ring_height == 0) and (physical.shape[1] % ring_width == 0)
  can_tile_transposed = (physical.shape[1] % ring_height == 0) and (physical.shape[0] % ring_width == 0)
  if not can_tile_normal and can_tile_transposed:
    physical = physical.T
  elif can_tile_normal and can_tile_transposed and physical.shape[0] > physical.shape[1]:
    physical = physical.T