Modified comprehensive_8dev_experiments.yaml from 20 num runs to 100

Ironwood/configs/attention/attention.yaml

@@ -0,0 +1,7 @@
+benchmarks:
+- benchmark_name: "tokamax_splash_attention"
+  benchmark_sweep_params:
+  - {batch_size: 1, q_seq_len: 4096, kv_seq_len: 4096, q_heads: 128, kv_heads: 128, qk_head_dim: 256, v_head_dim: 256, mode: ["fwd", "bwd"], causal: [true, false], num_runs: 20}
+  trace_dir: "../microbenchmarks/attention"
+  csv_path: "../microbenchmarks/attention"
+  xlml_metrics_dir: "../microbenchmarks/attention"

Ironwood/configs/bmm/multi_host_bmm.yaml

@@ -0,0 +1,75 @@
+benchmarks:
+- benchmark_name: "multi_host_bmm"
+  trace_dir: "../microbenchmarks/multi_host_bmm_bf16"
+  csv_path: "../microbenchmarks/multi_host_bmm_bf16"
+  xlml_metrics_dir: "../microbenchmarks/multi_host_bmm_bf16"
+  xla_dump_dir: "../microbenchmarks/multi_host_bmm_bf16/hlo_graphs"
+  benchmark_sweep_params:
+  - {b: 8, m: 128, k: 128, n: 128, num_runs: 100, dtype: 'bfloat16'}
+  - {b: 8, m: 256, k: 256, n: 256, num_runs: 100, dtype: 'bfloat16'}
+  - {b: 8, m: 512, k: 512, n: 512, num_runs: 100, dtype: 'bfloat16'}
+  - {b: 8, m: 1024, k: 1024, n: 1024, num_runs: 100, dtype: 'bfloat16'}
+  - {b: 8, m: 2048, k: 2048, n: 2048, num_runs: 100, dtype: 'bfloat16'}
+  - {b: 8, m: 4096, k: 4096, n: 4096, num_runs: 100, dtype: 'bfloat16'}
+  - {b: 8, m: 16384, k: 16384, n: 16384, num_runs: 100, dtype: 'bfloat16'}
+  - {b: 8, m: 32768, k: 32768, n: 32768, num_runs: 100, dtype: 'bfloat16'}
+
+- benchmark_name: "multi_host_bmm"
+  trace_dir: "../microbenchmarks/multi_host_bmm_f32"
+  csv_path: "../microbenchmarks/multi_host_bmm_f32"
+  xlml_metrics_dir: "../microbenchmarks/multi_host_bmm_f32"
+  xla_dump_dir: "../microbenchmarks/multi_host_bmm_f32/hlo_graphs"
+  benchmark_sweep_params:
+  - {b: 8, m: 128, k: 128, n: 128, num_runs: 100, dtype: 'float32'}
+  - {b: 8, m: 256, k: 256, n: 256, num_runs: 100, dtype: 'float32'}
+  - {b: 8, m: 512, k: 512, n: 512, num_runs: 100, dtype: 'float32'}
+  - {b: 8, m: 1024, k: 1024, n: 1024, num_runs: 100, dtype: 'float32'}
+  - {b: 8, m: 2048, k: 2048, n: 2048, num_runs: 100, dtype: 'float32'}
+  - {b: 8, m: 4096, k: 4096, n: 4096, num_runs: 100, dtype: 'float32'}
+  - {b: 8, m: 16384, k: 16384, n: 16384, num_runs: 100, dtype: 'float32'}
+  - {b: 8, m: 32768, k: 32768, n: 32768, num_runs: 100, dtype: 'float32'}
+
+- benchmark_name: "multi_host_bmm"
+  trace_dir: "../microbenchmarks/multi_host_bmm_fp16"
+  csv_path: "../microbenchmarks/multi_host_bmm_fp16"
+  xlml_metrics_dir: "../microbenchmarks/multi_host_bmm_fp16"
+  xla_dump_dir: "../microbenchmarks/multi_host_bmm_fp16/hlo_graphs"
+  benchmark_sweep_params:
+  - {b: 8, m: 128, k: 128, n: 128, num_runs: 100, dtype: 'float16'}
+  - {b: 8, m: 256, k: 256, n: 256, num_runs: 100, dtype: 'float16'}
+  - {b: 8, m: 512, k: 512, n: 512, num_runs: 100, dtype: 'float16'}
+  - {b: 8, m: 1024, k: 1024, n: 1024, num_runs: 100, dtype: 'float16'}
+  - {b: 8, m: 2048, k: 2048, n: 2048, num_runs: 100, dtype: 'float16'}
+  - {b: 8, m: 4096, k: 4096, n: 4096, num_runs: 100, dtype: 'float16'}
+  - {b: 8, m: 16384, k: 16384, n: 16384, num_runs: 100, dtype: 'float16'}
+  - {b: 8, m: 32768, k: 32768, n: 32768, num_runs: 100, dtype: 'float16'}
+
+- benchmark_name: "multi_host_bmm"
+  trace_dir: "../microbenchmarks/multi_host_bmm_fp8"
+  csv_path: "../microbenchmarks/multi_host_bmm_fp8"
+  xlml_metrics_dir: "../microbenchmarks/multi_host_bmm_fp8"
+  xla_dump_dir: "../microbenchmarks/multi_host_bmm_fp8/hlo_graphs"
+  benchmark_sweep_params:
+  - {b: 8, m: 128, k: 128, n: 128, num_runs: 100, dtype: 'float8'}
+  - {b: 8, m: 256, k: 256, n: 256, num_runs: 100, dtype: 'float8'}
+  - {b: 8, m: 512, k: 512, n: 512, num_runs: 100, dtype: 'float8'}
+  - {b: 8, m: 1024, k: 1024, n: 1024, num_runs: 100, dtype: 'float8'}
+  - {b: 8, m: 2048, k: 2048, n: 2048, num_runs: 100, dtype: 'float8'}
+  - {b: 8, m: 4096, k: 4096, n: 4096, num_runs: 100, dtype: 'float8'}
+  - {b: 8, m: 16384, k: 16384, n: 16384, num_runs: 100, dtype: 'float8'}
+  - {b: 8, m: 32768, k: 32768, n: 32768, num_runs: 100, dtype: 'float8'}
+
+- benchmark_name: "multi_host_bmm"
+  trace_dir: "../microbenchmarks/multi_host_bmm_fp4"
+  csv_path: "../microbenchmarks/multi_host_bmm_fp4"
+  xlml_metrics_dir: "../microbenchmarks/multi_host_bmm_fp4"
+  xla_dump_dir: "../microbenchmarks/multi_host_bmm_fp4/hlo_graphs"
+  benchmark_sweep_params:
+  - {b: 8, m: 128, k: 128, n: 128, num_runs: 100, dtype: 'float4'}
+  - {b: 8, m: 256, k: 256, n: 256, num_runs: 100, dtype: 'float4'}
+  - {b: 8, m: 512, k: 512, n: 512, num_runs: 100, dtype: 'float4'}
+  - {b: 8, m: 1024, k: 1024, n: 1024, num_runs: 100, dtype: 'float4'}
+  - {b: 8, m: 2048, k: 2048, n: 2048, num_runs: 100, dtype: 'float4'}
+  - {b: 8, m: 4096, k: 4096, n: 4096, num_runs: 100, dtype: 'float4'}
+  - {b: 8, m: 16384, k: 16384, n: 16384, num_runs: 100, dtype: 'float4'}
+  - {b: 8, m: 32768, k: 32768, n: 32768, num_runs: 100, dtype: 'float4'}

Ironwood/configs/bmm/single_device_bmm.yaml

@@ -0,0 +1,75 @@
+benchmarks:
+- benchmark_name: "single_device_bmm"
+  trace_dir: "../microbenchmarks/single_device_bmm_bf16"
+  csv_path: "../microbenchmarks/single_device_bmm_bf16"
+  xlml_metrics_dir: "../microbenchmarks/single_device_bmm_bf16"
+  xla_dump_dir: "../microbenchmarks/single_device_bmm_bf16/hlo_graphs"
+  benchmark_sweep_params:
+  - {b: 8, m: 128, k: 128, n: 128, num_runs: 100, dtype: 'bfloat16'}
+  - {b: 8, m: 256, k: 256, n: 256, num_runs: 100, dtype: 'bfloat16'}
+  - {b: 8, m: 512, k: 512, n: 512, num_runs: 100, dtype: 'bfloat16'}
+  - {b: 8, m: 1024, k: 1024, n: 1024, num_runs: 100, dtype: 'bfloat16'}
+  - {b: 8, m: 2048, k: 2048, n: 2048, num_runs: 100, dtype: 'bfloat16'}
+  - {b: 8, m: 4096, k: 4096, n: 4096, num_runs: 100, dtype: 'bfloat16'}
+  - {b: 8, m: 16384, k: 16384, n: 16384, num_runs: 100, dtype: 'bfloat16'}
+  - {b: 8, m: 32768, k: 32768, n: 32768, num_runs: 100, dtype: 'bfloat16'}
+
+- benchmark_name: "single_device_bmm"
+  trace_dir: "../microbenchmarks/single_device_bmm_f32"
+  csv_path: "../microbenchmarks/single_device_bmm_f32"
+  xlml_metrics_dir: "../microbenchmarks/single_device_bmm_f32"
+  xla_dump_dir: "../microbenchmarks/single_device_bmm_f32/hlo_graphs"
+  benchmark_sweep_params:
+  - {b: 8, m: 128, k: 128, n: 128, num_runs: 100, dtype: 'float32'}
+  - {b: 8, m: 256, k: 256, n: 256, num_runs: 100, dtype: 'float32'}
+  - {b: 8, m: 512, k: 512, n: 512, num_runs: 100, dtype: 'float32'}
+  - {b: 8, m: 1024, k: 1024, n: 1024, num_runs: 100, dtype: 'float32'}
+  - {b: 8, m: 2048, k: 2048, n: 2048, num_runs: 100, dtype: 'float32'}
+  - {b: 8, m: 4096, k: 4096, n: 4096, num_runs: 100, dtype: 'float32'}
+  - {b: 8, m: 16384, k: 16384, n: 16384, num_runs: 100, dtype: 'float32'}
+  - {b: 8, m: 32768, k: 32768, n: 32768, num_runs: 100, dtype: 'float32'}
+
+- benchmark_name: "single_device_bmm"
+  trace_dir: "../microbenchmarks/single_device_bmm_fp16"
+  csv_path: "../microbenchmarks/single_device_bmm_fp16"
+  xlml_metrics_dir: "../microbenchmarks/single_device_bmm_fp16"
+  xla_dump_dir: "../microbenchmarks/single_device_bmm_fp16/hlo_graphs"
+  benchmark_sweep_params:
+  - {b: 8, m: 128, k: 128, n: 128, num_runs: 100, dtype: 'float16'}
+  - {b: 8, m: 256, k: 256, n: 256, num_runs: 100, dtype: 'float16'}
+  - {b: 8, m: 512, k: 512, n: 512, num_runs: 100, dtype: 'float16'}
+  - {b: 8, m: 1024, k: 1024, n: 1024, num_runs: 100, dtype: 'float16'}
+  - {b: 8, m: 2048, k: 2048, n: 2048, num_runs: 100, dtype: 'float16'}
+  - {b: 8, m: 4096, k: 4096, n: 4096, num_runs: 100, dtype: 'float16'}
+  - {b: 8, m: 16384, k: 16384, n: 16384, num_runs: 100, dtype: 'float16'}
+  - {b: 8, m: 32768, k: 32768, n: 32768, num_runs: 100, dtype: 'float16'}
+
+- benchmark_name: "single_device_bmm"
+  trace_dir: "../microbenchmarks/single_device_bmm_fp8"
+  csv_path: "../microbenchmarks/single_device_bmm_fp8"
+  xlml_metrics_dir: "../microbenchmarks/single_device_bmm_fp8"
+  xla_dump_dir: "../microbenchmarks/single_device_bmm_fp8/hlo_graphs"
+  benchmark_sweep_params:
+  - {b: 8, m: 128, k: 128, n: 128, num_runs: 100, dtype: 'float8'}
+  - {b: 8, m: 256, k: 256, n: 256, num_runs: 100, dtype: 'float8'}
+  - {b: 8, m: 512, k: 512, n: 512, num_runs: 100, dtype: 'float8'}
+  - {b: 8, m: 1024, k: 1024, n: 1024, num_runs: 100, dtype: 'float8'}
+  - {b: 8, m: 2048, k: 2048, n: 2048, num_runs: 100, dtype: 'float8'}
+  - {b: 8, m: 4096, k: 4096, n: 4096, num_runs: 100, dtype: 'float8'}
+  - {b: 8, m: 16384, k: 16384, n: 16384, num_runs: 100, dtype: 'float8'}
+  - {b: 8, m: 32768, k: 32768, n: 32768, num_runs: 100, dtype: 'float8'}
+
+- benchmark_name: "single_device_bmm"
+  trace_dir: "../microbenchmarks/single_device_bmm_fp4"
+  csv_path: "../microbenchmarks/single_device_bmm_fp4"
+  xlml_metrics_dir: "../microbenchmarks/single_device_bmm_fp4"
+  xla_dump_dir: "../microbenchmarks/single_device_bmm_fp4/hlo_graphs"
+  benchmark_sweep_params:
+  - {b: 8, m: 128, k: 128, n: 128, num_runs: 100, dtype: 'float4'}
+  - {b: 8, m: 256, k: 256, n: 256, num_runs: 100, dtype: 'float4'}
+  - {b: 8, m: 512, k: 512, n: 512, num_runs: 100, dtype: 'float4'}
+  - {b: 8, m: 1024, k: 1024, n: 1024, num_runs: 100, dtype: 'float4'}
+  - {b: 8, m: 2048, k: 2048, n: 2048, num_runs: 100, dtype: 'float4'}
+  - {b: 8, m: 4096, k: 4096, n: 4096, num_runs: 100, dtype: 'float4'}
+  - {b: 8, m: 16384, k: 16384, n: 16384, num_runs: 100, dtype: 'float4'}
+  - {b: 8, m: 32768, k: 32768, n: 32768, num_runs: 100, dtype: 'float4'}

Ironwood/configs/collectives/all_gather_tpu7x_2x2x1.yaml

@@ -1,7 +1,9 @@
 benchmarks:
 - benchmark_name: all_gather
   benchmark_sweep_params:
-  - {matrix_dim_range: {start: 8, end: 16384, multiplier: 2}, dtype: "float32",  mesh_shape: "2x2x2", ici_size_range: 8, sharding_strategy: "2x2x1", op_dimension: 1, num_runs: 5}
+  - {matrix_dim_range: {start: 64, end: 2097152, multiplier: 2}, dtype: "float32",  mesh_shape: "2x2x2", ici_size_range: 8, sharding_strategy: "2x2x1", op_dimension: 1, num_runs: 20}
+  - {matrix_dim_range: {start: 64, end: 2097152, multiplier: 2}, dtype: "bfloat16",  mesh_shape: "2x2x2", ici_size_range: 8, sharding_strategy: "2x2x1", op_dimension: 1, num_runs: 20}
+  - {matrix_dim_range: {start: 64, end: 2097152, multiplier: 2}, dtype: "float8",  mesh_shape: "2x2x2", ici_size_range: 8, sharding_strategy: "2x2x1", op_dimension: 1, num_runs: 20}
   trace_dir: "../microbenchmarks/all_gather_tpu7x_2x2x1"
   csv_path: "../microbenchmarks/all_gather_tpu7x_2x2x1"
   xlml_metrics_dir: "../microbenchmarks/all_gather_tpu7x_2x2x1"

Ironwood/configs/collectives/all_gather_tpu7x_2x2x2.yaml

@@ -1,7 +1,9 @@
 benchmarks:
 - benchmark_name: all_gather
   benchmark_sweep_params:
-  - {matrix_dim_range: {start: 8, end: 16384, multiplier: 2}, dtype: "float32",  mesh_shape: "4x2x2", ici_size_range: 16, sharding_strategy: "4x2x1", op_dimension: 1, num_runs: 5}
+  - {matrix_dim_range: {start: 64, end: 2097152, multiplier: 2}, dtype: "float32",  mesh_shape: "4x2x2", ici_size_range: 16, sharding_strategy: "4x2x1", op_dimension: 1, num_runs: 20}
+  - {matrix_dim_range: {start: 64, end: 2097152, multiplier: 2}, dtype: "bfloat16",  mesh_shape: "4x2x2", ici_size_range: 16, sharding_strategy: "4x2x1", op_dimension: 1, num_runs: 20}
+  - {matrix_dim_range: {start: 64, end: 2097152, multiplier: 2}, dtype: "float8",  mesh_shape: "4x2x2", ici_size_range: 16, sharding_strategy: "4x2x1", op_dimension: 1, num_runs: 20}
   trace_dir: "../microbenchmarks/all_gather_tpu7x_2x2x2"
   csv_path: "../microbenchmarks/all_gather_tpu7x_2x2x2"
   xlml_metrics_dir: "../microbenchmarks/all_gather_tpu7x_2x2x2"

Ironwood/configs/collectives/all_gather_tpu7x_2x2x4.yaml

@@ -1,7 +1,9 @@
 benchmarks:
 - benchmark_name: all_gather
   benchmark_sweep_params:
-  - {matrix_dim_range: {start: 16, end: 16384, multiplier: 2}, dtype: "float32",  mesh_shape: "4x4x2", ici_size_range: 32, sharding_strategy: "4x4x1", op_dimension: 1, num_runs: 5}
+  - {matrix_dim_range: {start: 64, end: 2097152, multiplier: 2}, dtype: "float32",  mesh_shape: "4x4x2", ici_size_range: 32, sharding_strategy: "4x4x1", op_dimension: 1, num_runs: 20}
+  - {matrix_dim_range: {start: 64, end: 2097152, multiplier: 2}, dtype: "bfloat16",  mesh_shape: "4x4x2", ici_size_range: 32, sharding_strategy: "4x4x1", op_dimension: 1, num_runs: 20}
+  - {matrix_dim_range: {start: 64, end: 2097152, multiplier: 2}, dtype: "float8",  mesh_shape: "4x4x2", ici_size_range: 32, sharding_strategy: "4x4x1", op_dimension: 1, num_runs: 20}
   trace_dir: "../microbenchmarks/all_gather_tpu7x_2x2x4"
   csv_path: "../microbenchmarks/all_gather_tpu7x_2x2x4"
   xlml_metrics_dir: "../microbenchmarks/all_gather_tpu7x_2x2x4"

Ironwood/configs/collectives/all_gather_tpu7x_2x4x4.yaml

@@ -1,7 +1,9 @@
 benchmarks:
 - benchmark_name: all_gather
   benchmark_sweep_params:
-  - {matrix_dim_range: {start: 32, end: 16384, multiplier: 2}, dtype: "float32",  mesh_shape: "8x4x2", ici_size_range: 64, sharding_strategy: "8x4x1", op_dimension: 1, num_runs: 5}
+  - {matrix_dim_range: {start: 64, end: 2097152, multiplier: 2}, dtype: "float32",  mesh_shape: "8x4x2", ici_size_range: 64, sharding_strategy: "8x4x1", op_dimension: 1, num_runs: 20}
+  - {matrix_dim_range: {start: 64, end: 2097152, multiplier: 2}, dtype: "bfloat16",  mesh_shape: "8x4x2", ici_size_range: 64, sharding_strategy: "8x4x1", op_dimension: 1, num_runs: 20}
+  - {matrix_dim_range: {start: 64, end: 2097152, multiplier: 2}, dtype: "float8",  mesh_shape: "8x4x2", ici_size_range: 64, sharding_strategy: "8x4x1", op_dimension: 1, num_runs: 20}
   trace_dir: "../microbenchmarks/all_gather_tpu7x_2x4x4"
   csv_path: "../microbenchmarks/all_gather_tpu7x_2x4x4"
   xlml_metrics_dir: "../microbenchmarks/all_gather_tpu7x_2x4x4"

Ironwood/configs/collectives/all_gather_tpu7x_4x4x4.yaml

@@ -1,7 +1,9 @@
 benchmarks:
 - benchmark_name: all_gather
   benchmark_sweep_params:
-  - {matrix_dim_range: {start: 64, end: 16384, multiplier: 2}, dtype: "float32",  mesh_shape: "16x4x2", ici_size_range: 128, sharding_strategy: "16x4x1", op_dimension: 1, num_runs: 10}
+  - {matrix_dim_range: {start: 64, end: 16384, multiplier: 2}, dtype: "float32",  mesh_shape: "16x4x2", ici_size_range: 128, sharding_strategy: "16x4x1", op_dimension: 1, num_runs: 20}
+  - {matrix_dim_range: {start: 64, end: 16384, multiplier: 2}, dtype: "bfloat16",  mesh_shape: "16x4x2", ici_size_range: 128, sharding_strategy: "16x4x1", op_dimension: 1, num_runs: 20}
+  - {matrix_dim_range: {start: 64, end: 16384, multiplier: 2}, dtype: "float8",  mesh_shape: "16x4x2", ici_size_range: 128, sharding_strategy: "16x4x1", op_dimension: 1, num_runs: 20}
   trace_dir: "../microbenchmarks/all_gather_tpu7x_4x4x4"
   csv_path: "../microbenchmarks/all_gather_tpu7x_4x4x4"
   xlml_metrics_dir: "../microbenchmarks/all_gather_tpu7x_4x4x4"

Ironwood/configs/collectives/all_gather_tpu7x_4x4x8.yaml

@@ -1,7 +1,9 @@
 benchmarks:
 - benchmark_name: all_gather
   benchmark_sweep_params:
-  - {matrix_dim_range: {start: 128, end: 16384, multiplier: 2}, dtype: "float32",  mesh_shape: "32x4x2", ici_size_range: 256, sharding_strategy: "32x4x1", op_dimension: 1, num_runs: 5}
+  - {matrix_dim_range: {start: 64, end: 2097152, multiplier: 2}, dtype: "float32",  mesh_shape: "32x4x2", ici_size_range: 256, sharding_strategy: "32x4x1", op_dimension: 1, num_runs: 20}
+  - {matrix_dim_range: {start: 64, end: 2097152, multiplier: 2}, dtype: "bfloat16",  mesh_shape: "32x4x2", ici_size_range: 256, sharding_strategy: "32x4x1", op_dimension: 1, num_runs: 20}
+  - {matrix_dim_range: {start: 64, end: 2097152, multiplier: 2}, dtype: "float8",  mesh_shape: "32x4x2", ici_size_range: 256, sharding_strategy: "32x4x1", op_dimension: 1, num_runs: 20}
   trace_dir: "../microbenchmarks/all_gather_tpu7x_4x4x8"
   csv_path: "../microbenchmarks/all_gather_tpu7x_4x4x8"
   xlml_metrics_dir: "../microbenchmarks/all_gather_tpu7x_4x4x8"

Ironwood/configs/collectives/all_reduce_tpu7x_2x2x1.yaml

@@ -1,7 +1,9 @@
 benchmarks:
 - benchmark_name: psum
   benchmark_sweep_params:
-  - {matrix_dim_range: {start: 8, end: 16384, multiplier: 2}, dtype: "float32",  mesh_shape: "2x2x2", ici_size_range: 8, sharding_strategy: "2x2x1", op_dimension: 1, num_runs: 5}
+  - {matrix_dim_range: {start: 64, end: 2097152, multiplier: 2}, dtype: "float32",  mesh_shape: "2x2x2", ici_size_range: 8, sharding_strategy: "2x2x1", op_dimension: 1, num_runs: 20}
+  - {matrix_dim_range: {start: 64, end: 2097152, multiplier: 2}, dtype: "bfloat16",  mesh_shape: "2x2x2", ici_size_range: 8, sharding_strategy: "2x2x1", op_dimension: 1, num_runs: 20}
+  - {matrix_dim_range: {start: 64, end: 2097152, multiplier: 2}, dtype: "float8",  mesh_shape: "2x2x2", ici_size_range: 8, sharding_strategy: "2x2x1", op_dimension: 1, num_runs: 20}
   trace_dir: "../microbenchmarks/psum_tpu7x_2x2x1"
   csv_path: "../microbenchmarks/psum_tpu7x_2x2x1"
   xlml_metrics_dir: "../microbenchmarks/psum_tpu7x_2x2x1"

Ironwood/configs/collectives/all_reduce_tpu7x_2x2x2.yaml

@@ -1,7 +1,9 @@
 benchmarks:
 - benchmark_name: psum
   benchmark_sweep_params:
-  - {matrix_dim_range: {start: 8, end: 16384, multiplier: 2}, dtype: "float32",  mesh_shape: "4x2x2", ici_size_range: 16, sharding_strategy: "4x2x1", op_dimension: 1, num_runs: 5}
+  - {matrix_dim_range: {start: 64, end: 2097152, multiplier: 2}, dtype: "float32",  mesh_shape: "4x2x2", ici_size_range: 16, sharding_strategy: "4x2x1", op_dimension: 1, num_runs: 20}
+  - {matrix_dim_range: {start: 64, end: 2097152, multiplier: 2}, dtype: "bfloat16",  mesh_shape: "4x2x2", ici_size_range: 16, sharding_strategy: "4x2x1", op_dimension: 1, num_runs: 20}
+  - {matrix_dim_range: {start: 64, end: 2097152, multiplier: 2}, dtype: "float8",  mesh_shape: "4x2x2", ici_size_range: 16, sharding_strategy: "4x2x1", op_dimension: 1, num_runs: 20}
   trace_dir: "../microbenchmarks/psum_tpu7x_2x2x2"
   csv_path: "../microbenchmarks/psum_tpu7x_2x2x2"
   xlml_metrics_dir: "../microbenchmarks/psum_tpu7x_2x2x2"

Ironwood/configs/collectives/all_reduce_tpu7x_2x2x4.yaml

@@ -1,7 +1,9 @@
 benchmarks:
 - benchmark_name: psum
   benchmark_sweep_params:
-  - {matrix_dim_range: {start: 16, end: 16384, multiplier: 2}, dtype: "float32",  mesh_shape: "4x4x2", ici_size_range: 32, sharding_strategy: "4x4x1", op_dimension: 1, num_runs: 5}
+  - {matrix_dim_range: {start: 64, end: 2097152, multiplier: 2}, dtype: "float32",  mesh_shape: "4x4x2", ici_size_range: 32, sharding_strategy: "4x4x1", op_dimension: 1, num_runs: 20}
+  - {matrix_dim_range: {start: 64, end: 2097152, multiplier: 2}, dtype: "bfloat16",  mesh_shape: "4x4x2", ici_size_range: 32, sharding_strategy: "4x4x1", op_dimension: 1, num_runs: 20}
+  - {matrix_dim_range: {start: 64, end: 2097152, multiplier: 2}, dtype: "float8",  mesh_shape: "4x4x2", ici_size_range: 32, sharding_strategy: "4x4x1", op_dimension: 1, num_runs: 20}
   trace_dir: "../microbenchmarks/psum_tpu7x_2x2x4"
   csv_path: "../microbenchmarks/psum_tpu7x_2x2x4"
   xlml_metrics_dir: "../microbenchmarks/psum_tpu7x_2x2x4"

Ironwood/configs/collectives/all_reduce_tpu7x_2x4x4.yaml

@@ -1,7 +1,9 @@
 benchmarks:
 - benchmark_name: psum
   benchmark_sweep_params:
-  - {matrix_dim_range: {start: 32, end: 16384, multiplier: 2}, dtype: "float32",  mesh_shape: "8x4x2", ici_size_range: 64, sharding_strategy: "8x4x1", op_dimension: 1, num_runs: 5}
+  - {matrix_dim_range: {start: 64, end: 2097152, multiplier: 2}, dtype: "float32",  mesh_shape: "8x4x2", ici_size_range: 64, sharding_strategy: "8x4x1", op_dimension: 1, num_runs: 20}
+  - {matrix_dim_range: {start: 64, end: 2097152, multiplier: 2}, dtype: "bfloat16",  mesh_shape: "8x4x2", ici_size_range: 64, sharding_strategy: "8x4x1", op_dimension: 1, num_runs: 20}
+  - {matrix_dim_range: {start: 64, end: 2097152, multiplier: 2}, dtype: "float8",  mesh_shape: "8x4x2", ici_size_range: 64, sharding_strategy: "8x4x1", op_dimension: 1, num_runs: 20}
   trace_dir: "../microbenchmarks/psum_tpu7x_2x4x4"
   csv_path: "../microbenchmarks/psum_tpu7x_2x4x4"
   xlml_metrics_dir: "../microbenchmarks/psum_tpu7x_2x4x4"

Ironwood/configs/collectives/all_reduce_tpu7x_4x4x4.yaml

@@ -1,7 +1,9 @@
 benchmarks:
 - benchmark_name: psum
   benchmark_sweep_params:
-  - {matrix_dim_range: {start: 64, end: 16384, multiplier: 2}, dtype: "float32",  mesh_shape: "16x4x2", ici_size_range: 128, sharding_strategy: "16x4x1", op_dimension: 1, num_runs: 5}
+  - {matrix_dim_range: {start: 64, end: 2097152, multiplier: 2}, dtype: "float32",  mesh_shape: "16x4x2", ici_size_range: 128, sharding_strategy: "16x4x1", op_dimension: 1, num_runs: 20}
+  - {matrix_dim_range: {start: 64, end: 2097152, multiplier: 2}, dtype: "bfloat16",  mesh_shape: "16x4x2", ici_size_range: 128, sharding_strategy: "16x4x1", op_dimension: 1, num_runs: 20}
+  - {matrix_dim_range: {start: 64, end: 2097152, multiplier: 2}, dtype: "float8",  mesh_shape: "16x4x2", ici_size_range: 128, sharding_strategy: "16x4x1", op_dimension: 1, num_runs: 20}
   trace_dir: "../microbenchmarks/psum_tpu7x_4x4x4"
   csv_path: "../microbenchmarks/psum_tpu7x_4x4x4"
   xlml_metrics_dir: "../microbenchmarks/psum_tpu7x_4x4x4"