feat(throttling): add independent all-devices GEMM throttling benchmark & Max-TS telemetry

[linamy85]

Description

1. Summary

This PR introduces multi-device independent GEMM execution under the data_gen_once_noblock_stressed_no_sharding gap strategy to stress all on-host TPU cores continuously without sharding/collective communication barriers. It also adds a dedicated multi-device Max-TS metric extractor for accurate step timing and throughput aggregation across all active devices.

---

2. Key Changes

• Ironwood/src/benchmark_gemm_throttling.py:
  • Added data_gen_once_noblock_stressed_no_sharding strategy using ThreadPoolExecutor to launch un-sharded GEMMs on all local devices concurrently.
  • Normalized dtype handling for BF16 peak FLOPS calculation.
• Ironwood/src/benchmark_utils.py:
  • Added multiple_iteration_get_metrics_from_multithread_trace_throttling to extract per-step max(duration) across all TensorCores with hierarchical event filtering.
  • Preserved upstream multiple_iteration_get_metrics_from_trace and timing logic intact.
• Configs:
  • Added Ironwood/configs/throttling/throttling_all_devices_16k.yaml

---

3. Benchmark Telemetry & Physical Verification (1000 Iterations)

• Configuration: $M=N=K=16384$, dtype=bfloat16, $1000$ continuous iterations, 8 on-host TPU cores under stressed non-blocking execution.
• TPU v7x (Ironwood) — 1000 Iterations: https://screenshot-v2.corp.google.com/5pe4p7i7fch00
• TPU v6e (Viperlite) — 1000 Iterations: https://screenshot-v2.corp.google.com/0uhrcl5nlljd0

[gemini-code-assist]

Code Review

This pull request introduces a new multi-threaded gap strategy (data_gen_once_noblock_stressed_no_sharding) for the GEMM throttling benchmark, allowing execution across all devices without sharding, and adds a corresponding configuration file. The review feedback highlights several critical issues: a potential runtime crash when parsing float8 dtypes for peak FLOPS calculation, a shape mismatch bug in the multi-threaded run_device loop due to input chaining, a ValueError risk from changing the default gap_strategy to None, and the need to sort trace events chronologically while handling empty traces gracefully.

[gemini-code-assist]

The default dtype is jax.numpy.float8_e4m3fn. When converting this dtype to a string via jnp.dtype(dtype).name, it evaluates to "float8_e4m3fn". However, get_peak_flops_multiplier only supports "fp8", "bf16", "fp16", and "fp32", and will raise a RuntimeError for "float8_e4m3fn". We should map any float8-based dtypes to "fp8" to prevent runtime crashes during metric calculation.

    d_str = (dtype if isinstance(dtype, str) else jnp.dtype(dtype).name).lower()
    if "bf16" in d_str or "bfloat16" in d_str:
        in_dtype_str = "bf16"
    elif "float8" in d_str or "f8" in d_str:
        in_dtype_str = "fp8"
    else:
        in_dtype_str = d_str

[gemini-code-assist]

In run_device, the computation is chained as latest = compute_func(latest, rhs). Since compute_func performs a GEMM of shape (m, k) x (k, n) resulting in (m, n), passing latest (which now has shape (m, n)) back into the first argument of compute_func in the next iteration will cause a shape mismatch error if n != k. To support arbitrary dimensions, we should keep the inputs constant across iterations by using lhs and rhs directly.

Suggested change

	def run_device(dev_data):
	latest, rhs = dev_data[0], dev_data[1]
	for i in range(tries):
	with jax.profiler.StepTraceAnnotation(task, step_num=i):
	with jax.named_scope(f"{MARKER}_{i}"):
	latest = compute_func(latest, rhs)
	if latest is not None:
	jax.block_until_ready(latest)
	def run_device(dev_data):
	lhs, rhs = dev_data[0], dev_data[1]
	latest = None
	for i in range(tries):
	with jax.profiler.StepTraceAnnotation(task, step_num=i):
	with jax.named_scope(f"{MARKER}_{i}"):
	latest = compute_func(lhs, rhs)
	if latest is not None:
	jax.block_until_ready(latest)

[gemini-code-assist]

Changing the default value of gap_strategy to None will cause a ValueError at runtime if the caller does not explicitly provide a strategy, because multiple_iteration_timeit_from_trace_throttling does not handle None and raises an exception for unknown strategies. It is safer to keep the default as "data_gen_every_iter_block_every_iter".

Suggested change

	gap_strategy: str = None,
	gap_strategy: str = "data_gen_every_iter_block_every_iter",

[gemini-code-assist]

To ensure that the events are chronologically aligned across different PIDs (devices) and steps, we should sort marker_done_events by their timestamp (ts) before grouping them. Additionally, if no TPU execution events or markers are found in the trace, we should raise a clear RuntimeError rather than letting the code proceed and fail with a cryptic KeyError when calculating statistics.

Suggested change

	unique_pids = sorted(list(set([e["pid"] for e in marker_done_events])))
	print(f"TPU TensorCore PIDs found: {unique_pids}")
	marker_done_events.sort(key=lambda e: e.get("ts", 0))
	unique_pids = sorted(list(set([e["pid"] for e in marker_done_events])))
	if not unique_pids:
	raise RuntimeError(
	"No TPU execution events or markers were found in the trace. "
	"Please check if the profiling was successful and the trace is not empty."
	)
	print(f"TPU TensorCore PIDs found: {unique_pids}")