[simrankaurb]

Summary

Lets a benchmark initialize libtpu as a standalone single host on one VM of a multi-host TPU slice, without touching the slice's ICI routing or contacting tpunetd. Enabled per benchmark via a new config key single_host_tpu_init: true. Also bumps JAX/libtpu to pick up --deepsea_hal_test_allow_multichip_skip_slicebuilder.

Changes

File	Change
requirements.txt	jax[tpu] / jaxlib 0.10.1 → 0.11.2, libtpu 0.0.41 → 0.0.48
Ironwood/src/benchmark_utils.py	(a) New configure_single_host_tpu(): sets TPU_HOST_BOUNDS=1,1,1, TPU_SKIP_MDS_QUERY=true, TPU_WORKER_ID=0, TPU_WORKER_HOSTNAMES=localhost; appends --deepsea_hal_test_skip_slicebuilder=true --deepsea_hal_test_allow_multichip_skip_slicebuilder=true --xla_tpu_use_enhanced_launch_barrier=false to LIBTPU_INIT_ARGS (existing flags matched by exact name). is_single_host_tpu_init() exposes the state to benchmark modules. (b) libtpu 0.0.48 HLO dump format: rename_xla_dump() skips subdirectories now present in the dump dir, and extract_hlo_features_from_file() parses the new replica_groups=mesh[...]{...} syntax in addition to the explicit {{...}} list form.
Ironwood/src/run_benchmark.py	Calls configure_single_host_tpu() when the benchmark config has single_host_tpu_init: true. Runs after the benchmark module import (which sets LIBTPU_INIT_ARGS) and before the first JAX device call.
Ironwood/src/benchmark_gemm.py	(a) When single_host_tpu_init is active, generate inputs independently on each local device (jax.make_array_from_single_device_arrays) instead of device_put from device 0: SliceBuilder is skipped, so intra-host ICI is unprogrammed and a cross-chip copy would fault. Default path unchanged otherwise. (b) --xla_tpu_enable_async_collective_fusion=false, drop _fuse_all_gather — libtpu 0.0.48 rejects =true on v6e and 7x (FAILED_PRECONDITION: Continuation fusion for AllGather … not supported on platforms other than Viperlite); aligned with op_flags.yaml on main.

Usage

benchmarks:
- benchmark_name: gemm_multiple_run
  single_host_tpu_init: true        # benchmark-level, not a sweep param
  benchmark_sweep_params:
  - {m: 16384, k: 16384, n: 16384, num_runs: 100, dtype: bfloat16, run_on_local_node: True}

single_host_tpu_init controls libtpu process init; run_on_local_node (existing) controls the mesh. Do not set single_host_tpu_init for multi-host benchmarks (e.g. collectives).

Validation (CHS, 16-node pools, PR head 9d68aa1)

Pool	Flag	Run	Report
v6e 8x8	OFF	v6e-off3	https://storage.cloud.google.com/simranka/bills_of_health/20261008-144431_v6e-off3-f4tr.html
v6e 8x8	ON	v6e-on3	https://storage.cloud.google.com/simranka/bills_of_health/20261008-145803_v6e-on3-itlu.html
7x 4x4x4	OFF	7x-off3	https://storage.cloud.google.com/simranka/bills_of_health/20261008-144915_7x-off3-oyx6.html
7x 4x4x4	ON	7x-on3	https://storage.cloud.google.com/simranka/bills_of_health/20261008-145823_7x-on3-qkkd.html

Notes

• Other modules that still set --xla_tpu_enable_async_collective_fusion=true (benchmark_compute.py, benchmark_gemm_numerics.py, benchmark_gemm_throttling.py, benchmark_inference_compute.py) will need the same change on libtpu 0.0.48; left out to keep this PR minimal.
• Background on why skipping SliceBuilder is required: go/rca-tpu-hanging, go/disabling-softslicing, b/517071651.

[gemini-code-assist]

Code Review

This pull request introduces support for running benchmarks on a local node (single host without slicebuilder) and updates JAX/libtpu dependencies. Key feedback points out that appending TPU initialization flags using substring checks on LIBTPU_INIT_ARGS can cause silent bugs, recommending exact flag name matching instead. Additionally, in benchmark_gemm.py, several issues were identified regarding local data initialization, including incorrect indexing of seed_arr, shape mismatches with per_dev_seeds, and non-standard usage of jax.block_until_ready on a tuple. Finally, is_local_node_benchmark in run_benchmark.py should be expanded to check both benchmark_params and benchmark_sweep_params to ensure local node configurations are not missed.

[linamy85]

Some side note:

• jax 0.11.2 has the jax.device_get device-to-host regression (chs: is not using d2h yet)
• jax 0.11.1 the create_mesh has changed and will impact parallel 2x4x4 topology. (chs: non-parallel atm)