{
 "device": {
  "jax": "0.11.2",
  "backend": "gpu",
  "devices": [
   "cuda:0"
  ],
  "python": "3.14.8",
  "device_kind": "NVIDIA GeForce RTX 3090",
  "gpu": "NVIDIA GeForce RTX 3090",
  "driver": "615.78.08",
  "gpu_memory_gb": 24.0,
  "torch": "2.14.1+cu130",
  "torch_cuda": "13.0"
 },
 "config": {
  "sources": [
   {
    "file": "ab6_base_1.json",
    "code": "before",
    "note": "A/B rep 1, sorted pallas, base code (gpu-benchmarks)",
    "config": {
     "device": "gpu",
     "suites": [
      "layer"
     ],
     "dtypes": [
      "float32",
      "bfloat16"
     ],
     "tokens": [
      4096,
      32768,
      131072
     ],
     "experts": [
      8
     ],
     "top_k": [
      2
     ],
     "skew": [
      0.0
     ],
     "d_model": 512,
     "d_hidden": 1024,
     "impls": [
      "dropless:pallas:sorted"
     ],
     "block_m": 64,
     "block_ms": [
      32,
      64,
      128
     ],
     "chunks": [
      128,
      256,
      512,
      1024
     ],
     "dense_limit": 524288,
     "ragged_limit": 262144,
     "maxtok_start": 32768,
     "maxtok_stop": 2097152,
     "check": false
    }
   },
   {
    "file": "ab6_v6_1.json",
    "code": "final",
    "note": "A/B rep 1, sorted pallas, final code",
    "config": {
     "device": "gpu",
     "suites": [
      "layer"
     ],
     "dtypes": [
      "float32",
      "bfloat16"
     ],
     "tokens": [
      4096,
      32768,
      131072
     ],
     "experts": [
      8
     ],
     "top_k": [
      2
     ],
     "skew": [
      0.0
     ],
     "d_model": 512,
     "d_hidden": 1024,
     "impls": [
      "dropless:pallas:sorted"
     ],
     "block_m": 64,
     "block_ms": [
      32,
      64,
      128
     ],
     "chunks": [
      128,
      256,
      512,
      1024
     ],
     "dense_limit": 524288,
     "ragged_limit": 262144,
     "maxtok_start": 32768,
     "maxtok_stop": 2097152,
     "slab_elems": null,
     "check": false
    }
   },
   {
    "file": "ab6_base_2.json",
    "code": "before",
    "note": "A/B rep 2, base code",
    "config": {
     "device": "gpu",
     "suites": [
      "layer"
     ],
     "dtypes": [
      "float32",
      "bfloat16"
     ],
     "tokens": [
      4096,
      32768,
      131072
     ],
     "experts": [
      8
     ],
     "top_k": [
      2
     ],
     "skew": [
      0.0
     ],
     "d_model": 512,
     "d_hidden": 1024,
     "impls": [
      "dropless:pallas:sorted"
     ],
     "block_m": 64,
     "block_ms": [
      32,
      64,
      128
     ],
     "chunks": [
      128,
      256,
      512,
      1024
     ],
     "dense_limit": 524288,
     "ragged_limit": 262144,
     "maxtok_start": 32768,
     "maxtok_stop": 2097152,
     "check": false
    }
   },
   {
    "file": "ab6_v6_2.json",
    "code": "final",
    "note": "A/B rep 2, final code",
    "config": {
     "device": "gpu",
     "suites": [
      "layer"
     ],
     "dtypes": [
      "float32",
      "bfloat16"
     ],
     "tokens": [
      4096,
      32768,
      131072
     ],
     "experts": [
      8
     ],
     "top_k": [
      2
     ],
     "skew": [
      0.0
     ],
     "d_model": 512,
     "d_hidden": 1024,
     "impls": [
      "dropless:pallas:sorted"
     ],
     "block_m": 64,
     "block_ms": [
      32,
      64,
      128
     ],
     "chunks": [
      128,
      256,
      512,
      1024
     ],
     "dense_limit": 524288,
     "ragged_limit": 262144,
     "maxtok_start": 32768,
     "maxtok_stop": 2097152,
     "slab_elems": null,
     "check": false
    }
   },
   {
    "file": "f_gpu_a.json",
    "code": "final",
    "note": "all impls/dispatches 4k/32k fp32+bf16 with numerics vs fp32 dense (GPU shared: some 4k fp32 rows contended, see IQR)",
    "config": {
     "device": "gpu",
     "suites": [
      "layer"
     ],
     "dtypes": [
      "float32",
      "bfloat16"
     ],
     "tokens": [
      4096,
      32768
     ],
     "experts": [
      8
     ],
     "top_k": [
      2
     ],
     "skew": [
      0.0
     ],
     "d_model": 512,
     "d_hidden": 1024,
     "impls": null,
     "block_m": 64,
     "block_ms": [
      32,
      64,
      128
     ],
     "chunks": [
      128,
      256,
      512,
      1024
     ],
     "dense_limit": 524288,
     "ragged_limit": 262144,
     "maxtok_start": 32768,
     "maxtok_stop": 2097152,
     "slab_elems": null,
     "check": true
    }
   },
   {
    "file": "f_gpu_b.json",
    "code": "final",
    "note": "131k tokens",
    "config": {
     "device": "gpu",
     "suites": [
      "layer"
     ],
     "dtypes": [
      "float32",
      "bfloat16"
     ],
     "tokens": [
      131072
     ],
     "experts": [
      8
     ],
     "top_k": [
      2
     ],
     "skew": [
      0.0
     ],
     "d_model": 512,
     "d_hidden": 1024,
     "impls": [
      "dropless:pallas:sorted",
      "dropless:pallas:fused_gsum",
      "dropless:blocked:sorted",
      "capacity:auto:sorted"
     ],
     "block_m": 64,
     "block_ms": [
      32,
      64,
      128
     ],
     "chunks": [
      128,
      256,
      512,
      1024
     ],
     "dense_limit": 524288,
     "ragged_limit": 262144,
     "maxtok_start": 32768,
     "maxtok_stop": 2097152,
     "slab_elems": null,
     "check": false
    }
   },
   {
    "file": "f_gpu_c.json",
    "code": "final",
    "note": "64 experts top-4 d_hidden 256 bf16",
    "config": {
     "device": "gpu",
     "suites": [
      "layer"
     ],
     "dtypes": [
      "bfloat16"
     ],
     "tokens": [
      32768
     ],
     "experts": [
      64
     ],
     "top_k": [
      4
     ],
     "skew": [
      0.0
     ],
     "d_model": 512,
     "d_hidden": 256,
     "impls": null,
     "block_m": 64,
     "block_ms": [
      32,
      64,
      128
     ],
     "chunks": [
      128,
      256,
      512,
      1024
     ],
     "dense_limit": 524288,
     "ragged_limit": 262144,
     "maxtok_start": 32768,
     "maxtok_stop": 2097152,
     "slab_elems": null,
     "check": false
    }
   },
   {
    "file": "f_gpu_pieces.json",
    "code": "final",
    "note": "train-step cost split of the sorted path",
    "config": {
     "device": "gpu",
     "suites": [
      "pieces"
     ],
     "dtypes": [
      "float32",
      "bfloat16"
     ],
     "tokens": [
      4096,
      32768
     ],
     "experts": [
      8
     ],
     "top_k": [
      2
     ],
     "skew": [
      0.0
     ],
     "d_model": 512,
     "d_hidden": 1024,
     "impls": null,
     "block_m": 64,
     "block_ms": [
      32,
      64,
      128
     ],
     "chunks": [
      128,
      256,
      512,
      1024
     ],
     "dense_limit": 524288,
     "ragged_limit": 262144,
     "maxtok_start": 32768,
     "maxtok_stop": 2097152,
     "slab_elems": null,
     "check": true
    }
   },
   {
    "file": "f_gpu_maxtok.json",
    "code": "final",
    "note": "largest bf16 train step (peak_bytes_in_use, fresh subprocess)",
    "config": {
     "device": "gpu",
     "suites": [
      "maxtok"
     ],
     "dtypes": [
      "bfloat16"
     ],
     "tokens": [
      4096
     ],
     "experts": [
      8
     ],
     "top_k": [
      2
     ],
     "skew": [
      0.0
     ],
     "d_model": 512,
     "d_hidden": 1024,
     "impls": [
      "dropless:pallas:sorted",
      "dropless:pallas:fused_gsum"
     ],
     "block_m": 64,
     "block_ms": [
      32,
      64,
      128
     ],
     "chunks": [
      128,
      256,
      512,
      1024
     ],
     "dense_limit": 524288,
     "ragged_limit": 262144,
     "maxtok_start": 131072,
     "maxtok_stop": 2097152,
     "slab_elems": null,
     "check": true
    }
   },
   {
    "file": "gpu_maxtok_base.json",
    "code": "before",
    "note": "largest bf16 train step, base code",
    "config": {
     "device": "gpu",
     "suites": [
      "maxtok"
     ],
     "dtypes": [
      "bfloat16"
     ],
     "tokens": [
      4096
     ],
     "experts": [
      8
     ],
     "top_k": [
      2
     ],
     "skew": [
      0.0
     ],
     "d_model": 512,
     "d_hidden": 1024,
     "impls": [
      "dropless:pallas:sorted",
      "dropless:pallas:fused",
      "capacity:auto:sorted"
     ],
     "block_m": 64,
     "block_ms": [
      32,
      64,
      128
     ],
     "chunks": [
      128,
      256,
      512,
      1024
     ],
     "dense_limit": 524288,
     "ragged_limit": 262144,
     "maxtok_start": 131072,
     "maxtok_stop": 2097152,
     "check": true
    }
   },
   {
    "file": "gpu_maxtok_other.json",
    "code": "final",
    "note": "largest bf16 train step, blocked/ragged_dot/dense (code unchanged by this lane; GPU shared with other jobs)",
    "config": {
     "device": "gpu",
     "suites": [
      "maxtok"
     ],
     "dtypes": [
      "bfloat16"
     ],
     "tokens": [
      4096
     ],
     "experts": [
      8
     ],
     "top_k": [
      2
     ],
     "skew": [
      0.0
     ],
     "d_model": 512,
     "d_hidden": 1024,
     "impls": [
      "dropless:blocked:sorted",
      "dropless:ragged_dot:sorted",
      "dense:auto:sorted"
     ],
     "block_m": 64,
     "block_ms": [
      32,
      64,
      128
     ],
     "chunks": [
      128,
      256,
      512,
      1024
     ],
     "dense_limit": 524288,
     "ragged_limit": 262144,
     "maxtok_start": 32768,
     "maxtok_stop": 2097152,
     "check": true
    }
   },
   {
    "file": "gpu_layer_before.json",
    "code": "before",
    "note": "all impls, base code (contended run; prefer A/B rows for sorted)",
    "config": {
     "device": "gpu",
     "suites": [
      "layer"
     ],
     "dtypes": [
      "float32",
      "bfloat16"
     ],
     "tokens": [
      4096,
      32768
     ],
     "experts": [
      8
     ],
     "top_k": [
      2
     ],
     "skew": [
      0.0
     ],
     "d_model": 512,
     "d_hidden": 1024,
     "impls": null,
     "block_m": 64,
     "block_ms": [
      32,
      64,
      128
     ],
     "chunks": [
      128,
      256,
      512,
      1024
     ],
     "dense_limit": 524288,
     "ragged_limit": 262144,
     "maxtok_start": 32768,
     "maxtok_stop": 2097152,
     "check": true
    }
   },
   {
    "file": "gpu_pieces_before.json",
    "code": "before",
    "note": "cost split + block_m sweep (32/64/128), base code",
    "config": {
     "device": "gpu",
     "suites": [
      "pieces",
      "blockm"
     ],
     "dtypes": [
      "float32",
      "bfloat16"
     ],
     "tokens": [
      4096,
      32768
     ],
     "experts": [
      8
     ],
     "top_k": [
      2
     ],
     "skew": [
      0.0
     ],
     "d_model": 512,
     "d_hidden": 1024,
     "impls": null,
     "block_m": 64,
     "block_ms": [
      32,
      64,
      128
     ],
     "chunks": [
      128,
      256,
      512,
      1024
     ],
     "dense_limit": 524288,
     "ragged_limit": 262144,
     "maxtok_start": 32768,
     "maxtok_stop": 2097152,
     "check": true
    }
   },
   {
    "file": "gpu_slab_0.json",
    "code": "sweep",
    "note": "slab size sweep (intermediate scatter variant): one slab",
    "config": {
     "device": "gpu",
     "suites": [
      "layer"
     ],
     "dtypes": [
      "float32",
      "bfloat16"
     ],
     "tokens": [
      32768,
      131072
     ],
     "experts": [
      8
     ],
     "top_k": [
      2
     ],
     "skew": [
      0.0
     ],
     "d_model": 512,
     "d_hidden": 1024,
     "impls": [
      "dropless:pallas:sorted"
     ],
     "block_m": 64,
     "block_ms": [
      32,
      64,
      128
     ],
     "chunks": [
      128,
      256,
      512,
      1024
     ],
     "dense_limit": 524288,
     "ragged_limit": 262144,
     "maxtok_start": 32768,
     "maxtok_stop": 2097152,
     "slab_elems": 0,
     "check": false
    }
   },
   {
    "file": "gpu_slab_16777216.json",
    "code": "sweep",
    "note": "slab 2**24",
    "config": {
     "device": "gpu",
     "suites": [
      "layer"
     ],
     "dtypes": [
      "float32",
      "bfloat16"
     ],
     "tokens": [
      32768,
      131072
     ],
     "experts": [
      8
     ],
     "top_k": [
      2
     ],
     "skew": [
      0.0
     ],
     "d_model": 512,
     "d_hidden": 1024,
     "impls": [
      "dropless:pallas:sorted"
     ],
     "block_m": 64,
     "block_ms": [
      32,
      64,
      128
     ],
     "chunks": [
      128,
      256,
      512,
      1024
     ],
     "dense_limit": 524288,
     "ragged_limit": 262144,
     "maxtok_start": 32768,
     "maxtok_stop": 2097152,
     "slab_elems": 16777216,
     "check": false
    }
   },
   {
    "file": "gpu_slab_33554432.json",
    "code": "sweep",
    "note": "slab 2**25",
    "config": {
     "device": "gpu",
     "suites": [
      "layer"
     ],
     "dtypes": [
      "float32",
      "bfloat16"
     ],
     "tokens": [
      32768,
      131072
     ],
     "experts": [
      8
     ],
     "top_k": [
      2
     ],
     "skew": [
      0.0
     ],
     "d_model": 512,
     "d_hidden": 1024,
     "impls": [
      "dropless:pallas:sorted"
     ],
     "block_m": 64,
     "block_ms": [
      32,
      64,
      128
     ],
     "chunks": [
      128,
      256,
      512,
      1024
     ],
     "dense_limit": 524288,
     "ragged_limit": 262144,
     "maxtok_start": 32768,
     "maxtok_stop": 2097152,
     "slab_elems": 33554432,
     "check": false
    }
   },
   {
    "file": "gpu_slab_67108864.json",
    "code": "sweep",
    "note": "slab 2**26",
    "config": {
     "device": "gpu",
     "suites": [
      "layer"
     ],
     "dtypes": [
      "float32",
      "bfloat16"
     ],
     "tokens": [
      32768,
      131072
     ],
     "experts": [
      8
     ],
     "top_k": [
      2
     ],
     "skew": [
      0.0
     ],
     "d_model": 512,
     "d_hidden": 1024,
     "impls": [
      "dropless:pallas:sorted"
     ],
     "block_m": 64,
     "block_ms": [
      32,
      64,
      128
     ],
     "chunks": [
      128,
      256,
      512,
      1024
     ],
     "dense_limit": 524288,
     "ragged_limit": 262144,
     "maxtok_start": 32768,
     "maxtok_stop": 2097152,
     "slab_elems": 67108864,
     "check": false
    }
   },
   {
    "file": "gpu_C.json",
    "code": "intermediate",
    "note": "skew 0.9 (2**24 slab variant)",
    "config": {
     "device": "gpu",
     "suites": [
      "layer"
     ],
     "dtypes": [
      "bfloat16"
     ],
     "tokens": [
      32768
     ],
     "experts": [
      8
     ],
     "top_k": [
      2
     ],
     "skew": [
      0.9
     ],
     "d_model": 512,
     "d_hidden": 1024,
     "impls": [
      "dropless:pallas:sorted",
      "dropless:pallas:fused",
      "dropless:pallas:fused_gsum",
      "capacity:auto:sorted"
     ],
     "block_m": 64,
     "block_ms": [
      32,
      64,
      128
     ],
     "chunks": [
      128,
      256,
      512,
      1024
     ],
     "dense_limit": 524288,
     "ragged_limit": 262144,
     "maxtok_start": 32768,
     "maxtok_stop": 2097152,
     "slab_elems": null,
     "check": true
    }
   },
   {
    "file": "gpu_numerics.json",
    "code": "before",
    "note": "numerics sweep 1000 tokens d128 h256, E8/64, top1/4, skew 0/0.9, fp32/bf16, all impls vs fp32 dense",
    "config": {
     "device": "gpu",
     "suites": [
      "layer"
     ],
     "dtypes": [
      "float32",
      "bfloat16"
     ],
     "tokens": [
      1000
     ],
     "experts": [
      8,
      64
     ],
     "top_k": [
      1,
      4
     ],
     "skew": [
      0.0,
      0.9
     ],
     "d_model": 128,
     "d_hidden": 256,
     "impls": null,
     "block_m": 64,
     "block_ms": [
      32,
      64,
      128
     ],
     "chunks": [
      128,
      256,
      512,
      1024
     ],
     "dense_limit": 524288,
     "ragged_limit": 262144,
     "maxtok_start": 32768,
     "maxtok_stop": 2097152,
     "check": true
    }
   }
  ]
 },
 "results": [
  {
   "code": "before",
   "run": "ab6_base_1",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.135498046875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 85.15921783447266,
   "xla_arg_mb": 56.015625,
   "xla_out_mb": 8.0,
   "compile_s": 0.0032215419923886657,
   "inner": 36,
   "median_ms": 1.286565277647848,
   "q1_ms": 1.2819592916962899,
   "q3_ms": 1.2878219721541329
  },
  {
   "code": "before",
   "run": "ab6_base_1",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.135498046875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 136.19632720947266,
   "xla_arg_mb": 152.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.010725822998210788,
   "inner": 10,
   "median_ms": 4.645629400329199,
   "q1_ms": 4.6447499000350945,
   "q3_ms": 4.673444949730765
  },
  {
   "code": "before",
   "run": "ab6_base_1",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1292572021484375,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 646.2566299438477,
   "xla_arg_mb": 112.015625,
   "xla_out_mb": 64.0,
   "compile_s": 0.023673599993344396,
   "inner": 5,
   "median_ms": 8.635517000220716,
   "q1_ms": 8.629775898589287,
   "q3_ms": 8.66322060173843
  },
  {
   "code": "before",
   "run": "ab6_base_1",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1292572021484375,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 1033.6450576782227,
   "xla_arg_mb": 208.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.05966123100370169,
   "inner": 1,
   "median_ms": 31.87741999863647,
   "q1_ms": 31.823587502003647,
   "q3_ms": 31.886911994661205
  },
  {
   "code": "before",
   "run": "ab6_base_1",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 131072,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.12841415405273438,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 2570.0183486938477,
   "xla_arg_mb": 304.015625,
   "xla_out_mb": 256.0,
   "compile_s": 0.08912519400473684,
   "inner": 1,
   "median_ms": 33.944146998692304,
   "q1_ms": 33.77800699672662,
   "q3_ms": 34.04880100424634
  },
  {
   "code": "before",
   "run": "ab6_base_1",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 131072,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.12841415405273438,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 4110.535926818848,
   "xla_arg_mb": 400.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.2402062070032116,
   "inner": 1,
   "median_ms": 130.772338990937,
   "q1_ms": 130.42409699846758,
   "q3_ms": 131.86699500511168
  },
  {
   "code": "before",
   "run": "ab6_base_1",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1356201171875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 50.690467834472656,
   "xla_arg_mb": 52.015625,
   "xla_out_mb": 4.0,
   "compile_s": 0.002724748002947308,
   "inner": 51,
   "median_ms": 0.8741907059790238,
   "q1_ms": 0.8740151960427837,
   "q3_ms": 0.8754844508503619
  },
  {
   "code": "before",
   "run": "ab6_base_1",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1356201171875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 59.005409240722656,
   "xla_arg_mb": 148.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.005563603001064621,
   "inner": 16,
   "median_ms": 2.7398505617384217,
   "q1_ms": 2.7390627815293556,
   "q3_ms": 2.7417741871431645
  },
  {
   "code": "before",
   "run": "ab6_base_1",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1291961669921875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 324.00662994384766,
   "xla_arg_mb": 80.015625,
   "xla_out_mb": 32.0,
   "compile_s": 0.006510382998385467,
   "inner": 10,
   "median_ms": 4.884568600391503,
   "q1_ms": 4.875259399705101,
   "q3_ms": 4.893508600798668
  },
  {
   "code": "before",
   "run": "ab6_base_1",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1291961669921875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 524.0054092407227,
   "xla_arg_mb": 176.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.01837455900385976,
   "inner": 3,
   "median_ms": 16.982286668887053,
   "q1_ms": 16.88150333696588,
   "q3_ms": 17.008394332757838
  },
  {
   "code": "before",
   "run": "ab6_base_1",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 131072,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.12839126586914062,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 1288.5183486938477,
   "xla_arg_mb": 176.015625,
   "xla_out_mb": 128.0,
   "compile_s": 0.01922999399539549,
   "inner": 2,
   "median_ms": 18.648549499630462,
   "q1_ms": 18.64644825400319,
   "q3_ms": 18.75775725056883
  },
  {
   "code": "before",
   "run": "ab6_base_1",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 131072,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.12839126586914062,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 2060.0056533813477,
   "xla_arg_mb": 272.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.06774797799880616,
   "inner": 1,
   "median_ms": 68.9175910083577,
   "q1_ms": 68.76846249360824,
   "q3_ms": 69.43487899843603
  },
  {
   "code": "final",
   "run": "ab6_v6_1",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.135498046875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 85.15921783447266,
   "xla_arg_mb": 56.015625,
   "xla_out_mb": 8.0,
   "compile_s": 0.0034648449945962057,
   "inner": 35,
   "median_ms": 1.2847463429872212,
   "q1_ms": 1.2814011428937582,
   "q3_ms": 1.2873539713577233
  },
  {
   "code": "final",
   "run": "ab6_v6_1",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.135498046875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 136.19632720947266,
   "xla_arg_mb": 152.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.010347136005293578,
   "inner": 10,
   "median_ms": 4.647483599546831,
   "q1_ms": 4.6262454001407605,
   "q3_ms": 4.650281950307544
  },
  {
   "code": "final",
   "run": "ab6_v6_1",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1292572021484375,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 581.0122680664062,
   "xla_arg_mb": 112.015625,
   "xla_out_mb": 64.0,
   "compile_s": 0.024816681994707324,
   "inner": 5,
   "median_ms": 9.282120800344273,
   "q1_ms": 9.262449700327124,
   "q3_ms": 9.311618898937013
  },
  {
   "code": "final",
   "run": "ab6_v6_1",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1292572021484375,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 773.6470718383789,
   "xla_arg_mb": 208.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.03254191500309389,
   "inner": 1,
   "median_ms": 31.852801999775693,
   "q1_ms": 31.822410004679114,
   "q3_ms": 31.881046503258403
  },
  {
   "code": "final",
   "run": "ab6_v6_1",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 131072,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.12841415405273438,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 1339.0252075195312,
   "xla_arg_mb": 304.015625,
   "xla_out_mb": 256.0,
   "compile_s": 0.06398921200889163,
   "inner": 1,
   "median_ms": 36.68068700062577,
   "q1_ms": 36.61648449633503,
   "q3_ms": 36.82205949735362
  },
  {
   "code": "final",
   "run": "ab6_v6_1",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 131072,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.12841415405273438,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 1851.538429260254,
   "xla_arg_mb": 400.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.1324250450124964,
   "inner": 1,
   "median_ms": 131.4043159945868,
   "q1_ms": 131.22776650561718,
   "q3_ms": 131.6294414937147
  },
  {
   "code": "final",
   "run": "ab6_v6_1",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1356201171875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 50.690467834472656,
   "xla_arg_mb": 52.015625,
   "xla_out_mb": 4.0,
   "compile_s": 0.002620968996779993,
   "inner": 51,
   "median_ms": 0.8697468626504654,
   "q1_ms": 0.8656223823588032,
   "q3_ms": 0.8700695292574498
  },
  {
   "code": "final",
   "run": "ab6_v6_1",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1356201171875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 59.005409240722656,
   "xla_arg_mb": 148.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.005461431996081956,
   "inner": 16,
   "median_ms": 2.729736936998961,
   "q1_ms": 2.723985749980784,
   "q3_ms": 2.732449843279028
  },
  {
   "code": "final",
   "run": "ab6_v6_1",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1291961669921875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 307.26226806640625,
   "xla_arg_mb": 80.015625,
   "xla_out_mb": 32.0,
   "compile_s": 0.007013319001998752,
   "inner": 9,
   "median_ms": 5.001206222611169,
   "q1_ms": 4.996510221745767,
   "q3_ms": 5.01572866697744
  },
  {
   "code": "final",
   "run": "ab6_v6_1",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1291961669921875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 386.2601089477539,
   "xla_arg_mb": 176.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.018117245999746956,
   "inner": 3,
   "median_ms": 16.017293334395315,
   "q1_ms": 15.977912167727482,
   "q3_ms": 16.669076333831377
  },
  {
   "code": "final",
   "run": "ab6_v6_1",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 131072,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.12839126586914062,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 696.5252075195312,
   "xla_arg_mb": 176.015625,
   "xla_out_mb": 128.0,
   "compile_s": 0.020213372001308016,
   "inner": 2,
   "median_ms": 19.24347250314895,
   "q1_ms": 19.148790248436853,
   "q3_ms": 19.44157174875727
  },
  {
   "code": "final",
   "run": "ab6_v6_1",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 131072,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.12839126586914062,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 1032.117042541504,
   "xla_arg_mb": 272.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.06457873700128403,
   "inner": 1,
   "median_ms": 65.45954299508594,
   "q1_ms": 65.15349350229371,
   "q3_ms": 65.55349250265863
  },
  {
   "code": "before",
   "run": "ab6_base_2",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.135498046875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 85.15921783447266,
   "xla_arg_mb": 56.015625,
   "xla_out_mb": 8.0,
   "compile_s": 0.0033174259879160672,
   "inner": 35,
   "median_ms": 1.2814011711660507,
   "q1_ms": 1.2788119715488782,
   "q3_ms": 1.28767359981014
  },
  {
   "code": "before",
   "run": "ab6_base_2",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.135498046875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 136.19632720947266,
   "xla_arg_mb": 152.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.01096858199161943,
   "inner": 10,
   "median_ms": 4.645300499396399,
   "q1_ms": 4.64413025038084,
   "q3_ms": 4.650498849514406
  },
  {
   "code": "before",
   "run": "ab6_base_2",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1292572021484375,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 646.2566299438477,
   "xla_arg_mb": 112.015625,
   "xla_out_mb": 64.0,
   "compile_s": 0.025503160009975545,
   "inner": 5,
   "median_ms": 8.623270399402827,
   "q1_ms": 8.577624200552236,
   "q3_ms": 8.637392200762406
  },
  {
   "code": "before",
   "run": "ab6_base_2",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1292572021484375,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 1033.6450576782227,
   "xla_arg_mb": 208.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.060453763988334686,
   "inner": 1,
   "median_ms": 31.91056300420314,
   "q1_ms": 31.784428007085808,
   "q3_ms": 32.06977900117636
  },
  {
   "code": "before",
   "run": "ab6_base_2",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 131072,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.12841415405273438,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 2570.0183486938477,
   "xla_arg_mb": 304.015625,
   "xla_out_mb": 256.0,
   "compile_s": 0.0864492240070831,
   "inner": 1,
   "median_ms": 33.92118599731475,
   "q1_ms": 33.78949049510993,
   "q3_ms": 34.16411999933189
  },
  {
   "code": "before",
   "run": "ab6_base_2",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 131072,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.12841415405273438,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 4110.535926818848,
   "xla_arg_mb": 400.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.2360110320005333,
   "inner": 1,
   "median_ms": 131.36554398806766,
   "q1_ms": 130.3443540018634,
   "q3_ms": 131.9972310011508
  },
  {
   "code": "before",
   "run": "ab6_base_2",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1356201171875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 50.690467834472656,
   "xla_arg_mb": 52.015625,
   "xla_out_mb": 4.0,
   "compile_s": 0.0027797990042017773,
   "inner": 51,
   "median_ms": 0.870090058728523,
   "q1_ms": 0.8545311078812707,
   "q3_ms": 0.8713601175045101
  },
  {
   "code": "before",
   "run": "ab6_base_2",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1356201171875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 59.005409240722656,
   "xla_arg_mb": 148.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.005409072997281328,
   "inner": 16,
   "median_ms": 2.747234374510299,
   "q1_ms": 2.738723937454779,
   "q3_ms": 2.748029000031238
  },
  {
   "code": "before",
   "run": "ab6_base_2",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1291961669921875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 324.00662994384766,
   "xla_arg_mb": 80.015625,
   "xla_out_mb": 32.0,
   "compile_s": 0.006239875991013832,
   "inner": 10,
   "median_ms": 4.878326399193611,
   "q1_ms": 4.851743500330485,
   "q3_ms": 4.882617449766258
  },
  {
   "code": "before",
   "run": "ab6_base_2",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1291961669921875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 524.0054092407227,
   "xla_arg_mb": 176.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.01837943200371228,
   "inner": 3,
   "median_ms": 16.958150003726285,
   "q1_ms": 16.92697166921183,
   "q3_ms": 16.991476164548658
  },
  {
   "code": "before",
   "run": "ab6_base_2",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 131072,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.12839126586914062,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 1288.5183486938477,
   "xla_arg_mb": 176.015625,
   "xla_out_mb": 128.0,
   "compile_s": 0.019234268998843618,
   "inner": 2,
   "median_ms": 18.670683501113672,
   "q1_ms": 18.659343251783866,
   "q3_ms": 18.751779749436537
  },
  {
   "code": "before",
   "run": "ab6_base_2",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 131072,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.12839126586914062,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 2060.0056533813477,
   "xla_arg_mb": 272.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.06831712499842979,
   "inner": 1,
   "median_ms": 68.67243199667428,
   "q1_ms": 68.64810150000267,
   "q3_ms": 68.697714501468
  },
  {
   "code": "final",
   "run": "ab6_v6_2",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.135498046875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 85.15921783447266,
   "xla_arg_mb": 56.015625,
   "xla_out_mb": 8.0,
   "compile_s": 0.0034090479894075543,
   "inner": 35,
   "median_ms": 1.2928499715469246,
   "q1_ms": 1.2891489856493925,
   "q3_ms": 1.2952044144705204
  },
  {
   "code": "final",
   "run": "ab6_v6_2",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.135498046875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 136.19632720947266,
   "xla_arg_mb": 152.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.010678507998818532,
   "inner": 11,
   "median_ms": 4.644223545752043,
   "q1_ms": 4.631203954473181,
   "q3_ms": 4.648261136687572
  },
  {
   "code": "final",
   "run": "ab6_v6_2",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1292572021484375,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 581.0122680664062,
   "xla_arg_mb": 112.015625,
   "xla_out_mb": 64.0,
   "compile_s": 0.026169809003476985,
   "inner": 5,
   "median_ms": 9.234536800067872,
   "q1_ms": 9.215672301070299,
   "q3_ms": 9.273880401451606
  },
  {
   "code": "final",
   "run": "ab6_v6_2",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1292572021484375,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 773.6470718383789,
   "xla_arg_mb": 208.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.03281907699420117,
   "inner": 1,
   "median_ms": 31.69436899770517,
   "q1_ms": 31.689921495853923,
   "q3_ms": 31.85352050786605
  },
  {
   "code": "final",
   "run": "ab6_v6_2",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 131072,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.12841415405273438,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 1339.0252075195312,
   "xla_arg_mb": 304.015625,
   "xla_out_mb": 256.0,
   "compile_s": 0.06487772498803679,
   "inner": 1,
   "median_ms": 36.48624299967196,
   "q1_ms": 36.46537899476243,
   "q3_ms": 36.5221809988725
  },
  {
   "code": "final",
   "run": "ab6_v6_2",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 131072,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.12841415405273438,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 1851.538429260254,
   "xla_arg_mb": 400.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.13442561400006525,
   "inner": 1,
   "median_ms": 131.00303000828717,
   "q1_ms": 130.76019100117264,
   "q3_ms": 131.72209049662342
  },
  {
   "code": "final",
   "run": "ab6_v6_2",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1356201171875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 50.690467834472656,
   "xla_arg_mb": 52.015625,
   "xla_out_mb": 4.0,
   "compile_s": 0.0026441109948791564,
   "inner": 51,
   "median_ms": 0.8515685684887655,
   "q1_ms": 0.8495156666065347,
   "q3_ms": 0.8700596373790748
  },
  {
   "code": "final",
   "run": "ab6_v6_2",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1356201171875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 59.005409240722656,
   "xla_arg_mb": 148.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.005644124990794808,
   "inner": 16,
   "median_ms": 2.7441699376140605,
   "q1_ms": 2.7384544373489916,
   "q3_ms": 2.7468890934869705
  },
  {
   "code": "final",
   "run": "ab6_v6_2",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1291961669921875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 307.26226806640625,
   "xla_arg_mb": 80.015625,
   "xla_out_mb": 32.0,
   "compile_s": 0.006762842007447034,
   "inner": 9,
   "median_ms": 5.013278109901067,
   "q1_ms": 5.004586611499286,
   "q3_ms": 5.015429388246654
  },
  {
   "code": "final",
   "run": "ab6_v6_2",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1291961669921875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 386.2601089477539,
   "xla_arg_mb": 176.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.01777212599699851,
   "inner": 3,
   "median_ms": 16.03705033388299,
   "q1_ms": 15.962157332978677,
   "q3_ms": 16.073384334352642
  },
  {
   "code": "final",
   "run": "ab6_v6_2",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 131072,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.12839126586914062,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 696.5252075195312,
   "xla_arg_mb": 176.015625,
   "xla_out_mb": 128.0,
   "compile_s": 0.020264459992176853,
   "inner": 2,
   "median_ms": 19.311356998514384,
   "q1_ms": 19.289645748358453,
   "q3_ms": 19.315938003273914
  },
  {
   "code": "final",
   "run": "ab6_v6_2",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 131072,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.12839126586914062,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 1032.117042541504,
   "xla_arg_mb": 272.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.0649371479958063,
   "inner": 1,
   "median_ms": 65.0605219998397,
   "q1_ms": 65.00652400427498,
   "q3_ms": 65.22019299882231
  },
  {
   "code": "final",
   "run": "f_gpu_a",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.135498046875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 85.15921783447266,
   "xla_arg_mb": 56.015625,
   "xla_out_mb": 8.0,
   "compile_s": 0.08292652299860492,
   "inner": 13,
   "median_ms": 3.044169307731164,
   "q1_ms": 2.9527878840991226,
   "q3_ms": 3.060675807668863,
   "rel_err_out": 0.01140456501241541,
   "rel_err_grad_params": 0.010704822355157323,
   "rel_err_grad_x": 0.00843557025022219
  },
  {
   "code": "final",
   "run": "f_gpu_a",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.135498046875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 136.19632720947266,
   "xla_arg_mb": 152.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.017775547996279784,
   "inner": 9,
   "median_ms": 5.036389112420793,
   "q1_ms": 4.807465278847506,
   "q3_ms": 5.378594055299699,
   "rel_err_out": 0.01140456501241541,
   "rel_err_grad_params": 0.010704822355157323,
   "rel_err_grad_x": 0.00843557025022219
  },
  {
   "code": "final",
   "run": "f_gpu_a",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "gather",
   "dtype": "float32",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.135498046875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 117.16336822509766,
   "xla_arg_mb": 56.015625,
   "xla_out_mb": 8.0,
   "compile_s": 0.0057558560074539855,
   "inner": 33,
   "median_ms": 1.5551401209615339,
   "q1_ms": 1.5496004244513726,
   "q3_ms": 1.583486818135222,
   "rel_err_out": 0.01140456501241541,
   "rel_err_grad_params": 0.010704822355157323,
   "rel_err_grad_x": 0.008434855028591981
  },
  {
   "code": "final",
   "run": "f_gpu_a",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "gather",
   "dtype": "float32",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.135498046875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 202.0703582763672,
   "xla_arg_mb": 152.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.007695186999626458,
   "inner": 10,
   "median_ms": 5.147058799047954,
   "q1_ms": 4.882383300719084,
   "q3_ms": 5.504070399911143,
   "rel_err_out": 0.01140456501241541,
   "rel_err_grad_params": 0.010704822355157323,
   "rel_err_grad_x": 0.008434855028591981
  },
  {
   "code": "final",
   "run": "f_gpu_a",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused",
   "dtype": "float32",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.135498046875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 102.10208892822266,
   "xla_arg_mb": 56.015625,
   "xla_out_mb": 8.0,
   "compile_s": 0.007119019006495364,
   "inner": 20,
   "median_ms": 1.814727400051197,
   "q1_ms": 1.740861774669611,
   "q3_ms": 1.8643285748112248,
   "rel_err_out": 0.01140456501241541,
   "rel_err_grad_params": 0.010704822355157323,
   "rel_err_grad_x": 0.00843485500352305
  },
  {
   "code": "final",
   "run": "f_gpu_a",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused",
   "dtype": "float32",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.135498046875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 170.0408172607422,
   "xla_arg_mb": 152.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.007243935004225932,
   "inner": 10,
   "median_ms": 5.055467000056524,
   "q1_ms": 4.938792199391173,
   "q3_ms": 5.172816050617257,
   "rel_err_out": 0.01140456501241541,
   "rel_err_grad_params": 0.010704822355157323,
   "rel_err_grad_x": 0.00843485500352305
  },
  {
   "code": "final",
   "run": "f_gpu_a",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused_gsum",
   "dtype": "float32",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.135498046875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 102.16141510009766,
   "xla_arg_mb": 56.015625,
   "xla_out_mb": 8.0,
   "compile_s": 0.010180500001297332,
   "inner": 33,
   "median_ms": 1.5096472727071326,
   "q1_ms": 1.396134787875891,
   "q3_ms": 1.5299152274206844,
   "rel_err_out": 0.01140456501241541,
   "rel_err_grad_params": 0.010704822355157323,
   "rel_err_grad_x": 0.00843485500352305
  },
  {
   "code": "final",
   "run": "f_gpu_a",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused_gsum",
   "dtype": "float32",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.135498046875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 170.0410614013672,
   "xla_arg_mb": 152.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.006862905996968038,
   "inner": 11,
   "median_ms": 5.033600908725269,
   "q1_ms": 4.813760409226895,
   "q3_ms": 5.1607704544122415,
   "rel_err_out": 0.01140456501241541,
   "rel_err_grad_params": 0.010704822355157323,
   "rel_err_grad_x": 0.00843485500352305
  },
  {
   "code": "final",
   "run": "f_gpu_a",
   "suite": "layer",
   "routing": "dropless",
   "impl": "blocked",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.135498046875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 486.15872955322266,
   "xla_arg_mb": 56.015625,
   "xla_out_mb": 8.0,
   "compile_s": 0.006183913996210322,
   "inner": 18,
   "median_ms": 2.610955222432191,
   "q1_ms": 2.36881777770274,
   "q3_ms": 2.7833614997992604,
   "rel_err_out": 0.01140456501241541,
   "rel_err_grad_params": 0.010704833542337283,
   "rel_err_grad_x": 0.0084355704322865
  },
  {
   "code": "final",
   "run": "f_gpu_a",
   "suite": "layer",
   "routing": "dropless",
   "impl": "blocked",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.135498046875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 504.25467681884766,
   "xla_arg_mb": 152.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.008819398994091898,
   "inner": 7,
   "median_ms": 7.134214141322965,
   "q1_ms": 6.891721927136782,
   "q3_ms": 7.187357929069549,
   "rel_err_out": 0.01140456501241541,
   "rel_err_grad_params": 0.010704833542337283,
   "rel_err_grad_x": 0.0084355704322865
  },
  {
   "code": "final",
   "run": "f_gpu_a",
   "suite": "layer",
   "routing": "dropless",
   "impl": "ragged_dot",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.135498046875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 288.09498596191406,
   "xla_arg_mb": 56.015625,
   "xla_out_mb": 8.0,
   "compile_s": 0.008054235993768089,
   "inner": 6,
   "median_ms": 6.607968498428818,
   "q1_ms": 6.589531082378623,
   "q3_ms": 6.732340916641988,
   "rel_err_out": 0.011404565034145914,
   "rel_err_grad_params": 0.010704815714799832,
   "rel_err_grad_x": 0.008435547878322612
  },
  {
   "code": "final",
   "run": "f_gpu_a",
   "suite": "layer",
   "routing": "dropless",
   "impl": "ragged_dot",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.135498046875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 768.0039367675781,
   "xla_arg_mb": 152.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.020875426998827606,
   "inner": 2,
   "median_ms": 21.421117999125272,
   "q1_ms": 21.074793497973587,
   "q3_ms": 21.793005496874684,
   "rel_err_out": 0.011404565034145914,
   "rel_err_grad_params": 0.010704815714799832,
   "rel_err_grad_x": 0.008435547878322612
  },
  {
   "code": "final",
   "run": "f_gpu_a",
   "suite": "layer",
   "routing": "capacity",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.135498046875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 120.15675354003906,
   "xla_arg_mb": 56.015625,
   "xla_out_mb": 8.0,
   "compile_s": 0.04734614500193857,
   "inner": 12,
   "median_ms": 3.360633834139056,
   "q1_ms": 2.9154443758064494,
   "q3_ms": 3.4734812088572653
  },
  {
   "code": "final",
   "run": "f_gpu_a",
   "suite": "layer",
   "routing": "capacity",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.135498046875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 200.0024642944336,
   "xla_arg_mb": 152.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.0133912190067349,
   "inner": 11,
   "median_ms": 3.7565339999590917,
   "q1_ms": 3.71890786416109,
   "q3_ms": 3.7777380454248157
  },
  {
   "code": "final",
   "run": "f_gpu_a",
   "suite": "layer",
   "routing": "dense",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.135498046875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 256.00001525878906,
   "xla_arg_mb": 56.015625,
   "xla_out_mb": 8.0,
   "compile_s": 0.0046783279976807535,
   "inner": 14,
   "median_ms": 3.7692865007556975,
   "q1_ms": 3.7289835355685295,
   "q3_ms": 3.8428340355624515,
   "rel_err_out": 0.01140456501241541,
   "rel_err_grad_params": 0.010704805168015577,
   "rel_err_grad_x": 0.00843557463968868
  },
  {
   "code": "final",
   "run": "f_gpu_a",
   "suite": "layer",
   "routing": "dense",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.135498046875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 448.0022201538086,
   "xla_arg_mb": 152.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.011347177001880482,
   "inner": 4,
   "median_ms": 11.084725749242352,
   "q1_ms": 9.820819248488988,
   "q3_ms": 11.212172752493643,
   "rel_err_out": 0.01140456501241541,
   "rel_err_grad_params": 0.010704805168015577,
   "rel_err_grad_x": 0.00843557463968868
  },
  {
   "code": "final",
   "run": "f_gpu_a",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1292572021484375,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 581.0122680664062,
   "xla_arg_mb": 112.015625,
   "xla_out_mb": 64.0,
   "compile_s": 0.0175547629914945,
   "inner": 5,
   "median_ms": 11.57039419922512,
   "q1_ms": 9.299754700623453,
   "q3_ms": 12.371740098751616,
   "rel_err_out": 0.01101133789638197,
   "rel_err_grad_params": 0.01021944625795938,
   "rel_err_grad_x": 0.013931359767753112
  },
  {
   "code": "final",
   "run": "f_gpu_a",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1292572021484375,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 773.6470718383789,
   "xla_arg_mb": 208.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.03335735799919348,
   "inner": 1,
   "median_ms": 32.30411700496916,
   "q1_ms": 32.14505749929231,
   "q3_ms": 33.68257649708539,
   "rel_err_out": 0.01101133789638197,
   "rel_err_grad_params": 0.01021944625795938,
   "rel_err_grad_x": 0.013931359767753112
  },
  {
   "code": "final",
   "run": "f_gpu_a",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "gather",
   "dtype": "float32",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1292572021484375,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 776.0098037719727,
   "xla_arg_mb": 112.015625,
   "xla_out_mb": 64.0,
   "compile_s": 0.011212301003979519,
   "inner": 5,
   "median_ms": 9.14976540079806,
   "q1_ms": 9.087017901765648,
   "q3_ms": 9.21113329968648,
   "rel_err_out": 0.01101133789638197,
   "rel_err_grad_params": 0.010216445921312082,
   "rel_err_grad_x": 0.013930634470649978
  },
  {
   "code": "final",
   "run": "f_gpu_a",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "gather",
   "dtype": "float32",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1292572021484375,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 1419.5195770263672,
   "xla_arg_mb": 208.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.03771233498991933,
   "inner": 1,
   "median_ms": 37.71038100239821,
   "q1_ms": 36.78305350331357,
   "q3_ms": 40.131133995600976,
   "rel_err_out": 0.01101133789638197,
   "rel_err_grad_params": 0.010216445921312082,
   "rel_err_grad_x": 0.013930634470649978
  },
  {
   "code": "final",
   "run": "f_gpu_a",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused",
   "dtype": "float32",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1292572021484375,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 774.7615127563477,
   "xla_arg_mb": 112.015625,
   "xla_out_mb": 64.0,
   "compile_s": 0.01272402799804695,
   "inner": 4,
   "median_ms": 11.389561750547728,
   "q1_ms": 11.069567124650348,
   "q3_ms": 11.57173212413909,
   "rel_err_out": 0.01101133789638197,
   "rel_err_grad_params": 0.010216445921312082,
   "rel_err_grad_x": 0.013930634460644907
  },
  {
   "code": "final",
   "run": "f_gpu_a",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused",
   "dtype": "float32",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1292572021484375,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 1290.0193328857422,
   "xla_arg_mb": 208.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.04046998400008306,
   "inner": 1,
   "median_ms": 37.013975001173094,
   "q1_ms": 36.670837507699616,
   "q3_ms": 37.49052599596325,
   "rel_err_out": 0.01101133789638197,
   "rel_err_grad_params": 0.010216445921312082,
   "rel_err_grad_x": 0.013930634460644907
  },
  {
   "code": "final",
   "run": "f_gpu_a",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused_gsum",
   "dtype": "float32",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1292572021484375,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 775.0078506469727,
   "xla_arg_mb": 112.015625,
   "xla_out_mb": 64.0,
   "compile_s": 0.010043883987236768,
   "inner": 5,
   "median_ms": 8.635223400779068,
   "q1_ms": 8.604051399743184,
   "q3_ms": 8.773502499388997,
   "rel_err_out": 0.01101133789638197,
   "rel_err_grad_params": 0.010216445921312082,
   "rel_err_grad_x": 0.013930634460644907
  },
  {
   "code": "final",
   "run": "f_gpu_a",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused_gsum",
   "dtype": "float32",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1292572021484375,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 1290.0195770263672,
   "xla_arg_mb": 208.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.03432314199744724,
   "inner": 1,
   "median_ms": 33.37158399517648,
   "q1_ms": 33.34340599394636,
   "q3_ms": 33.49624999827938,
   "rel_err_out": 0.01101133789638197,
   "rel_err_grad_params": 0.010216445921312082,
   "rel_err_grad_x": 0.013930634460644907
  },
  {
   "code": "final",
   "run": "f_gpu_a",
   "suite": "layer",
   "routing": "dropless",
   "impl": "blocked",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1292572021484375,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 3511.2527236938477,
   "xla_arg_mb": 112.015625,
   "xla_out_mb": 64.0,
   "compile_s": 0.017434416004107334,
   "inner": 3,
   "median_ms": 15.90392533398699,
   "q1_ms": 15.832299667332942,
   "q3_ms": 16.083262167133704,
   "rel_err_out": 0.01101133789638197,
   "rel_err_grad_params": 0.01021944625795938,
   "rel_err_grad_x": 0.013931359603212637
  },
  {
   "code": "final",
   "run": "f_gpu_a",
   "suite": "layer",
   "routing": "dropless",
   "impl": "blocked",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1292572021484375,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 3641.5044326782227,
   "xla_arg_mb": 208.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.03856451399042271,
   "inner": 1,
   "median_ms": 45.56837599375285,
   "q1_ms": 42.24125049950089,
   "q3_ms": 45.633929505129345,
   "rel_err_out": 0.01101133789638197,
   "rel_err_grad_params": 0.01021944625795938,
   "rel_err_grad_x": 0.013931359603212637
  },
  {
   "code": "final",
   "run": "f_gpu_a",
   "suite": "layer",
   "routing": "dropless",
   "impl": "ragged_dot",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1292572021484375,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 2304.7514877319336,
   "xla_arg_mb": 112.015625,
   "xla_out_mb": 64.0,
   "compile_s": 0.04920065800251905,
   "inner": 1,
   "median_ms": 49.88638299982995,
   "q1_ms": 49.353943002643064,
   "q3_ms": 51.16156200529076,
   "rel_err_out": 0.01101133789638197,
   "rel_err_grad_params": 0.01021944625795938,
   "rel_err_grad_x": 0.013931359603212637
  },
  {
   "code": "final",
   "run": "f_gpu_a",
   "suite": "layer",
   "routing": "dropless",
   "impl": "ragged_dot",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1292572021484375,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 6145.644317626953,
   "xla_arg_mb": 208.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.15985014200850856,
   "inner": 1,
   "median_ms": 158.63840399833862,
   "q1_ms": 156.39561050193151,
   "q3_ms": 171.86443450918887,
   "rel_err_out": 0.01101133789638197,
   "rel_err_grad_params": 0.01021944625795938,
   "rel_err_grad_x": 0.013931359603212637
  },
  {
   "code": "final",
   "run": "f_gpu_a",
   "suite": "layer",
   "routing": "capacity",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1292572021484375,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 800.7502670288086,
   "xla_arg_mb": 112.015625,
   "xla_out_mb": 64.0,
   "compile_s": 0.011717205998138525,
   "inner": 5,
   "median_ms": 10.244876198703423,
   "q1_ms": 10.018160899926443,
   "q3_ms": 11.103905100026168
  },
  {
   "code": "final",
   "run": "f_gpu_a",
   "suite": "layer",
   "routing": "capacity",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1292572021484375,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 1504.1274642944336,
   "xla_arg_mb": 208.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.03581774499616586,
   "inner": 1,
   "median_ms": 26.77482200670056,
   "q1_ms": 26.452718499058392,
   "q3_ms": 26.818835009180475
  },
  {
   "code": "final",
   "run": "f_gpu_a",
   "suite": "layer",
   "routing": "dense",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1292572021484375,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 2048.000015258789,
   "xla_arg_mb": 112.015625,
   "xla_out_mb": 64.0,
   "compile_s": 0.028434848005417734,
   "inner": 1,
   "median_ms": 29.013470993959345,
   "q1_ms": 28.544675995362923,
   "q3_ms": 29.09434900357155,
   "rel_err_out": 0.01101133789638197,
   "rel_err_grad_params": 0.01021944625795938,
   "rel_err_grad_x": 0.013931359603212637
  },
  {
   "code": "final",
   "run": "f_gpu_a",
   "suite": "layer",
   "routing": "dense",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1292572021484375,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 4163.767845153809,
   "xla_arg_mb": 208.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.0818868670030497,
   "inner": 1,
   "median_ms": 76.45969500299543,
   "q1_ms": 75.6009475007886,
   "q3_ms": 77.449844997318,
   "rel_err_out": 0.01101133789638197,
   "rel_err_grad_params": 0.01021944625795938,
   "rel_err_grad_x": 0.013931359603212637
  },
  {
   "code": "final",
   "run": "f_gpu_a",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1356201171875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 50.690467834472656,
   "xla_arg_mb": 52.015625,
   "xla_out_mb": 4.0,
   "compile_s": 0.0027807139995275065,
   "inner": 50,
   "median_ms": 0.9507511201081797,
   "q1_ms": 0.9257891699962784,
   "q3_ms": 0.9609288499632385,
   "rel_err_out": 0.020131099618709833,
   "rel_err_grad_params": 0.03448199141601909,
   "rel_err_grad_x": 0.029022456572827038
  },
  {
   "code": "final",
   "run": "f_gpu_a",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1356201171875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 59.005409240722656,
   "xla_arg_mb": 148.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.006106325992732309,
   "inner": 16,
   "median_ms": 2.8183631247884477,
   "q1_ms": 2.8022498127029394,
   "q3_ms": 2.899610968597699,
   "rel_err_out": 0.020131099618709833,
   "rel_err_grad_params": 0.03448199141601909,
   "rel_err_grad_x": 0.029022456572827038
  },
  {
   "code": "final",
   "run": "f_gpu_a",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "gather",
   "dtype": "bfloat16",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1356201171875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 58.663368225097656,
   "xla_arg_mb": 52.015625,
   "xla_out_mb": 4.0,
   "compile_s": 0.002908495007432066,
   "inner": 45,
   "median_ms": 0.9077952000855779,
   "q1_ms": 0.8954595554516548,
   "q3_ms": 0.9204105778028153,
   "rel_err_out": 0.020131099618709833,
   "rel_err_grad_params": 0.03445727698235376,
   "rel_err_grad_x": 0.028982459890591516
  },
  {
   "code": "final",
   "run": "f_gpu_a",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "gather",
   "dtype": "bfloat16",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1356201171875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 102.12992858886719,
   "xla_arg_mb": 148.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.005276268013403751,
   "inner": 18,
   "median_ms": 2.614986833779969,
   "q1_ms": 2.610445722287598,
   "q3_ms": 2.783796000181206,
   "rel_err_out": 0.020131099618709833,
   "rel_err_grad_params": 0.03445727698235376,
   "rel_err_grad_x": 0.028982459890591516
  },
  {
   "code": "final",
   "run": "f_gpu_a",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused",
   "dtype": "bfloat16",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1356201171875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 51.102088928222656,
   "xla_arg_mb": 52.015625,
   "xla_out_mb": 4.0,
   "compile_s": 0.0027312800084473565,
   "inner": 46,
   "median_ms": 0.9591114781948779,
   "q1_ms": 0.9492310762099172,
   "q3_ms": 0.9837804781968968,
   "rel_err_out": 0.020063212378551402,
   "rel_err_grad_params": 0.0344319496824839,
   "rel_err_grad_x": 0.02899075704691895
  },
  {
   "code": "final",
   "run": "f_gpu_a",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused",
   "dtype": "bfloat16",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1356201171875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 102.16264343261719,
   "xla_arg_mb": 148.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.005208068992942572,
   "inner": 18,
   "median_ms": 2.749719666705156,
   "q1_ms": 2.708055972309214,
   "q3_ms": 2.7657496664485537,
   "rel_err_out": 0.020063212378551402,
   "rel_err_grad_params": 0.0344319496824839,
   "rel_err_grad_x": 0.02899075704691895
  },
  {
   "code": "final",
   "run": "f_gpu_a",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused_gsum",
   "dtype": "bfloat16",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1356201171875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 51.192665100097656,
   "xla_arg_mb": 52.015625,
   "xla_out_mb": 4.0,
   "compile_s": 0.002750996995018795,
   "inner": 53,
   "median_ms": 0.8953930565731529,
   "q1_ms": 0.8646201980432038,
   "q3_ms": 0.9262416981828381,
   "rel_err_out": 0.020131099618709833,
   "rel_err_grad_params": 0.03445727698235376,
   "rel_err_grad_x": 0.028982459890591516
  },
  {
   "code": "final",
   "run": "f_gpu_a",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused_gsum",
   "dtype": "bfloat16",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1356201171875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 102.19413757324219,
   "xla_arg_mb": 148.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.005532972005312331,
   "inner": 18,
   "median_ms": 2.956606611001512,
   "q1_ms": 2.8024770832012615,
   "q3_ms": 3.011267666604706,
   "rel_err_out": 0.020131099618709833,
   "rel_err_grad_params": 0.03445727698235376,
   "rel_err_grad_x": 0.028982459890591516
  },
  {
   "code": "final",
   "run": "f_gpu_a",
   "suite": "layer",
   "routing": "dropless",
   "impl": "blocked",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1356201171875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 486.15872955322266,
   "xla_arg_mb": 52.015625,
   "xla_out_mb": 4.0,
   "compile_s": 0.004145882005104795,
   "inner": 18,
   "median_ms": 2.376373221826119,
   "q1_ms": 2.355175666808565,
   "q3_ms": 2.3939130272386437,
   "rel_err_out": 0.019465065167553218,
   "rel_err_grad_params": 0.03442669881708805,
   "rel_err_grad_x": 0.028680775702455974
  },
  {
   "code": "final",
   "run": "f_gpu_a",
   "suite": "layer",
   "routing": "dropless",
   "impl": "blocked",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1356201171875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 504.12943267822266,
   "xla_arg_mb": 148.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.008176697010640055,
   "inner": 8,
   "median_ms": 5.818356748932274,
   "q1_ms": 5.709188124455977,
   "q3_ms": 5.977200000415905,
   "rel_err_out": 0.019465065167553218,
   "rel_err_grad_params": 0.03442669881708805,
   "rel_err_grad_x": 0.028680775702455974
  },
  {
   "code": "final",
   "run": "f_gpu_a",
   "suite": "layer",
   "routing": "dropless",
   "impl": "ragged_dot",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1356201171875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 288.09474182128906,
   "xla_arg_mb": 52.015625,
   "xla_out_mb": 4.0,
   "compile_s": 0.007595078001031652,
   "inner": 7,
   "median_ms": 6.8517068575602025,
   "q1_ms": 6.603023499857434,
   "q3_ms": 7.002400072191708,
   "rel_err_out": 0.019465065167553218,
   "rel_err_grad_params": 0.03442669881708805,
   "rel_err_grad_x": 0.02868077401564277
  },
  {
   "code": "final",
   "run": "f_gpu_a",
   "suite": "layer",
   "routing": "dropless",
   "impl": "ragged_dot",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1356201171875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 768.0036926269531,
   "xla_arg_mb": 148.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.020669359000748955,
   "inner": 2,
   "median_ms": 19.026318499527406,
   "q1_ms": 18.872660752094816,
   "q3_ms": 21.49033125169808,
   "rel_err_out": 0.019465065167553218,
   "rel_err_grad_params": 0.03442669881708805,
   "rel_err_grad_x": 0.02868077401564277
  },
  {
   "code": "final",
   "run": "f_gpu_a",
   "suite": "layer",
   "routing": "capacity",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1356201171875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 60.15675354003906,
   "xla_arg_mb": 52.015625,
   "xla_out_mb": 4.0,
   "compile_s": 0.0023667509958613664,
   "inner": 52,
   "median_ms": 0.8811505963072145,
   "q1_ms": 0.8338856250002908,
   "q3_ms": 0.9360363942855305
  },
  {
   "code": "final",
   "run": "f_gpu_a",
   "suite": "layer",
   "routing": "capacity",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1356201171875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 104.1975326538086,
   "xla_arg_mb": 148.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.00473580599646084,
   "inner": 21,
   "median_ms": 2.384366333163122,
   "q1_ms": 2.1932276663608805,
   "q3_ms": 2.5870962379572338
  },
  {
   "code": "final",
   "run": "f_gpu_a",
   "suite": "layer",
   "routing": "dense",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1356201171875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 128.12525939941406,
   "xla_arg_mb": 52.015625,
   "xla_out_mb": 8.0,
   "compile_s": 0.0029100969986757264,
   "inner": 24,
   "median_ms": 2.3162544997224663,
   "q1_ms": 2.132461083419912,
   "q3_ms": 2.341661270899446,
   "rel_err_out": 0.01985194671361912,
   "rel_err_grad_params": 0.034447798045072635,
   "rel_err_grad_x": 0.02890855175770541
  },
  {
   "code": "final",
   "run": "f_gpu_a",
   "suite": "layer",
   "routing": "dense",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1356201171875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 232.0019760131836,
   "xla_arg_mb": 148.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.006729433007421903,
   "inner": 9,
   "median_ms": 5.182525110689716,
   "q1_ms": 5.154727000141671,
   "q3_ms": 5.205471222224232,
   "rel_err_out": 0.01985194671361912,
   "rel_err_grad_params": 0.034447798045072635,
   "rel_err_grad_x": 0.02890855175770541
  },
  {
   "code": "final",
   "run": "f_gpu_a",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1291961669921875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 307.26226806640625,
   "xla_arg_mb": 80.015625,
   "xla_out_mb": 32.0,
   "compile_s": 0.010793932000524364,
   "inner": 9,
   "median_ms": 5.044297223018172,
   "q1_ms": 5.040708777313638,
   "q3_ms": 5.089238721767389,
   "rel_err_out": 0.0319747751514326,
   "rel_err_grad_params": 0.02832304765654251,
   "rel_err_grad_x": 0.04002731502646485
  },
  {
   "code": "final",
   "run": "f_gpu_a",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1291961669921875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 386.2601089477539,
   "xla_arg_mb": 176.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.017987561004702002,
   "inner": 3,
   "median_ms": 16.011087335452128,
   "q1_ms": 15.950608333999602,
   "q3_ms": 16.09647383641762,
   "rel_err_out": 0.0319747751514326,
   "rel_err_grad_params": 0.02832304765654251,
   "rel_err_grad_x": 0.04002731502646485
  },
  {
   "code": "final",
   "run": "f_gpu_a",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "gather",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1291961669921875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 389.00980377197266,
   "xla_arg_mb": 80.015625,
   "xla_out_mb": 32.0,
   "compile_s": 0.007354993998887949,
   "inner": 9,
   "median_ms": 5.122019555225658,
   "q1_ms": 5.113335389372272,
   "q3_ms": 5.2757606665383925,
   "rel_err_out": 0.0319747751514326,
   "rel_err_grad_params": 0.02830920487332268,
   "rel_err_grad_x": 0.04000338068614109
  },
  {
   "code": "final",
   "run": "f_gpu_a",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "gather",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1291961669921875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 774.0046844482422,
   "xla_arg_mb": 176.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.022497882004245184,
   "inner": 2,
   "median_ms": 19.60137500282144,
   "q1_ms": 19.538496500899782,
   "q3_ms": 19.852577250276227,
   "rel_err_out": 0.0319747751514326,
   "rel_err_grad_params": 0.02830920487332268,
   "rel_err_grad_x": 0.04000338068614109
  },
  {
   "code": "final",
   "run": "f_gpu_a",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1291961669921875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 387.76151275634766,
   "xla_arg_mb": 80.015625,
   "xla_out_mb": 32.0,
   "compile_s": 0.0069248820072971284,
   "inner": 9,
   "median_ms": 5.46664622217779,
   "q1_ms": 5.391874444386405,
   "q3_ms": 5.482245722507489,
   "rel_err_out": 0.0319327827765743,
   "rel_err_grad_params": 0.028317132192730844,
   "rel_err_grad_x": 0.04000904859273607
  },
  {
   "code": "final",
   "run": "f_gpu_a",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1291961669921875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 677.2715301513672,
   "xla_arg_mb": 176.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.020879534000414424,
   "inner": 2,
   "median_ms": 19.29745099914726,
   "q1_ms": 19.229157001973363,
   "q3_ms": 19.446131751465145,
   "rel_err_out": 0.0319327827765743,
   "rel_err_grad_params": 0.028317132192730844,
   "rel_err_grad_x": 0.04000904859273607
  },
  {
   "code": "final",
   "run": "f_gpu_a",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused_gsum",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1291961669921875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 388.00785064697266,
   "xla_arg_mb": 80.015625,
   "xla_out_mb": 32.0,
   "compile_s": 0.006232273997738957,
   "inner": 11,
   "median_ms": 4.591810091038828,
   "q1_ms": 4.581675090858797,
   "q3_ms": 4.654951226918704,
   "rel_err_out": 0.0319747751514326,
   "rel_err_grad_params": 0.02830920487332268,
   "rel_err_grad_x": 0.04000338068614109
  },
  {
   "code": "final",
   "run": "f_gpu_a",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused_gsum",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1291961669921875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 677.2717742919922,
   "xla_arg_mb": 176.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.020230488997185603,
   "inner": 2,
   "median_ms": 20.68502749898471,
   "q1_ms": 20.648054996854626,
   "q3_ms": 20.82335375234834,
   "rel_err_out": 0.0319747751514326,
   "rel_err_grad_params": 0.02830920487332268,
   "rel_err_grad_x": 0.04000338068614109
  },
  {
   "code": "final",
   "run": "f_gpu_a",
   "suite": "layer",
   "routing": "dropless",
   "impl": "blocked",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1291961669921875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 3511.2527236938477,
   "xla_arg_mb": 80.015625,
   "xla_out_mb": 32.0,
   "compile_s": 0.01822264499787707,
   "inner": 3,
   "median_ms": 15.591140666704936,
   "q1_ms": 15.588809665738761,
   "q3_ms": 15.679312498832587,
   "rel_err_out": 0.031560396855348956,
   "rel_err_grad_params": 0.02828702150706022,
   "rel_err_grad_x": 0.03977877787376183
  },
  {
   "code": "final",
   "run": "f_gpu_a",
   "suite": "layer",
   "routing": "dropless",
   "impl": "blocked",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1291961669921875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 3641.5044326782227,
   "xla_arg_mb": 176.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.03798791499866638,
   "inner": 1,
   "median_ms": 36.12887300550938,
   "q1_ms": 36.05123150191503,
   "q3_ms": 36.37997300393181,
   "rel_err_out": 0.031560396855348956,
   "rel_err_grad_params": 0.02828702150706022,
   "rel_err_grad_x": 0.03977877787376183
  },
  {
   "code": "final",
   "run": "f_gpu_a",
   "suite": "layer",
   "routing": "dropless",
   "impl": "ragged_dot",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1291961669921875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 2304.7514877319336,
   "xla_arg_mb": 80.015625,
   "xla_out_mb": 32.0,
   "compile_s": 0.05590897100046277,
   "inner": 1,
   "median_ms": 55.05125998752192,
   "q1_ms": 51.8827345003956,
   "q3_ms": 55.124248494394124,
   "rel_err_out": 0.031560396855348956,
   "rel_err_grad_params": 0.028286975495102697,
   "rel_err_grad_x": 0.03977877862442696
  },
  {
   "code": "final",
   "run": "f_gpu_a",
   "suite": "layer",
   "routing": "dropless",
   "impl": "ragged_dot",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1291961669921875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 6145.644317626953,
   "xla_arg_mb": 176.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.1489165469974978,
   "inner": 1,
   "median_ms": 150.63394099706784,
   "q1_ms": 146.36338900163537,
   "q3_ms": 157.91976049513323,
   "rel_err_out": 0.031560396855348956,
   "rel_err_grad_params": 0.028286975495102697,
   "rel_err_grad_x": 0.03977877862442696
  },
  {
   "code": "final",
   "run": "f_gpu_a",
   "suite": "layer",
   "routing": "capacity",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1291961669921875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 481.7509994506836,
   "xla_arg_mb": 80.015625,
   "xla_out_mb": 32.0,
   "compile_s": 0.006477536997408606,
   "inner": 8,
   "median_ms": 6.452036501286784,
   "q1_ms": 6.219927936399472,
   "q3_ms": 6.600918312869908
  },
  {
   "code": "final",
   "run": "f_gpu_a",
   "suite": "layer",
   "routing": "capacity",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1291961669921875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 800.0039291381836,
   "xla_arg_mb": 176.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.014788097003474832,
   "inner": 3,
   "median_ms": 14.049625334640345,
   "q1_ms": 13.738111832935829,
   "q3_ms": 14.31821116663438
  },
  {
   "code": "final",
   "run": "f_gpu_a",
   "suite": "layer",
   "routing": "dense",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1291961669921875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 1025.000259399414,
   "xla_arg_mb": 80.015625,
   "xla_out_mb": 64.0,
   "compile_s": 0.01481795399740804,
   "inner": 3,
   "median_ms": 15.345663666569939,
   "q1_ms": 14.955209332886927,
   "q3_ms": 17.0887466641337,
   "rel_err_out": 0.03180156440897449,
   "rel_err_grad_params": 0.028307605822455884,
   "rel_err_grad_x": 0.039942760953268215
  },
  {
   "code": "final",
   "run": "f_gpu_a",
   "suite": "layer",
   "routing": "dense",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1291961669921875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 2115.7683334350586,
   "xla_arg_mb": 176.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.03766462400380988,
   "inner": 1,
   "median_ms": 38.00077999767382,
   "q1_ms": 37.537589501880575,
   "q3_ms": 38.20970949891489,
   "rel_err_out": 0.03180156440897449,
   "rel_err_grad_params": 0.028307605822455884,
   "rel_err_grad_x": 0.039942760953268215
  },
  {
   "code": "final",
   "run": "f_gpu_b",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 131072,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.12841415405273438,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 1339.0252075195312,
   "xla_arg_mb": 304.015625,
   "xla_out_mb": 256.0,
   "compile_s": 0.06572303200664464,
   "inner": 1,
   "median_ms": 36.941147001925856,
   "q1_ms": 36.82614550052676,
   "q3_ms": 36.964947496016975
  },
  {
   "code": "final",
   "run": "f_gpu_b",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 131072,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.12841415405273438,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 1851.538429260254,
   "xla_arg_mb": 400.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.1314108060032595,
   "inner": 1,
   "median_ms": 132.312700006878,
   "q1_ms": 131.1977340010344,
   "q3_ms": 136.83226299326634
  },
  {
   "code": "final",
   "run": "f_gpu_b",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused_gsum",
   "dtype": "float32",
   "tokens": 131072,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.12841415405273438,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 3082.0198135375977,
   "xla_arg_mb": 304.015625,
   "xla_out_mb": 256.0,
   "compile_s": 0.0897929339989787,
   "inner": 1,
   "median_ms": 33.60401299141813,
   "q1_ms": 33.53526350838365,
   "q3_ms": 33.76314450724749
  },
  {
   "code": "final",
   "run": "f_gpu_b",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused_gsum",
   "dtype": "float32",
   "tokens": 131072,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.12841415405273438,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 5130.035690307617,
   "xla_arg_mb": 400.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.2510207779996563,
   "inner": 1,
   "median_ms": 142.171564992168,
   "q1_ms": 141.9407885041437,
   "q3_ms": 143.19744799286127
  },
  {
   "code": "final",
   "run": "f_gpu_b",
   "suite": "layer",
   "routing": "dropless",
   "impl": "blocked",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 131072,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.12841415405273438,
   "empty_experts": 0,
   "op": "forward",
   "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 13.56GiB with allocator GPU_0_bfc on device 0. [executable_name='jit_predict'] [tf-allocator-allocation-error='']"
  },
  {
   "code": "final",
   "run": "f_gpu_b",
   "suite": "layer",
   "routing": "dropless",
   "impl": "blocked",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 131072,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.12841415405273438,
   "empty_experts": 0,
   "op": "train_step",
   "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 14.06GiB with allocator GPU_0_bfc on device 0. [tf-allocator-allocation-error=''] [executable_name='jit_train_step']"
  },
  {
   "code": "final",
   "run": "f_gpu_b",
   "suite": "layer",
   "routing": "capacity",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 131072,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.12841415405273438,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 3203.0005111694336,
   "xla_arg_mb": 304.015625,
   "xla_out_mb": 256.0,
   "compile_s": 0.03934257499349769,
   "inner": 1,
   "median_ms": 39.96783601178322,
   "q1_ms": 39.895629008242395,
   "q3_ms": 40.04383450228488
  },
  {
   "code": "final",
   "run": "f_gpu_b",
   "suite": "layer",
   "routing": "capacity",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 131072,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.12841415405273438,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 6016.753440856934,
   "xla_arg_mb": 400.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.10754655298660509,
   "inner": 1,
   "median_ms": 110.10972499207128,
   "q1_ms": 109.87831750389887,
   "q3_ms": 113.63896699913312
  },
  {
   "code": "final",
   "run": "f_gpu_b",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 131072,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.12839126586914062,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 696.5252075195312,
   "xla_arg_mb": 176.015625,
   "xla_out_mb": 128.0,
   "compile_s": 0.020600727992132306,
   "inner": 2,
   "median_ms": 19.325498498801608,
   "q1_ms": 19.30241000081878,
   "q3_ms": 19.328867751028156
  },
  {
   "code": "final",
   "run": "f_gpu_b",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 131072,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.12839126586914062,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 1032.117042541504,
   "xla_arg_mb": 272.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.06559703299717512,
   "inner": 1,
   "median_ms": 65.53116900613531,
   "q1_ms": 65.28445798903704,
   "q3_ms": 70.71929999801796
  },
  {
   "code": "final",
   "run": "f_gpu_b",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused_gsum",
   "dtype": "bfloat16",
   "tokens": 131072,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.12839126586914062,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 1543.0198135375977,
   "xla_arg_mb": 176.015625,
   "xla_out_mb": 128.0,
   "compile_s": 0.019333413991262205,
   "inner": 2,
   "median_ms": 18.093241000315174,
   "q1_ms": 17.939856250450248,
   "q3_ms": 18.481800001609372
  },
  {
   "code": "final",
   "run": "f_gpu_b",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused_gsum",
   "dtype": "bfloat16",
   "tokens": 131072,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.12839126586914062,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 2589.035690307617,
   "xla_arg_mb": 272.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.07738376100314781,
   "inner": 1,
   "median_ms": 75.70134199340828,
   "q1_ms": 75.54422500834335,
   "q3_ms": 76.21714149718173
  },
  {
   "code": "final",
   "run": "f_gpu_b",
   "suite": "layer",
   "routing": "dropless",
   "impl": "blocked",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 131072,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.12839126586914062,
   "empty_experts": 0,
   "op": "forward",
   "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 13.56GiB with allocator GPU_0_bfc on device 0. [executable_name='jit_predict'] [tf-allocator-allocation-error='']"
  },
  {
   "code": "final",
   "run": "f_gpu_b",
   "suite": "layer",
   "routing": "dropless",
   "impl": "blocked",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 131072,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.12839126586914062,
   "empty_experts": 0,
   "op": "train_step",
   "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 14.06GiB with allocator GPU_0_bfc on device 0. [tf-allocator-allocation-error=''] [executable_name='jit_train_step']"
  },
  {
   "code": "final",
   "run": "f_gpu_b",
   "suite": "layer",
   "routing": "capacity",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 131072,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.12839126586914062,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 1927.0007553100586,
   "xla_arg_mb": 176.015625,
   "xla_out_mb": 128.0,
   "compile_s": 0.02128140399872791,
   "inner": 2,
   "median_ms": 21.603282002615742,
   "q1_ms": 21.540899499086663,
   "q3_ms": 21.728621246438706
  },
  {
   "code": "final",
   "run": "f_gpu_b",
   "suite": "layer",
   "routing": "capacity",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 131072,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.12839126586914062,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 3009.5034408569336,
   "xla_arg_mb": 272.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.05274265899788588,
   "inner": 1,
   "median_ms": 53.80375499953516,
   "q1_ms": 53.624374500941485,
   "q3_ms": 53.906633002043236
  },
  {
   "code": "final",
   "run": "f_gpu_c",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.018524169921875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 266.88553619384766,
   "xla_arg_mb": 128.125,
   "xla_out_mb": 32.0,
   "compile_s": 0.006353703996865079,
   "inner": 12,
   "median_ms": 3.78857816758682,
   "q1_ms": 3.7653611252608243,
   "q3_ms": 3.78960383325951
  },
  {
   "code": "final",
   "run": "f_gpu_c",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.018524169921875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 724.8335342407227,
   "xla_arg_mb": 320.37500381469727,
   "xla_out_mb": 288.37511444091797,
   "compile_s": 0.028505472000688314,
   "inner": 3,
   "median_ms": 13.368729332190318,
   "q1_ms": 13.35802500155599,
   "q3_ms": 13.37737266779489
  },
  {
   "code": "final",
   "run": "f_gpu_c",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "gather",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.018524169921875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 339.88455963134766,
   "xla_arg_mb": 128.125,
   "xla_out_mb": 32.0,
   "compile_s": 0.005993250000756234,
   "inner": 11,
   "median_ms": 4.242331182467751,
   "q1_ms": 4.235081499659414,
   "q3_ms": 4.269417681577828
  },
  {
   "code": "final",
   "run": "f_gpu_c",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "gather",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.018524169921875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 593.7236785888672,
   "xla_arg_mb": 320.37500381469727,
   "xla_out_mb": 288.37511444091797,
   "compile_s": 0.015477824010304175,
   "inner": 3,
   "median_ms": 13.327219669008628,
   "q1_ms": 13.271921835742734,
   "q3_ms": 13.34864800204135
  },
  {
   "code": "final",
   "run": "f_gpu_c",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.018524169921875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 207.94681549072266,
   "xla_arg_mb": 128.125,
   "xla_out_mb": 32.0,
   "compile_s": 0.006435219998820685,
   "inner": 10,
   "median_ms": 4.614628599665593,
   "q1_ms": 4.579510799521813,
   "q3_ms": 4.698916349298088
  },
  {
   "code": "final",
   "run": "f_gpu_c",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.018524169921875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 330.3481903076172,
   "xla_arg_mb": 320.37500381469727,
   "xla_out_mb": 288.37511444091797,
   "compile_s": 0.01455666299443692,
   "inner": 4,
   "median_ms": 12.613330251042498,
   "q1_ms": 12.521423252110253,
   "q3_ms": 12.646369250433054
  },
  {
   "code": "final",
   "run": "f_gpu_c",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused_gsum",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.018524169921875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 215.93167877197266,
   "xla_arg_mb": 128.125,
   "xla_out_mb": 32.0,
   "compile_s": 0.0045755050086881965,
   "inner": 16,
   "median_ms": 2.9884173127356917,
   "q1_ms": 2.9812054685862677,
   "q3_ms": 2.9893027499383606
  },
  {
   "code": "final",
   "run": "f_gpu_c",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused_gsum",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.018524169921875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 329.8484344482422,
   "xla_arg_mb": 320.37500381469727,
   "xla_out_mb": 288.37511444091797,
   "compile_s": 0.013057023999863304,
   "inner": 4,
   "median_ms": 11.01987149741035,
   "q1_ms": 10.995950626238482,
   "q3_ms": 11.06828537376714
  },
  {
   "code": "final",
   "run": "f_gpu_c",
   "suite": "layer",
   "routing": "dropless",
   "impl": "blocked",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.018524169921875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 2042.5027236938477,
   "xla_arg_mb": 128.125,
   "xla_out_mb": 32.0,
   "compile_s": 0.03553454199573025,
   "inner": 5,
   "median_ms": 8.111647400073707,
   "q1_ms": 8.079887002531905,
   "q3_ms": 8.126270199136343
  },
  {
   "code": "final",
   "run": "f_gpu_c",
   "suite": "layer",
   "routing": "dropless",
   "impl": "blocked",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.018524169921875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 2312.5083389282227,
   "xla_arg_mb": 320.37500381469727,
   "xla_out_mb": 288.37511444091797,
   "compile_s": 0.07702206300746184,
   "inner": 2,
   "median_ms": 21.314350495231338,
   "q1_ms": 21.289660999173066,
   "q3_ms": 21.343701497244183
  },
  {
   "code": "final",
   "run": "f_gpu_c",
   "suite": "layer",
   "routing": "dropless",
   "impl": "ragged_dot",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.018524169921875,
   "empty_experts": 0,
   "op": "forward",
   "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 16.02GiB with allocator GPU_0_bfc on device 0. [tf-allocator-allocation-error='']"
  },
  {
   "code": "final",
   "run": "f_gpu_c",
   "suite": "layer",
   "routing": "dropless",
   "impl": "ragged_dot",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.018524169921875,
   "empty_experts": 0,
   "op": "train_step",
   "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Failed to get configs for: 5 out of 33 instructions. See logs for all failures. Example failure: "
  },
  {
   "code": "final",
   "run": "f_gpu_c",
   "suite": "layer",
   "routing": "capacity",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.018524169921875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 321.7507553100586,
   "xla_arg_mb": 128.125,
   "xla_out_mb": 32.0,
   "compile_s": 0.005468643008498475,
   "inner": 11,
   "median_ms": 4.540256544714794,
   "q1_ms": 4.422434364228552,
   "q3_ms": 4.546457318197513
  },
  {
   "code": "final",
   "run": "f_gpu_c",
   "suite": "layer",
   "routing": "capacity",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.018524169921875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 752.0027084350586,
   "xla_arg_mb": 320.37500381469727,
   "xla_out_mb": 288.37511444091797,
   "compile_s": 0.01270412300073076,
   "inner": 4,
   "median_ms": 11.160542249854188,
   "q1_ms": 11.158528997839312,
   "q3_ms": 11.182880500200554
  },
  {
   "code": "final",
   "run": "f_gpu_pieces",
   "suite": "pieces",
   "piece": "route",
   "dtype": "float32",
   "op": "fwd+bwd",
   "block_m": 64,
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "xla_temp_mb": 0.2812652587890625,
   "xla_arg_mb": 8.015625,
   "xla_out_mb": 8.015640258789062,
   "compile_s": 0.000857332008308731,
   "inner": 200,
   "median_ms": 0.049753130006138235,
   "q1_ms": 0.04972099246515427,
   "q3_ms": 0.05016355498810299
  },
  {
   "code": "final",
   "run": "f_gpu_pieces",
   "suite": "pieces",
   "piece": "sort_permute",
   "dtype": "float32",
   "op": "fwd+bwd",
   "block_m": 64,
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "xla_temp_mb": 0.0,
   "xla_arg_mb": 8.0,
   "xla_out_mb": 8.000019073486328,
   "compile_s": 0.0006296310020843521,
   "inner": 200,
   "median_ms": 0.04131992500333581,
   "q1_ms": 0.04130129003897309,
   "q3_ms": 0.04134356997383293
  },
  {
   "code": "final",
   "run": "f_gpu_pieces",
   "suite": "pieces",
   "piece": "experts_grouped_swiglu",
   "dtype": "float32",
   "op": "fwd+bwd",
   "block_m": 64,
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "xla_temp_mb": 136.07278442382812,
   "xla_arg_mb": 64.0,
   "xla_out_mb": 64.00004196166992,
   "compile_s": 0.009938161994796246,
   "inner": 10,
   "median_ms": 4.748922599537764,
   "q1_ms": 4.746567599795526,
   "q3_ms": 4.789776799589163
  },
  {
   "code": "final",
   "run": "f_gpu_pieces",
   "suite": "pieces",
   "piece": "unpermute_combine",
   "dtype": "float32",
   "op": "fwd+bwd",
   "block_m": 64,
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "xla_temp_mb": 8.000259399414062,
   "xla_arg_mb": 16.03125,
   "xla_out_mb": 16.03127670288086,
   "compile_s": 0.0007643459975952283,
   "inner": 200,
   "median_ms": 0.10256629502691794,
   "q1_ms": 0.10252394000417553,
   "q3_ms": 0.1026400099726743
  },
  {
   "code": "final",
   "run": "f_gpu_pieces",
   "suite": "pieces",
   "piece": "route",
   "dtype": "float32",
   "op": "fwd+bwd",
   "block_m": 64,
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "xla_temp_mb": 2.2500152587890625,
   "xla_arg_mb": 64.015625,
   "xla_out_mb": 64.01564025878906,
   "compile_s": 0.0011557960096979514,
   "inner": 125,
   "median_ms": 0.29938777606002986,
   "q1_ms": 0.2992659840383567,
   "q3_ms": 0.2994089759886265
  },
  {
   "code": "final",
   "run": "f_gpu_pieces",
   "suite": "pieces",
   "piece": "sort_permute",
   "dtype": "float32",
   "op": "fwd+bwd",
   "block_m": 64,
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "xla_temp_mb": 0.7502670288085938,
   "xla_arg_mb": 64.0,
   "xla_out_mb": 64.00001907348633,
   "compile_s": 0.0011127150064567104,
   "inner": 106,
   "median_ms": 0.3758046698079989,
   "q1_ms": 0.37510740566056416,
   "q3_ms": 0.3758701226066684
  },
  {
   "code": "final",
   "run": "f_gpu_pieces",
   "suite": "pieces",
   "piece": "experts_grouped_swiglu",
   "dtype": "float32",
   "op": "fwd+bwd",
   "block_m": 64,
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "xla_temp_mb": 822.5132751464844,
   "xla_arg_mb": 176.0,
   "xla_out_mb": 176.00004196166992,
   "compile_s": 0.049888173991348594,
   "inner": 1,
   "median_ms": 36.14800899231341,
   "q1_ms": 36.01676050311653,
   "q3_ms": 36.37692900520051
  },
  {
   "code": "final",
   "run": "f_gpu_pieces",
   "suite": "pieces",
   "piece": "unpermute_combine",
   "dtype": "float32",
   "op": "fwd+bwd",
   "block_m": 64,
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "xla_temp_mb": 64.00025939941406,
   "xla_arg_mb": 128.25,
   "xla_out_mb": 128.25002670288086,
   "compile_s": 0.00148300600994844,
   "inner": 55,
   "median_ms": 0.8035527999957346,
   "q1_ms": 0.8034098909278824,
   "q3_ms": 0.8036206545181233
  },
  {
   "code": "final",
   "run": "f_gpu_pieces",
   "suite": "pieces",
   "piece": "route",
   "dtype": "bfloat16",
   "op": "fwd+bwd",
   "block_m": 64,
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "xla_temp_mb": 0.2812652587890625,
   "xla_arg_mb": 4.015625,
   "xla_out_mb": 4.0156402587890625,
   "compile_s": 0.0008870989986462519,
   "inner": 200,
   "median_ms": 0.05666691002261359,
   "q1_ms": 0.056494657474104315,
   "q3_ms": 0.056929082529677544
  },
  {
   "code": "final",
   "run": "f_gpu_pieces",
   "suite": "pieces",
   "piece": "sort_permute",
   "dtype": "bfloat16",
   "op": "fwd+bwd",
   "block_m": 64,
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "xla_temp_mb": 0.0,
   "xla_arg_mb": 4.0,
   "xla_out_mb": 4.000019073486328,
   "compile_s": 0.0006405419990187511,
   "inner": 200,
   "median_ms": 0.05022171499149408,
   "q1_ms": 0.049768332501116674,
   "q3_ms": 0.05028451003454393
  },
  {
   "code": "final",
   "run": "f_gpu_pieces",
   "suite": "pieces",
   "piece": "experts_grouped_swiglu",
   "dtype": "bfloat16",
   "op": "fwd+bwd",
   "block_m": 64,
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "xla_temp_mb": 68.07278442382812,
   "xla_arg_mb": 32.0,
   "xla_out_mb": 32.00004196166992,
   "compile_s": 0.00412517499353271,
   "inner": 19,
   "median_ms": 2.4336047369472094,
   "q1_ms": 2.420971815713289,
   "q3_ms": 2.4372173157692147
  },
  {
   "code": "final",
   "run": "f_gpu_pieces",
   "suite": "pieces",
   "piece": "unpermute_combine",
   "dtype": "bfloat16",
   "op": "fwd+bwd",
   "block_m": 64,
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "xla_temp_mb": 1.52587890625e-05,
   "xla_arg_mb": 8.03125,
   "xla_out_mb": 8.03127670288086,
   "compile_s": 0.0006941129977349192,
   "inner": 200,
   "median_ms": 0.057161344957421534,
   "q1_ms": 0.057116634998237714,
   "q3_ms": 0.05800436501885997
  },
  {
   "code": "final",
   "run": "f_gpu_pieces",
   "suite": "pieces",
   "piece": "route",
   "dtype": "bfloat16",
   "op": "fwd+bwd",
   "block_m": 64,
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "xla_temp_mb": 2.2500152587890625,
   "xla_arg_mb": 32.015625,
   "xla_out_mb": 32.01564025878906,
   "compile_s": 0.0010156319913221523,
   "inner": 180,
   "median_ms": 0.18415869441620697,
   "q1_ms": 0.1840895027675692,
   "q3_ms": 0.1842052499543772
  },
  {
   "code": "final",
   "run": "f_gpu_pieces",
   "suite": "pieces",
   "piece": "sort_permute",
   "dtype": "bfloat16",
   "op": "fwd+bwd",
   "block_m": 64,
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "xla_temp_mb": 0.7502670288085938,
   "xla_arg_mb": 32.0,
   "xla_out_mb": 32.00001907348633,
   "compile_s": 0.0009730109886731952,
   "inner": 151,
   "median_ms": 0.2205331257809669,
   "q1_ms": 0.2202170629525366,
   "q3_ms": 0.22060262249969254
  },
  {
   "code": "final",
   "run": "f_gpu_pieces",
   "suite": "pieces",
   "piece": "experts_grouped_swiglu",
   "dtype": "bfloat16",
   "op": "fwd+bwd",
   "block_m": 64,
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "xla_temp_mb": 379.0132751464844,
   "xla_arg_mb": 88.0,
   "xla_out_mb": 88.00004196166992,
   "compile_s": 0.018771845003357157,
   "inner": 2,
   "median_ms": 17.82167200144613,
   "q1_ms": 17.807665251893923,
   "q3_ms": 17.823012254666537
  },
  {
   "code": "final",
   "run": "f_gpu_pieces",
   "suite": "pieces",
   "piece": "unpermute_combine",
   "dtype": "bfloat16",
   "op": "fwd+bwd",
   "block_m": 64,
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "xla_temp_mb": 64.00025939941406,
   "xla_arg_mb": 64.25,
   "xla_out_mb": 64.25002670288086,
   "compile_s": 0.0012550940009532496,
   "inner": 76,
   "median_ms": 0.5527435132535174,
   "q1_ms": 0.552666782919244,
   "q3_ms": 0.5528535197969943
  },
  {
   "code": "final",
   "run": "f_gpu_maxtok",
   "suite": "maxtok",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "op": "train_step",
   "tokens": 131072,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "peak_mb": 1448.213623046875
  },
  {
   "code": "final",
   "run": "f_gpu_maxtok",
   "suite": "maxtok",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "op": "train_step",
   "tokens": 262144,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "peak_mb": 2605.166748046875
  },
  {
   "code": "final",
   "run": "f_gpu_maxtok",
   "suite": "maxtok",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "op": "train_step",
   "tokens": 524288,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "peak_mb": 4909.5419921875
  },
  {
   "code": "final",
   "run": "f_gpu_maxtok",
   "suite": "maxtok",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "op": "train_step",
   "tokens": 1048576,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "peak_mb": 9558.4794921875
  },
  {
   "code": "final",
   "run": "f_gpu_maxtok",
   "suite": "maxtok",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "op": "train_step",
   "tokens": 2097152,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "error": "jax.errors.JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 16.09GiB with allocator GPU_0_bfc on device 0. [executable_name='jit_train_step'] [tf-allocator-allocation-error='']"
  },
  {
   "code": "final",
   "run": "f_gpu_maxtok",
   "suite": "maxtok_summary",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "max_tokens": 1048576
  },
  {
   "code": "final",
   "run": "f_gpu_maxtok",
   "suite": "maxtok",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused_gsum",
   "dtype": "bfloat16",
   "op": "train_step",
   "tokens": 131072,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "peak_mb": 3005.13232421875
  },
  {
   "code": "final",
   "run": "f_gpu_maxtok",
   "suite": "maxtok",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused_gsum",
   "dtype": "bfloat16",
   "op": "train_step",
   "tokens": 262144,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "peak_mb": 5701.13232421875
  },
  {
   "code": "final",
   "run": "f_gpu_maxtok",
   "suite": "maxtok",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused_gsum",
   "dtype": "bfloat16",
   "op": "train_step",
   "tokens": 524288,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "peak_mb": 11093.16357421875
  },
  {
   "code": "final",
   "run": "f_gpu_maxtok",
   "suite": "maxtok",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused_gsum",
   "dtype": "bfloat16",
   "op": "train_step",
   "tokens": 1048576,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "error": "jax.errors.JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 20.08GiB with allocator GPU_0_bfc on device 0. [tf-allocator-allocation-error=''] [executable_name='jit_train_step']"
  },
  {
   "code": "final",
   "run": "f_gpu_maxtok",
   "suite": "maxtok_summary",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused_gsum",
   "dtype": "bfloat16",
   "max_tokens": 524288
  },
  {
   "code": "before",
   "run": "gpu_maxtok_base",
   "suite": "maxtok",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "op": "train_step",
   "tokens": 131072,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "peak_mb": 2476.102294921875
  },
  {
   "code": "before",
   "run": "gpu_maxtok_base",
   "suite": "maxtok",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "op": "train_step",
   "tokens": 262144,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "peak_mb": 4652.102294921875
  },
  {
   "code": "before",
   "run": "gpu_maxtok_base",
   "suite": "maxtok",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "op": "train_step",
   "tokens": 524288,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "peak_mb": 9004.102294921875
  },
  {
   "code": "before",
   "run": "gpu_maxtok_base",
   "suite": "maxtok",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "op": "train_step",
   "tokens": 1048576,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "error": "jax.errors.JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 16.01GiB with allocator GPU_0_bfc on device 0. [tf-allocator-allocation-error=''] [executable_name='jit_train_step']"
  },
  {
   "code": "before",
   "run": "gpu_maxtok_base",
   "suite": "maxtok_summary",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "max_tokens": 524288
  },
  {
   "code": "before",
   "run": "gpu_maxtok_base",
   "suite": "maxtok",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused",
   "dtype": "bfloat16",
   "op": "train_step",
   "tokens": 131072,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "peak_mb": 3004.1318359375
  },
  {
   "code": "before",
   "run": "gpu_maxtok_base",
   "suite": "maxtok",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused",
   "dtype": "bfloat16",
   "op": "train_step",
   "tokens": 262144,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "peak_mb": 5699.1318359375
  },
  {
   "code": "before",
   "run": "gpu_maxtok_base",
   "suite": "maxtok",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused",
   "dtype": "bfloat16",
   "op": "train_step",
   "tokens": 524288,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "peak_mb": 11089.1630859375
  },
  {
   "code": "before",
   "run": "gpu_maxtok_base",
   "suite": "maxtok",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused",
   "dtype": "bfloat16",
   "op": "train_step",
   "tokens": 1048576,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "error": "jax.errors.JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 20.08GiB with allocator GPU_0_bfc on device 0. [executable_name='jit_train_step'] [tf-allocator-allocation-error='']"
  },
  {
   "code": "before",
   "run": "gpu_maxtok_base",
   "suite": "maxtok_summary",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused",
   "dtype": "bfloat16",
   "max_tokens": 524288
  },
  {
   "code": "before",
   "run": "gpu_maxtok_base",
   "suite": "maxtok",
   "routing": "capacity",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "op": "train_step",
   "tokens": 131072,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "peak_mb": 3425.599853515625
  },
  {
   "code": "before",
   "run": "gpu_maxtok_base",
   "suite": "maxtok",
   "routing": "capacity",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "op": "train_step",
   "tokens": 262144,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "peak_mb": 6567.10009765625
  },
  {
   "code": "before",
   "run": "gpu_maxtok_base",
   "suite": "maxtok",
   "routing": "capacity",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "op": "train_step",
   "tokens": 524288,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "error": "jax.errors.JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 11.76GiB with allocator GPU_0_bfc on device 0. [tf-allocator-allocation-error=''] [executable_name='jit_train_step']"
  },
  {
   "code": "before",
   "run": "gpu_maxtok_base",
   "suite": "maxtok_summary",
   "routing": "capacity",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "max_tokens": 262144
  },
  {
   "code": "final",
   "run": "gpu_maxtok_other",
   "suite": "maxtok",
   "routing": "dropless",
   "impl": "blocked",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "op": "train_step",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "peak_mb": 3961.60107421875
  },
  {
   "code": "final",
   "run": "gpu_maxtok_other",
   "suite": "maxtok",
   "routing": "dropless",
   "impl": "blocked",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "op": "train_step",
   "tokens": 65536,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "peak_mb": 7579.101806640625
  },
  {
   "code": "final",
   "run": "gpu_maxtok_other",
   "suite": "maxtok",
   "routing": "dropless",
   "impl": "blocked",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "op": "train_step",
   "tokens": 131072,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "error": "jax.errors.JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 14.06GiB with allocator GPU_0_bfc on device 0. [executable_name='jit_train_step'] [tf-allocator-allocation-error='']"
  },
  {
   "code": "final",
   "run": "gpu_maxtok_other",
   "suite": "maxtok_summary",
   "routing": "dropless",
   "impl": "blocked",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "max_tokens": 65536
  },
  {
   "code": "final",
   "run": "gpu_maxtok_other",
   "suite": "maxtok",
   "routing": "dropless",
   "impl": "ragged_dot",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "op": "train_step",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "peak_mb": 6465.740966796875
  },
  {
   "code": "final",
   "run": "gpu_maxtok_other",
   "suite": "maxtok",
   "routing": "dropless",
   "impl": "ragged_dot",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "op": "train_step",
   "tokens": 65536,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "error": "For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these."
  },
  {
   "code": "final",
   "run": "gpu_maxtok_other",
   "suite": "maxtok_summary",
   "routing": "dropless",
   "impl": "ragged_dot",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "max_tokens": 32768
  },
  {
   "code": "final",
   "run": "gpu_maxtok_other",
   "suite": "maxtok",
   "routing": "dense",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "op": "train_step",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "peak_mb": 2435.864501953125
  },
  {
   "code": "final",
   "run": "gpu_maxtok_other",
   "suite": "maxtok",
   "routing": "dense",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "op": "train_step",
   "tokens": 65536,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "peak_mb": 4583.61474609375
  },
  {
   "code": "final",
   "run": "gpu_maxtok_other",
   "suite": "maxtok",
   "routing": "dense",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "op": "train_step",
   "tokens": 131072,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "error": "jax.errors.JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 8.26GiB with allocator GPU_0_bfc on device 0. [executable_name='jit_train_step'] [tf-allocator-allocation-error='']"
  },
  {
   "code": "final",
   "run": "gpu_maxtok_other",
   "suite": "maxtok_summary",
   "routing": "dense",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "max_tokens": 65536
  },
  {
   "code": "before",
   "run": "gpu_layer_before",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.135498046875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 85.15921783447266,
   "xla_arg_mb": 56.015625,
   "xla_out_mb": 8.0,
   "compile_s": 0.0038574580030399375,
   "inner": 35,
   "median_ms": 1.3928792857249002,
   "q1_ms": 1.3673781858025387,
   "q3_ms": 1.4385926428076345,
   "rel_err_out": 0.01140456501241541,
   "rel_err_grad_params": 0.010704822355157323,
   "rel_err_grad_x": 0.00843557025022219
  },
  {
   "code": "before",
   "run": "gpu_layer_before",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.135498046875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 136.19632720947266,
   "xla_arg_mb": 152.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.014268345999880694,
   "inner": 9,
   "median_ms": 5.176025222075016,
   "q1_ms": 5.073723277746467,
   "q3_ms": 5.320447833784985,
   "rel_err_out": 0.01140456501241541,
   "rel_err_grad_params": 0.010704822355157323,
   "rel_err_grad_x": 0.00843557025022219
  },
  {
   "code": "before",
   "run": "gpu_layer_before",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "gather",
   "dtype": "float32",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.135498046875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 117.16336822509766,
   "xla_arg_mb": 56.015625,
   "xla_out_mb": 8.0,
   "compile_s": 0.0032953550035017543,
   "inner": 30,
   "median_ms": 1.5807256667661327,
   "q1_ms": 1.5276187332347035,
   "q3_ms": 1.595466916599738,
   "rel_err_out": 0.01140456501241541,
   "rel_err_grad_params": 0.010704822355157323,
   "rel_err_grad_x": 0.008434855028541565
  },
  {
   "code": "before",
   "run": "gpu_layer_before",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "gather",
   "dtype": "float32",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.135498046875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 202.0703582763672,
   "xla_arg_mb": 152.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.0072280950043932535,
   "inner": 10,
   "median_ms": 5.41605399994296,
   "q1_ms": 5.261301749851555,
   "q3_ms": 5.4245481001999,
   "rel_err_out": 0.01140456501241541,
   "rel_err_grad_params": 0.010704822355157323,
   "rel_err_grad_x": 0.008434855028541565
  },
  {
   "code": "before",
   "run": "gpu_layer_before",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused",
   "dtype": "float32",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.135498046875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 102.10208892822266,
   "xla_arg_mb": 56.015625,
   "xla_out_mb": 8.0,
   "compile_s": 0.003742480999790132,
   "inner": 27,
   "median_ms": 1.8267977036346026,
   "q1_ms": 1.7451744074437388,
   "q3_ms": 1.8933479074837156,
   "rel_err_out": 0.01140456501241541,
   "rel_err_grad_params": 0.010704822355157323,
   "rel_err_grad_x": 0.00843485500352305
  },
  {
   "code": "before",
   "run": "gpu_layer_before",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused",
   "dtype": "float32",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.135498046875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 170.0408172607422,
   "xla_arg_mb": 152.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.008317095001984853,
   "inner": 9,
   "median_ms": 5.104878222078292,
   "q1_ms": 5.062296611059638,
   "q3_ms": 5.456125611090101,
   "rel_err_out": 0.01140456501241541,
   "rel_err_grad_params": 0.010704822355157323,
   "rel_err_grad_x": 0.00843485500352305
  },
  {
   "code": "before",
   "run": "gpu_layer_before",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused_gsum",
   "dtype": "float32",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.135498046875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 102.16141510009766,
   "xla_arg_mb": 56.015625,
   "xla_out_mb": 8.0,
   "compile_s": 0.0032402110009570606,
   "inner": 33,
   "median_ms": 1.4375883635666193,
   "q1_ms": 1.3994167423912915,
   "q3_ms": 1.4843140302356472,
   "rel_err_out": 0.01140456501241541,
   "rel_err_grad_params": 0.010704822355157323,
   "rel_err_grad_x": 0.00843485500352305
  },
  {
   "code": "before",
   "run": "gpu_layer_before",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused_gsum",
   "dtype": "float32",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.135498046875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 170.0410614013672,
   "xla_arg_mb": 152.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.0070448380065499805,
   "inner": 10,
   "median_ms": 5.1114025998685975,
   "q1_ms": 4.8941275505058,
   "q3_ms": 5.376901949421153,
   "rel_err_out": 0.01140456501241541,
   "rel_err_grad_params": 0.010704822355157323,
   "rel_err_grad_x": 0.00843485500352305
  },
  {
   "code": "before",
   "run": "gpu_layer_before",
   "suite": "layer",
   "routing": "dropless",
   "impl": "blocked",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.135498046875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 486.15872955322266,
   "xla_arg_mb": 56.015625,
   "xla_out_mb": 8.0,
   "compile_s": 0.004340271007094998,
   "inner": 20,
   "median_ms": 2.6130288002605084,
   "q1_ms": 2.527103974898637,
   "q3_ms": 2.7821449999464676,
   "rel_err_out": 0.01140456501241541,
   "rel_err_grad_params": 0.010704833542363809,
   "rel_err_grad_x": 0.0084355704322865
  },
  {
   "code": "before",
   "run": "gpu_layer_before",
   "suite": "layer",
   "routing": "dropless",
   "impl": "blocked",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.135498046875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 504.25467681884766,
   "xla_arg_mb": 152.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.009870003996184096,
   "inner": 7,
   "median_ms": 6.185711857272379,
   "q1_ms": 6.1170003573352005,
   "q3_ms": 7.160327928431798,
   "rel_err_out": 0.01140456501241541,
   "rel_err_grad_params": 0.010704833542363809,
   "rel_err_grad_x": 0.0084355704322865
  },
  {
   "code": "before",
   "run": "gpu_layer_before",
   "suite": "layer",
   "routing": "dropless",
   "impl": "ragged_dot",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.135498046875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 288.09498596191406,
   "xla_arg_mb": 56.015625,
   "xla_out_mb": 8.0,
   "compile_s": 0.014720630999363493,
   "inner": 7,
   "median_ms": 7.807836428193176,
   "q1_ms": 7.6522104999250065,
   "q3_ms": 8.056819785581736,
   "rel_err_out": 0.01140456501241541,
   "rel_err_grad_params": 0.010704822589761478,
   "rel_err_grad_x": 0.0084355704322865
  },
  {
   "code": "before",
   "run": "gpu_layer_before",
   "suite": "layer",
   "routing": "dropless",
   "impl": "ragged_dot",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.135498046875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 768.0034484863281,
   "xla_arg_mb": 152.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.028356850001728162,
   "inner": 2,
   "median_ms": 21.970593003061367,
   "q1_ms": 21.77698875129863,
   "q3_ms": 22.34441749897087,
   "rel_err_out": 0.01140456501241541,
   "rel_err_grad_params": 0.010704822589761478,
   "rel_err_grad_x": 0.0084355704322865
  },
  {
   "code": "before",
   "run": "gpu_layer_before",
   "suite": "layer",
   "routing": "capacity",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.135498046875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 120.15675354003906,
   "xla_arg_mb": 56.015625,
   "xla_out_mb": 8.0,
   "compile_s": 0.003254506998928264,
   "inner": 28,
   "median_ms": 1.5597424285910424,
   "q1_ms": 1.5177520893173226,
   "q3_ms": 1.5709990536509264
  },
  {
   "code": "before",
   "run": "gpu_layer_before",
   "suite": "layer",
   "routing": "capacity",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.135498046875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 200.0024642944336,
   "xla_arg_mb": 152.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.0059631420008372515,
   "inner": 11,
   "median_ms": 4.175903454200703,
   "q1_ms": 3.980625818340658,
   "q3_ms": 4.307982363653455
  },
  {
   "code": "before",
   "run": "gpu_layer_before",
   "suite": "layer",
   "routing": "dense",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.135498046875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 256.00001525878906,
   "xla_arg_mb": 56.015625,
   "xla_out_mb": 8.0,
   "compile_s": 0.004395596006361302,
   "inner": 14,
   "median_ms": 3.96858771429314,
   "q1_ms": 3.717392857132446,
   "q3_ms": 4.183351928596884,
   "rel_err_out": 0.01140456501241541,
   "rel_err_grad_params": 0.010704805168015577,
   "rel_err_grad_x": 0.00843557463968868
  },
  {
   "code": "before",
   "run": "gpu_layer_before",
   "suite": "layer",
   "routing": "dense",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.135498046875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 448.0022201538086,
   "xla_arg_mb": 152.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.011006243999872822,
   "inner": 5,
   "median_ms": 10.534803400514647,
   "q1_ms": 9.789206899586134,
   "q3_ms": 11.010056699888082,
   "rel_err_out": 0.01140456501241541,
   "rel_err_grad_params": 0.010704805168015577,
   "rel_err_grad_x": 0.00843557463968868
  },
  {
   "code": "before",
   "run": "gpu_layer_before",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1292572021484375,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 646.2566299438477,
   "xla_arg_mb": 112.015625,
   "xla_out_mb": 64.0,
   "compile_s": 0.010381792002590373,
   "inner": 5,
   "median_ms": 8.609276400238741,
   "q1_ms": 8.598757599975215,
   "q3_ms": 8.65843749998021,
   "rel_err_out": 0.01101133789638197,
   "rel_err_grad_params": 0.010219446957116794,
   "rel_err_grad_x": 0.013931359767753112
  },
  {
   "code": "before",
   "run": "gpu_layer_before",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1292572021484375,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 1033.6450576782227,
   "xla_arg_mb": 208.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.0621792110032402,
   "inner": 1,
   "median_ms": 32.18574400671059,
   "q1_ms": 31.927360498229973,
   "q3_ms": 32.279651997669134,
   "rel_err_out": 0.01101133789638197,
   "rel_err_grad_params": 0.010219446957116794,
   "rel_err_grad_x": 0.013931359767753112
  },
  {
   "code": "before",
   "run": "gpu_layer_before",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "gather",
   "dtype": "float32",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1292572021484375,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 776.0098037719727,
   "xla_arg_mb": 112.015625,
   "xla_out_mb": 64.0,
   "compile_s": 0.010361453001678456,
   "inner": 5,
   "median_ms": 8.908923401031643,
   "q1_ms": 8.88659720076248,
   "q3_ms": 8.951790500577772,
   "rel_err_out": 0.01101133789638197,
   "rel_err_grad_params": 0.010216446286068267,
   "rel_err_grad_x": 0.013930634470649978
  },
  {
   "code": "before",
   "run": "gpu_layer_before",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "gather",
   "dtype": "float32",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1292572021484375,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 1419.5195770263672,
   "xla_arg_mb": 208.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.03641749999951571,
   "inner": 1,
   "median_ms": 35.34981099801371,
   "q1_ms": 35.188099998777034,
   "q3_ms": 35.49306200511637,
   "rel_err_out": 0.01101133789638197,
   "rel_err_grad_params": 0.010216446286068267,
   "rel_err_grad_x": 0.013930634470649978
  },
  {
   "code": "before",
   "run": "gpu_layer_before",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused",
   "dtype": "float32",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1292572021484375,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 774.7615127563477,
   "xla_arg_mb": 112.015625,
   "xla_out_mb": 64.0,
   "compile_s": 0.012626718998944853,
   "inner": 4,
   "median_ms": 10.924353749942384,
   "q1_ms": 10.916209375864128,
   "q3_ms": 11.004838375811232,
   "rel_err_out": 0.01101133789638197,
   "rel_err_grad_params": 0.010216446286068267,
   "rel_err_grad_x": 0.013930634460644907
  },
  {
   "code": "before",
   "run": "gpu_layer_before",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused",
   "dtype": "float32",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1292572021484375,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 1290.0193328857422,
   "xla_arg_mb": 208.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.036704183999972884,
   "inner": 1,
   "median_ms": 35.63924899935955,
   "q1_ms": 35.55245900133741,
   "q3_ms": 36.78639849749743,
   "rel_err_out": 0.01101133789638197,
   "rel_err_grad_params": 0.010216446286068267,
   "rel_err_grad_x": 0.013930634460644907
  },
  {
   "code": "before",
   "run": "gpu_layer_before",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused_gsum",
   "dtype": "float32",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1292572021484375,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 775.0078506469727,
   "xla_arg_mb": 112.015625,
   "xla_out_mb": 64.0,
   "compile_s": 0.010263759002555162,
   "inner": 5,
   "median_ms": 9.138052799971774,
   "q1_ms": 8.860096999706002,
   "q3_ms": 9.572270399803529,
   "rel_err_out": 0.01101133789638197,
   "rel_err_grad_params": 0.010216446286068267,
   "rel_err_grad_x": 0.013930634460644907
  },
  {
   "code": "before",
   "run": "gpu_layer_before",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused_gsum",
   "dtype": "float32",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1292572021484375,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 1290.0195770263672,
   "xla_arg_mb": 208.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.03701896800339455,
   "inner": 1,
   "median_ms": 40.88775899435859,
   "q1_ms": 40.0373845041031,
   "q3_ms": 41.287975502200425,
   "rel_err_out": 0.01101133789638197,
   "rel_err_grad_params": 0.010216446286068267,
   "rel_err_grad_x": 0.013930634460644907
  },
  {
   "code": "before",
   "run": "gpu_layer_before",
   "suite": "layer",
   "routing": "dropless",
   "impl": "blocked",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1292572021484375,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 3511.2527236938477,
   "xla_arg_mb": 112.015625,
   "xla_out_mb": 64.0,
   "compile_s": 0.017512433005322237,
   "inner": 2,
   "median_ms": 18.376262501988094,
   "q1_ms": 16.992223501802073,
   "q3_ms": 18.594733001009445,
   "rel_err_out": 0.01101133789638197,
   "rel_err_grad_params": 0.010219446957116794,
   "rel_err_grad_x": 0.013931359603212637
  },
  {
   "code": "before",
   "run": "gpu_layer_before",
   "suite": "layer",
   "routing": "dropless",
   "impl": "blocked",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1292572021484375,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 3641.5041885375977,
   "xla_arg_mb": 208.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.0442654990038136,
   "inner": 1,
   "median_ms": 42.457277995708864,
   "q1_ms": 41.85076050271164,
   "q3_ms": 45.00331449889927,
   "rel_err_out": 0.01101133789638197,
   "rel_err_grad_params": 0.010219446957116794,
   "rel_err_grad_x": 0.013931359603212637
  },
  {
   "code": "before",
   "run": "gpu_layer_before",
   "suite": "layer",
   "routing": "dropless",
   "impl": "ragged_dot",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1292572021484375,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 2304.7514877319336,
   "xla_arg_mb": 112.015625,
   "xla_out_mb": 64.0,
   "compile_s": 0.05793050300417235,
   "inner": 1,
   "median_ms": 55.73713299963856,
   "q1_ms": 51.98300050324178,
   "q3_ms": 57.58542949843104,
   "rel_err_out": 0.01101133789638197,
   "rel_err_grad_params": 0.010219446957116794,
   "rel_err_grad_x": 0.013931359603212637
  },
  {
   "code": "before",
   "run": "gpu_layer_before",
   "suite": "layer",
   "routing": "dropless",
   "impl": "ragged_dot",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1292572021484375,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 6145.644561767578,
   "xla_arg_mb": 208.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.16556128700176487,
   "inner": 1,
   "median_ms": 176.65569700329797,
   "q1_ms": 171.58187750101206,
   "q3_ms": 181.62443800247274,
   "rel_err_out": 0.01101133789638197,
   "rel_err_grad_params": 0.010219446957116794,
   "rel_err_grad_x": 0.013931359603212637
  },
  {
   "code": "before",
   "run": "gpu_layer_before",
   "suite": "layer",
   "routing": "capacity",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1292572021484375,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 801.7505111694336,
   "xla_arg_mb": 112.015625,
   "xla_out_mb": 64.0,
   "compile_s": 0.014874180997139774,
   "inner": 4,
   "median_ms": 11.265741250099381,
   "q1_ms": 10.997067874995992,
   "q3_ms": 11.632635374553502
  },
  {
   "code": "before",
   "run": "gpu_layer_before",
   "suite": "layer",
   "routing": "capacity",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1292572021484375,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 1504.1909408569336,
   "xla_arg_mb": 208.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.036596707002900075,
   "inner": 1,
   "median_ms": 28.40983799978858,
   "q1_ms": 28.163819002656965,
   "q3_ms": 29.97587650315836
  },
  {
   "code": "before",
   "run": "gpu_layer_before",
   "suite": "layer",
   "routing": "dense",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1292572021484375,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 2048.000259399414,
   "xla_arg_mb": 112.015625,
   "xla_out_mb": 64.0,
   "compile_s": 0.03667874300299445,
   "inner": 1,
   "median_ms": 30.94256999611389,
   "q1_ms": 30.485250001220265,
   "q3_ms": 32.08103450015187,
   "rel_err_out": 0.01101133789638197,
   "rel_err_grad_params": 0.010219446957116794,
   "rel_err_grad_x": 0.013931359603212637
  },
  {
   "code": "before",
   "run": "gpu_layer_before",
   "suite": "layer",
   "routing": "dense",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1292572021484375,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 4163.768089294434,
   "xla_arg_mb": 208.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.08115319500211626,
   "inner": 1,
   "median_ms": 86.15760200336808,
   "q1_ms": 81.01642049587099,
   "q3_ms": 89.07914399969741,
   "rel_err_out": 0.01101133789638197,
   "rel_err_grad_params": 0.010219446957116794,
   "rel_err_grad_x": 0.013931359603212637
  },
  {
   "code": "before",
   "run": "gpu_layer_before",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1356201171875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 50.690467834472656,
   "xla_arg_mb": 52.015625,
   "xla_out_mb": 4.0,
   "compile_s": 0.006631516000197735,
   "inner": 47,
   "median_ms": 0.9915452765392517,
   "q1_ms": 0.9693462447141932,
   "q3_ms": 1.0114266063263788,
   "rel_err_out": 0.020131099618709833,
   "rel_err_grad_params": 0.0344819908194601,
   "rel_err_grad_x": 0.029022456572827038
  },
  {
   "code": "before",
   "run": "gpu_layer_before",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1356201171875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 59.005409240722656,
   "xla_arg_mb": 148.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.00600106299680192,
   "inner": 15,
   "median_ms": 3.0787005996292764,
   "q1_ms": 2.974966700033595,
   "q3_ms": 3.147558166771584,
   "rel_err_out": 0.020131099618709833,
   "rel_err_grad_params": 0.0344819908194601,
   "rel_err_grad_x": 0.029022456572827038
  },
  {
   "code": "before",
   "run": "gpu_layer_before",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "gather",
   "dtype": "bfloat16",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1356201171875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 58.663368225097656,
   "xla_arg_mb": 52.015625,
   "xla_out_mb": 4.0,
   "compile_s": 0.0028766920004272833,
   "inner": 42,
   "median_ms": 0.9862177381971076,
   "q1_ms": 0.9268316904732881,
   "q3_ms": 1.005202821466983,
   "rel_err_out": 0.020131099618709833,
   "rel_err_grad_params": 0.03445727541303919,
   "rel_err_grad_x": 0.028982459890591516
  },
  {
   "code": "before",
   "run": "gpu_layer_before",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "gather",
   "dtype": "bfloat16",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1356201171875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 102.12992858886719,
   "xla_arg_mb": 148.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.005386110002291389,
   "inner": 18,
   "median_ms": 2.774898166535422,
   "q1_ms": 2.7269986666295964,
   "q3_ms": 2.784831416445539,
   "rel_err_out": 0.020131099618709833,
   "rel_err_grad_params": 0.03445727541303919,
   "rel_err_grad_x": 0.028982459890591516
  },
  {
   "code": "before",
   "run": "gpu_layer_before",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused",
   "dtype": "bfloat16",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1356201171875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 51.102088928222656,
   "xla_arg_mb": 52.015625,
   "xla_out_mb": 4.0,
   "compile_s": 0.0031733640062157065,
   "inner": 36,
   "median_ms": 1.1019848611289893,
   "q1_ms": 1.0501716805164405,
   "q3_ms": 1.1133271389351447,
   "rel_err_out": 0.020063212378551402,
   "rel_err_grad_params": 0.034431948875563535,
   "rel_err_grad_x": 0.02899075704691895
  },
  {
   "code": "before",
   "run": "gpu_layer_before",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused",
   "dtype": "bfloat16",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1356201171875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 102.16264343261719,
   "xla_arg_mb": 148.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.020820420002564788,
   "inner": 8,
   "median_ms": 2.9974433746247087,
   "q1_ms": 2.9555948126471776,
   "q3_ms": 3.0879026244292618,
   "rel_err_out": 0.020063212378551402,
   "rel_err_grad_params": 0.034431948875563535,
   "rel_err_grad_x": 0.02899075704691895
  },
  {
   "code": "before",
   "run": "gpu_layer_before",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused_gsum",
   "dtype": "bfloat16",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1356201171875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 51.192665100097656,
   "xla_arg_mb": 52.015625,
   "xla_out_mb": 4.0,
   "compile_s": 0.01125541500368854,
   "inner": 42,
   "median_ms": 0.8533642380609221,
   "q1_ms": 0.8456589524333442,
   "q3_ms": 0.8612653929005665,
   "rel_err_out": 0.020131099618709833,
   "rel_err_grad_params": 0.03445727541303919,
   "rel_err_grad_x": 0.028982459890591516
  },
  {
   "code": "before",
   "run": "gpu_layer_before",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused_gsum",
   "dtype": "bfloat16",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1356201171875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 102.19413757324219,
   "xla_arg_mb": 148.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.005966347998764832,
   "inner": 18,
   "median_ms": 2.7369613890186884,
   "q1_ms": 2.716995833098837,
   "q3_ms": 2.778257861185314,
   "rel_err_out": 0.020131099618709833,
   "rel_err_grad_params": 0.03445727541303919,
   "rel_err_grad_x": 0.028982459890591516
  },
  {
   "code": "before",
   "run": "gpu_layer_before",
   "suite": "layer",
   "routing": "dropless",
   "impl": "blocked",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1356201171875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 486.15872955322266,
   "xla_arg_mb": 52.015625,
   "xla_out_mb": 4.0,
   "compile_s": 0.0053095550028956495,
   "inner": 16,
   "median_ms": 2.5115428747994883,
   "q1_ms": 2.45563993757969,
   "q3_ms": 2.5993561562245304,
   "rel_err_out": 0.019465065167553218,
   "rel_err_grad_params": 0.03442669863608144,
   "rel_err_grad_x": 0.028680775702455974
  },
  {
   "code": "before",
   "run": "gpu_layer_before",
   "suite": "layer",
   "routing": "dropless",
   "impl": "blocked",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1356201171875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 504.12943267822266,
   "xla_arg_mb": 148.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.014924937000614591,
   "inner": 7,
   "median_ms": 6.63033542930082,
   "q1_ms": 6.5532294289109165,
   "q3_ms": 6.650816428840958,
   "rel_err_out": 0.019465065167553218,
   "rel_err_grad_params": 0.03442669863608144,
   "rel_err_grad_x": 0.028680775702455974
  },
  {
   "code": "before",
   "run": "gpu_layer_before",
   "suite": "layer",
   "routing": "dropless",
   "impl": "ragged_dot",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1356201171875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 288.09498596191406,
   "xla_arg_mb": 52.015625,
   "xla_out_mb": 4.0,
   "compile_s": 0.007754941994789988,
   "inner": 7,
   "median_ms": 6.99414371462938,
   "q1_ms": 6.931771356903482,
   "q3_ms": 7.140053285443823,
   "rel_err_out": 0.019465065167553218,
   "rel_err_grad_params": 0.03442669863608144,
   "rel_err_grad_x": 0.028680775702455974
  },
  {
   "code": "before",
   "run": "gpu_layer_before",
   "suite": "layer",
   "routing": "dropless",
   "impl": "ragged_dot",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1356201171875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 768.0034484863281,
   "xla_arg_mb": 148.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.02254013399942778,
   "inner": 2,
   "median_ms": 19.433125999057665,
   "q1_ms": 18.75920074962778,
   "q3_ms": 20.88977600033104,
   "rel_err_out": 0.019465065167553218,
   "rel_err_grad_params": 0.03442669863608144,
   "rel_err_grad_x": 0.028680775702455974
  },
  {
   "code": "before",
   "run": "gpu_layer_before",
   "suite": "layer",
   "routing": "capacity",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1356201171875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 60.15675354003906,
   "xla_arg_mb": 52.015625,
   "xla_out_mb": 4.0,
   "compile_s": 0.0023103109997464344,
   "inner": 53,
   "median_ms": 0.862164905691247,
   "q1_ms": 0.848144896186774,
   "q3_ms": 0.878156188744744
  },
  {
   "code": "before",
   "run": "gpu_layer_before",
   "suite": "layer",
   "routing": "capacity",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1356201171875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 104.1975326538086,
   "xla_arg_mb": 148.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.005027781997341663,
   "inner": 21,
   "median_ms": 2.260429857068673,
   "q1_ms": 2.2111506666641656,
   "q3_ms": 2.3340994524741765
  },
  {
   "code": "before",
   "run": "gpu_layer_before",
   "suite": "layer",
   "routing": "dense",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1356201171875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 128.12525939941406,
   "xla_arg_mb": 52.015625,
   "xla_out_mb": 8.0,
   "compile_s": 0.004279255997971632,
   "inner": 23,
   "median_ms": 2.1136492609714526,
   "q1_ms": 2.10822380430055,
   "q3_ms": 2.2630416522430177,
   "rel_err_out": 0.01985194671361912,
   "rel_err_grad_params": 0.03444779553851689,
   "rel_err_grad_x": 0.02890855175770541
  },
  {
   "code": "before",
   "run": "gpu_layer_before",
   "suite": "layer",
   "routing": "dense",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1356201171875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 232.0019760131836,
   "xla_arg_mb": 148.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.007640463998541236,
   "inner": 9,
   "median_ms": 5.261959555759353,
   "q1_ms": 5.216002666808587,
   "q3_ms": 5.44476988862799,
   "rel_err_out": 0.01985194671361912,
   "rel_err_grad_params": 0.03444779553851689,
   "rel_err_grad_x": 0.02890855175770541
  },
  {
   "code": "before",
   "run": "gpu_layer_before",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1291961669921875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 324.00662994384766,
   "xla_arg_mb": 80.015625,
   "xla_out_mb": 32.0,
   "compile_s": 0.00637587300298037,
   "inner": 10,
   "median_ms": 5.526079799892614,
   "q1_ms": 5.424249150019023,
   "q3_ms": 5.721544149855617,
   "rel_err_out": 0.0319747751514326,
   "rel_err_grad_params": 0.02832304717759384,
   "rel_err_grad_x": 0.04002731502646485
  },
  {
   "code": "before",
   "run": "gpu_layer_before",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1291961669921875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 524.0054092407227,
   "xla_arg_mb": 176.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.022014108995790593,
   "inner": 2,
   "median_ms": 19.661893500597216,
   "q1_ms": 19.236787002228084,
   "q3_ms": 19.79759049936547,
   "rel_err_out": 0.0319747751514326,
   "rel_err_grad_params": 0.02832304717759384,
   "rel_err_grad_x": 0.04002731502646485
  },
  {
   "code": "before",
   "run": "gpu_layer_before",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "gather",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1291961669921875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 389.00980377197266,
   "xla_arg_mb": 80.015625,
   "xla_out_mb": 32.0,
   "compile_s": 0.006744419995811768,
   "inner": 9,
   "median_ms": 5.284613443816649,
   "q1_ms": 5.225319166735668,
   "q3_ms": 5.520766444129144,
   "rel_err_out": 0.0319747751514326,
   "rel_err_grad_params": 0.028309205338670283,
   "rel_err_grad_x": 0.04000338068614109
  },
  {
   "code": "before",
   "run": "gpu_layer_before",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "gather",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1291961669921875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 774.0046844482422,
   "xla_arg_mb": 176.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.022491912000987213,
   "inner": 2,
   "median_ms": 20.7098409991886,
   "q1_ms": 20.082045750314137,
   "q3_ms": 21.731491749960696,
   "rel_err_out": 0.0319747751514326,
   "rel_err_grad_params": 0.028309205338670283,
   "rel_err_grad_x": 0.04000338068614109
  },
  {
   "code": "before",
   "run": "gpu_layer_before",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1291961669921875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 387.76151275634766,
   "xla_arg_mb": 80.015625,
   "xla_out_mb": 32.0,
   "compile_s": 0.01042695600335719,
   "inner": 9,
   "median_ms": 5.4839306661354685,
   "q1_ms": 5.44957955531168,
   "q3_ms": 6.166605555335991,
   "rel_err_out": 0.0319327827765743,
   "rel_err_grad_params": 0.028317131254175693,
   "rel_err_grad_x": 0.04000904859273607
  },
  {
   "code": "before",
   "run": "gpu_layer_before",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1291961669921875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 677.2715301513672,
   "xla_arg_mb": 176.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.02325936400302453,
   "inner": 2,
   "median_ms": 19.432714998401934,
   "q1_ms": 19.263527250586776,
   "q3_ms": 19.80490425012249,
   "rel_err_out": 0.0319327827765743,
   "rel_err_grad_params": 0.028317131254175693,
   "rel_err_grad_x": 0.04000904859273607
  },
  {
   "code": "before",
   "run": "gpu_layer_before",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused_gsum",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1291961669921875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 388.00785064697266,
   "xla_arg_mb": 80.015625,
   "xla_out_mb": 32.0,
   "compile_s": 0.007633091998286545,
   "inner": 9,
   "median_ms": 4.938566666775539,
   "q1_ms": 4.770694055575101,
   "q3_ms": 5.016771277749083,
   "rel_err_out": 0.0319747751514326,
   "rel_err_grad_params": 0.028309205338670283,
   "rel_err_grad_x": 0.04000338068614109
  },
  {
   "code": "before",
   "run": "gpu_layer_before",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused_gsum",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1291961669921875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 677.2717742919922,
   "xla_arg_mb": 176.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.02054163299908396,
   "inner": 2,
   "median_ms": 18.836061499314383,
   "q1_ms": 18.673914000828518,
   "q3_ms": 20.078043000467005,
   "rel_err_out": 0.0319747751514326,
   "rel_err_grad_params": 0.028309205338670283,
   "rel_err_grad_x": 0.04000338068614109
  },
  {
   "code": "before",
   "run": "gpu_layer_before",
   "suite": "layer",
   "routing": "dropless",
   "impl": "blocked",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1291961669921875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 3511.2527236938477,
   "xla_arg_mb": 80.015625,
   "xla_out_mb": 32.0,
   "compile_s": 0.018321142997592688,
   "inner": 3,
   "median_ms": 16.2261783334543,
   "q1_ms": 15.992558332072804,
   "q3_ms": 16.99083499988774,
   "rel_err_out": 0.031560396855348956,
   "rel_err_grad_params": 0.02828702133814283,
   "rel_err_grad_x": 0.03977877787376183
  },
  {
   "code": "before",
   "run": "gpu_layer_before",
   "suite": "layer",
   "routing": "dropless",
   "impl": "blocked",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1291961669921875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 3641.5041885375977,
   "xla_arg_mb": 176.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.040429880005831365,
   "inner": 1,
   "median_ms": 41.255971998907626,
   "q1_ms": 38.368544501281576,
   "q3_ms": 41.945762997784186,
   "rel_err_out": 0.031560396855348956,
   "rel_err_grad_params": 0.02828702133814283,
   "rel_err_grad_x": 0.03977877787376183
  },
  {
   "code": "before",
   "run": "gpu_layer_before",
   "suite": "layer",
   "routing": "dropless",
   "impl": "ragged_dot",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1291961669921875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 2304.7517318725586,
   "xla_arg_mb": 80.015625,
   "xla_out_mb": 32.0,
   "compile_s": 0.05229185099597089,
   "inner": 1,
   "median_ms": 55.85842900472926,
   "q1_ms": 53.126565500861034,
   "q3_ms": 56.521883998357225,
   "rel_err_out": 0.031560396855348956,
   "rel_err_grad_params": 0.028286975459827633,
   "rel_err_grad_x": 0.03977877862442696
  },
  {
   "code": "before",
   "run": "gpu_layer_before",
   "suite": "layer",
   "routing": "dropless",
   "impl": "ragged_dot",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1291961669921875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 6145.644805908203,
   "xla_arg_mb": 176.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.14887038899905747,
   "inner": 1,
   "median_ms": 165.6325109943282,
   "q1_ms": 164.4945435000409,
   "q3_ms": 166.20483900260297,
   "rel_err_out": 0.031560396855348956,
   "rel_err_grad_params": 0.028286975459827633,
   "rel_err_grad_x": 0.03977877862442696
  },
  {
   "code": "before",
   "run": "gpu_layer_before",
   "suite": "layer",
   "routing": "capacity",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1291961669921875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 481.7507553100586,
   "xla_arg_mb": 80.015625,
   "xla_out_mb": 32.0,
   "compile_s": 0.008671886003867257,
   "inner": 7,
   "median_ms": 6.390408713611707,
   "q1_ms": 5.756029214743259,
   "q3_ms": 6.467573857142789
  },
  {
   "code": "before",
   "run": "gpu_layer_before",
   "suite": "layer",
   "routing": "capacity",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1291961669921875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 752.3784408569336,
   "xla_arg_mb": 176.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.018924773998151068,
   "inner": 3,
   "median_ms": 15.146168334467802,
   "q1_ms": 15.13026683339073,
   "q3_ms": 15.555332499692062
  },
  {
   "code": "before",
   "run": "gpu_layer_before",
   "suite": "layer",
   "routing": "dense",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1291961669921875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 1025.000259399414,
   "xla_arg_mb": 80.015625,
   "xla_out_mb": 64.0,
   "compile_s": 0.019180637995305005,
   "inner": 3,
   "median_ms": 17.237985998993583,
   "q1_ms": 16.91200750065036,
   "q3_ms": 17.405259334433744,
   "rel_err_out": 0.03180156440897449,
   "rel_err_grad_params": 0.02830780452356545,
   "rel_err_grad_x": 0.03994272001523532
  },
  {
   "code": "before",
   "run": "gpu_layer_before",
   "suite": "layer",
   "routing": "dense",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1291961669921875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 2115.7680892944336,
   "xla_arg_mb": 176.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.0405629710003268,
   "inner": 1,
   "median_ms": 40.9404260062729,
   "q1_ms": 39.64181849732995,
   "q3_ms": 42.410707003000425,
   "rel_err_out": 0.03180156440897449,
   "rel_err_grad_params": 0.02830780452356545,
   "rel_err_grad_x": 0.03994272001523532
  },
  {
   "code": "before",
   "run": "gpu_pieces_before",
   "suite": "pieces",
   "piece": "route",
   "dtype": "float32",
   "op": "fwd+bwd",
   "block_m": 64,
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "xla_temp_mb": 0.2812652587890625,
   "xla_arg_mb": 8.015625,
   "xla_out_mb": 8.015640258789062,
   "compile_s": 0.0009124310017796233,
   "inner": 200,
   "median_ms": 0.061225350000313476,
   "q1_ms": 0.05838153250806499,
   "q3_ms": 0.06449676000556792
  },
  {
   "code": "before",
   "run": "gpu_pieces_before",
   "suite": "pieces",
   "piece": "sort_permute",
   "dtype": "float32",
   "op": "fwd+bwd",
   "block_m": 64,
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "xla_temp_mb": 0.0,
   "xla_arg_mb": 8.0,
   "xla_out_mb": 8.000019073486328,
   "compile_s": 0.0008090750052360818,
   "inner": 200,
   "median_ms": 0.051125089994457085,
   "q1_ms": 0.046729002497158945,
   "q3_ms": 0.06021836748914211
  },
  {
   "code": "before",
   "run": "gpu_pieces_before",
   "suite": "pieces",
   "piece": "experts_grouped_swiglu",
   "dtype": "float32",
   "op": "fwd+bwd",
   "block_m": 64,
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "xla_temp_mb": 136.07278442382812,
   "xla_arg_mb": 64.0,
   "xla_out_mb": 64.00004196166992,
   "compile_s": 0.011619033000897616,
   "inner": 10,
   "median_ms": 4.776782100088894,
   "q1_ms": 4.76193045033142,
   "q3_ms": 4.804494149720995
  },
  {
   "code": "before",
   "run": "gpu_pieces_before",
   "suite": "pieces",
   "piece": "unpermute_combine",
   "dtype": "float32",
   "op": "fwd+bwd",
   "block_m": 64,
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "xla_temp_mb": 8.000259399414062,
   "xla_arg_mb": 16.03125,
   "xla_out_mb": 16.03127670288086,
   "compile_s": 0.0009218190025421791,
   "inner": 200,
   "median_ms": 0.10333894500945462,
   "q1_ms": 0.1028186875009851,
   "q3_ms": 0.10410872999273124
  },
  {
   "code": "before",
   "run": "gpu_pieces_before",
   "suite": "pieces",
   "piece": "route",
   "dtype": "float32",
   "op": "fwd+bwd",
   "block_m": 64,
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "xla_temp_mb": 2.2500152587890625,
   "xla_arg_mb": 64.015625,
   "xla_out_mb": 64.01564025878906,
   "compile_s": 0.0014073600032133982,
   "inner": 123,
   "median_ms": 0.2981661382364109,
   "q1_ms": 0.2944280772126371,
   "q3_ms": 0.29913407724518215
  },
  {
   "code": "before",
   "run": "gpu_pieces_before",
   "suite": "pieces",
   "piece": "sort_permute",
   "dtype": "float32",
   "op": "fwd+bwd",
   "block_m": 64,
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "xla_temp_mb": 0.7502670288085938,
   "xla_arg_mb": 64.0,
   "xla_out_mb": 64.00001907348633,
   "compile_s": 0.0011841369996545836,
   "inner": 107,
   "median_ms": 0.3755362616676759,
   "q1_ms": 0.375360752310099,
   "q3_ms": 0.3756891682289245
  },
  {
   "code": "before",
   "run": "gpu_pieces_before",
   "suite": "pieces",
   "piece": "experts_grouped_swiglu",
   "dtype": "float32",
   "op": "fwd+bwd",
   "block_m": 64,
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "xla_temp_mb": 1032.2715148925781,
   "xla_arg_mb": 176.0,
   "xla_out_mb": 176.00004196166992,
   "compile_s": 0.06525702099315822,
   "inner": 1,
   "median_ms": 35.556545000872575,
   "q1_ms": 35.48979849801981,
   "q3_ms": 35.58694749881397
  },
  {
   "code": "before",
   "run": "gpu_pieces_before",
   "suite": "pieces",
   "piece": "unpermute_combine",
   "dtype": "float32",
   "op": "fwd+bwd",
   "block_m": 64,
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "xla_temp_mb": 64.00025939941406,
   "xla_arg_mb": 128.25,
   "xla_out_mb": 128.25002670288086,
   "compile_s": 0.0016494000010425225,
   "inner": 54,
   "median_ms": 0.8066432407455674,
   "q1_ms": 0.8064181111640883,
   "q3_ms": 0.8076310555393489
  },
  {
   "code": "before",
   "run": "gpu_pieces_before",
   "suite": "pieces",
   "piece": "route",
   "dtype": "bfloat16",
   "op": "fwd+bwd",
   "block_m": 64,
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "xla_temp_mb": 0.2812652587890625,
   "xla_arg_mb": 4.015625,
   "xla_out_mb": 4.0156402587890625,
   "compile_s": 0.0008212079992517829,
   "inner": 200,
   "median_ms": 0.055999084979703184,
   "q1_ms": 0.05590030001258128,
   "q3_ms": 0.056575482503831154
  },
  {
   "code": "before",
   "run": "gpu_pieces_before",
   "suite": "pieces",
   "piece": "sort_permute",
   "dtype": "bfloat16",
   "op": "fwd+bwd",
   "block_m": 64,
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "xla_temp_mb": 0.0,
   "xla_arg_mb": 4.0,
   "xla_out_mb": 4.000019073486328,
   "compile_s": 0.0006298949956544675,
   "inner": 200,
   "median_ms": 0.04225395998219028,
   "q1_ms": 0.04177379998509423,
   "q3_ms": 0.04368312249425799
  },
  {
   "code": "before",
   "run": "gpu_pieces_before",
   "suite": "pieces",
   "piece": "experts_grouped_swiglu",
   "dtype": "bfloat16",
   "op": "fwd+bwd",
   "block_m": 64,
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "xla_temp_mb": 68.07278442382812,
   "xla_arg_mb": 32.0,
   "xla_out_mb": 32.00004196166992,
   "compile_s": 0.003978407999966294,
   "inner": 20,
   "median_ms": 2.4347585000214167,
   "q1_ms": 2.4323537250893423,
   "q3_ms": 2.4353994751436403
  },
  {
   "code": "before",
   "run": "gpu_pieces_before",
   "suite": "pieces",
   "piece": "unpermute_combine",
   "dtype": "bfloat16",
   "op": "fwd+bwd",
   "block_m": 64,
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "xla_temp_mb": 1.52587890625e-05,
   "xla_arg_mb": 8.03125,
   "xla_out_mb": 8.03127670288086,
   "compile_s": 0.0007029129992588423,
   "inner": 200,
   "median_ms": 0.046860395013936795,
   "q1_ms": 0.04677513748902129,
   "q3_ms": 0.05061757750809193
  },
  {
   "code": "before",
   "run": "gpu_pieces_before",
   "suite": "pieces",
   "piece": "route",
   "dtype": "bfloat16",
   "op": "fwd+bwd",
   "block_m": 64,
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "xla_temp_mb": 2.2500152587890625,
   "xla_arg_mb": 32.015625,
   "xla_out_mb": 32.01564025878906,
   "compile_s": 0.0009647199985920452,
   "inner": 191,
   "median_ms": 0.17905136127986687,
   "q1_ms": 0.17875584555669144,
   "q3_ms": 0.18249631936921548
  },
  {
   "code": "before",
   "run": "gpu_pieces_before",
   "suite": "pieces",
   "piece": "sort_permute",
   "dtype": "bfloat16",
   "op": "fwd+bwd",
   "block_m": 64,
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "xla_temp_mb": 0.7502670288085938,
   "xla_arg_mb": 32.0,
   "xla_out_mb": 32.00001907348633,
   "compile_s": 0.0009744990020408295,
   "inner": 148,
   "median_ms": 0.22033705406372334,
   "q1_ms": 0.2202573141649073,
   "q3_ms": 0.2203900641836446
  },
  {
   "code": "before",
   "run": "gpu_pieces_before",
   "suite": "pieces",
   "piece": "experts_grouped_swiglu",
   "dtype": "bfloat16",
   "op": "fwd+bwd",
   "block_m": 64,
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "xla_temp_mb": 516.2715148925781,
   "xla_arg_mb": 88.0,
   "xla_out_mb": 88.00004196166992,
   "compile_s": 0.018710747004661243,
   "inner": 2,
   "median_ms": 18.341135499213124,
   "q1_ms": 18.30343400069978,
   "q3_ms": 18.40064299904043
  },
  {
   "code": "before",
   "run": "gpu_pieces_before",
   "suite": "pieces",
   "piece": "unpermute_combine",
   "dtype": "bfloat16",
   "op": "fwd+bwd",
   "block_m": 64,
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "xla_temp_mb": 64.00025939941406,
   "xla_arg_mb": 64.25,
   "xla_out_mb": 64.25002670288086,
   "compile_s": 0.00127532899932703,
   "inner": 74,
   "median_ms": 0.5526168919262804,
   "q1_ms": 0.5526090405289102,
   "q3_ms": 0.5526858040448623
  },
  {
   "code": "before",
   "run": "gpu_pieces_before",
   "suite": "blockm",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "max_expert_frac": 0.135498046875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 82.65970611572266,
   "xla_arg_mb": 56.015625,
   "xla_out_mb": 8.0,
   "compile_s": 0.003049616003409028,
   "inner": 34,
   "median_ms": 1.4013078824277996,
   "q1_ms": 1.3932303382342286,
   "q3_ms": 1.4028587354593638,
   "block_m": 32
  },
  {
   "code": "before",
   "run": "gpu_pieces_before",
   "suite": "blockm",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "max_expert_frac": 0.135498046875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 132.19632720947266,
   "xla_arg_mb": 152.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.00693328299530549,
   "inner": 10,
   "median_ms": 4.646821600181283,
   "q1_ms": 4.644318300415762,
   "q3_ms": 4.65391504985746,
   "block_m": 32
  },
  {
   "code": "before",
   "run": "gpu_pieces_before",
   "suite": "blockm",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "max_expert_frac": 0.135498046875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 85.15921783447266,
   "xla_arg_mb": 56.015625,
   "xla_out_mb": 8.0,
   "compile_s": 0.0030980079973232932,
   "inner": 35,
   "median_ms": 1.2874562000173944,
   "q1_ms": 1.2860813999265832,
   "q3_ms": 1.2886480143476677,
   "block_m": 64
  },
  {
   "code": "before",
   "run": "gpu_pieces_before",
   "suite": "blockm",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "max_expert_frac": 0.135498046875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 136.19632720947266,
   "xla_arg_mb": 152.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.006913484998221975,
   "inner": 10,
   "median_ms": 4.663350999908289,
   "q1_ms": 4.64680410004803,
   "q3_ms": 4.672036400006618,
   "block_m": 64
  },
  {
   "code": "before",
   "run": "gpu_pieces_before",
   "suite": "blockm",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "max_expert_frac": 0.135498046875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 90.15897369384766,
   "xla_arg_mb": 56.015625,
   "xla_out_mb": 8.0,
   "compile_s": 0.0030027360044186935,
   "inner": 32,
   "median_ms": 1.4124235938197671,
   "q1_ms": 1.380653328169501,
   "q3_ms": 1.4433049999524883,
   "block_m": 128
  },
  {
   "code": "before",
   "run": "gpu_pieces_before",
   "suite": "blockm",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "max_expert_frac": 0.135498046875,
   "empty_experts": 0,
   "op": "train_step",
   "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Shared memory size limit exceeded: requested 131072, available: 101376, context: []",
   "block_m": 128
  },
  {
   "code": "before",
   "run": "gpu_pieces_before",
   "suite": "blockm",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "max_expert_frac": 0.1292572021484375,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 643.7605361938477,
   "xla_arg_mb": 112.015625,
   "xla_out_mb": 64.0,
   "compile_s": 0.011043040998629294,
   "inner": 5,
   "median_ms": 9.896290401229635,
   "q1_ms": 9.893079200264765,
   "q3_ms": 9.910316199966474,
   "block_m": 32
  },
  {
   "code": "before",
   "run": "gpu_pieces_before",
   "suite": "blockm",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "max_expert_frac": 0.1292572021484375,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 1029.6450576782227,
   "xla_arg_mb": 208.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.03372348799894098,
   "inner": 1,
   "median_ms": 33.60316999896895,
   "q1_ms": 33.55514249778935,
   "q3_ms": 33.67847650224576,
   "block_m": 32
  },
  {
   "code": "before",
   "run": "gpu_pieces_before",
   "suite": "blockm",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "max_expert_frac": 0.1292572021484375,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 646.2566299438477,
   "xla_arg_mb": 112.015625,
   "xla_out_mb": 64.0,
   "compile_s": 0.009813627002586145,
   "inner": 5,
   "median_ms": 8.632298599695787,
   "q1_ms": 8.624167999369092,
   "q3_ms": 8.661048900830792,
   "block_m": 64
  },
  {
   "code": "before",
   "run": "gpu_pieces_before",
   "suite": "blockm",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "max_expert_frac": 0.1292572021484375,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 1033.6450576782227,
   "xla_arg_mb": 208.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.03293934999965131,
   "inner": 1,
   "median_ms": 31.818123999983072,
   "q1_ms": 31.771119000040926,
   "q3_ms": 31.94490749956458,
   "block_m": 64
  },
  {
   "code": "before",
   "run": "gpu_pieces_before",
   "suite": "blockm",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "max_expert_frac": 0.1292572021484375,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 651.2546768188477,
   "xla_arg_mb": 112.015625,
   "xla_out_mb": 64.0,
   "compile_s": 0.010143263003556058,
   "inner": 5,
   "median_ms": 8.61605740064988,
   "q1_ms": 8.554265499697067,
   "q3_ms": 8.639962899906095,
   "block_m": 128
  },
  {
   "code": "before",
   "run": "gpu_pieces_before",
   "suite": "blockm",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "max_expert_frac": 0.1292572021484375,
   "empty_experts": 0,
   "op": "train_step",
   "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Shared memory size limit exceeded: requested 131072, available: 101376, context: []",
   "block_m": 128
  },
  {
   "code": "before",
   "run": "gpu_pieces_before",
   "suite": "blockm",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "max_expert_frac": 0.1356201171875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 49.440956115722656,
   "xla_arg_mb": 52.015625,
   "xla_out_mb": 4.0,
   "compile_s": 0.0027995909986202605,
   "inner": 50,
   "median_ms": 0.8821641799295321,
   "q1_ms": 0.8814040199649753,
   "q3_ms": 0.8823454199591652,
   "block_m": 32
  },
  {
   "code": "before",
   "run": "gpu_pieces_before",
   "suite": "blockm",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "max_expert_frac": 0.1356201171875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 57.505409240722656,
   "xla_arg_mb": 148.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.005451371995150112,
   "inner": 16,
   "median_ms": 2.793261000078928,
   "q1_ms": 2.788872374821949,
   "q3_ms": 2.79763868752525,
   "block_m": 32
  },
  {
   "code": "before",
   "run": "gpu_pieces_before",
   "suite": "blockm",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "max_expert_frac": 0.1356201171875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 50.690467834472656,
   "xla_arg_mb": 52.015625,
   "xla_out_mb": 4.0,
   "compile_s": 0.0024641639975016005,
   "inner": 52,
   "median_ms": 0.8692945385309241,
   "q1_ms": 0.8688127595982783,
   "q3_ms": 0.8693277788132456,
   "block_m": 64
  },
  {
   "code": "before",
   "run": "gpu_pieces_before",
   "suite": "blockm",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "max_expert_frac": 0.1356201171875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 59.005409240722656,
   "xla_arg_mb": 148.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.005409453005995601,
   "inner": 16,
   "median_ms": 2.719245625030453,
   "q1_ms": 2.7180091874470236,
   "q3_ms": 2.7372021563678572,
   "block_m": 64
  },
  {
   "code": "before",
   "run": "gpu_pieces_before",
   "suite": "blockm",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "max_expert_frac": 0.1356201171875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 53.190223693847656,
   "xla_arg_mb": 52.015625,
   "xla_out_mb": 4.0,
   "compile_s": 0.0026887099957093596,
   "inner": 42,
   "median_ms": 1.084045118962725,
   "q1_ms": 1.0832628095591819,
   "q3_ms": 1.0843822142522273,
   "block_m": 128
  },
  {
   "code": "before",
   "run": "gpu_pieces_before",
   "suite": "blockm",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 4096,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "max_expert_frac": 0.1356201171875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 62.005409240722656,
   "xla_arg_mb": 148.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.006105372005549725,
   "inner": 13,
   "median_ms": 3.334030461421487,
   "q1_ms": 3.3213082309745143,
   "q3_ms": 3.359036692507708,
   "block_m": 128
  },
  {
   "code": "before",
   "run": "gpu_pieces_before",
   "suite": "blockm",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "max_expert_frac": 0.1291961669921875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 322.76053619384766,
   "xla_arg_mb": 80.015625,
   "xla_out_mb": 32.0,
   "compile_s": 0.006880874003400095,
   "inner": 9,
   "median_ms": 5.536698333647412,
   "q1_ms": 5.525888833896412,
   "q3_ms": 5.552149722158598,
   "block_m": 32
  },
  {
   "code": "before",
   "run": "gpu_pieces_before",
   "suite": "blockm",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "max_expert_frac": 0.1291961669921875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 522.0054092407227,
   "xla_arg_mb": 176.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.01864667499467032,
   "inner": 3,
   "median_ms": 17.982063999321934,
   "q1_ms": 17.928452498987703,
   "q3_ms": 17.99988616645957,
   "block_m": 32
  },
  {
   "code": "before",
   "run": "gpu_pieces_before",
   "suite": "blockm",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "max_expert_frac": 0.1291961669921875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 324.00662994384766,
   "xla_arg_mb": 80.015625,
   "xla_out_mb": 32.0,
   "compile_s": 0.006190394000441302,
   "inner": 10,
   "median_ms": 4.8785752995172516,
   "q1_ms": 4.876044900447596,
   "q3_ms": 4.8916565498075215,
   "block_m": 64
  },
  {
   "code": "before",
   "run": "gpu_pieces_before",
   "suite": "blockm",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "max_expert_frac": 0.1291961669921875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 524.0054092407227,
   "xla_arg_mb": 176.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.01841888299532002,
   "inner": 3,
   "median_ms": 16.945648332087632,
   "q1_ms": 16.924046666341987,
   "q3_ms": 16.996364166213123,
   "block_m": 64
  },
  {
   "code": "before",
   "run": "gpu_pieces_before",
   "suite": "blockm",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "max_expert_frac": 0.1291961669921875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 326.50467681884766,
   "xla_arg_mb": 80.015625,
   "xla_out_mb": 32.0,
   "compile_s": 0.00795223599561723,
   "inner": 7,
   "median_ms": 6.5296977149306,
   "q1_ms": 6.526955928815628,
   "q3_ms": 6.538585857405062,
   "block_m": 128
  },
  {
   "code": "before",
   "run": "gpu_pieces_before",
   "suite": "blockm",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "max_expert_frac": 0.1291961669921875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 528.0054092407227,
   "xla_arg_mb": 176.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.0218724150035996,
   "inner": 2,
   "median_ms": 20.295647998864297,
   "q1_ms": 20.262483001715736,
   "q3_ms": 20.424545249625226,
   "block_m": 128
  },
  {
   "code": "sweep",
   "run": "gpu_slab_0",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1292572021484375,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 646.2566299438477,
   "xla_arg_mb": 112.015625,
   "xla_out_mb": 64.0,
   "compile_s": 0.024975460997666232,
   "inner": 5,
   "median_ms": 8.628804399631917,
   "q1_ms": 8.618918899446726,
   "q3_ms": 8.635910200246144
  },
  {
   "code": "sweep",
   "run": "gpu_slab_0",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1292572021484375,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 1033.6450576782227,
   "xla_arg_mb": 208.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.05929093099985039,
   "inner": 1,
   "median_ms": 31.88553300424246,
   "q1_ms": 31.858938000368653,
   "q3_ms": 31.96794950054027
  },
  {
   "code": "sweep",
   "run": "gpu_slab_0",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 131072,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.12841415405273438,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 2570.0183486938477,
   "xla_arg_mb": 304.015625,
   "xla_out_mb": 256.0,
   "compile_s": 0.0873739509988809,
   "inner": 1,
   "median_ms": 33.99238900601631,
   "q1_ms": 33.88030250062002,
   "q3_ms": 34.03641800105106
  },
  {
   "code": "sweep",
   "run": "gpu_slab_0",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 131072,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.12841415405273438,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 4110.535926818848,
   "xla_arg_mb": 400.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.2309832320024725,
   "inner": 1,
   "median_ms": 130.7909609968192,
   "q1_ms": 130.5024105022312,
   "q3_ms": 131.18102850057767
  },
  {
   "code": "sweep",
   "run": "gpu_slab_0",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1291961669921875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 324.00662994384766,
   "xla_arg_mb": 80.015625,
   "xla_out_mb": 32.0,
   "compile_s": 0.006308114003331866,
   "inner": 10,
   "median_ms": 4.8791597997478675,
   "q1_ms": 4.876009300278383,
   "q3_ms": 4.881140450015664
  },
  {
   "code": "sweep",
   "run": "gpu_slab_0",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1291961669921875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 524.0054092407227,
   "xla_arg_mb": 176.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.01838689599389909,
   "inner": 3,
   "median_ms": 17.007610001504265,
   "q1_ms": 16.95503516748431,
   "q3_ms": 17.10336383378793
  },
  {
   "code": "sweep",
   "run": "gpu_slab_0",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 131072,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.12839126586914062,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 1288.5183486938477,
   "xla_arg_mb": 176.015625,
   "xla_out_mb": 128.0,
   "compile_s": 0.019375045005290303,
   "inner": 2,
   "median_ms": 18.666790001589106,
   "q1_ms": 18.647872251676745,
   "q3_ms": 18.75331650080625
  },
  {
   "code": "sweep",
   "run": "gpu_slab_0",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 131072,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.12839126586914062,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 2060.0056533813477,
   "xla_arg_mb": 272.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.06819986699701985,
   "inner": 1,
   "median_ms": 68.6648150003748,
   "q1_ms": 68.60499849790358,
   "q3_ms": 68.8528744976793
  },
  {
   "code": "sweep",
   "run": "gpu_slab_16777216",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1292572021484375,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 334.63555908203125,
   "xla_arg_mb": 112.015625,
   "xla_out_mb": 64.0,
   "compile_s": 0.019445190999249462,
   "inner": 5,
   "median_ms": 9.524235600838438,
   "q1_ms": 9.510661799868103,
   "q3_ms": 9.555366299900925
  },
  {
   "code": "sweep",
   "run": "gpu_slab_16777216",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1292572021484375,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 464.6465835571289,
   "xla_arg_mb": 208.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.03484054200089304,
   "inner": 1,
   "median_ms": 34.380672994302586,
   "q1_ms": 34.24637300122413,
   "q3_ms": 34.38221550459275
  },
  {
   "code": "sweep",
   "run": "gpu_slab_16777216",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 131072,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.12841415405273438,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 1118.7742309570312,
   "xla_arg_mb": 304.015625,
   "xla_out_mb": 256.0,
   "compile_s": 0.06437037800060352,
   "inner": 1,
   "median_ms": 37.106148003658745,
   "q1_ms": 36.89716100052465,
   "q3_ms": 37.245695501042064
  },
  {
   "code": "sweep",
   "run": "gpu_slab_16777216",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 131072,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.12841415405273438,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 1798.553810119629,
   "xla_arg_mb": 400.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.1334789800021099,
   "inner": 1,
   "median_ms": 134.50366799952462,
   "q1_ms": 134.15289999829838,
   "q3_ms": 134.72615349746775
  },
  {
   "code": "sweep",
   "run": "gpu_slab_16777216",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1291961669921875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 191.82305908203125,
   "xla_arg_mb": 80.015625,
   "xla_out_mb": 32.0,
   "compile_s": 0.007191292999777943,
   "inner": 9,
   "median_ms": 5.378283111591979,
   "q1_ms": 5.363333333257794,
   "q3_ms": 5.407859611295862
  },
  {
   "code": "sweep",
   "run": "gpu_slab_16777216",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1291961669921875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 257.9002456665039,
   "xla_arg_mb": 176.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.019464097000309266,
   "inner": 2,
   "median_ms": 17.685438000626164,
   "q1_ms": 17.671266501565697,
   "q3_ms": 17.75852574974124
  },
  {
   "code": "sweep",
   "run": "gpu_slab_16777216",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 131072,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.12839126586914062,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 585.3992309570312,
   "xla_arg_mb": 176.015625,
   "xla_out_mb": 128.0,
   "compile_s": 0.021636426004988607,
   "inner": 2,
   "median_ms": 20.456674497836502,
   "q1_ms": 20.355939002911327,
   "q3_ms": 20.541782998407143
  },
  {
   "code": "sweep",
   "run": "gpu_slab_16777216",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 131072,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.12839126586914062,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 1031.558448791504,
   "xla_arg_mb": 272.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.07345181599521311,
   "inner": 1,
   "median_ms": 72.0042180037126,
   "q1_ms": 71.97582200024044,
   "q3_ms": 72.3025504994439
  },
  {
   "code": "sweep",
   "run": "gpu_slab_33554432",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1292572021484375,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 429.01031494140625,
   "xla_arg_mb": 112.015625,
   "xla_out_mb": 64.0,
   "compile_s": 0.01878368299367139,
   "inner": 5,
   "median_ms": 9.288543999718968,
   "q1_ms": 9.256305400049314,
   "q3_ms": 9.305222499824595
  },
  {
   "code": "sweep",
   "run": "gpu_slab_33554432",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1292572021484375,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 601.6465835571289,
   "xla_arg_mb": 208.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.047249407005438115,
   "inner": 1,
   "median_ms": 33.17832500033546,
   "q1_ms": 33.03427950231708,
   "q3_ms": 33.25555600167718
  },
  {
   "code": "sweep",
   "run": "gpu_slab_33554432",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 131072,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.12841415405273438,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 1199.0227661132812,
   "xla_arg_mb": 304.015625,
   "xla_out_mb": 256.0,
   "compile_s": 0.0643902030060417,
   "inner": 1,
   "median_ms": 36.7289920031908,
   "q1_ms": 36.609225997381145,
   "q3_ms": 36.812867496337276
  },
  {
   "code": "sweep",
   "run": "gpu_slab_33554432",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 131072,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.12841415405273438,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 1798.553810119629,
   "xla_arg_mb": 400.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.13114345500071067,
   "inner": 1,
   "median_ms": 131.02977999369614,
   "q1_ms": 130.71195150041603,
   "q3_ms": 131.93069249973632
  },
  {
   "code": "sweep",
   "run": "gpu_slab_33554432",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1291961669921875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 239.01031494140625,
   "xla_arg_mb": 80.015625,
   "xla_out_mb": 32.0,
   "compile_s": 0.006993425995460711,
   "inner": 9,
   "median_ms": 5.237188777957474,
   "q1_ms": 5.227682000016406,
   "q3_ms": 5.250118944483498
  },
  {
   "code": "sweep",
   "run": "gpu_slab_33554432",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1291961669921875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 300.2596206665039,
   "xla_arg_mb": 176.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.018376724001427647,
   "inner": 3,
   "median_ms": 16.730773999976616,
   "q1_ms": 16.686745834401034,
   "q3_ms": 16.735767000985412
  },
  {
   "code": "sweep",
   "run": "gpu_slab_33554432",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 131072,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.12839126586914062,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 625.5227661132812,
   "xla_arg_mb": 176.015625,
   "xla_out_mb": 128.0,
   "compile_s": 0.02162552799563855,
   "inner": 2,
   "median_ms": 20.191294999676757,
   "q1_ms": 20.182380998448934,
   "q3_ms": 20.23025900052744
  },
  {
   "code": "sweep",
   "run": "gpu_slab_33554432",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 131072,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.12839126586914062,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 1030.554054260254,
   "xla_arg_mb": 272.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.0677149199982523,
   "inner": 1,
   "median_ms": 68.60781899740687,
   "q1_ms": 68.36133699835045,
   "q3_ms": 68.85494249945623
  },
  {
   "code": "sweep",
   "run": "gpu_slab_67108864",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1292572021484375,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 579.5110473632812,
   "xla_arg_mb": 112.015625,
   "xla_out_mb": 64.0,
   "compile_s": 0.025291241996455938,
   "inner": 5,
   "median_ms": 9.311306400923058,
   "q1_ms": 9.277655099867843,
   "q3_ms": 9.351908099779394
  },
  {
   "code": "sweep",
   "run": "gpu_slab_67108864",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1292572021484375,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 773.6465835571289,
   "xla_arg_mb": 208.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.0324987339990912,
   "inner": 1,
   "median_ms": 31.844326003920287,
   "q1_ms": 31.80845850147307,
   "q3_ms": 31.998190501326462
  },
  {
   "code": "sweep",
   "run": "gpu_slab_67108864",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 131072,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.12841415405273438,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 1335.8992309570312,
   "xla_arg_mb": 304.015625,
   "xla_out_mb": 256.0,
   "compile_s": 0.063436593998631,
   "inner": 1,
   "median_ms": 36.9691879968741,
   "q1_ms": 36.724713500007056,
   "q3_ms": 37.01100500256871
  },
  {
   "code": "sweep",
   "run": "gpu_slab_67108864",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 131072,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.12841415405273438,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 1851.553810119629,
   "xla_arg_mb": 400.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.13131889999931445,
   "inner": 1,
   "median_ms": 131.54044799739495,
   "q1_ms": 130.92242999846349,
   "q3_ms": 132.23293450209894
  },
  {
   "code": "sweep",
   "run": "gpu_slab_67108864",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1291961669921875,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 306.26104736328125,
   "xla_arg_mb": 80.015625,
   "xla_out_mb": 32.0,
   "compile_s": 0.0068985469988547266,
   "inner": 9,
   "median_ms": 5.2278896667606505,
   "q1_ms": 5.208924221935579,
   "q3_ms": 5.243936000220453
  },
  {
   "code": "sweep",
   "run": "gpu_slab_67108864",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.1291961669921875,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 386.2596206665039,
   "xla_arg_mb": 176.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.018061107999528758,
   "inner": 3,
   "median_ms": 16.223890333397623,
   "q1_ms": 16.214715332656247,
   "q3_ms": 16.33107649953066
  },
  {
   "code": "sweep",
   "run": "gpu_slab_67108864",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 131072,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.12839126586914062,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 693.9617309570312,
   "xla_arg_mb": 176.015625,
   "xla_out_mb": 128.0,
   "compile_s": 0.021114583003509324,
   "inner": 2,
   "median_ms": 20.17521050220239,
   "q1_ms": 20.114968250709353,
   "q3_ms": 20.21776300171041
  },
  {
   "code": "sweep",
   "run": "gpu_slab_67108864",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 131072,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.0,
   "max_expert_frac": 0.12839126586914062,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 1030.554054260254,
   "xla_arg_mb": 272.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.06698547799896915,
   "inner": 1,
   "median_ms": 66.33425399922999,
   "q1_ms": 66.02848200054723,
   "q3_ms": 66.54852249994292
  },
  {
   "code": "intermediate",
   "run": "gpu_C",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.9,
   "max_expert_frac": 0.471588134765625,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 191.82305908203125,
   "xla_arg_mb": 80.015625,
   "xla_out_mb": 32.0,
   "compile_s": 0.008855177002260461,
   "inner": 9,
   "median_ms": 5.414120000043315,
   "q1_ms": 5.401436777901836,
   "q3_ms": 5.449126333486573,
   "rel_err_out": 0.00813548201925598,
   "rel_err_grad_params": 0.0763183976119977,
   "rel_err_grad_x": 0.04581131524433544
  },
  {
   "code": "intermediate",
   "run": "gpu_C",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.9,
   "max_expert_frac": 0.471588134765625,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 257.9002456665039,
   "xla_arg_mb": 176.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.02637604200572241,
   "inner": 2,
   "median_ms": 17.87362249888247,
   "q1_ms": 17.865854750198196,
   "q3_ms": 17.901342749610194,
   "rel_err_out": 0.00813548201925598,
   "rel_err_grad_params": 0.0763183976119977,
   "rel_err_grad_x": 0.04581131524433544
  },
  {
   "code": "intermediate",
   "run": "gpu_C",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.9,
   "max_expert_frac": 0.471588134765625,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 387.76151275634766,
   "xla_arg_mb": 80.015625,
   "xla_out_mb": 32.0,
   "compile_s": 0.006812663996242918,
   "inner": 9,
   "median_ms": 5.393781666094179,
   "q1_ms": 5.374857444419629,
   "q3_ms": 5.400858444772894,
   "rel_err_out": 0.0081210528913697,
   "rel_err_grad_params": 0.07628494532457057,
   "rel_err_grad_x": 0.045809509906336796
  },
  {
   "code": "intermediate",
   "run": "gpu_C",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.9,
   "max_expert_frac": 0.471588134765625,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 677.2715301513672,
   "xla_arg_mb": 176.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.022085135999077465,
   "inner": 2,
   "median_ms": 20.48625300085405,
   "q1_ms": 20.458831249925424,
   "q3_ms": 20.685824749307358,
   "rel_err_out": 0.0081210528913697,
   "rel_err_grad_params": 0.07628494532457057,
   "rel_err_grad_x": 0.045809509906336796
  },
  {
   "code": "intermediate",
   "run": "gpu_C",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused_gsum",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.9,
   "max_expert_frac": 0.471588134765625,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 388.00785064697266,
   "xla_arg_mb": 80.015625,
   "xla_out_mb": 32.0,
   "compile_s": 0.006101138002122752,
   "inner": 11,
   "median_ms": 4.598593090618538,
   "q1_ms": 4.590066954485586,
   "q3_ms": 4.611150000114735,
   "rel_err_out": 0.00813548201925598,
   "rel_err_grad_params": 0.07631728577023915,
   "rel_err_grad_x": 0.045811245476744256
  },
  {
   "code": "intermediate",
   "run": "gpu_C",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused_gsum",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.9,
   "max_expert_frac": 0.471588134765625,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 677.2717742919922,
   "xla_arg_mb": 176.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.021336531004635617,
   "inner": 2,
   "median_ms": 19.68440650307457,
   "q1_ms": 19.636762750451453,
   "q3_ms": 19.94492599988007,
   "rel_err_out": 0.00813548201925598,
   "rel_err_grad_params": 0.07631728577023915,
   "rel_err_grad_x": 0.045811245476744256
  },
  {
   "code": "intermediate",
   "run": "gpu_C",
   "suite": "layer",
   "routing": "capacity",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.9,
   "max_expert_frac": 0.471588134765625,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 481.7509994506836,
   "xla_arg_mb": 80.015625,
   "xla_out_mb": 32.0,
   "compile_s": 0.006353855998895597,
   "inner": 9,
   "median_ms": 5.374410555911406,
   "q1_ms": 5.372809166551128,
   "q3_ms": 5.3786703336097865
  },
  {
   "code": "intermediate",
   "run": "gpu_C",
   "suite": "layer",
   "routing": "capacity",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 32768,
   "d_model": 512,
   "d_hidden": 1024,
   "experts": 8,
   "top_k": 2,
   "skew": 0.9,
   "max_expert_frac": 0.471588134765625,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 752.3786849975586,
   "xla_arg_mb": 176.04687881469727,
   "xla_out_mb": 144.04698944091797,
   "compile_s": 0.014661196997622028,
   "inner": 4,
   "median_ms": 12.948655999935,
   "q1_ms": 12.88125437531562,
   "q3_ms": 13.011749874749512
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 1,
   "skew": 0.0,
   "max_expert_frac": 0.15900000929832458,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 3.7986221313476562,
   "xla_arg_mb": 3.4921875,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.011597837001318112,
   "inner": 24,
   "median_ms": 0.6435993330645337,
   "q1_ms": 0.6145063334770384,
   "q3_ms": 0.7990897708320215,
   "rel_err_out": 0.0005186846986736415,
   "rel_err_grad_params": 0.000619610548708075,
   "rel_err_grad_x": 0.0005381070236430502
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 1,
   "skew": 0.0,
   "max_expert_frac": 0.15900000929832458,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 4.505165100097656,
   "xla_arg_mb": 9.500003814697266,
   "xla_out_mb": 9.011833190917969,
   "compile_s": 0.004808053003216628,
   "inner": 80,
   "median_ms": 0.35881357498510624,
   "q1_ms": 0.34640859375940636,
   "q3_ms": 0.3641531937319087,
   "rel_err_out": 0.0005186846986736415,
   "rel_err_grad_params": 0.000619610548708075,
   "rel_err_grad_x": 0.0005381070236430502
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "gather",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 1,
   "skew": 0.0,
   "max_expert_frac": 0.15900000929832458,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 5.776161193847656,
   "xla_arg_mb": 3.4921875,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.002142015000572428,
   "inner": 166,
   "median_ms": 0.18980314460525527,
   "q1_ms": 0.1847632680788176,
   "q3_ms": 0.19788630725745765,
   "rel_err_out": 0.0005186846986736415,
   "rel_err_grad_params": 0.000619610548708075,
   "rel_err_grad_x": 0.0005381222076587467
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "gather",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 1,
   "skew": 0.0,
   "max_expert_frac": 0.15900000929832458,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 9.019088745117188,
   "xla_arg_mb": 9.500003814697266,
   "xla_out_mb": 9.011833190917969,
   "compile_s": 0.004094953001185786,
   "inner": 88,
   "median_ms": 0.365999977310077,
   "q1_ms": 0.3591752954641627,
   "q3_ms": 0.37064115337092185,
   "rel_err_out": 0.0005186846986736415,
   "rel_err_grad_params": 0.000619610548708075,
   "rel_err_grad_x": 0.0005381222076587467
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 1,
   "skew": 0.0,
   "max_expert_frac": 0.15900000929832458,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 5.019569396972656,
   "xla_arg_mb": 3.4921875,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.0024277060001622885,
   "inner": 147,
   "median_ms": 0.21355621769927682,
   "q1_ms": 0.20675463947926043,
   "q3_ms": 0.24241221768026822,
   "rel_err_out": 0.0005186846986736415,
   "rel_err_grad_params": 0.000619610548708075,
   "rel_err_grad_x": 0.0005381222076587467
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 1,
   "skew": 0.0,
   "max_expert_frac": 0.15900000929832458,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 9.024459838867188,
   "xla_arg_mb": 9.500003814697266,
   "xla_out_mb": 9.011833190917969,
   "compile_s": 0.0036990230000810698,
   "inner": 78,
   "median_ms": 0.4119473589339461,
   "q1_ms": 0.3991677179631855,
   "q3_ms": 0.451564929533514,
   "rel_err_out": 0.0005186846986736415,
   "rel_err_grad_params": 0.000619610548708075,
   "rel_err_grad_x": 0.0005381222076587467
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused_gsum",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 1,
   "skew": 0.0,
   "max_expert_frac": 0.15900000929832458,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 5.030067443847656,
   "xla_arg_mb": 3.4921875,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.0034242729961988516,
   "inner": 90,
   "median_ms": 0.23222611109506236,
   "q1_ms": 0.21099502219941413,
   "q3_ms": 0.2453074888762785,
   "rel_err_out": 0.0005186846986736415,
   "rel_err_grad_params": 0.000619610548708075,
   "rel_err_grad_x": 0.0005381222076587467
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused_gsum",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 1,
   "skew": 0.0,
   "max_expert_frac": 0.15900000929832458,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 9.028610229492188,
   "xla_arg_mb": 9.500003814697266,
   "xla_out_mb": 9.011833190917969,
   "compile_s": 0.004558460997941438,
   "inner": 78,
   "median_ms": 0.403906230740876,
   "q1_ms": 0.38288240382546346,
   "q3_ms": 0.4335350256438594,
   "rel_err_out": 0.0005186846986736415,
   "rel_err_grad_params": 0.000619610548708075,
   "rel_err_grad_x": 0.0005381222076587467
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "blocked",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 1,
   "skew": 0.0,
   "max_expert_frac": 0.15900000929832458,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 11.048622131347656,
   "xla_arg_mb": 3.4921875,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.003281151999544818,
   "inner": 24,
   "median_ms": 0.19573516662300486,
   "q1_ms": 0.18295324995657816,
   "q3_ms": 0.21541637503711777,
   "rel_err_out": 0.0005186846986736415,
   "rel_err_grad_params": 0.0006196100250229748,
   "rel_err_grad_x": 0.0005381076783174353
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "blocked",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 1,
   "skew": 0.0,
   "max_expert_frac": 0.15900000929832458,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 11.054237365722656,
   "xla_arg_mb": 9.500003814697266,
   "xla_out_mb": 9.011833190917969,
   "compile_s": 0.005166221002582461,
   "inner": 55,
   "median_ms": 0.4979984726312316,
   "q1_ms": 0.39881579994636757,
   "q3_ms": 0.5816208909726066,
   "rel_err_out": 0.0005186846986736415,
   "rel_err_grad_params": 0.0006196100250229748,
   "rel_err_grad_x": 0.0005381076783174353
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "ragged_dot",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 1,
   "skew": 0.0,
   "max_expert_frac": 0.15900000929832458,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 9.804458618164062,
   "xla_arg_mb": 3.4921875,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.0014217620046110824,
   "inner": 200,
   "median_ms": 0.1276771250195452,
   "q1_ms": 0.12658646248382865,
   "q3_ms": 0.1292860950161412,
   "rel_err_out": 0.0005186846986736415,
   "rel_err_grad_params": 0.0006196101919023643,
   "rel_err_grad_x": 0.0005381076783174353
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "ragged_dot",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 1,
   "skew": 0.0,
   "max_expert_frac": 0.15900000929832458,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 23.487823486328125,
   "xla_arg_mb": 9.500003814697266,
   "xla_out_mb": 9.011833190917969,
   "compile_s": 0.012266222998732701,
   "inner": 23,
   "median_ms": 0.6881790000277207,
   "q1_ms": 0.6309680652718627,
   "q3_ms": 0.7868490652269784,
   "rel_err_out": 0.0005186846986736415,
   "rel_err_grad_params": 0.0006196101919023643,
   "rel_err_grad_x": 0.0005381076783174353
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "capacity",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 1,
   "skew": 0.0,
   "max_expert_frac": 0.15900000929832458,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 3.7451324462890625,
   "xla_arg_mb": 3.4921875,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.001213898001878988,
   "inner": 200,
   "median_ms": 0.0734147949697217,
   "q1_ms": 0.0680245125113288,
   "q3_ms": 0.07671456500247587
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "capacity",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 1,
   "skew": 0.0,
   "max_expert_frac": 0.15900000929832458,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 6.653587341308594,
   "xla_arg_mb": 9.500003814697266,
   "xla_out_mb": 9.011833190917969,
   "compile_s": 0.07549837799888337,
   "inner": 20,
   "median_ms": 0.3755039502721047,
   "q1_ms": 0.3745361000255798,
   "q3_ms": 0.3767495500142104
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dense",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 1,
   "skew": 0.0,
   "max_expert_frac": 0.15900000929832458,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 23.437515258789062,
   "xla_arg_mb": 3.4921875,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.0011310510017210618,
   "inner": 200,
   "median_ms": 0.10866267501114635,
   "q1_ms": 0.10744787748990348,
   "q3_ms": 0.10871973750909092,
   "rel_err_out": 0.0005186846986736415,
   "rel_err_grad_params": 0.0006196128820785577,
   "rel_err_grad_x": 0.0005381076783174353
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dense",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 1,
   "skew": 0.0,
   "max_expert_frac": 0.15900000929832458,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 31.251976013183594,
   "xla_arg_mb": 9.500003814697266,
   "xla_out_mb": 9.011833190917969,
   "compile_s": 0.07389545399928465,
   "inner": 19,
   "median_ms": 0.8405751053176477,
   "q1_ms": 0.836700947394619,
   "q3_ms": 0.8423049473908337,
   "rel_err_out": 0.0005186846986736415,
   "rel_err_grad_params": 0.0006196128820785577,
   "rel_err_grad_x": 0.0005381076783174353
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 1,
   "skew": 0.9,
   "max_expert_frac": 0.9390000700950623,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 3.7986221313476562,
   "xla_arg_mb": 3.4921875,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.0017359560006298125,
   "inner": 186,
   "median_ms": 0.16845016130333848,
   "q1_ms": 0.1677210887181147,
   "q3_ms": 0.1768951774025599,
   "rel_err_out": 0.0005147988291311027,
   "rel_err_grad_params": 0.00043578268584146956,
   "rel_err_grad_x": 0.0005386833074921734
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 1,
   "skew": 0.9,
   "max_expert_frac": 0.9390000700950623,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 4.505165100097656,
   "xla_arg_mb": 9.500003814697266,
   "xla_out_mb": 9.011833190917969,
   "compile_s": 0.0034514239960117266,
   "inner": 93,
   "median_ms": 0.3119375699410035,
   "q1_ms": 0.31153374197122996,
   "q3_ms": 0.32096821505742856,
   "rel_err_out": 0.0005147988291311027,
   "rel_err_grad_params": 0.00043578268584146956,
   "rel_err_grad_x": 0.0005386833074921734
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "gather",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 1,
   "skew": 0.9,
   "max_expert_frac": 0.9390000700950623,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 5.776161193847656,
   "xla_arg_mb": 3.4921875,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.049561742998776026,
   "inner": 20,
   "median_ms": 0.3730913496838184,
   "q1_ms": 0.37214457497611875,
   "q3_ms": 0.37461745014297776,
   "rel_err_out": 0.0005147988291311027,
   "rel_err_grad_params": 0.00043578268584146956,
   "rel_err_grad_x": 0.0005386714737035058
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "gather",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 1,
   "skew": 0.9,
   "max_expert_frac": 0.9390000700950623,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 9.019088745117188,
   "xla_arg_mb": 9.500003814697266,
   "xla_out_mb": 9.011833190917969,
   "compile_s": 0.0030030859998078085,
   "inner": 98,
   "median_ms": 0.3010890204120874,
   "q1_ms": 0.29854675512422324,
   "q3_ms": 0.30521346937404115,
   "rel_err_out": 0.0005147988291311027,
   "rel_err_grad_params": 0.00043578268584146956,
   "rel_err_grad_x": 0.0005386714737035058
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 1,
   "skew": 0.9,
   "max_expert_frac": 0.9390000700950623,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 5.019569396972656,
   "xla_arg_mb": 3.4921875,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.0020037939975736663,
   "inner": 157,
   "median_ms": 0.23008299361029047,
   "q1_ms": 0.18648421336695772,
   "q3_ms": 0.34650950636981637,
   "rel_err_out": 0.0005147988291311027,
   "rel_err_grad_params": 0.00043578268584146956,
   "rel_err_grad_x": 0.0005386714737035058
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 1,
   "skew": 0.9,
   "max_expert_frac": 0.9390000700950623,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 9.024459838867188,
   "xla_arg_mb": 9.500003814697266,
   "xla_out_mb": 9.011833190917969,
   "compile_s": 0.0032881450024433434,
   "inner": 95,
   "median_ms": 0.33025104203261435,
   "q1_ms": 0.32916867371807856,
   "q3_ms": 0.33434500529485944,
   "rel_err_out": 0.0005147988291311027,
   "rel_err_grad_params": 0.00043578268584146956,
   "rel_err_grad_x": 0.0005386714737035058
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused_gsum",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 1,
   "skew": 0.9,
   "max_expert_frac": 0.9390000700950623,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 5.030067443847656,
   "xla_arg_mb": 3.4921875,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.08354661599878455,
   "inner": 19,
   "median_ms": 0.4228658420432016,
   "q1_ms": 0.41981771075944907,
   "q3_ms": 0.42404573679576923,
   "rel_err_out": 0.0005147988291311027,
   "rel_err_grad_params": 0.00043578268584146956,
   "rel_err_grad_x": 0.0005386714737035058
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused_gsum",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 1,
   "skew": 0.9,
   "max_expert_frac": 0.9390000700950623,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 9.028610229492188,
   "xla_arg_mb": 9.500003814697266,
   "xla_out_mb": 9.011833190917969,
   "compile_s": 0.003152268000121694,
   "inner": 91,
   "median_ms": 0.3274637143316955,
   "q1_ms": 0.32699898348745104,
   "q3_ms": 0.3425049890410974,
   "rel_err_out": 0.0005147988291311027,
   "rel_err_grad_params": 0.00043578268584146956,
   "rel_err_grad_x": 0.0005386714737035058
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "blocked",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 1,
   "skew": 0.9,
   "max_expert_frac": 0.9390000700950623,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 11.048622131347656,
   "xla_arg_mb": 3.4921875,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.001790729998901952,
   "inner": 192,
   "median_ms": 0.17147996876095326,
   "q1_ms": 0.17064416666319931,
   "q3_ms": 0.17796941926917498,
   "rel_err_out": 0.0005147988291311027,
   "rel_err_grad_params": 0.0004357963191464059,
   "rel_err_grad_x": 0.0005386826579973071
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "blocked",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 1,
   "skew": 0.9,
   "max_expert_frac": 0.9390000700950623,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 11.054237365722656,
   "xla_arg_mb": 9.500003814697266,
   "xla_out_mb": 9.011833190917969,
   "compile_s": 0.004695749994425569,
   "inner": 100,
   "median_ms": 0.2792976899945643,
   "q1_ms": 0.2754367200031993,
   "q3_ms": 0.2829923850003979,
   "rel_err_out": 0.0005147988291311027,
   "rel_err_grad_params": 0.0004357963191464059,
   "rel_err_grad_x": 0.0005386826579973071
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "ragged_dot",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 1,
   "skew": 0.9,
   "max_expert_frac": 0.9390000700950623,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 9.804458618164062,
   "xla_arg_mb": 3.4921875,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.0013491439967765473,
   "inner": 200,
   "median_ms": 0.12776730000041425,
   "q1_ms": 0.12595073500051512,
   "q3_ms": 0.12978815000678878,
   "rel_err_out": 0.0005147988291311027,
   "rel_err_grad_params": 0.00043578268063009965,
   "rel_err_grad_x": 0.0005386826579973071
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "ragged_dot",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 1,
   "skew": 0.9,
   "max_expert_frac": 0.9390000700950623,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 23.487823486328125,
   "xla_arg_mb": 9.500003814697266,
   "xla_out_mb": 9.011833190917969,
   "compile_s": 0.0026058430012199096,
   "inner": 83,
   "median_ms": 0.35122326508245755,
   "q1_ms": 0.33791584340737657,
   "q3_ms": 0.3520732530142343,
   "rel_err_out": 0.0005147988291311027,
   "rel_err_grad_params": 0.00043578268063009965,
   "rel_err_grad_x": 0.0005386826579973071
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "capacity",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 1,
   "skew": 0.9,
   "max_expert_frac": 0.9390000700950623,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 3.7451324462890625,
   "xla_arg_mb": 3.4921875,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.0013850519972038455,
   "inner": 200,
   "median_ms": 0.06139378499938176,
   "q1_ms": 0.06072672000300372,
   "q3_ms": 0.0647342325100908
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "capacity",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 1,
   "skew": 0.9,
   "max_expert_frac": 0.9390000700950623,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 6.653587341308594,
   "xla_arg_mb": 9.500003814697266,
   "xla_out_mb": 9.011833190917969,
   "compile_s": 0.0023085610009729862,
   "inner": 167,
   "median_ms": 0.1687109700409684,
   "q1_ms": 0.16624735331777651,
   "q3_ms": 0.1704590419273542
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dense",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 1,
   "skew": 0.9,
   "max_expert_frac": 0.9390000700950623,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 23.437515258789062,
   "xla_arg_mb": 3.4921875,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.0012215030001243576,
   "inner": 194,
   "median_ms": 0.1113570618791837,
   "q1_ms": 0.11129034018520861,
   "q3_ms": 0.11146213658968318,
   "rel_err_out": 0.0005147988291311027,
   "rel_err_grad_params": 0.0004357816547894153,
   "rel_err_grad_x": 0.0005386826579973071
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dense",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 1,
   "skew": 0.9,
   "max_expert_frac": 0.9390000700950623,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 31.251976013183594,
   "xla_arg_mb": 9.500003814697266,
   "xla_out_mb": 9.011833190917969,
   "compile_s": 0.005363415002648253,
   "inner": 34,
   "median_ms": 0.44864411756861955,
   "q1_ms": 0.4154854410872384,
   "q3_ms": 0.4831934852785814,
   "rel_err_out": 0.0005147988291311027,
   "rel_err_grad_params": 0.0004357816547894153,
   "rel_err_grad_x": 0.0005386826579973071
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.12849999964237213,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 11.173377990722656,
   "xla_arg_mb": 3.4921875,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.006394987001840491,
   "inner": 29,
   "median_ms": 0.3080691034295436,
   "q1_ms": 0.2941524309852835,
   "q3_ms": 0.35981691376193714,
   "rel_err_out": 0.01218432611866739,
   "rel_err_grad_params": 0.01099811138607451,
   "rel_err_grad_x": 0.010645545593928019
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.12849999964237213,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 17.755409240722656,
   "xla_arg_mb": 9.500003814697266,
   "xla_out_mb": 9.011833190917969,
   "compile_s": 0.004275814993889071,
   "inner": 61,
   "median_ms": 0.48686963943554634,
   "q1_ms": 0.4811135082401633,
   "q3_ms": 0.4911222049038949,
   "rel_err_out": 0.01218432611866739,
   "rel_err_grad_params": 0.01099811138607451,
   "rel_err_grad_x": 0.010645545593928019
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "gather",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.12849999964237213,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 13.396278381347656,
   "xla_arg_mb": 3.4921875,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.010201242999755777,
   "inner": 51,
   "median_ms": 0.5952821176798155,
   "q1_ms": 0.39690272544748056,
   "q3_ms": 0.83475040196836,
   "rel_err_out": 0.012184325960935733,
   "rel_err_grad_params": 0.010998106251530089,
   "rel_err_grad_x": 0.010644399048593357
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "gather",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.12849999964237213,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 24.458053588867188,
   "xla_arg_mb": 9.500003814697266,
   "xla_out_mb": 9.011833190917969,
   "compile_s": 0.005035704001784325,
   "inner": 69,
   "median_ms": 0.4640190870133296,
   "q1_ms": 0.4545238695847228,
   "q3_ms": 0.5157441231911408,
   "rel_err_out": 0.012184325960935733,
   "rel_err_grad_params": 0.010998106251530089,
   "rel_err_grad_x": 0.010644399048593357
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.12849999964237213,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 13.366737365722656,
   "xla_arg_mb": 3.4921875,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.0077903990022605285,
   "inner": 73,
   "median_ms": 0.6180255752964239,
   "q1_ms": 0.5800256095762837,
   "q3_ms": 0.8277754863079843,
   "rel_err_out": 0.012184325992167103,
   "rel_err_grad_params": 0.010998105919565978,
   "rel_err_grad_x": 0.010644396772357527
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.12849999964237213,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 22.209762573242188,
   "xla_arg_mb": 9.500003814697266,
   "xla_out_mb": 9.011833190917969,
   "compile_s": 0.006457005001720972,
   "inner": 30,
   "median_ms": 0.6207896668153505,
   "q1_ms": 0.6153226500221839,
   "q3_ms": 0.6861544166895328,
   "rel_err_out": 0.012184325992167103,
   "rel_err_grad_params": 0.010998105919565978,
   "rel_err_grad_x": 0.010644396772357527
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused_gsum",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.12849999964237213,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 13.394325256347656,
   "xla_arg_mb": 3.4921875,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.0028829079965362325,
   "inner": 42,
   "median_ms": 0.31041959525845453,
   "q1_ms": 0.28623716666299426,
   "q3_ms": 0.38838151184401276,
   "rel_err_out": 0.01218432611866739,
   "rel_err_grad_params": 0.01099810651694539,
   "rel_err_grad_x": 0.010644397495676404
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused_gsum",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.12849999964237213,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 22.210006713867188,
   "xla_arg_mb": 9.500003814697266,
   "xla_out_mb": 9.011833190917969,
   "compile_s": 0.005727333002141677,
   "inner": 44,
   "median_ms": 0.4574157954431774,
   "q1_ms": 0.45131548871392047,
   "q3_ms": 0.4699860567350036,
   "rel_err_out": 0.01218432611866739,
   "rel_err_grad_params": 0.01099810651694539,
   "rel_err_grad_x": 0.010644397495676404
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "blocked",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.12849999964237213,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 27.579383850097656,
   "xla_arg_mb": 3.4921875,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.007933451001008507,
   "inner": 83,
   "median_ms": 0.6136901928069833,
   "q1_ms": 0.5593502651044471,
   "q3_ms": 0.6495964216700958,
   "rel_err_out": 0.01218432611866739,
   "rel_err_grad_params": 0.010998115983172916,
   "rel_err_grad_x": 0.010645545550711373
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "blocked",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.12849999964237213,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 30.066444396972656,
   "xla_arg_mb": 9.500003814697266,
   "xla_out_mb": 9.011833190917969,
   "compile_s": 0.0031910320030874573,
   "inner": 73,
   "median_ms": 0.478587054714484,
   "q1_ms": 0.47199443156698606,
   "q3_ms": 0.48761757537012934,
   "rel_err_out": 0.01218432611866739,
   "rel_err_grad_params": 0.010998115983172916,
   "rel_err_grad_x": 0.010645545550711373
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "ragged_dot",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.12849999964237213,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 35.20314025878906,
   "xla_arg_mb": 3.4921875,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.0035812100031762384,
   "inner": 64,
   "median_ms": 0.4303306718611566,
   "q1_ms": 0.41544006251115206,
   "q3_ms": 0.4323152031702193,
   "rel_err_out": 0.01218432611866739,
   "rel_err_grad_params": 0.010998111200330993,
   "rel_err_grad_x": 0.010645545550711373
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "ragged_dot",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.12849999964237213,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 93.75369262695312,
   "xla_arg_mb": 9.500003814697266,
   "xla_out_mb": 9.011833190917969,
   "compile_s": 0.005987075004668441,
   "inner": 29,
   "median_ms": 1.148895517338854,
   "q1_ms": 1.1429285172951118,
   "q3_ms": 1.1517279309718387,
   "rel_err_out": 0.01218432611866739,
   "rel_err_grad_params": 0.010998111200330993,
   "rel_err_grad_x": 0.010645545550711373
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "capacity",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.12849999964237213,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 14.694839477539062,
   "xla_arg_mb": 3.4921875,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.0015627989996573888,
   "inner": 186,
   "median_ms": 0.16342521504616686,
   "q1_ms": 0.15961321236885537,
   "q3_ms": 0.1645134301118863
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "capacity",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.12849999964237213,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 24.416526794433594,
   "xla_arg_mb": 9.500003814697266,
   "xla_out_mb": 9.011833190917969,
   "compile_s": 0.002781827002763748,
   "inner": 95,
   "median_ms": 0.3452347789884937,
   "q1_ms": 0.3441949684393445,
   "q3_ms": 0.3744721842323136
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dense",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.12849999964237213,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 23.437515258789062,
   "xla_arg_mb": 3.4921875,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.00509967499965569,
   "inner": 94,
   "median_ms": 0.2434017872648373,
   "q1_ms": 0.2166669627455276,
   "q3_ms": 0.28867041487017636,
   "rel_err_out": 0.012184325958399817,
   "rel_err_grad_params": 0.010998109264901535,
   "rel_err_grad_x": 0.010645544481103969
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dense",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.12849999964237213,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 31.251976013183594,
   "xla_arg_mb": 9.500003814697266,
   "xla_out_mb": 9.011833190917969,
   "compile_s": 0.002872979996027425,
   "inner": 82,
   "median_ms": 0.4082811463421222,
   "q1_ms": 0.3997894085531276,
   "q3_ms": 0.4138819389836721,
   "rel_err_out": 0.012184325958399817,
   "rel_err_grad_params": 0.010998109264901535,
   "rel_err_grad_x": 0.010645544481103969
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 4,
   "skew": 0.9,
   "max_expert_frac": 0.23800000548362732,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 11.173377990722656,
   "xla_arg_mb": 3.4921875,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.0037962970018270425,
   "inner": 64,
   "median_ms": 0.284523687582805,
   "q1_ms": 0.2660956640738732,
   "q3_ms": 0.29232188279593174,
   "rel_err_out": 0.0005197545384898878,
   "rel_err_grad_params": 0.001096224029899918,
   "rel_err_grad_x": 0.0006835645892495535
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 4,
   "skew": 0.9,
   "max_expert_frac": 0.23800000548362732,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 17.755409240722656,
   "xla_arg_mb": 9.500003814697266,
   "xla_out_mb": 9.011833190917969,
   "compile_s": 0.006325626003672369,
   "inner": 51,
   "median_ms": 0.5761700587897269,
   "q1_ms": 0.5268040489788879,
   "q3_ms": 0.5845783236116975,
   "rel_err_out": 0.0005197545384898878,
   "rel_err_grad_params": 0.001096224029899918,
   "rel_err_grad_x": 0.0006835645892495535
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "gather",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 4,
   "skew": 0.9,
   "max_expert_frac": 0.23800000548362732,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 13.396278381347656,
   "xla_arg_mb": 3.4921875,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.003732505996595137,
   "inner": 117,
   "median_ms": 0.2140579145684297,
   "q1_ms": 0.21245209828544503,
   "q3_ms": 0.22804593589902122,
   "rel_err_out": 0.0005197545553623271,
   "rel_err_grad_params": 0.0011294922053488084,
   "rel_err_grad_x": 0.0006958447603475682
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "gather",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 4,
   "skew": 0.9,
   "max_expert_frac": 0.23800000548362732,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 24.458053588867188,
   "xla_arg_mb": 9.500003814697266,
   "xla_out_mb": 9.011833190917969,
   "compile_s": 0.003673835002700798,
   "inner": 75,
   "median_ms": 0.44371553328043467,
   "q1_ms": 0.4432364332994136,
   "q3_ms": 0.4470989467032875,
   "rel_err_out": 0.0005197545553623271,
   "rel_err_grad_params": 0.0011294922053488084,
   "rel_err_grad_x": 0.0006958447603475682
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 4,
   "skew": 0.9,
   "max_expert_frac": 0.23800000548362732,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 13.366737365722656,
   "xla_arg_mb": 3.4921875,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.004702423000708222,
   "inner": 94,
   "median_ms": 0.6242600744720203,
   "q1_ms": 0.3018490850932876,
   "q3_ms": 0.7477141808773554,
   "rel_err_out": 0.0005197545262774343,
   "rel_err_grad_params": 0.0011292347766239618,
   "rel_err_grad_x": 0.0006958410488883556
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 4,
   "skew": 0.9,
   "max_expert_frac": 0.23800000548362732,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 22.209762573242188,
   "xla_arg_mb": 9.500003814697266,
   "xla_out_mb": 9.011833190917969,
   "compile_s": 0.02407592199597275,
   "inner": 5,
   "median_ms": 3.4603914013132453,
   "q1_ms": 2.357743800530443,
   "q3_ms": 3.9695224004390184,
   "rel_err_out": 0.0005197545262774343,
   "rel_err_grad_params": 0.0011292347766239618,
   "rel_err_grad_x": 0.0006958410488883556
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused_gsum",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 4,
   "skew": 0.9,
   "max_expert_frac": 0.23800000548362732,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 13.394325256347656,
   "xla_arg_mb": 3.4921875,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.009143741001025774,
   "inner": 74,
   "median_ms": 0.27484400001216075,
   "q1_ms": 0.2690462702417249,
   "q3_ms": 0.28543168239569217,
   "rel_err_out": 0.0005197545384898878,
   "rel_err_grad_params": 0.0011292347766239618,
   "rel_err_grad_x": 0.0006958413158504303
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused_gsum",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 4,
   "skew": 0.9,
   "max_expert_frac": 0.23800000548362732,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 22.210006713867188,
   "xla_arg_mb": 9.500003814697266,
   "xla_out_mb": 9.011833190917969,
   "compile_s": 0.01721647500380641,
   "inner": 39,
   "median_ms": 0.6800279999226452,
   "q1_ms": 0.6162807564042174,
   "q3_ms": 0.9313758974969507,
   "rel_err_out": 0.0005197545384898878,
   "rel_err_grad_params": 0.0011292347766239618,
   "rel_err_grad_x": 0.0006958413158504303
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "blocked",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 4,
   "skew": 0.9,
   "max_expert_frac": 0.23800000548362732,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 27.579383850097656,
   "xla_arg_mb": 3.4921875,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.003490358001727145,
   "inner": 80,
   "median_ms": 0.3204088625352597,
   "q1_ms": 0.30301319370664714,
   "q3_ms": 0.3489840312340675,
   "rel_err_out": 0.0005197545384898878,
   "rel_err_grad_params": 0.001096224029899918,
   "rel_err_grad_x": 0.0006835640971531726
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "blocked",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 4,
   "skew": 0.9,
   "max_expert_frac": 0.23800000548362732,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 30.066444396972656,
   "xla_arg_mb": 9.500003814697266,
   "xla_out_mb": 9.011833190917969,
   "compile_s": 0.006719001001329161,
   "inner": 40,
   "median_ms": 0.48392270000476856,
   "q1_ms": 0.47206987501340336,
   "q3_ms": 0.49914977498701774,
   "rel_err_out": 0.0005197545384898878,
   "rel_err_grad_params": 0.001096224029899918,
   "rel_err_grad_x": 0.0006835640971531726
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "ragged_dot",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 4,
   "skew": 0.9,
   "max_expert_frac": 0.23800000548362732,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 35.20314025878906,
   "xla_arg_mb": 3.4921875,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.0022725219969288446,
   "inner": 76,
   "median_ms": 0.38333918420999236,
   "q1_ms": 0.3703613881378643,
   "q3_ms": 0.4071946118289672,
   "rel_err_out": 0.0005197545384898878,
   "rel_err_grad_params": 0.001096224029899918,
   "rel_err_grad_x": 0.0006835640971531726
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "ragged_dot",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 4,
   "skew": 0.9,
   "max_expert_frac": 0.23800000548362732,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 93.75369262695312,
   "xla_arg_mb": 9.500003814697266,
   "xla_out_mb": 9.011833190917969,
   "compile_s": 0.007614686997840181,
   "inner": 19,
   "median_ms": 1.0883573685484742,
   "q1_ms": 1.0776563159245263,
   "q3_ms": 1.1232083685938146,
   "rel_err_out": 0.0005197545384898878,
   "rel_err_grad_params": 0.001096224029899918,
   "rel_err_grad_x": 0.0006835640971531726
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "capacity",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 4,
   "skew": 0.9,
   "max_expert_frac": 0.23800000548362732,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 14.694839477539062,
   "xla_arg_mb": 3.4921875,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.006096452001656871,
   "inner": 17,
   "median_ms": 0.3138718234635342,
   "q1_ms": 0.2620973823735874,
   "q3_ms": 0.3268467940763324
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "capacity",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 4,
   "skew": 0.9,
   "max_expert_frac": 0.23800000548362732,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 24.416526794433594,
   "xla_arg_mb": 9.500003814697266,
   "xla_out_mb": 9.011833190917969,
   "compile_s": 0.006023423993610777,
   "inner": 59,
   "median_ms": 0.3585726101892622,
   "q1_ms": 0.3418509407048863,
   "q3_ms": 0.3682405762971277
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dense",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 4,
   "skew": 0.9,
   "max_expert_frac": 0.23800000548362732,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 23.437515258789062,
   "xla_arg_mb": 3.4921875,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.002336062003450934,
   "inner": 150,
   "median_ms": 0.11465166666312143,
   "q1_ms": 0.11156698999305566,
   "q3_ms": 0.1164837333150596,
   "rel_err_out": 0.0005197545523893959,
   "rel_err_grad_params": 0.001096143756134782,
   "rel_err_grad_x": 0.0006835634501909138
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dense",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 4,
   "skew": 0.9,
   "max_expert_frac": 0.23800000548362732,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 31.251976013183594,
   "xla_arg_mb": 9.500003814697266,
   "xla_out_mb": 9.011833190917969,
   "compile_s": 0.004403496997838374,
   "inner": 80,
   "median_ms": 0.3491058249892376,
   "q1_ms": 0.3430196812587383,
   "q3_ms": 0.3639802375346335,
   "rel_err_out": 0.0005197545523893959,
   "rel_err_grad_params": 0.001096143756134782,
   "rel_err_grad_x": 0.0006835634501909138
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 1,
   "skew": 0.0,
   "max_expert_frac": 0.02800000086426735,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 12.606483459472656,
   "xla_arg_mb": 24.51953125,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.0037676729989470914,
   "inner": 79,
   "median_ms": 0.40266805058021005,
   "q1_ms": 0.38884105065968655,
   "q3_ms": 0.4282020949913319,
   "rel_err_out": 0.0005194084077871895,
   "rel_err_grad_params": 0.0007197537143998042,
   "rel_err_grad_x": 0.0005395474572417026
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 1,
   "skew": 0.0,
   "max_expert_frac": 0.02800000086426735,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 0.02176666259765625,
   "xla_arg_mb": 72.58203506469727,
   "xla_out_mb": 72.09386444091797,
   "compile_s": 0.007585322004160844,
   "inner": 27,
   "median_ms": 0.8503095183793145,
   "q1_ms": 0.8426767037437659,
   "q3_ms": 1.0852133147838887,
   "rel_err_out": 0.0005194084077871895,
   "rel_err_grad_params": 0.0007197537143998042,
   "rel_err_grad_x": 0.0005395474572417026
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "gather",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 1,
   "skew": 0.0,
   "max_expert_frac": 0.02800000086426735,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 28.397010803222656,
   "xla_arg_mb": 24.51953125,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.004333894998126198,
   "inner": 24,
   "median_ms": 0.4601799998151061,
   "q1_ms": 0.4342740833938782,
   "q3_ms": 0.5568361041999499,
   "rel_err_out": 0.0005194084077871895,
   "rel_err_grad_params": 0.0007197537143998042,
   "rel_err_grad_x": 0.0005395619945495329
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "gather",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 1,
   "skew": 0.0,
   "max_expert_frac": 0.02800000086426735,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 36.49125671386719,
   "xla_arg_mb": 72.58203506469727,
   "xla_out_mb": 72.09386444091797,
   "compile_s": 0.00639266300277086,
   "inner": 33,
   "median_ms": 0.8908411515276936,
   "q1_ms": 0.8853093181228773,
   "q3_ms": 0.9028086364692586,
   "rel_err_out": 0.0005194084077871895,
   "rel_err_grad_params": 0.0007197537143998042,
   "rel_err_grad_x": 0.0005395619945495329
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 1,
   "skew": 0.0,
   "max_expert_frac": 0.02800000086426735,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 25.935096740722656,
   "xla_arg_mb": 24.51953125,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.003614251996623352,
   "inner": 49,
   "median_ms": 0.4712092652750601,
   "q1_ms": 0.4642235917901639,
   "q3_ms": 0.4817014796200817,
   "rel_err_out": 0.0005194084077871895,
   "rel_err_grad_params": 0.0007197537143998042,
   "rel_err_grad_x": 0.0005395619945495329
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 1,
   "skew": 0.0,
   "max_expert_frac": 0.02800000086426735,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 36.31278991699219,
   "xla_arg_mb": 72.58203506469727,
   "xla_out_mb": 72.09386444091797,
   "compile_s": 0.007982632996572647,
   "inner": 21,
   "median_ms": 1.185854190442201,
   "q1_ms": 1.1763698095732944,
   "q3_ms": 1.2900412142958624,
   "rel_err_out": 0.0005194084077871895,
   "rel_err_grad_params": 0.0007197537143998042,
   "rel_err_grad_x": 0.0005395619945495329
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused_gsum",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 1,
   "skew": 0.0,
   "max_expert_frac": 0.02800000086426735,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 25.932167053222656,
   "xla_arg_mb": 24.51953125,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.005952480001724325,
   "inner": 34,
   "median_ms": 0.46541849994408785,
   "q1_ms": 0.45951836770708565,
   "q3_ms": 0.49909652946499755,
   "rel_err_out": 0.0005194084077871895,
   "rel_err_grad_params": 0.0007197537143998042,
   "rel_err_grad_x": 0.0005395619945495329
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused_gsum",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 1,
   "skew": 0.0,
   "max_expert_frac": 0.02800000086426735,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 36.31694030761719,
   "xla_arg_mb": 72.58203506469727,
   "xla_out_mb": 72.09386444091797,
   "compile_s": 0.008165859995642677,
   "inner": 26,
   "median_ms": 1.1298333078089995,
   "q1_ms": 1.1144956153237637,
   "q3_ms": 1.1484057691772112,
   "rel_err_out": 0.0005194084077871895,
   "rel_err_grad_params": 0.0007197537143998042,
   "rel_err_grad_x": 0.0005395619945495329
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "blocked",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 1,
   "skew": 0.0,
   "max_expert_frac": 0.02800000086426735,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 49.762489318847656,
   "xla_arg_mb": 24.51953125,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.0067656289975275286,
   "inner": 11,
   "median_ms": 0.703871727058537,
   "q1_ms": 0.6250547273587224,
   "q3_ms": 0.776231000186685,
   "rel_err_out": 0.0005194084077871895,
   "rel_err_grad_params": 0.0007197537143998042,
   "rel_err_grad_x": 0.0005395488216186534
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "blocked",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 1,
   "skew": 0.0,
   "max_expert_frac": 0.02800000086426735,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 45.021034240722656,
   "xla_arg_mb": 72.58203506469727,
   "xla_out_mb": 72.09386444091797,
   "compile_s": 0.026529967995884363,
   "inner": 25,
   "median_ms": 1.734574839938432,
   "q1_ms": 1.5568781600450166,
   "q3_ms": 2.2912300999450963,
   "rel_err_out": 0.0005194084077871895,
   "rel_err_grad_params": 0.0007197537143998042,
   "rel_err_grad_x": 0.0005395488216186534
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "ragged_dot",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 1,
   "skew": 0.0,
   "max_expert_frac": 0.02800000086426735,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 66.65895080566406,
   "xla_arg_mb": 24.51953125,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.003487481997581199,
   "inner": 52,
   "median_ms": 0.6774757307889209,
   "q1_ms": 0.6748399615529566,
   "q3_ms": 0.6832641826034747,
   "rel_err_out": 0.0005194084077871895,
   "rel_err_grad_params": 0.0007197538619694328,
   "rel_err_grad_x": 0.0005395488216186534
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "ragged_dot",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 1,
   "skew": 0.0,
   "max_expert_frac": 0.02800000086426735,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 156.25857543945312,
   "xla_arg_mb": 72.58203506469727,
   "xla_out_mb": 72.09386444091797,
   "compile_s": 0.0071094800005084835,
   "inner": 16,
   "median_ms": 2.271901312724367,
   "q1_ms": 2.2183186251822917,
   "q3_ms": 2.352303843736081,
   "rel_err_out": 0.0005194084077871895,
   "rel_err_grad_params": 0.0007197538619694328,
   "rel_err_grad_x": 0.0005395488216186534
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "capacity",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 1,
   "skew": 0.0,
   "max_expert_frac": 0.02800000086426735,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 4.2426910400390625,
   "xla_arg_mb": 24.51953125,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.0031842079988564365,
   "inner": 154,
   "median_ms": 0.13397820133265326,
   "q1_ms": 0.1084661363909783,
   "q3_ms": 0.15221331168858898
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "capacity",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 1,
   "skew": 0.0,
   "max_expert_frac": 0.02800000086426735,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 17.603782653808594,
   "xla_arg_mb": 72.58203506469727,
   "xla_out_mb": 72.09386444091797,
   "compile_s": 0.020357512003101874,
   "inner": 43,
   "median_ms": 0.4701738837675928,
   "q1_ms": 0.40591315117253124,
   "q3_ms": 0.49009532550942736
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dense",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 1,
   "skew": 0.0,
   "max_expert_frac": 0.02800000086426735,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 125.00001525878906,
   "xla_arg_mb": 24.51953125,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.002362172002904117,
   "inner": 53,
   "median_ms": 0.7546147359600993,
   "q1_ms": 0.7513962736307152,
   "q3_ms": 0.7593848868569728,
   "rel_err_out": 0.0005194084077871895,
   "rel_err_grad_params": 0.0007197553677978324,
   "rel_err_grad_x": 0.0005395488216186534
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dense",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 1,
   "skew": 0.0,
   "max_expert_frac": 0.02800000086426735,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 218.7527084350586,
   "xla_arg_mb": 72.58203506469727,
   "xla_out_mb": 72.09386444091797,
   "compile_s": 0.004520258000411559,
   "inner": 20,
   "median_ms": 2.2333122498821467,
   "q1_ms": 2.2066417750465916,
   "q3_ms": 2.2401629748856067,
   "rel_err_out": 0.0005194084077871895,
   "rel_err_grad_params": 0.0007197553677978324,
   "rel_err_grad_x": 0.0005395488216186534
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 1,
   "skew": 0.9,
   "max_expert_frac": 0.9320000410079956,
   "empty_experts": 16,
   "op": "forward",
   "xla_temp_mb": 12.606483459472656,
   "xla_arg_mb": 24.51953125,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.0020757699967361987,
   "inner": 110,
   "median_ms": 0.2616289727236356,
   "q1_ms": 0.25149934091049625,
   "q3_ms": 0.27847937275179324,
   "rel_err_out": 0.0005147767136513409,
   "rel_err_grad_params": 0.00043615885679524176,
   "rel_err_grad_x": 0.0005389297406673184
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 1,
   "skew": 0.9,
   "max_expert_frac": 0.9320000410079956,
   "empty_experts": 16,
   "op": "train_step",
   "xla_temp_mb": 0.02176666259765625,
   "xla_arg_mb": 72.58203506469727,
   "xla_out_mb": 72.09386444091797,
   "compile_s": 0.003712618999998085,
   "inner": 50,
   "median_ms": 0.7474316800653469,
   "q1_ms": 0.7457114199496573,
   "q3_ms": 0.7523844599927543,
   "rel_err_out": 0.0005147767136513409,
   "rel_err_grad_params": 0.00043615885679524176,
   "rel_err_grad_x": 0.0005389297406673184
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "gather",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 1,
   "skew": 0.9,
   "max_expert_frac": 0.9320000410079956,
   "empty_experts": 16,
   "op": "forward",
   "xla_temp_mb": 28.397010803222656,
   "xla_arg_mb": 24.51953125,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.0028572099981829524,
   "inner": 86,
   "median_ms": 0.42011876745206345,
   "q1_ms": 0.37127897671463106,
   "q3_ms": 0.5283486686049972,
   "rel_err_out": 0.0005147767136513409,
   "rel_err_grad_params": 0.00043615885679524176,
   "rel_err_grad_x": 0.0005389166678085416
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "gather",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 1,
   "skew": 0.9,
   "max_expert_frac": 0.9320000410079956,
   "empty_experts": 16,
   "op": "train_step",
   "xla_temp_mb": 36.49125671386719,
   "xla_arg_mb": 72.58203506469727,
   "xla_out_mb": 72.09386444091797,
   "compile_s": 0.004327572998590767,
   "inner": 32,
   "median_ms": 0.9157086562936456,
   "q1_ms": 0.9029806718672262,
   "q3_ms": 1.0104311405711996,
   "rel_err_out": 0.0005147767136513409,
   "rel_err_grad_params": 0.00043615885679524176,
   "rel_err_grad_x": 0.0005389166678085416
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 1,
   "skew": 0.9,
   "max_expert_frac": 0.9320000410079956,
   "empty_experts": 16,
   "op": "forward",
   "xla_temp_mb": 25.935096740722656,
   "xla_arg_mb": 24.51953125,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.010747597996669356,
   "inner": 14,
   "median_ms": 1.089396142820728,
   "q1_ms": 0.972197749953401,
   "q3_ms": 1.1866015715245988,
   "rel_err_out": 0.0005147767136513409,
   "rel_err_grad_params": 0.00043615885679524176,
   "rel_err_grad_x": 0.0005389166678085416
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 1,
   "skew": 0.9,
   "max_expert_frac": 0.9320000410079956,
   "empty_experts": 16,
   "op": "train_step",
   "xla_temp_mb": 36.31278991699219,
   "xla_arg_mb": 72.58203506469727,
   "xla_out_mb": 72.09386444091797,
   "compile_s": 0.029096418002154678,
   "inner": 24,
   "median_ms": 1.288176333522036,
   "q1_ms": 1.1944587290599884,
   "q3_ms": 1.3721378125713577,
   "rel_err_out": 0.0005147767136513409,
   "rel_err_grad_params": 0.00043615885679524176,
   "rel_err_grad_x": 0.0005389166678085416
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused_gsum",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 1,
   "skew": 0.9,
   "max_expert_frac": 0.9320000410079956,
   "empty_experts": 16,
   "op": "forward",
   "xla_temp_mb": 25.932167053222656,
   "xla_arg_mb": 24.51953125,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.08791379399917787,
   "inner": 18,
   "median_ms": 0.3752743889183168,
   "q1_ms": 0.37300227773408146,
   "q3_ms": 0.3924326944494775,
   "rel_err_out": 0.0005147767136513409,
   "rel_err_grad_params": 0.00043615885679524176,
   "rel_err_grad_x": 0.0005389166678085416
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused_gsum",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 1,
   "skew": 0.9,
   "max_expert_frac": 0.9320000410079956,
   "empty_experts": 16,
   "op": "train_step",
   "xla_temp_mb": 36.31694030761719,
   "xla_arg_mb": 72.58203506469727,
   "xla_out_mb": 72.09386444091797,
   "compile_s": 0.006126888998551294,
   "inner": 30,
   "median_ms": 1.6102242332029466,
   "q1_ms": 1.3766041166187883,
   "q3_ms": 1.6966236666727732,
   "rel_err_out": 0.0005147767136513409,
   "rel_err_grad_params": 0.00043615885679524176,
   "rel_err_grad_x": 0.0005389166678085416
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "blocked",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 1,
   "skew": 0.9,
   "max_expert_frac": 0.9320000410079956,
   "empty_experts": 16,
   "op": "forward",
   "xla_temp_mb": 49.762489318847656,
   "xla_arg_mb": 24.51953125,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.0022266770029091276,
   "inner": 77,
   "median_ms": 0.3502997403432216,
   "q1_ms": 0.349477928584629,
   "q3_ms": 0.3711618506689798,
   "rel_err_out": 0.0005147767136513409,
   "rel_err_grad_params": 0.0004361772768670316,
   "rel_err_grad_x": 0.000538929596000033
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "blocked",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 1,
   "skew": 0.9,
   "max_expert_frac": 0.9320000410079956,
   "empty_experts": 16,
   "op": "train_step",
   "xla_temp_mb": 45.021034240722656,
   "xla_arg_mb": 72.58203506469727,
   "xla_out_mb": 72.09386444091797,
   "compile_s": 0.004364574997453019,
   "inner": 41,
   "median_ms": 0.9253484877706618,
   "q1_ms": 0.9191980609945704,
   "q3_ms": 1.0186810975095042,
   "rel_err_out": 0.0005147767136513409,
   "rel_err_grad_params": 0.0004361772768670316,
   "rel_err_grad_x": 0.000538929596000033
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "ragged_dot",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 1,
   "skew": 0.9,
   "max_expert_frac": 0.9320000410079956,
   "empty_experts": 16,
   "op": "forward",
   "xla_temp_mb": 66.65895080566406,
   "xla_arg_mb": 24.51953125,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.002330493000044953,
   "inner": 66,
   "median_ms": 0.7007632878164507,
   "q1_ms": 0.6483019620425308,
   "q3_ms": 0.8176211060455182,
   "rel_err_out": 0.0005147767136513409,
   "rel_err_grad_params": 0.00043615885702718096,
   "rel_err_grad_x": 0.000538929596000033
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "ragged_dot",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 1,
   "skew": 0.9,
   "max_expert_frac": 0.9320000410079956,
   "empty_experts": 16,
   "op": "train_step",
   "xla_temp_mb": 156.25857543945312,
   "xla_arg_mb": 72.58203506469727,
   "xla_out_mb": 72.09386444091797,
   "compile_s": 0.0046129260008456185,
   "inner": 20,
   "median_ms": 2.2100803500507027,
   "q1_ms": 2.0992221498090657,
   "q3_ms": 3.0134194250422297,
   "rel_err_out": 0.0005147767136513409,
   "rel_err_grad_params": 0.00043615885702718096,
   "rel_err_grad_x": 0.000538929596000033
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "capacity",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 1,
   "skew": 0.9,
   "max_expert_frac": 0.9320000410079956,
   "empty_experts": 16,
   "op": "forward",
   "xla_temp_mb": 4.2426910400390625,
   "xla_arg_mb": 24.51953125,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.0015443450029124506,
   "inner": 200,
   "median_ms": 0.11411886000132654,
   "q1_ms": 0.10636297500241199,
   "q3_ms": 0.12807882500055712
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "capacity",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 1,
   "skew": 0.9,
   "max_expert_frac": 0.9320000410079956,
   "empty_experts": 16,
   "op": "train_step",
   "xla_temp_mb": 17.603782653808594,
   "xla_arg_mb": 72.58203506469727,
   "xla_out_mb": 72.09386444091797,
   "compile_s": 0.002917876001447439,
   "inner": 75,
   "median_ms": 0.46281646665496134,
   "q1_ms": 0.46079076001963887,
   "q3_ms": 0.6762380866469659
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dense",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 1,
   "skew": 0.9,
   "max_expert_frac": 0.9320000410079956,
   "empty_experts": 16,
   "op": "forward",
   "xla_temp_mb": 125.00001525878906,
   "xla_arg_mb": 24.51953125,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.006968195993977133,
   "inner": 25,
   "median_ms": 0.757262039987836,
   "q1_ms": 0.7494045999192167,
   "q3_ms": 0.759569020010531,
   "rel_err_out": 0.0005147767136513409,
   "rel_err_grad_params": 0.0004361549381115259,
   "rel_err_grad_x": 0.000538929596000033
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dense",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 1,
   "skew": 0.9,
   "max_expert_frac": 0.9320000410079956,
   "empty_experts": 16,
   "op": "train_step",
   "xla_temp_mb": 218.7527084350586,
   "xla_arg_mb": 72.58203506469727,
   "xla_out_mb": 72.09386444091797,
   "compile_s": 0.004061281004396733,
   "inner": 21,
   "median_ms": 2.5032404761683815,
   "q1_ms": 2.5005244046735156,
   "q3_ms": 2.510822976113386,
   "rel_err_out": 0.0005147767136513409,
   "rel_err_grad_params": 0.0004361549381115259,
   "rel_err_grad_x": 0.000538929596000033
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.022750001400709152,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 19.767372131347656,
   "xla_arg_mb": 24.51953125,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.002511806997063104,
   "inner": 87,
   "median_ms": 0.4539988390885925,
   "q1_ms": 0.41608410918698435,
   "q3_ms": 0.6527696896646167,
   "rel_err_out": 0.0005772428273623252,
   "rel_err_grad_params": 0.0008223208181534585,
   "rel_err_grad_x": 0.0006922784117970971
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.022750001400709152,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 0.31595611572265625,
   "xla_arg_mb": 72.58203506469727,
   "xla_out_mb": 72.09386444091797,
   "compile_s": 0.054658329994708765,
   "inner": 32,
   "median_ms": 0.9716250312976626,
   "q1_ms": 0.9572637657129235,
   "q3_ms": 1.0007166717969085,
   "rel_err_out": 0.0005772428273623252,
   "rel_err_grad_params": 0.0008223208181534585,
   "rel_err_grad_x": 0.0006922784117970971
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "gather",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.022750001400709152,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 35.798133850097656,
   "xla_arg_mb": 24.51953125,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.002525572999729775,
   "inner": 82,
   "median_ms": 0.5060685975100879,
   "q1_ms": 0.49091350606098916,
   "q3_ms": 0.7648557012300177,
   "rel_err_out": 0.0005772430447566889,
   "rel_err_grad_params": 0.0008014748130307203,
   "rel_err_grad_x": 0.0006809472686841025
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "gather",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.022750001400709152,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 32.32109069824219,
   "xla_arg_mb": 72.58203506469727,
   "xla_out_mb": 72.09386444091797,
   "compile_s": 0.013073818998236675,
   "inner": 17,
   "median_ms": 1.3025319413453176,
   "q1_ms": 1.2133466176130563,
   "q3_ms": 1.4486283531261828,
   "rel_err_out": 0.0005772430447566889,
   "rel_err_grad_params": 0.0008014748130307203,
   "rel_err_grad_x": 0.0006809472686841025
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.022750001400709152,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 31.844520568847656,
   "xla_arg_mb": 24.51953125,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.0023628850030945614,
   "inner": 76,
   "median_ms": 0.4558831973745806,
   "q1_ms": 0.4479270789344906,
   "q3_ms": 0.4687277829170374,
   "rel_err_out": 0.0005772429845821953,
   "rel_err_grad_params": 0.0008014693977111677,
   "rel_err_grad_x": 0.0006809453631386828
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.022750001400709152,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 47.90287780761719,
   "xla_arg_mb": 72.58203506469727,
   "xla_out_mb": 72.09386444091797,
   "compile_s": 0.003997751002316363,
   "inner": 39,
   "median_ms": 1.1380048461619596,
   "q1_ms": 1.120504269168193,
   "q3_ms": 1.17767933326645,
   "rel_err_out": 0.0005772429845821953,
   "rel_err_grad_params": 0.0008014693977111677,
   "rel_err_grad_x": 0.0006809453631386828
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused_gsum",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.022750001400709152,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 31.876014709472656,
   "xla_arg_mb": 24.51953125,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.002352846000576392,
   "inner": 89,
   "median_ms": 0.43942013479444836,
   "q1_ms": 0.42187083147765725,
   "q3_ms": 0.46770955055619834,
   "rel_err_out": 0.0005772428273623252,
   "rel_err_grad_params": 0.0008014644200040841,
   "rel_err_grad_x": 0.0006809489699858832
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused_gsum",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.022750001400709152,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 47.91850280761719,
   "xla_arg_mb": 72.58203506469727,
   "xla_out_mb": 72.09386444091797,
   "compile_s": 0.004269224999006838,
   "inner": 32,
   "median_ms": 1.2342470624844282,
   "q1_ms": 1.149491859337104,
   "q3_ms": 1.2459415311241173,
   "rel_err_out": 0.0005772428273623252,
   "rel_err_grad_params": 0.0008014644200040841,
   "rel_err_grad_x": 0.0006809489699858832
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "blocked",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.022750001400709152,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 65.39212799072266,
   "xla_arg_mb": 24.51953125,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.0021716629998991266,
   "inner": 78,
   "median_ms": 0.4901186794785855,
   "q1_ms": 0.4869209551543463,
   "q3_ms": 0.49854418592496164,
   "rel_err_out": 0.0005772428273623252,
   "rel_err_grad_params": 0.0008223208181534585,
   "rel_err_grad_x": 0.0006922692195913684
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "blocked",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.022750001400709152,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 59.918251037597656,
   "xla_arg_mb": 72.58203506469727,
   "xla_out_mb": 72.09386444091797,
   "compile_s": 0.005171683995286003,
   "inner": 32,
   "median_ms": 1.1883022500569496,
   "q1_ms": 1.1803969375705492,
   "q3_ms": 1.1890981094211384,
   "rel_err_out": 0.0005772428273623252,
   "rel_err_grad_params": 0.0008223208181534585,
   "rel_err_grad_x": 0.0006922692195913684
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "ragged_dot",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.022750001400709152,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 261.85963439941406,
   "xla_arg_mb": 24.51953125,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.0036425279977265745,
   "inner": 21,
   "median_ms": 2.155883190252574,
   "q1_ms": 2.127756357367616,
   "q3_ms": 2.172158785903039,
   "rel_err_out": 0.0005772426021365553,
   "rel_err_grad_params": 0.0008223208181534585,
   "rel_err_grad_x": 0.0006922601994012503
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "ragged_dot",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.022750001400709152,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 625.0044250488281,
   "xla_arg_mb": 72.58203506469727,
   "xla_out_mb": 72.09386444091797,
   "compile_s": 0.010312439000699669,
   "inner": 6,
   "median_ms": 7.957688166546481,
   "q1_ms": 7.759668250097699,
   "q3_ms": 8.349375749579242,
   "rel_err_out": 0.0005772426021365553,
   "rel_err_grad_params": 0.0008223208181534585,
   "rel_err_grad_x": 0.0006922601994012503
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "capacity",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.022750001400709152,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 15.072769165039062,
   "xla_arg_mb": 24.51953125,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.001580062002176419,
   "inner": 123,
   "median_ms": 0.20516113005673164,
   "q1_ms": 0.20280474797719372,
   "q3_ms": 0.20580954878823832
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "capacity",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.022750001400709152,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 16.742698669433594,
   "xla_arg_mb": 72.58203506469727,
   "xla_out_mb": 72.09386444091797,
   "compile_s": 0.0029895810002926737,
   "inner": 62,
   "median_ms": 0.5724325967622139,
   "q1_ms": 0.5673019032626818,
   "q3_ms": 0.5906058870623969
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dense",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.022750001400709152,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 125.00001525878906,
   "xla_arg_mb": 24.51953125,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.001795340001990553,
   "inner": 60,
   "median_ms": 0.8491161833565759,
   "q1_ms": 0.7669430583215823,
   "q3_ms": 0.8499657833453966,
   "rel_err_out": 0.0005772429011971193,
   "rel_err_grad_params": 0.0008215619959479718,
   "rel_err_grad_x": 0.0006918220272142461
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dense",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.022750001400709152,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 218.7522201538086,
   "xla_arg_mb": 72.58203506469727,
   "xla_out_mb": 72.09386444091797,
   "compile_s": 0.0044256409964873455,
   "inner": 21,
   "median_ms": 2.198431857063302,
   "q1_ms": 2.18988949997029,
   "q3_ms": 2.2053566192750753,
   "rel_err_out": 0.0005772429011971193,
   "rel_err_grad_params": 0.0008215619959479718,
   "rel_err_grad_x": 0.0006918220272142461
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.9,
   "max_expert_frac": 0.234250009059906,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 19.767372131347656,
   "xla_arg_mb": 24.51953125,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.002160550997359678,
   "inner": 84,
   "median_ms": 0.3715167737633844,
   "q1_ms": 0.3693004940768809,
   "q3_ms": 0.38932883931660656,
   "rel_err_out": 0.0005186876225241922,
   "rel_err_grad_params": 0.0010479485133974365,
   "rel_err_grad_x": 0.0006176730986909898
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.9,
   "max_expert_frac": 0.234250009059906,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 0.31595611572265625,
   "xla_arg_mb": 72.58203506469727,
   "xla_out_mb": 72.09386444091797,
   "compile_s": 0.004167942999629304,
   "inner": 40,
   "median_ms": 0.9721307249492384,
   "q1_ms": 0.9666847749940644,
   "q3_ms": 1.0203088374510116,
   "rel_err_out": 0.0005186876225241922,
   "rel_err_grad_params": 0.0010479485133974365,
   "rel_err_grad_x": 0.0006176730986909898
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "gather",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.9,
   "max_expert_frac": 0.234250009059906,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 35.798133850097656,
   "xla_arg_mb": 24.51953125,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.013126377998560201,
   "inner": 38,
   "median_ms": 0.7587792895624021,
   "q1_ms": 0.6451009999421473,
   "q3_ms": 0.7882288946826808,
   "rel_err_out": 0.000518687627554293,
   "rel_err_grad_params": 0.0010363673402051827,
   "rel_err_grad_x": 0.0006181328979627797
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "gather",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.9,
   "max_expert_frac": 0.234250009059906,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 32.32109069824219,
   "xla_arg_mb": 72.58203506469727,
   "xla_out_mb": 72.09386444091797,
   "compile_s": 0.008584447001339868,
   "inner": 25,
   "median_ms": 1.6983743201126345,
   "q1_ms": 1.5658615400025155,
   "q3_ms": 1.8737794800836127,
   "rel_err_out": 0.000518687627554293,
   "rel_err_grad_params": 0.0010363673402051827,
   "rel_err_grad_x": 0.0006181328979627797
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.9,
   "max_expert_frac": 0.234250009059906,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 31.844520568847656,
   "xla_arg_mb": 24.51953125,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.0022881529948790558,
   "inner": 81,
   "median_ms": 0.5019499629629591,
   "q1_ms": 0.4628789259489471,
   "q3_ms": 0.5074414012024473,
   "rel_err_out": 0.0005186877701922546,
   "rel_err_grad_params": 0.0010363673402051827,
   "rel_err_grad_x": 0.0006181321170116436
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.9,
   "max_expert_frac": 0.234250009059906,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 47.90287780761719,
   "xla_arg_mb": 72.58203506469727,
   "xla_out_mb": 72.09386444091797,
   "compile_s": 0.004200043993478175,
   "inner": 35,
   "median_ms": 1.1263890857143062,
   "q1_ms": 1.1234392286236732,
   "q3_ms": 1.137574714292506,
   "rel_err_out": 0.0005186877701922546,
   "rel_err_grad_params": 0.0010363673402051827,
   "rel_err_grad_x": 0.0006181321170116436
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused_gsum",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.9,
   "max_expert_frac": 0.234250009059906,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 31.876014709472656,
   "xla_arg_mb": 24.51953125,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.0022311549982987344,
   "inner": 73,
   "median_ms": 0.4616791095579847,
   "q1_ms": 0.4307352260495887,
   "q3_ms": 0.4899687671352956,
   "rel_err_out": 0.0005186876225241922,
   "rel_err_grad_params": 0.0010363673402051827,
   "rel_err_grad_x": 0.0006181331462347699
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused_gsum",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.9,
   "max_expert_frac": 0.234250009059906,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 47.91850280761719,
   "xla_arg_mb": 72.58203506469727,
   "xla_out_mb": 72.09386444091797,
   "compile_s": 0.0037110470002517104,
   "inner": 41,
   "median_ms": 1.0443103415238466,
   "q1_ms": 1.0390269268189114,
   "q3_ms": 1.0525333536639754,
   "rel_err_out": 0.0005186876225241922,
   "rel_err_grad_params": 0.0010363673402051827,
   "rel_err_grad_x": 0.0006181331462347699
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "blocked",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.9,
   "max_expert_frac": 0.234250009059906,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 65.39212799072266,
   "xla_arg_mb": 24.51953125,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.0020926630022586323,
   "inner": 81,
   "median_ms": 0.5142866543287204,
   "q1_ms": 0.4941479568021159,
   "q3_ms": 0.6937260247165493,
   "rel_err_out": 0.0005186876225241922,
   "rel_err_grad_params": 0.0010479485133974365,
   "rel_err_grad_x": 0.0006176728188322955
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "blocked",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.9,
   "max_expert_frac": 0.234250009059906,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 59.918251037597656,
   "xla_arg_mb": 72.58203506469727,
   "xla_out_mb": 72.09386444091797,
   "compile_s": 0.010571500002697576,
   "inner": 33,
   "median_ms": 1.1688693030560953,
   "q1_ms": 1.15407340920435,
   "q3_ms": 1.260647500065247,
   "rel_err_out": 0.0005186876225241922,
   "rel_err_grad_params": 0.0010479485133974365,
   "rel_err_grad_x": 0.0006176728188322955
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "ragged_dot",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.9,
   "max_expert_frac": 0.234250009059906,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 261.85963439941406,
   "xla_arg_mb": 24.51953125,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.0046060329987085424,
   "inner": 18,
   "median_ms": 2.323988833167176,
   "q1_ms": 2.170729166917024,
   "q3_ms": 2.334878610983853,
   "rel_err_out": 0.0005186877596714431,
   "rel_err_grad_params": 0.0010479485133974365,
   "rel_err_grad_x": 0.0006176723503605931
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "ragged_dot",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.9,
   "max_expert_frac": 0.234250009059906,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 625.0044250488281,
   "xla_arg_mb": 72.58203506469727,
   "xla_out_mb": 72.09386444091797,
   "compile_s": 0.010308150995115284,
   "inner": 5,
   "median_ms": 8.857784399879165,
   "q1_ms": 7.94913869976881,
   "q3_ms": 11.408665199269308,
   "rel_err_out": 0.0005186877596714431,
   "rel_err_grad_params": 0.0010479485133974365,
   "rel_err_grad_x": 0.0006176723503605931
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "capacity",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.9,
   "max_expert_frac": 0.234250009059906,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 15.072769165039062,
   "xla_arg_mb": 24.51953125,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.0016173429976333864,
   "inner": 177,
   "median_ms": 0.17996040676190325,
   "q1_ms": 0.16781675706516902,
   "q3_ms": 0.18113975705086297
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "capacity",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.9,
   "max_expert_frac": 0.234250009059906,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 16.742698669433594,
   "xla_arg_mb": 72.58203506469727,
   "xla_out_mb": 72.09386444091797,
   "compile_s": 0.004377910001494456,
   "inner": 52,
   "median_ms": 0.642617153901553,
   "q1_ms": 0.5873520096009508,
   "q3_ms": 0.6824540287897645
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dense",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.9,
   "max_expert_frac": 0.234250009059906,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 125.00001525878906,
   "xla_arg_mb": 24.51953125,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.0023440489967470057,
   "inner": 60,
   "median_ms": 0.7934691333502997,
   "q1_ms": 0.7533518333730171,
   "q3_ms": 0.7984857166472162,
   "rel_err_out": 0.0005186876418448801,
   "rel_err_grad_params": 0.0010481054925347798,
   "rel_err_grad_x": 0.0006176689063939197
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dense",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "float32",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.9,
   "max_expert_frac": 0.234250009059906,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 218.7522201538086,
   "xla_arg_mb": 72.58203506469727,
   "xla_out_mb": 72.09386444091797,
   "compile_s": 0.004963579005561769,
   "inner": 20,
   "median_ms": 2.301344700026675,
   "q1_ms": 2.2767693750211038,
   "q3_ms": 2.516362875030609,
   "rel_err_out": 0.0005186876418448801,
   "rel_err_grad_params": 0.0010481054925347798,
   "rel_err_grad_x": 0.0006176689063939197
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 1,
   "skew": 0.0,
   "max_expert_frac": 0.15900000929832458,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 2.4236221313476562,
   "xla_arg_mb": 3.248046875,
   "xla_out_mb": 0.244140625,
   "compile_s": 0.002001270004257094,
   "inner": 174,
   "median_ms": 0.19280812068308448,
   "q1_ms": 0.18252424712860876,
   "q3_ms": 0.21282165230906852,
   "rel_err_out": 0.05459915645864793,
   "rel_err_grad_params": 0.053237652261722695,
   "rel_err_grad_x": 0.053695951445510724
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 1,
   "skew": 0.0,
   "max_expert_frac": 0.15900000929832458,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 1.3801651000976562,
   "xla_arg_mb": 9.255863189697266,
   "xla_out_mb": 9.011833190917969,
   "compile_s": 0.018573574001493398,
   "inner": 32,
   "median_ms": 0.4114238438432949,
   "q1_ms": 0.34600226547354396,
   "q3_ms": 0.45120946867882594,
   "rel_err_out": 0.05459915645864793,
   "rel_err_grad_params": 0.053237652261722695,
   "rel_err_grad_x": 0.053695951445510724
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "gather",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 1,
   "skew": 0.0,
   "max_expert_frac": 0.15900000929832458,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 2.9011611938476562,
   "xla_arg_mb": 3.248046875,
   "xla_out_mb": 0.244140625,
   "compile_s": 0.007381086994428188,
   "inner": 110,
   "median_ms": 0.18940988176671616,
   "q1_ms": 0.18865972726251296,
   "q3_ms": 0.21257146816589573,
   "rel_err_out": 0.05459915645864793,
   "rel_err_grad_params": 0.053237652261722695,
   "rel_err_grad_x": 0.05364085470253975
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "gather",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 1,
   "skew": 0.0,
   "max_expert_frac": 0.15900000929832458,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 3.0710906982421875,
   "xla_arg_mb": 9.255863189697266,
   "xla_out_mb": 9.011833190917969,
   "compile_s": 0.009279543002776336,
   "inner": 94,
   "median_ms": 0.3150682659520015,
   "q1_ms": 0.30195181913787816,
   "q3_ms": 0.3224442446848348,
   "rel_err_out": 0.05459915645864793,
   "rel_err_grad_params": 0.053237652261722695,
   "rel_err_grad_x": 0.05364085470253975
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 1,
   "skew": 0.0,
   "max_expert_frac": 0.15900000929832458,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 2.5195693969726562,
   "xla_arg_mb": 3.248046875,
   "xla_out_mb": 0.244140625,
   "compile_s": 0.0031309389960370027,
   "inner": 125,
   "median_ms": 0.25045069598127156,
   "q1_ms": 0.24419379999744706,
   "q3_ms": 0.2578949879971333,
   "rel_err_out": 0.054619682930089404,
   "rel_err_grad_params": 0.05322258081450044,
   "rel_err_grad_x": 0.053617261975707625
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 1,
   "skew": 0.0,
   "max_expert_frac": 0.15900000929832458,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 3.0879364013671875,
   "xla_arg_mb": 9.255863189697266,
   "xla_out_mb": 9.011833190917969,
   "compile_s": 0.003235427000618074,
   "inner": 112,
   "median_ms": 0.3011807142943975,
   "q1_ms": 0.29453966070442605,
   "q3_ms": 0.3031213169444007,
   "rel_err_out": 0.054619682930089404,
   "rel_err_grad_params": 0.05322258081450044,
   "rel_err_grad_x": 0.053617261975707625
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused_gsum",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 1,
   "skew": 0.0,
   "max_expert_frac": 0.15900000929832458,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 2.5339736938476562,
   "xla_arg_mb": 3.248046875,
   "xla_out_mb": 0.244140625,
   "compile_s": 0.009754161997989286,
   "inner": 50,
   "median_ms": 0.5295109598955605,
   "q1_ms": 0.42066713998792693,
   "q3_ms": 0.6043436800246127,
   "rel_err_out": 0.05459915645864793,
   "rel_err_grad_params": 0.053237652261722695,
   "rel_err_grad_x": 0.05364085470253975
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused_gsum",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 1,
   "skew": 0.0,
   "max_expert_frac": 0.15900000929832458,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 3.0920867919921875,
   "xla_arg_mb": 9.255863189697266,
   "xla_out_mb": 9.011833190917969,
   "compile_s": 0.006127954002295155,
   "inner": 67,
   "median_ms": 0.5762970298502609,
   "q1_ms": 0.5117354477662767,
   "q3_ms": 0.735728738822109,
   "rel_err_out": 0.05459915645864793,
   "rel_err_grad_params": 0.053237652261722695,
   "rel_err_grad_x": 0.05364085470253975
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "blocked",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 1,
   "skew": 0.0,
   "max_expert_frac": 0.15900000929832458,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 11.048622131347656,
   "xla_arg_mb": 3.248046875,
   "xla_out_mb": 0.244140625,
   "compile_s": 0.0016998690043692477,
   "inner": 183,
   "median_ms": 0.1747817595621642,
   "q1_ms": 0.1668420519194403,
   "q3_ms": 0.18212437979820378,
   "rel_err_out": 0.05435781463005756,
   "rel_err_grad_params": 0.05286217657152205,
   "rel_err_grad_x": 0.053236651698532644
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "blocked",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 1,
   "skew": 0.0,
   "max_expert_frac": 0.15900000929832458,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 11.082801818847656,
   "xla_arg_mb": 9.255863189697266,
   "xla_out_mb": 9.011833190917969,
   "compile_s": 0.002905332003138028,
   "inner": 111,
   "median_ms": 0.30081345044760255,
   "q1_ms": 0.29524745948472153,
   "q3_ms": 0.30992290542984324,
   "rel_err_out": 0.05435781463005756,
   "rel_err_grad_params": 0.05286217657152205,
   "rel_err_grad_x": 0.053236651698532644
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "ragged_dot",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 1,
   "skew": 0.0,
   "max_expert_frac": 0.15900000929832458,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 9.804458618164062,
   "xla_arg_mb": 3.248046875,
   "xla_out_mb": 0.244140625,
   "compile_s": 0.0013371020031627268,
   "inner": 200,
   "median_ms": 0.12145320502895629,
   "q1_ms": 0.11687485750371707,
   "q3_ms": 0.12338070248006261,
   "rel_err_out": 0.05435781463005756,
   "rel_err_grad_params": 0.05286217732762523,
   "rel_err_grad_x": 0.053236651698532644
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "ragged_dot",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 1,
   "skew": 0.0,
   "max_expert_frac": 0.15900000929832458,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 23.487823486328125,
   "xla_arg_mb": 9.255863189697266,
   "xla_out_mb": 9.011833190917969,
   "compile_s": 0.00252701599674765,
   "inner": 106,
   "median_ms": 0.32460008487588604,
   "q1_ms": 0.3218317359017608,
   "q3_ms": 0.32553431131264177,
   "rel_err_out": 0.05435781463005756,
   "rel_err_grad_params": 0.05286217732762523,
   "rel_err_grad_x": 0.053236651698532644
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "capacity",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 1,
   "skew": 0.0,
   "max_expert_frac": 0.15900000929832458,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 1.9052886962890625,
   "xla_arg_mb": 3.248046875,
   "xla_out_mb": 0.244140625,
   "compile_s": 0.03358034299890278,
   "inner": 21,
   "median_ms": 0.16084504750324413,
   "q1_ms": 0.1597050714259413,
   "q3_ms": 0.16271383321422728
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "capacity",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 1,
   "skew": 0.0,
   "max_expert_frac": 0.15900000929832458,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 1.8403549194335938,
   "xla_arg_mb": 9.255863189697266,
   "xla_out_mb": 9.011833190917969,
   "compile_s": 0.002205666998634115,
   "inner": 182,
   "median_ms": 0.17847936816486673,
   "q1_ms": 0.16653078021660733,
   "q3_ms": 0.18934849724026426
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dense",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 1,
   "skew": 0.0,
   "max_expert_frac": 0.15900000929832458,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 11.718765258789062,
   "xla_arg_mb": 3.248046875,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.0017978139949264005,
   "inner": 15,
   "median_ms": 0.16650326676123467,
   "q1_ms": 0.1137460666844466,
   "q3_ms": 0.1814432667742949,
   "rel_err_out": 0.05451225366514498,
   "rel_err_grad_params": 0.0531317679628266,
   "rel_err_grad_x": 0.05356507358813334
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dense",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 1,
   "skew": 0.0,
   "max_expert_frac": 0.15900000929832458,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 17.673851013183594,
   "xla_arg_mb": 9.255863189697266,
   "xla_out_mb": 9.011833190917969,
   "compile_s": 0.00418586700106971,
   "inner": 81,
   "median_ms": 0.2893276666585689,
   "q1_ms": 0.28696374692465465,
   "q3_ms": 0.29711401239856733,
   "rel_err_out": 0.05451225366514498,
   "rel_err_grad_params": 0.0531317679628266,
   "rel_err_grad_x": 0.05356507358813334
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 1,
   "skew": 0.9,
   "max_expert_frac": 0.9390000700950623,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 2.4236221313476562,
   "xla_arg_mb": 3.248046875,
   "xla_out_mb": 0.244140625,
   "compile_s": 0.0018575479989522137,
   "inner": 133,
   "median_ms": 0.25901254138487384,
   "q1_ms": 0.228568902252333,
   "q3_ms": 0.2933240150041525,
   "rel_err_out": 0.00600335473842787,
   "rel_err_grad_params": 0.005180517114491812,
   "rel_err_grad_x": 0.006858158553468144
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 1,
   "skew": 0.9,
   "max_expert_frac": 0.9390000700950623,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 1.3801651000976562,
   "xla_arg_mb": 9.255863189697266,
   "xla_out_mb": 9.011833190917969,
   "compile_s": 0.035067939999862574,
   "inner": 88,
   "median_ms": 0.2834078408837244,
   "q1_ms": 0.28275489768764767,
   "q3_ms": 0.2990631761266823,
   "rel_err_out": 0.00600335473842787,
   "rel_err_grad_params": 0.005180517114491812,
   "rel_err_grad_x": 0.006858158553468144
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "gather",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 1,
   "skew": 0.9,
   "max_expert_frac": 0.9390000700950623,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 2.9011611938476562,
   "xla_arg_mb": 3.248046875,
   "xla_out_mb": 0.244140625,
   "compile_s": 0.006406763000995852,
   "inner": 14,
   "median_ms": 0.4932919998412087,
   "q1_ms": 0.3415317857746102,
   "q3_ms": 0.5083974286078177,
   "rel_err_out": 0.00600335473842787,
   "rel_err_grad_params": 0.005180517114491812,
   "rel_err_grad_x": 0.006640401213473802
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "gather",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 1,
   "skew": 0.9,
   "max_expert_frac": 0.9390000700950623,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 3.0710906982421875,
   "xla_arg_mb": 9.255863189697266,
   "xla_out_mb": 9.011833190917969,
   "compile_s": 0.0032476100022904575,
   "inner": 113,
   "median_ms": 0.2969619203760676,
   "q1_ms": 0.2928776902829676,
   "q3_ms": 0.30620953095695663,
   "rel_err_out": 0.00600335473842787,
   "rel_err_grad_params": 0.005180517114491812,
   "rel_err_grad_x": 0.006640401213473802
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 1,
   "skew": 0.9,
   "max_expert_frac": 0.9390000700950623,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 2.5195693969726562,
   "xla_arg_mb": 3.248046875,
   "xla_out_mb": 0.244140625,
   "compile_s": 0.0019356060001882724,
   "inner": 182,
   "median_ms": 0.1633518736375813,
   "q1_ms": 0.16095494506621583,
   "q3_ms": 0.17054214012015298,
   "rel_err_out": 0.006006652903119468,
   "rel_err_grad_params": 0.005190957323630938,
   "rel_err_grad_x": 0.006642145985342667
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 1,
   "skew": 0.9,
   "max_expert_frac": 0.9390000700950623,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 3.0879364013671875,
   "xla_arg_mb": 9.255863189697266,
   "xla_out_mb": 9.011833190917969,
   "compile_s": 0.004424107995873783,
   "inner": 86,
   "median_ms": 0.30524083719450745,
   "q1_ms": 0.2972847384182049,
   "q3_ms": 0.30644923259218515,
   "rel_err_out": 0.006006652903119468,
   "rel_err_grad_params": 0.005190957323630938,
   "rel_err_grad_x": 0.006642145985342667
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused_gsum",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 1,
   "skew": 0.9,
   "max_expert_frac": 0.9390000700950623,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 2.5339736938476562,
   "xla_arg_mb": 3.248046875,
   "xla_out_mb": 0.244140625,
   "compile_s": 0.0020358859983389266,
   "inner": 193,
   "median_ms": 0.2097616424854034,
   "q1_ms": 0.1686592201806045,
   "q3_ms": 0.40292926422963304,
   "rel_err_out": 0.00600335473842787,
   "rel_err_grad_params": 0.005180517114491812,
   "rel_err_grad_x": 0.006640401213473802
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused_gsum",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 1,
   "skew": 0.9,
   "max_expert_frac": 0.9390000700950623,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 3.0920867919921875,
   "xla_arg_mb": 9.255863189697266,
   "xla_out_mb": 9.011833190917969,
   "compile_s": 0.003664318996015936,
   "inner": 101,
   "median_ms": 0.4725795048882604,
   "q1_ms": 0.36573150988476866,
   "q3_ms": 0.5459356980281882,
   "rel_err_out": 0.00600335473842787,
   "rel_err_grad_params": 0.005180517114491812,
   "rel_err_grad_x": 0.006640401213473802
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "blocked",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 1,
   "skew": 0.9,
   "max_expert_frac": 0.9390000700950623,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 11.048622131347656,
   "xla_arg_mb": 3.248046875,
   "xla_out_mb": 0.244140625,
   "compile_s": 0.002290276999701746,
   "inner": 184,
   "median_ms": 0.16909484783473486,
   "q1_ms": 0.16387120107854414,
   "q3_ms": 0.18268095108515298,
   "rel_err_out": 0.0029485431029799315,
   "rel_err_grad_params": 0.002117893929549733,
   "rel_err_grad_x": 0.0031752336273649573
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "blocked",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 1,
   "skew": 0.9,
   "max_expert_frac": 0.9390000700950623,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 11.082801818847656,
   "xla_arg_mb": 9.255863189697266,
   "xla_out_mb": 9.011833190917969,
   "compile_s": 0.0029652649973286316,
   "inner": 115,
   "median_ms": 0.3122166521675156,
   "q1_ms": 0.2863775912808943,
   "q3_ms": 0.3198363000157532,
   "rel_err_out": 0.0029485431029799315,
   "rel_err_grad_params": 0.002117893929549733,
   "rel_err_grad_x": 0.0031752336273649573
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "ragged_dot",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 1,
   "skew": 0.9,
   "max_expert_frac": 0.9390000700950623,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 9.804458618164062,
   "xla_arg_mb": 3.248046875,
   "xla_out_mb": 0.244140625,
   "compile_s": 0.0014553860019077547,
   "inner": 200,
   "median_ms": 0.12692786996922223,
   "q1_ms": 0.12542084250526386,
   "q3_ms": 0.12964161000127206,
   "rel_err_out": 0.0029485431029799315,
   "rel_err_grad_params": 0.002118012025276932,
   "rel_err_grad_x": 0.0031752336273649573
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "ragged_dot",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 1,
   "skew": 0.9,
   "max_expert_frac": 0.9390000700950623,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 23.487823486328125,
   "xla_arg_mb": 9.255863189697266,
   "xla_out_mb": 9.011833190917969,
   "compile_s": 0.0025159650031127967,
   "inner": 98,
   "median_ms": 0.31785376535369825,
   "q1_ms": 0.3142622244922677,
   "q3_ms": 0.32742531632776467,
   "rel_err_out": 0.0029485431029799315,
   "rel_err_grad_params": 0.002118012025276932,
   "rel_err_grad_x": 0.0031752336273649573
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "capacity",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 1,
   "skew": 0.9,
   "max_expert_frac": 0.9390000700950623,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 1.9052886962890625,
   "xla_arg_mb": 3.248046875,
   "xla_out_mb": 0.244140625,
   "compile_s": 0.009271287002775352,
   "inner": 155,
   "median_ms": 0.07732152261005173,
   "q1_ms": 0.07294940643928824,
   "q3_ms": 0.08514065484632738
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "capacity",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 1,
   "skew": 0.9,
   "max_expert_frac": 0.9390000700950623,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 1.8403549194335938,
   "xla_arg_mb": 9.255863189697266,
   "xla_out_mb": 9.011833190917969,
   "compile_s": 0.002431283995974809,
   "inner": 172,
   "median_ms": 0.16923484885367987,
   "q1_ms": 0.16725322096368303,
   "q3_ms": 0.17038772674886574
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dense",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 1,
   "skew": 0.9,
   "max_expert_frac": 0.9390000700950623,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 11.718765258789062,
   "xla_arg_mb": 3.248046875,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.0012293979962123558,
   "inner": 200,
   "median_ms": 0.08216484999138629,
   "q1_ms": 0.07679132500925334,
   "q3_ms": 0.09668392000094173,
   "rel_err_out": 0.005180683864492286,
   "rel_err_grad_params": 0.004625526567562788,
   "rel_err_grad_x": 0.006141626831136807
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dense",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 1,
   "skew": 0.9,
   "max_expert_frac": 0.9390000700950623,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 17.673851013183594,
   "xla_arg_mb": 9.255863189697266,
   "xla_out_mb": 9.011833190917969,
   "compile_s": 0.012579882997670211,
   "inner": 17,
   "median_ms": 0.6960758824682082,
   "q1_ms": 0.642405088429394,
   "q3_ms": 0.7598507059703801,
   "rel_err_out": 0.005180683864492286,
   "rel_err_grad_params": 0.004625526567562788,
   "rel_err_grad_x": 0.006141626831136807
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.12849999964237213,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 6.141883850097656,
   "xla_arg_mb": 3.248046875,
   "xla_out_mb": 0.244140625,
   "compile_s": 0.001993625999602955,
   "inner": 169,
   "median_ms": 0.18612063316264624,
   "q1_ms": 0.1832520650669081,
   "q3_ms": 0.1875174763273536,
   "rel_err_out": 0.026973097927077614,
   "rel_err_grad_params": 0.03351844096768902,
   "rel_err_grad_x": 0.031955345857099984
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.12849999964237213,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 9.380165100097656,
   "xla_arg_mb": 9.255863189697266,
   "xla_out_mb": 9.011833190917969,
   "compile_s": 0.0034475590000511147,
   "inner": 94,
   "median_ms": 0.3233604468132785,
   "q1_ms": 0.3156272393445103,
   "q3_ms": 0.3242964787127579,
   "rel_err_out": 0.026973097927077614,
   "rel_err_grad_params": 0.03351844096768902,
   "rel_err_grad_x": 0.031955345857099984
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "gather",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.12849999964237213,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 6.740028381347656,
   "xla_arg_mb": 3.248046875,
   "xla_out_mb": 0.244140625,
   "compile_s": 0.0025687439992907457,
   "inner": 137,
   "median_ms": 0.18584362042551167,
   "q1_ms": 0.18454616422846531,
   "q3_ms": 0.18913174819423864,
   "rel_err_out": 0.026973097927077614,
   "rel_err_grad_params": 0.033537575188625775,
   "rel_err_grad_x": 0.03194101985758647
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "gather",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.12849999964237213,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 12.318893432617188,
   "xla_arg_mb": 9.255863189697266,
   "xla_out_mb": 9.011833190917969,
   "compile_s": 0.0033349560035276227,
   "inner": 92,
   "median_ms": 0.3133381738418551,
   "q1_ms": 0.30973749455923727,
   "q3_ms": 0.3265273804256298,
   "rel_err_out": 0.026973097927077614,
   "rel_err_grad_params": 0.033537575188625775,
   "rel_err_grad_x": 0.03194101985758647
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.12849999964237213,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 6.710487365722656,
   "xla_arg_mb": 3.248046875,
   "xla_out_mb": 0.244140625,
   "compile_s": 0.002151284003048204,
   "inner": 142,
   "median_ms": 0.22196500001780192,
   "q1_ms": 0.21384744012349574,
   "q3_ms": 0.23181473940420683,
   "rel_err_out": 0.02694082156944546,
   "rel_err_grad_params": 0.03348378185247052,
   "rel_err_grad_x": 0.031949466720892654
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.12849999964237213,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 11.226364135742188,
   "xla_arg_mb": 9.255863189697266,
   "xla_out_mb": 9.011833190917969,
   "compile_s": 0.002923956999438815,
   "inner": 106,
   "median_ms": 0.33235128305606404,
   "q1_ms": 0.32885796226421593,
   "q3_ms": 0.3363267782985514,
   "rel_err_out": 0.02694082156944546,
   "rel_err_grad_params": 0.03348378185247052,
   "rel_err_grad_x": 0.031949466720892654
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused_gsum",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.12849999964237213,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 6.753456115722656,
   "xla_arg_mb": 3.248046875,
   "xla_out_mb": 0.244140625,
   "compile_s": 0.0018991260003531352,
   "inner": 168,
   "median_ms": 0.197119666674033,
   "q1_ms": 0.18429651189373025,
   "q3_ms": 0.21330352380270295,
   "rel_err_out": 0.026973097927077614,
   "rel_err_grad_params": 0.033537575188625775,
   "rel_err_grad_x": 0.03194101985758647
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused_gsum",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.12849999964237213,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 11.241989135742188,
   "xla_arg_mb": 9.255863189697266,
   "xla_out_mb": 9.011833190917969,
   "compile_s": 0.0035765130014624447,
   "inner": 90,
   "median_ms": 0.33294718887191266,
   "q1_ms": 0.3286621833669616,
   "q3_ms": 0.3544555333468856,
   "rel_err_out": 0.026973097927077614,
   "rel_err_grad_params": 0.033537575188625775,
   "rel_err_grad_x": 0.03194101985758647
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "blocked",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.12849999964237213,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 27.579383850097656,
   "xla_arg_mb": 3.248046875,
   "xla_out_mb": 0.244140625,
   "compile_s": 0.0018132539989892393,
   "inner": 142,
   "median_ms": 0.22238120421296653,
   "q1_ms": 0.22008382395245168,
   "q3_ms": 0.22301522183394334,
   "rel_err_out": 0.026502545735319634,
   "rel_err_grad_params": 0.03324057827419623,
   "rel_err_grad_x": 0.03164159613364224
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "blocked",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.12849999964237213,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 30.066444396972656,
   "xla_arg_mb": 9.255863189697266,
   "xla_out_mb": 9.011833190917969,
   "compile_s": 0.0034602320010890253,
   "inner": 64,
   "median_ms": 0.4490250313438082,
   "q1_ms": 0.44316001566357954,
   "q3_ms": 0.45216308586759624,
   "rel_err_out": 0.026502545735319634,
   "rel_err_grad_params": 0.03324057827419623,
   "rel_err_grad_x": 0.03164159613364224
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "ragged_dot",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.12849999964237213,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 35.20314025878906,
   "xla_arg_mb": 3.248046875,
   "xla_out_mb": 0.244140625,
   "compile_s": 0.0032147680030902848,
   "inner": 60,
   "median_ms": 0.4008280667524862,
   "q1_ms": 0.40060607498162426,
   "q3_ms": 0.4027563000515026,
   "rel_err_out": 0.026502545735319634,
   "rel_err_grad_params": 0.03324058089707135,
   "rel_err_grad_x": 0.03164159616562729
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "ragged_dot",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.12849999964237213,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 89.94143676757812,
   "xla_arg_mb": 9.255863189697266,
   "xla_out_mb": 9.011833190917969,
   "compile_s": 0.015371048997621983,
   "inner": 30,
   "median_ms": 1.4392065000720322,
   "q1_ms": 1.3504848167940509,
   "q3_ms": 1.6003850834143425,
   "rel_err_out": 0.026502545735319634,
   "rel_err_grad_params": 0.03324058089707135,
   "rel_err_grad_x": 0.03164159616562729
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "capacity",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.12849999964237213,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 7.3706207275390625,
   "xla_arg_mb": 3.248046875,
   "xla_out_mb": 0.244140625,
   "compile_s": 0.0389511940011289,
   "inner": 21,
   "median_ms": 0.3465596188713486,
   "q1_ms": 0.344944666548621,
   "q3_ms": 0.3480388331809081
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "capacity",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.12849999964237213,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 16.131126403808594,
   "xla_arg_mb": 9.255863189697266,
   "xla_out_mb": 9.011833190917969,
   "compile_s": 0.0026457400017534383,
   "inner": 97,
   "median_ms": 0.2816372783404545,
   "q1_ms": 0.26575637112381223,
   "q3_ms": 0.284712474202718
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dense",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.12849999964237213,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 11.718765258789062,
   "xla_arg_mb": 3.248046875,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.019098388002021238,
   "inner": 21,
   "median_ms": 0.17817790475895717,
   "q1_ms": 0.17809966655996895,
   "q3_ms": 0.1782728570452448,
   "rel_err_out": 0.026805889095744845,
   "rel_err_grad_params": 0.03325374097576774,
   "rel_err_grad_x": 0.031827543861478576
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dense",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.12849999964237213,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 19.579856872558594,
   "xla_arg_mb": 9.255863189697266,
   "xla_out_mb": 9.011833190917969,
   "compile_s": 0.0023653300013393164,
   "inner": 132,
   "median_ms": 0.21896420455022686,
   "q1_ms": 0.2185249280349429,
   "q3_ms": 0.22309362121273862,
   "rel_err_out": 0.026805889095744845,
   "rel_err_grad_params": 0.03325374097576774,
   "rel_err_grad_x": 0.031827543861478576
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 4,
   "skew": 0.9,
   "max_expert_frac": 0.23800000548362732,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 6.141883850097656,
   "xla_arg_mb": 3.248046875,
   "xla_out_mb": 0.244140625,
   "compile_s": 0.001873658002295997,
   "inner": 173,
   "median_ms": 0.17592877454415093,
   "q1_ms": 0.1744995028843075,
   "q3_ms": 0.17645065896441847,
   "rel_err_out": 0.0060341307957337935,
   "rel_err_grad_params": 0.007934434783747368,
   "rel_err_grad_x": 0.007275428712151981
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 4,
   "skew": 0.9,
   "max_expert_frac": 0.23800000548362732,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 9.380165100097656,
   "xla_arg_mb": 9.255863189697266,
   "xla_out_mb": 9.011833190917969,
   "compile_s": 0.0038204650045372546,
   "inner": 93,
   "median_ms": 0.34642841935130236,
   "q1_ms": 0.33990783870854835,
   "q3_ms": 0.3764938440825415,
   "rel_err_out": 0.0060341307957337935,
   "rel_err_grad_params": 0.007934434783747368,
   "rel_err_grad_x": 0.007275428712151981
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "gather",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 4,
   "skew": 0.9,
   "max_expert_frac": 0.23800000548362732,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 6.740028381347656,
   "xla_arg_mb": 3.248046875,
   "xla_out_mb": 0.244140625,
   "compile_s": 0.0027331160017638467,
   "inner": 139,
   "median_ms": 0.23697812950637245,
   "q1_ms": 0.22431955035866352,
   "q3_ms": 0.2912086402574846,
   "rel_err_out": 0.0060341307957337935,
   "rel_err_grad_params": 0.007940135103541253,
   "rel_err_grad_x": 0.007071549260159952
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "gather",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 4,
   "skew": 0.9,
   "max_expert_frac": 0.23800000548362732,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 12.318893432617188,
   "xla_arg_mb": 9.255863189697266,
   "xla_out_mb": 9.011833190917969,
   "compile_s": 0.1195821689980221,
   "inner": 19,
   "median_ms": 0.6827287369033959,
   "q1_ms": 0.6798556316521411,
   "q3_ms": 0.6872375000200823,
   "rel_err_out": 0.0060341307957337935,
   "rel_err_grad_params": 0.007940135103541253,
   "rel_err_grad_x": 0.007071549260159952
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 4,
   "skew": 0.9,
   "max_expert_frac": 0.23800000548362732,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 6.710487365722656,
   "xla_arg_mb": 3.248046875,
   "xla_out_mb": 0.244140625,
   "compile_s": 0.005877430005057249,
   "inner": 20,
   "median_ms": 0.5701391997718019,
   "q1_ms": 0.5620420999548514,
   "q3_ms": 0.5787472751762834,
   "rel_err_out": 0.006024652734461824,
   "rel_err_grad_params": 0.007794950741587546,
   "rel_err_grad_x": 0.007060407646734258
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 4,
   "skew": 0.9,
   "max_expert_frac": 0.23800000548362732,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 11.226364135742188,
   "xla_arg_mb": 9.255863189697266,
   "xla_out_mb": 9.011833190917969,
   "compile_s": 0.19700401699810755,
   "inner": 18,
   "median_ms": 0.8662159444712517,
   "q1_ms": 0.8161711666616611,
   "q3_ms": 0.9227583054275277,
   "rel_err_out": 0.006024652734461824,
   "rel_err_grad_params": 0.007794950741587546,
   "rel_err_grad_x": 0.007060407646734258
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused_gsum",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 4,
   "skew": 0.9,
   "max_expert_frac": 0.23800000548362732,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 6.753456115722656,
   "xla_arg_mb": 3.248046875,
   "xla_out_mb": 0.244140625,
   "compile_s": 0.1517343180021271,
   "inner": 9,
   "median_ms": 0.4859813335416321,
   "q1_ms": 0.45768961071088493,
   "q3_ms": 0.6042476666657925,
   "rel_err_out": 0.0060341307957337935,
   "rel_err_grad_params": 0.007940135103541253,
   "rel_err_grad_x": 0.007071549260159952
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused_gsum",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 4,
   "skew": 0.9,
   "max_expert_frac": 0.23800000548362732,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 11.241989135742188,
   "xla_arg_mb": 9.255863189697266,
   "xla_out_mb": 9.011833190917969,
   "compile_s": 0.13721815699682338,
   "inner": 54,
   "median_ms": 0.776786777822525,
   "q1_ms": 0.7503097963139966,
   "q3_ms": 0.8142588610806141,
   "rel_err_out": 0.0060341307957337935,
   "rel_err_grad_params": 0.007940135103541253,
   "rel_err_grad_x": 0.007071549260159952
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "blocked",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 4,
   "skew": 0.9,
   "max_expert_frac": 0.23800000548362732,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 27.579383850097656,
   "xla_arg_mb": 3.248046875,
   "xla_out_mb": 0.244140625,
   "compile_s": 0.09781567900063237,
   "inner": 19,
   "median_ms": 0.444390526114914,
   "q1_ms": 0.3582566052117679,
   "q3_ms": 0.615408078853139,
   "rel_err_out": 0.003012837494843347,
   "rel_err_grad_params": 0.007487031689549278,
   "rel_err_grad_x": 0.004435323056190156
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "blocked",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 4,
   "skew": 0.9,
   "max_expert_frac": 0.23800000548362732,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 30.066444396972656,
   "xla_arg_mb": 9.255863189697266,
   "xla_out_mb": 9.011833190917969,
   "compile_s": 0.0033355470004607923,
   "inner": 77,
   "median_ms": 0.4564214935496014,
   "q1_ms": 0.45515603896883594,
   "q3_ms": 0.46141041561246054,
   "rel_err_out": 0.003012837494843347,
   "rel_err_grad_params": 0.007487031689549278,
   "rel_err_grad_x": 0.004435323056190156
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "ragged_dot",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 4,
   "skew": 0.9,
   "max_expert_frac": 0.23800000548362732,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 35.20314025878906,
   "xla_arg_mb": 3.248046875,
   "xla_out_mb": 0.244140625,
   "compile_s": 0.0018263479942106642,
   "inner": 92,
   "median_ms": 0.38108810867529624,
   "q1_ms": 0.3770256304525664,
   "q3_ms": 0.40525315216655144,
   "rel_err_out": 0.003012837494843347,
   "rel_err_grad_params": 0.007487239840121623,
   "rel_err_grad_x": 0.00443532935361362
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "ragged_dot",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 4,
   "skew": 0.9,
   "max_expert_frac": 0.23800000548362732,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 89.94143676757812,
   "xla_arg_mb": 9.255863189697266,
   "xla_out_mb": 9.011833190917969,
   "compile_s": 0.003653308995126281,
   "inner": 35,
   "median_ms": 1.0880436286762623,
   "q1_ms": 1.0824553285991507,
   "q3_ms": 1.0891299572124677,
   "rel_err_out": 0.003012837494843347,
   "rel_err_grad_params": 0.007487239840121623,
   "rel_err_grad_x": 0.00443532935361362
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "capacity",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 4,
   "skew": 0.9,
   "max_expert_frac": 0.23800000548362732,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 7.3706207275390625,
   "xla_arg_mb": 3.248046875,
   "xla_out_mb": 0.244140625,
   "compile_s": 0.0019929229965782724,
   "inner": 184,
   "median_ms": 0.13936985867432033,
   "q1_ms": 0.1361339999794322,
   "q3_ms": 0.14214219837689912
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "capacity",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 4,
   "skew": 0.9,
   "max_expert_frac": 0.23800000548362732,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 16.131126403808594,
   "xla_arg_mb": 9.255863189697266,
   "xla_out_mb": 9.011833190917969,
   "compile_s": 0.006396051998308394,
   "inner": 25,
   "median_ms": 0.5802568802027963,
   "q1_ms": 0.42173307985649444,
   "q3_ms": 0.677051179955015
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dense",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 4,
   "skew": 0.9,
   "max_expert_frac": 0.23800000548362732,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 11.718765258789062,
   "xla_arg_mb": 3.248046875,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.0016714550001779571,
   "inner": 199,
   "median_ms": 0.08302537186759239,
   "q1_ms": 0.08044231155929366,
   "q3_ms": 0.08339189196192079,
   "rel_err_out": 0.005200601389917184,
   "rel_err_grad_params": 0.008011331866720603,
   "rel_err_grad_x": 0.006660321374786306
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dense",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 8,
   "top_k": 4,
   "skew": 0.9,
   "max_expert_frac": 0.23800000548362732,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 19.579856872558594,
   "xla_arg_mb": 9.255863189697266,
   "xla_out_mb": 9.011833190917969,
   "compile_s": 0.012015303997031879,
   "inner": 13,
   "median_ms": 0.758639000056204,
   "q1_ms": 0.6905866539455019,
   "q3_ms": 0.8846095000850396,
   "rel_err_out": 0.005200601389917184,
   "rel_err_grad_params": 0.008011331866720603,
   "rel_err_grad_x": 0.006660321374786306
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 1,
   "skew": 0.0,
   "max_expert_frac": 0.02800000086426735,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 10.434608459472656,
   "xla_arg_mb": 24.275390625,
   "xla_out_mb": 0.244140625,
   "compile_s": 0.07563307999953395,
   "inner": 30,
   "median_ms": 0.73011699981483,
   "q1_ms": 0.6122282333914579,
   "q3_ms": 1.054663433266493,
   "rel_err_out": 0.07857382434818175,
   "rel_err_grad_params": 0.07247188532470779,
   "rel_err_grad_x": 0.06583717647501273
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 1,
   "skew": 0.0,
   "max_expert_frac": 0.02800000086426735,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 4.005409240722656,
   "xla_arg_mb": 72.33789443969727,
   "xla_out_mb": 72.09386444091797,
   "compile_s": 0.10858193700551055,
   "inner": 17,
   "median_ms": 0.5993214705513845,
   "q1_ms": 0.5807291764700921,
   "q3_ms": 0.8904192057825104,
   "rel_err_out": 0.07857382434818175,
   "rel_err_grad_params": 0.07247188532470779,
   "rel_err_grad_x": 0.06583717647501273
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "gather",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 1,
   "skew": 0.0,
   "max_expert_frac": 0.02800000086426735,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 14.225135803222656,
   "xla_arg_mb": 24.275390625,
   "xla_out_mb": 0.244140625,
   "compile_s": 0.003559951001079753,
   "inner": 34,
   "median_ms": 0.4112829705905717,
   "q1_ms": 0.38524323538106675,
   "q3_ms": 0.46529202935583747,
   "rel_err_out": 0.07857382434818175,
   "rel_err_grad_params": 0.07247188532470779,
   "rel_err_grad_x": 0.06581046068503288
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "gather",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 1,
   "skew": 0.0,
   "max_expert_frac": 0.02800000086426735,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 5.7405242919921875,
   "xla_arg_mb": 72.33789443969727,
   "xla_out_mb": 72.09386444091797,
   "compile_s": 0.016003015996830072,
   "inner": 23,
   "median_ms": 0.8443899132544175,
   "q1_ms": 0.7779297175938669,
   "q3_ms": 1.4160344566651317,
   "rel_err_out": 0.07857382434818175,
   "rel_err_grad_params": 0.07247188532470779,
   "rel_err_grad_x": 0.06581046068503288
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 1,
   "skew": 0.0,
   "max_expert_frac": 0.02800000086426735,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 12.997596740722656,
   "xla_arg_mb": 24.275390625,
   "xla_out_mb": 0.244140625,
   "compile_s": 0.028771416000381578,
   "inner": 12,
   "median_ms": 1.2240530832059449,
   "q1_ms": 0.5981602498650318,
   "q3_ms": 1.6233763332517508,
   "rel_err_out": 0.07858246364680722,
   "rel_err_grad_params": 0.07247059702324513,
   "rel_err_grad_x": 0.06577247932666894
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 1,
   "skew": 0.0,
   "max_expert_frac": 0.02800000086426735,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 4.7695770263671875,
   "xla_arg_mb": 72.33789443969727,
   "xla_out_mb": 72.09386444091797,
   "compile_s": 0.1318514939994202,
   "inner": 15,
   "median_ms": 1.0956028000994897,
   "q1_ms": 0.9645793334736178,
   "q3_ms": 1.126707067063156,
   "rel_err_out": 0.07858246364680722,
   "rel_err_grad_params": 0.07247059702324513,
   "rel_err_grad_x": 0.06577247932666894
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused_gsum",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 1,
   "skew": 0.0,
   "max_expert_frac": 0.02800000086426735,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 12.998573303222656,
   "xla_arg_mb": 24.275390625,
   "xla_out_mb": 0.244140625,
   "compile_s": 0.0022980719950282946,
   "inner": 95,
   "median_ms": 0.3530337052468799,
   "q1_ms": 0.33925773685424293,
   "q3_ms": 0.3992809315892182,
   "rel_err_out": 0.07857382434818175,
   "rel_err_grad_params": 0.07247188532470779,
   "rel_err_grad_x": 0.06581046068503288
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused_gsum",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 1,
   "skew": 0.0,
   "max_expert_frac": 0.02800000086426735,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 4.7737274169921875,
   "xla_arg_mb": 72.33789443969727,
   "xla_out_mb": 72.09386444091797,
   "compile_s": 0.004278633001376875,
   "inner": 43,
   "median_ms": 1.1134676976727735,
   "q1_ms": 0.9686688372735368,
   "q3_ms": 1.82435418613496,
   "rel_err_out": 0.07857382434818175,
   "rel_err_grad_params": 0.07247188532470779,
   "rel_err_grad_x": 0.06581046068503288
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "blocked",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 1,
   "skew": 0.0,
   "max_expert_frac": 0.02800000086426735,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 49.762489318847656,
   "xla_arg_mb": 24.275390625,
   "xla_out_mb": 0.244140625,
   "compile_s": 0.002104655999573879,
   "inner": 103,
   "median_ms": 0.3505449125953516,
   "q1_ms": 0.3440448883197551,
   "q3_ms": 0.36368406791967856,
   "rel_err_out": 0.07838134961625355,
   "rel_err_grad_params": 0.07215071372527429,
   "rel_err_grad_x": 0.06551706507759424
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "blocked",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 1,
   "skew": 0.0,
   "max_expert_frac": 0.02800000086426735,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 45.263221740722656,
   "xla_arg_mb": 72.33789443969727,
   "xla_out_mb": 72.09386444091797,
   "compile_s": 0.003985647999797948,
   "inner": 35,
   "median_ms": 1.0248273428130363,
   "q1_ms": 0.97415452847989,
   "q3_ms": 1.0493406142424129,
   "rel_err_out": 0.07838134961625355,
   "rel_err_grad_params": 0.07215071372527429,
   "rel_err_grad_x": 0.06551706507759424
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "ragged_dot",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 1,
   "skew": 0.0,
   "max_expert_frac": 0.02800000086426735,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 66.65895080566406,
   "xla_arg_mb": 24.275390625,
   "xla_out_mb": 0.244140625,
   "compile_s": 0.006661514998995699,
   "inner": 25,
   "median_ms": 0.6797729199752212,
   "q1_ms": 0.6398673600051552,
   "q3_ms": 0.7454432800295763,
   "rel_err_out": 0.07838134961625355,
   "rel_err_grad_params": 0.07215071378428076,
   "rel_err_grad_x": 0.06551706507759424
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "ragged_dot",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 1,
   "skew": 0.0,
   "max_expert_frac": 0.02800000086426735,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 156.25686645507812,
   "xla_arg_mb": 72.33789443969727,
   "xla_out_mb": 72.09386444091797,
   "compile_s": 0.005315176000294741,
   "inner": 21,
   "median_ms": 2.134939285627167,
   "q1_ms": 2.086310476269795,
   "q3_ms": 3.5800468569851502,
   "rel_err_out": 0.07838134961625355,
   "rel_err_grad_params": 0.07215071378428076,
   "rel_err_grad_x": 0.06551706507759424
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "capacity",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 1,
   "skew": 0.0,
   "max_expert_frac": 0.02800000086426735,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 2.3676910400390625,
   "xla_arg_mb": 24.275390625,
   "xla_out_mb": 0.244140625,
   "compile_s": 0.0014002220050315373,
   "inner": 200,
   "median_ms": 0.0866619649968925,
   "q1_ms": 0.08540881748558604,
   "q3_ms": 0.0894942800005083
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "capacity",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 1,
   "skew": 0.0,
   "max_expert_frac": 0.02800000086426735,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 5.048118591308594,
   "xla_arg_mb": 72.33789443969727,
   "xla_out_mb": 72.09386444091797,
   "compile_s": 0.003164132001984399,
   "inner": 87,
   "median_ms": 0.3699624597164535,
   "q1_ms": 0.36813036777679914,
   "q3_ms": 0.37682286213451843
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dense",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 1,
   "skew": 0.0,
   "max_expert_frac": 0.02800000086426735,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 93.99415588378906,
   "xla_arg_mb": 24.275390625,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.001570184002048336,
   "inner": 90,
   "median_ms": 0.42374112219679066,
   "q1_ms": 0.42336518886764274,
   "q3_ms": 0.4246062111229144,
   "rel_err_out": 0.07851990329991068,
   "rel_err_grad_params": 0.07251122514007487,
   "rel_err_grad_x": 0.06578765287277401
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dense",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 1,
   "skew": 0.0,
   "max_expert_frac": 0.02800000086426735,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 125.0027084350586,
   "xla_arg_mb": 72.33789443969727,
   "xla_out_mb": 72.09386444091797,
   "compile_s": 0.003175755002303049,
   "inner": 36,
   "median_ms": 1.2121238056958343,
   "q1_ms": 1.2114520972722351,
   "q3_ms": 1.2171874166395154,
   "rel_err_out": 0.07851990329991068,
   "rel_err_grad_params": 0.07251122514007487,
   "rel_err_grad_x": 0.06578765287277401
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 1,
   "skew": 0.9,
   "max_expert_frac": 0.9320000410079956,
   "empty_experts": 16,
   "op": "forward",
   "xla_temp_mb": 10.434608459472656,
   "xla_arg_mb": 24.275390625,
   "xla_out_mb": 0.244140625,
   "compile_s": 0.01396958500117762,
   "inner": 57,
   "median_ms": 0.2579087544163166,
   "q1_ms": 0.24895179829623532,
   "q3_ms": 0.2928162719303815,
   "rel_err_out": 0.04420499081319327,
   "rel_err_grad_params": 0.026246257998996186,
   "rel_err_grad_x": 0.038298252075406075
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 1,
   "skew": 0.9,
   "max_expert_frac": 0.9320000410079956,
   "empty_experts": 16,
   "op": "train_step",
   "xla_temp_mb": 4.005409240722656,
   "xla_arg_mb": 72.33789443969727,
   "xla_out_mb": 72.09386444091797,
   "compile_s": 0.006302724999841303,
   "inner": 42,
   "median_ms": 0.6312625952150916,
   "q1_ms": 0.6307573571359084,
   "q3_ms": 0.6746162499829162,
   "rel_err_out": 0.04420499081319327,
   "rel_err_grad_params": 0.026246257998996186,
   "rel_err_grad_x": 0.038298252075406075
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "gather",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 1,
   "skew": 0.9,
   "max_expert_frac": 0.9320000410079956,
   "empty_experts": 16,
   "op": "forward",
   "xla_temp_mb": 14.225135803222656,
   "xla_arg_mb": 24.275390625,
   "xla_out_mb": 0.244140625,
   "compile_s": 0.0022234710049815476,
   "inner": 105,
   "median_ms": 0.32588591427719665,
   "q1_ms": 0.31949810474456863,
   "q3_ms": 0.33607944760227665,
   "rel_err_out": 0.04420499081319327,
   "rel_err_grad_params": 0.026246257998996186,
   "rel_err_grad_x": 0.03827221471232962
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "gather",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 1,
   "skew": 0.9,
   "max_expert_frac": 0.9320000410079956,
   "empty_experts": 16,
   "op": "train_step",
   "xla_temp_mb": 5.7405242919921875,
   "xla_arg_mb": 72.33789443969727,
   "xla_out_mb": 72.09386444091797,
   "compile_s": 0.11702494599740021,
   "inner": 17,
   "median_ms": 1.4715738819658701,
   "q1_ms": 1.4299217351328801,
   "q3_ms": 1.5104185883554986,
   "rel_err_out": 0.04420499081319327,
   "rel_err_grad_params": 0.026246257998996186,
   "rel_err_grad_x": 0.03827221471232962
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 1,
   "skew": 0.9,
   "max_expert_frac": 0.9320000410079956,
   "empty_experts": 16,
   "op": "forward",
   "xla_temp_mb": 12.997596740722656,
   "xla_arg_mb": 24.275390625,
   "xla_out_mb": 0.244140625,
   "compile_s": 0.00207246500212932,
   "inner": 113,
   "median_ms": 0.35563800003600055,
   "q1_ms": 0.33264787168014354,
   "q3_ms": 0.575276022090317,
   "rel_err_out": 0.04422687964398338,
   "rel_err_grad_params": 0.026253371064965374,
   "rel_err_grad_x": 0.03827644091395787
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 1,
   "skew": 0.9,
   "max_expert_frac": 0.9320000410079956,
   "empty_experts": 16,
   "op": "train_step",
   "xla_temp_mb": 4.7695770263671875,
   "xla_arg_mb": 72.33789443969727,
   "xla_out_mb": 72.09386444091797,
   "compile_s": 0.003966572003264446,
   "inner": 43,
   "median_ms": 0.932066720980761,
   "q1_ms": 0.9022734418469427,
   "q3_ms": 1.0485070000191459,
   "rel_err_out": 0.04422687964398338,
   "rel_err_grad_params": 0.026253371064965374,
   "rel_err_grad_x": 0.03827644091395787
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused_gsum",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 1,
   "skew": 0.9,
   "max_expert_frac": 0.9320000410079956,
   "empty_experts": 16,
   "op": "forward",
   "xla_temp_mb": 12.998573303222656,
   "xla_arg_mb": 24.275390625,
   "xla_out_mb": 0.244140625,
   "compile_s": 0.00348265500360867,
   "inner": 17,
   "median_ms": 0.4601885294130839,
   "q1_ms": 0.446919029320056,
   "q3_ms": 0.7847702941302593,
   "rel_err_out": 0.04420499081319327,
   "rel_err_grad_params": 0.026246257998996186,
   "rel_err_grad_x": 0.03827221471232962
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused_gsum",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 1,
   "skew": 0.9,
   "max_expert_frac": 0.9320000410079956,
   "empty_experts": 16,
   "op": "train_step",
   "xla_temp_mb": 4.7737274169921875,
   "xla_arg_mb": 72.33789443969727,
   "xla_out_mb": 72.09386444091797,
   "compile_s": 0.003614936002122704,
   "inner": 45,
   "median_ms": 0.8919090666394267,
   "q1_ms": 0.8897602111877253,
   "q3_ms": 0.8999352556000102,
   "rel_err_out": 0.04420499081319327,
   "rel_err_grad_params": 0.026246257998996186,
   "rel_err_grad_x": 0.03827221471232962
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "blocked",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 1,
   "skew": 0.9,
   "max_expert_frac": 0.9320000410079956,
   "empty_experts": 16,
   "op": "forward",
   "xla_temp_mb": 49.762489318847656,
   "xla_arg_mb": 24.275390625,
   "xla_out_mb": 0.244140625,
   "compile_s": 0.001924374002555851,
   "inner": 109,
   "median_ms": 0.3436378257000956,
   "q1_ms": 0.3345961651592323,
   "q3_ms": 0.35597789907665633,
   "rel_err_out": 0.043863832847361124,
   "rel_err_grad_params": 0.02573532681117683,
   "rel_err_grad_x": 0.03771309372743574
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "blocked",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 1,
   "skew": 0.9,
   "max_expert_frac": 0.9320000410079956,
   "empty_experts": 16,
   "op": "train_step",
   "xla_temp_mb": 45.263221740722656,
   "xla_arg_mb": 72.33789443969727,
   "xla_out_mb": 72.09386444091797,
   "compile_s": 0.003692523001518566,
   "inner": 42,
   "median_ms": 0.9622734999408878,
   "q1_ms": 0.9099260356985421,
   "q3_ms": 1.8834377142844634,
   "rel_err_out": 0.043863832847361124,
   "rel_err_grad_params": 0.02573532681117683,
   "rel_err_grad_x": 0.03771309372743574
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "ragged_dot",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 1,
   "skew": 0.9,
   "max_expert_frac": 0.9320000410079956,
   "empty_experts": 16,
   "op": "forward",
   "xla_temp_mb": 66.65895080566406,
   "xla_arg_mb": 24.275390625,
   "xla_out_mb": 0.244140625,
   "compile_s": 0.03918532799434615,
   "inner": 14,
   "median_ms": 1.5405480714564743,
   "q1_ms": 1.4908851428897054,
   "q3_ms": 1.5628711069959016,
   "rel_err_out": 0.043863832847361124,
   "rel_err_grad_params": 0.02573533619817596,
   "rel_err_grad_x": 0.03771309372743574
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "ragged_dot",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 1,
   "skew": 0.9,
   "max_expert_frac": 0.9320000410079956,
   "empty_experts": 16,
   "op": "train_step",
   "xla_temp_mb": 156.25686645507812,
   "xla_arg_mb": 72.33789443969727,
   "xla_out_mb": 72.09386444091797,
   "compile_s": 0.008611928999016527,
   "inner": 12,
   "median_ms": 2.0830230829839516,
   "q1_ms": 2.062840250194616,
   "q3_ms": 2.100441041875456,
   "rel_err_out": 0.043863832847361124,
   "rel_err_grad_params": 0.02573533619817596,
   "rel_err_grad_x": 0.03771309372743574
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "capacity",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 1,
   "skew": 0.9,
   "max_expert_frac": 0.9320000410079956,
   "empty_experts": 16,
   "op": "forward",
   "xla_temp_mb": 2.3676910400390625,
   "xla_arg_mb": 24.275390625,
   "xla_out_mb": 0.244140625,
   "compile_s": 0.004562580004858319,
   "inner": 21,
   "median_ms": 0.13590261923603802,
   "q1_ms": 0.1245735237760181,
   "q3_ms": 0.18491519048915334
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "capacity",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 1,
   "skew": 0.9,
   "max_expert_frac": 0.9320000410079956,
   "empty_experts": 16,
   "op": "train_step",
   "xla_temp_mb": 5.048118591308594,
   "xla_arg_mb": 72.33789443969727,
   "xla_out_mb": 72.09386444091797,
   "compile_s": 0.002606095004011877,
   "inner": 87,
   "median_ms": 0.4343049884733648,
   "q1_ms": 0.4307351896543345,
   "q3_ms": 0.6138953218375729
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dense",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 1,
   "skew": 0.9,
   "max_expert_frac": 0.9320000410079956,
   "empty_experts": 16,
   "op": "forward",
   "xla_temp_mb": 93.99415588378906,
   "xla_arg_mb": 24.275390625,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.001717252001981251,
   "inner": 98,
   "median_ms": 0.4294375305887245,
   "q1_ms": 0.42779355606732994,
   "q3_ms": 0.6472622346930677,
   "rel_err_out": 0.04413179723476757,
   "rel_err_grad_params": 0.026145094041341476,
   "rel_err_grad_x": 0.0381980653414495
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dense",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 1,
   "skew": 0.9,
   "max_expert_frac": 0.9320000410079956,
   "empty_experts": 16,
   "op": "train_step",
   "xla_temp_mb": 125.0027084350586,
   "xla_arg_mb": 72.33789443969727,
   "xla_out_mb": 72.09386444091797,
   "compile_s": 0.003244544001063332,
   "inner": 36,
   "median_ms": 1.2382154165809174,
   "q1_ms": 1.2303153194403016,
   "q3_ms": 1.2456480971094503,
   "rel_err_out": 0.04413179723476757,
   "rel_err_grad_params": 0.026145094041341476,
   "rel_err_grad_x": 0.0381980653414495
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.022750001400709152,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 13.923622131347656,
   "xla_arg_mb": 24.275390625,
   "xla_out_mb": 0.244140625,
   "compile_s": 0.002341934996366035,
   "inner": 99,
   "median_ms": 0.7084390808354283,
   "q1_ms": 0.6976334191888166,
   "q3_ms": 0.7452486919132155,
   "rel_err_out": 0.05223494452238976,
   "rel_err_grad_params": 0.06823323440901599,
   "rel_err_grad_x": 0.053014643510888144
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.022750001400709152,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 4.005653381347656,
   "xla_arg_mb": 72.33789443969727,
   "xla_out_mb": 72.09386444091797,
   "compile_s": 0.004069907998200506,
   "inner": 49,
   "median_ms": 0.7075994082177723,
   "q1_ms": 0.703380612136169,
   "q3_ms": 0.7130638367350318,
   "rel_err_out": 0.05223494452238976,
   "rel_err_grad_params": 0.06823323440901599,
   "rel_err_grad_x": 0.053014643510888144
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "gather",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.022750001400709152,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 17.954383850097656,
   "xla_arg_mb": 24.275390625,
   "xla_out_mb": 0.244140625,
   "compile_s": 0.003730734999408014,
   "inner": 64,
   "median_ms": 0.3879836406213144,
   "q1_ms": 0.3787724531321146,
   "q3_ms": 0.3956101171525006,
   "rel_err_out": 0.05223494452417147,
   "rel_err_grad_params": 0.06823522896579659,
   "rel_err_grad_x": 0.052988981831444945
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "gather",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.022750001400709152,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 5.9746551513671875,
   "xla_arg_mb": 72.33789443969727,
   "xla_out_mb": 72.09386444091797,
   "compile_s": 0.006324686997686513,
   "inner": 24,
   "median_ms": 0.7733177914512149,
   "q1_ms": 0.7688163541388349,
   "q3_ms": 1.0957663542588612,
   "rel_err_out": 0.05223494452417147,
   "rel_err_grad_params": 0.06823522896579659,
   "rel_err_grad_x": 0.052988981831444945
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.022750001400709152,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 15.969520568847656,
   "xla_arg_mb": 24.275390625,
   "xla_out_mb": 0.244140625,
   "compile_s": 0.06451342299988028,
   "inner": 19,
   "median_ms": 0.8632167894240027,
   "q1_ms": 0.8629128159985779,
   "q3_ms": 0.8663733158180932,
   "rel_err_out": 0.05222141089255765,
   "rel_err_grad_params": 0.06822191610426952,
   "rel_err_grad_x": 0.05302113502777241
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.022750001400709152,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 4.3418426513671875,
   "xla_arg_mb": 72.33789443969727,
   "xla_out_mb": 72.09386444091797,
   "compile_s": 0.003987129995948635,
   "inner": 37,
   "median_ms": 0.92335618919114,
   "q1_ms": 0.9203775675937446,
   "q3_ms": 0.9291567432621771,
   "rel_err_out": 0.05222141089255765,
   "rel_err_grad_params": 0.06822191610426952,
   "rel_err_grad_x": 0.05302113502777241
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused_gsum",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.022750001400709152,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 16.016395568847656,
   "xla_arg_mb": 24.275390625,
   "xla_out_mb": 0.244140625,
   "compile_s": 0.0023551100020995364,
   "inner": 102,
   "median_ms": 0.4145995588652452,
   "q1_ms": 0.38624001960367116,
   "q3_ms": 0.7036641323508115,
   "rel_err_out": 0.05223494452238976,
   "rel_err_grad_params": 0.06823523247917898,
   "rel_err_grad_x": 0.052988982309311435
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused_gsum",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.022750001400709152,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 4.3420867919921875,
   "xla_arg_mb": 72.33789443969727,
   "xla_out_mb": 72.09386444091797,
   "compile_s": 0.0991928959992947,
   "inner": 15,
   "median_ms": 1.9815739331534135,
   "q1_ms": 1.5773758331003287,
   "q3_ms": 2.0149628665725077,
   "rel_err_out": 0.05223494452238976,
   "rel_err_grad_params": 0.06823523247917898,
   "rel_err_grad_x": 0.052988982309311435
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "blocked",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.022750001400709152,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 65.39212799072266,
   "xla_arg_mb": 24.275390625,
   "xla_out_mb": 0.244140625,
   "compile_s": 0.0039259549957932904,
   "inner": 56,
   "median_ms": 0.4616145536471907,
   "q1_ms": 0.45920265178470543,
   "q3_ms": 0.46833968748874866,
   "rel_err_out": 0.05200831227565055,
   "rel_err_grad_params": 0.06810486323365883,
   "rel_err_grad_x": 0.05279453728454358
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "blocked",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.022750001400709152,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 59.918251037597656,
   "xla_arg_mb": 72.33789443969727,
   "xla_out_mb": 72.09386444091797,
   "compile_s": 0.004014643003756646,
   "inner": 36,
   "median_ms": 1.1804771945915287,
   "q1_ms": 1.1654165972585258,
   "q3_ms": 1.2856072638290548,
   "rel_err_out": 0.05200831227565055,
   "rel_err_grad_params": 0.06810486323365883,
   "rel_err_grad_x": 0.05279453728454358
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "ragged_dot",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.022750001400709152,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 261.85963439941406,
   "xla_arg_mb": 24.275390625,
   "xla_out_mb": 0.244140625,
   "compile_s": 0.0033105990005424246,
   "inner": 22,
   "median_ms": 2.151967591172169,
   "q1_ms": 2.1362188864176543,
   "q3_ms": 2.2988318409757498,
   "rel_err_out": 0.052008308622547,
   "rel_err_grad_params": 0.0681047202520355,
   "rel_err_grad_x": 0.0527945168370264
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "ragged_dot",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.022750001400709152,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 625.0044250488281,
   "xla_arg_mb": 72.33789443969727,
   "xla_out_mb": 72.09386444091797,
   "compile_s": 0.011556304998521227,
   "inner": 6,
   "median_ms": 7.828655499906745,
   "q1_ms": 7.703081000120923,
   "q3_ms": 7.9008113328503295,
   "rel_err_out": 0.052008308622547,
   "rel_err_grad_params": 0.0681047202520355,
   "rel_err_grad_x": 0.0527945168370264
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "capacity",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.022750001400709152,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 7.6665191650390625,
   "xla_arg_mb": 24.275390625,
   "xla_out_mb": 0.244140625,
   "compile_s": 0.0016433720011264086,
   "inner": 160,
   "median_ms": 0.16800966873233847,
   "q1_ms": 0.16190363435271138,
   "q3_ms": 0.17821236560848774
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "capacity",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.022750001400709152,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 4.246360778808594,
   "xla_arg_mb": 72.33789443969727,
   "xla_out_mb": 72.09386444091797,
   "compile_s": 0.0038121600009617396,
   "inner": 72,
   "median_ms": 0.47753948618062875,
   "q1_ms": 0.47429291670495227,
   "q3_ms": 0.5085598194354033
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dense",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.022750001400709152,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 93.99415588378906,
   "xla_arg_mb": 24.275390625,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.002236284999526106,
   "inner": 86,
   "median_ms": 0.4398473721203329,
   "q1_ms": 0.43529480235308654,
   "q3_ms": 0.44804196511185457,
   "rel_err_out": 0.052127411047239225,
   "rel_err_grad_params": 0.06821531727580259,
   "rel_err_grad_x": 0.052954075268389
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dense",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.0,
   "max_expert_frac": 0.022750001400709152,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 125.0022201538086,
   "xla_arg_mb": 72.33789443969727,
   "xla_out_mb": 72.09386444091797,
   "compile_s": 0.00423664299887605,
   "inner": 28,
   "median_ms": 1.2376981784265288,
   "q1_ms": 1.2308471607736595,
   "q3_ms": 1.348279732285716,
   "rel_err_out": 0.052127411047239225,
   "rel_err_grad_params": 0.06821531727580259,
   "rel_err_grad_x": 0.052954075268389
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.9,
   "max_expert_frac": 0.234250009059906,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 13.923622131347656,
   "xla_arg_mb": 24.275390625,
   "xla_out_mb": 0.244140625,
   "compile_s": 0.0023959879981703125,
   "inner": 92,
   "median_ms": 0.35669364130982384,
   "q1_ms": 0.35073017391803896,
   "q3_ms": 0.3615361739202458,
   "rel_err_out": 0.007580264863689277,
   "rel_err_grad_params": 0.033201628913450965,
   "rel_err_grad_x": 0.007410366951859016
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.9,
   "max_expert_frac": 0.234250009059906,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 4.005653381347656,
   "xla_arg_mb": 72.33789443969727,
   "xla_out_mb": 72.09386444091797,
   "compile_s": 0.013373576002777554,
   "inner": 36,
   "median_ms": 0.8368848612008151,
   "q1_ms": 0.8111840833306714,
   "q3_ms": 0.8483567221951995,
   "rel_err_out": 0.007580264863689277,
   "rel_err_grad_params": 0.033201628913450965,
   "rel_err_grad_x": 0.007410366951859016
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "gather",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.9,
   "max_expert_frac": 0.234250009059906,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 17.954383850097656,
   "xla_arg_mb": 24.275390625,
   "xla_out_mb": 0.244140625,
   "compile_s": 0.0023387589972116984,
   "inner": 87,
   "median_ms": 0.4227612988525552,
   "q1_ms": 0.4214811723896077,
   "q3_ms": 0.47264748848218674,
   "rel_err_out": 0.007580264863689277,
   "rel_err_grad_params": 0.03319284959362698,
   "rel_err_grad_x": 0.007214262340817582
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "gather",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.9,
   "max_expert_frac": 0.234250009059906,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 5.9746551513671875,
   "xla_arg_mb": 72.33789443969727,
   "xla_out_mb": 72.09386444091797,
   "compile_s": 0.008637416998681147,
   "inner": 22,
   "median_ms": 0.9114050910945728,
   "q1_ms": 0.8397516364559785,
   "q3_ms": 0.9290814999655017,
   "rel_err_out": 0.007580264863689277,
   "rel_err_grad_params": 0.03319284959362698,
   "rel_err_grad_x": 0.007214262340817582
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.9,
   "max_expert_frac": 0.234250009059906,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 15.969520568847656,
   "xla_arg_mb": 24.275390625,
   "xla_out_mb": 0.244140625,
   "compile_s": 0.0024943430034909397,
   "inner": 90,
   "median_ms": 0.3854830110665514,
   "q1_ms": 0.38161635550042977,
   "q3_ms": 0.43621285557391176,
   "rel_err_out": 0.0075718151968306715,
   "rel_err_grad_params": 0.03313253577310485,
   "rel_err_grad_x": 0.007168786280517155
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.9,
   "max_expert_frac": 0.234250009059906,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 4.3418426513671875,
   "xla_arg_mb": 72.33789443969727,
   "xla_out_mb": 72.09386444091797,
   "compile_s": 0.0038938340003369376,
   "inner": 38,
   "median_ms": 0.9428934999481211,
   "q1_ms": 0.9398352104701764,
   "q3_ms": 0.952936197408187,
   "rel_err_out": 0.0075718151968306715,
   "rel_err_grad_params": 0.03313253577310485,
   "rel_err_grad_x": 0.007168786280517155
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused_gsum",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.9,
   "max_expert_frac": 0.234250009059906,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 16.016395568847656,
   "xla_arg_mb": 24.275390625,
   "xla_out_mb": 0.244140625,
   "compile_s": 0.0023113779971026815,
   "inner": 106,
   "median_ms": 0.40637873579091216,
   "q1_ms": 0.3729687169944331,
   "q3_ms": 0.424142702852479,
   "rel_err_out": 0.007580264863689277,
   "rel_err_grad_params": 0.03319284959362698,
   "rel_err_grad_x": 0.007214262340817582
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "pallas",
   "dispatch": "fused_gsum",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.9,
   "max_expert_frac": 0.234250009059906,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 4.3420867919921875,
   "xla_arg_mb": 72.33789443969727,
   "xla_out_mb": 72.09386444091797,
   "compile_s": 0.004062583997438196,
   "inner": 46,
   "median_ms": 0.8912188478343128,
   "q1_ms": 0.8810683043323123,
   "q3_ms": 0.8996358587667245,
   "rel_err_out": 0.007580264863689277,
   "rel_err_grad_params": 0.03319284959362698,
   "rel_err_grad_x": 0.007214262340817582
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "blocked",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.9,
   "max_expert_frac": 0.234250009059906,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 65.39212799072266,
   "xla_arg_mb": 24.275390625,
   "xla_out_mb": 0.244140625,
   "compile_s": 0.0020840869983658195,
   "inner": 83,
   "median_ms": 0.47200192767125276,
   "q1_ms": 0.4664192469990448,
   "q3_ms": 0.4863790361332028,
   "rel_err_out": 0.0054851148612827015,
   "rel_err_grad_params": 0.033032700363583536,
   "rel_err_grad_x": 0.0044514993700453705
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "blocked",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.9,
   "max_expert_frac": 0.234250009059906,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 59.918251037597656,
   "xla_arg_mb": 72.33789443969727,
   "xla_out_mb": 72.09386444091797,
   "compile_s": 0.003815225005382672,
   "inner": 38,
   "median_ms": 1.1744984210389176,
   "q1_ms": 1.161789486918439,
   "q3_ms": 1.1813523156697077,
   "rel_err_out": 0.0054851148612827015,
   "rel_err_grad_params": 0.033032700363583536,
   "rel_err_grad_x": 0.0044514993700453705
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "ragged_dot",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.9,
   "max_expert_frac": 0.234250009059906,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 261.85963439941406,
   "xla_arg_mb": 24.275390625,
   "xla_out_mb": 0.244140625,
   "compile_s": 0.0035891869993065484,
   "inner": 19,
   "median_ms": 2.281818052563865,
   "q1_ms": 2.2405011052461834,
   "q3_ms": 2.298606684474324,
   "rel_err_out": 0.0054851148612827015,
   "rel_err_grad_params": 0.03303270035575194,
   "rel_err_grad_x": 0.0044514993700453705
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dropless",
   "impl": "ragged_dot",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.9,
   "max_expert_frac": 0.234250009059906,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 625.0044250488281,
   "xla_arg_mb": 72.33789443969727,
   "xla_out_mb": 72.09386444091797,
   "compile_s": 0.015343596998718567,
   "inner": 6,
   "median_ms": 7.820837500427539,
   "q1_ms": 7.7900583334364155,
   "q3_ms": 8.299387750109116,
   "rel_err_out": 0.0054851148612827015,
   "rel_err_grad_params": 0.03303270035575194,
   "rel_err_grad_x": 0.0044514993700453705
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "capacity",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.9,
   "max_expert_frac": 0.234250009059906,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 7.6665191650390625,
   "xla_arg_mb": 24.275390625,
   "xla_out_mb": 0.244140625,
   "compile_s": 0.0015482719973078929,
   "inner": 200,
   "median_ms": 0.2119160650181584,
   "q1_ms": 0.15838561997952638,
   "q3_ms": 0.2224503400066169
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "capacity",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.9,
   "max_expert_frac": 0.234250009059906,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 4.246360778808594,
   "xla_arg_mb": 72.33789443969727,
   "xla_out_mb": 72.09386444091797,
   "compile_s": 0.0028599959987332113,
   "inner": 80,
   "median_ms": 0.46526836249540793,
   "q1_ms": 0.46258752504400036,
   "q3_ms": 0.47173446873784997
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dense",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.9,
   "max_expert_frac": 0.234250009059906,
   "empty_experts": 0,
   "op": "forward",
   "xla_temp_mb": 93.99415588378906,
   "xla_arg_mb": 24.275390625,
   "xla_out_mb": 0.48828125,
   "compile_s": 0.0014387940027518198,
   "inner": 95,
   "median_ms": 0.4390003052417581,
   "q1_ms": 0.4309245315916518,
   "q3_ms": 0.4584878947322355,
   "rel_err_out": 0.006932838339270058,
   "rel_err_grad_params": 0.03305448320884851,
   "rel_err_grad_x": 0.0067378918824145375
  },
  {
   "code": "before",
   "run": "gpu_numerics",
   "suite": "layer",
   "routing": "dense",
   "impl": "auto",
   "dispatch": "sorted",
   "dtype": "bfloat16",
   "tokens": 1000,
   "d_model": 128,
   "d_hidden": 256,
   "experts": 64,
   "top_k": 4,
   "skew": 0.9,
   "max_expert_frac": 0.234250009059906,
   "empty_experts": 0,
   "op": "train_step",
   "xla_temp_mb": 125.0022201538086,
   "xla_arg_mb": 72.33789443969727,
   "xla_out_mb": 72.09386444091797,
   "compile_s": 0.0033557950009708293,
   "inner": 33,
   "median_ms": 1.2622853030210757,
   "q1_ms": 1.2395008030359873,
   "q3_ms": 1.3965257423864608,
   "rel_err_out": 0.006932838339270058,
   "rel_err_grad_params": 0.03305448320884851,
   "rel_err_grad_x": 0.0067378918824145375
  }
 ]
}