{
  "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": {
    "device": "gpu",
    "suite": "dispatch",
    "model_dim": 256,
    "output": "/home/manuel/src/probjax-wt-attn-gpu/docs/benchmarks/kernel_attn_gpu_dispatch.json",
    "append": true,
    "tag": "after",
    "batch": 64,
    "heads": 8,
    "head_dims": [
      32,
      64
    ],
    "seq_lens": [
      16,
      32,
      64,
      128,
      256,
      512
    ],
    "dtypes": [
      "float32",
      "bfloat16"
    ],
    "modes": [
      "fwd",
      "fwd_bwd"
    ],
    "causal": [
      0,
      1
    ],
    "impls": [
      "probjax_default",
      "probjax_flex",
      "jax_cudnn"
    ],
    "find_max": false,
    "max_seq": 262144,
    "child": null
  },
  "results": [
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 16,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0005622011376544833,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.2033943790011108,
      "inner": 200,
      "median_ms": 0.053284935020201374,
      "q1_ms": 0.047861154980637366,
      "q3_ms": 0.05386084751080489,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 16,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0005622011376544833,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.14383930299663916,
      "inner": 200,
      "median_ms": 0.052807885003858246,
      "q1_ms": 0.0514645175098849,
      "q3_ms": 0.058090767488465644,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 16,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0005374412285163999,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.11455835800006753,
      "inner": 200,
      "median_ms": 0.04689049499575049,
      "q1_ms": 0.04573846748826327,
      "q3_ms": 0.04983213000741671,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 16,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0005374412285163999,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.10972127399872988,
      "inner": 200,
      "median_ms": 0.04735157002869528,
      "q1_ms": 0.04696182748375577,
      "q3_ms": 0.0482355874919449,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 16,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0005622011376544833,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.20689488600328332,
      "inner": 17,
      "median_ms": 0.0895604706088574,
      "q1_ms": 0.061869352922418755,
      "q3_ms": 0.11726049984538214,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 16,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006108639063313603,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 4.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.36217675099760527,
      "inner": 200,
      "median_ms": 0.06960373499168782,
      "q1_ms": 0.06555332500283839,
      "q3_ms": 0.07090121251167147,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 16,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006108639063313603,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 4.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.24357885499921395,
      "inner": 200,
      "median_ms": 0.06779361501685344,
      "q1_ms": 0.06740760501998011,
      "q3_ms": 0.07019918752121157,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 16,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006692991591989994,
      "xla_temp_mb": 1.062744140625,
      "xla_arg_mb": 4.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.2123465229960857,
      "inner": 200,
      "median_ms": 0.05699210498278262,
      "q1_ms": 0.056899852515925886,
      "q3_ms": 0.05833461749716662,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 16,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006692991591989994,
      "xla_temp_mb": 1.062744140625,
      "xla_arg_mb": 4.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.20359132200246677,
      "inner": 200,
      "median_ms": 0.04868577500019455,
      "q1_ms": 0.04749133999212063,
      "q3_ms": 0.05516874749446288,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 16,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006108639063313603,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 4.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.21250322500418406,
      "inner": 200,
      "median_ms": 0.061097219986550044,
      "q1_ms": 0.05671091750627966,
      "q3_ms": 0.061998355013201945,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 16,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0004711545188911259,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.48925819500436774,
      "inner": 200,
      "median_ms": 0.10781061500892974,
      "q1_ms": 0.09303772249040776,
      "q3_ms": 0.1090376124921022,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 16,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0004711545188911259,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.4254937619989505,
      "inner": 121,
      "median_ms": 0.08628518183038317,
      "q1_ms": 0.08004899172199897,
      "q3_ms": 0.0946124752118514,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 16,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0005859605153091252,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.16579996300424682,
      "inner": 186,
      "median_ms": 0.04479588706970155,
      "q1_ms": 0.04437856182693228,
      "q3_ms": 0.047837099461298015,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 16,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0005859605153091252,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.13607438399776584,
      "inner": 200,
      "median_ms": 0.05716738502087537,
      "q1_ms": 0.04560727249554475,
      "q3_ms": 0.08299505749164382,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 16,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0004711545188911259,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.2614271910060779,
      "inner": 200,
      "median_ms": 0.043857384989678394,
      "q1_ms": 0.043511002495506546,
      "q3_ms": 0.04817648499738425,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 16,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006589461117982864,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 4.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.6680269120042794,
      "inner": 200,
      "median_ms": 0.07075195502693532,
      "q1_ms": 0.06604574000448338,
      "q3_ms": 0.07127274000595207,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 16,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006589461117982864,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 4.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.8160310829989612,
      "inner": 43,
      "median_ms": 0.18146076746593598,
      "q1_ms": 0.17752959294824017,
      "q3_ms": 0.1880369418358896,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 16,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007603430422022939,
      "xla_temp_mb": 1.062744140625,
      "xla_arg_mb": 4.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.2820703739998862,
      "inner": 21,
      "median_ms": 0.13773171425176164,
      "q1_ms": 0.1375671189245103,
      "q3_ms": 0.13791178571409546,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 16,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007603430422022939,
      "xla_temp_mb": 1.062744140625,
      "xla_arg_mb": 4.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.2748842489963863,
      "inner": 156,
      "median_ms": 0.07889134618520845,
      "q1_ms": 0.07122099036794866,
      "q3_ms": 0.08172202885758359,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 16,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006589461117982864,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 4.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.632838029996492,
      "inner": 21,
      "median_ms": 0.15401590471377685,
      "q1_ms": 0.1535371189848298,
      "q3_ms": 0.15467045228606802,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 32,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.000846682523842901,
      "xla_temp_mb": 4.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.40880971500155283,
      "inner": 200,
      "median_ms": 0.0461293349871994,
      "q1_ms": 0.044813342501583975,
      "q3_ms": 0.0508231700041506,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 32,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.000846682523842901,
      "xla_temp_mb": 4.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.1584868869977072,
      "inner": 200,
      "median_ms": 0.046044474984228145,
      "q1_ms": 0.0437420399794064,
      "q3_ms": 0.050535177488200134,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 32,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006542886258102953,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.21674274100223556,
      "inner": 200,
      "median_ms": 0.05447504001494963,
      "q1_ms": 0.052168844995321706,
      "q3_ms": 0.05543518000195036,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 32,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006542886258102953,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.1431590780048282,
      "inner": 21,
      "median_ms": 0.12641504768648051,
      "q1_ms": 0.09152307150846657,
      "q3_ms": 0.12659900000601607,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 32,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.000846682523842901,
      "xla_temp_mb": 4.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.15878793799492996,
      "inner": 21,
      "median_ms": 0.14867766661891002,
      "q1_ms": 0.14862947621926043,
      "q3_ms": 0.1487962141427921,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 32,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007412488339468837,
      "xla_temp_mb": 4.0,
      "xla_arg_mb": 8.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 1.0187194150057621,
      "inner": 14,
      "median_ms": 0.21210371466752673,
      "q1_ms": 0.14264007144707389,
      "q3_ms": 0.2342563568942881,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 32,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007412488339468837,
      "xla_temp_mb": 4.0,
      "xla_arg_mb": 8.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.492622495003161,
      "inner": 133,
      "median_ms": 0.0807893458868489,
      "q1_ms": 0.08077349245051077,
      "q3_ms": 0.09276344737605213,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 32,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007395902648568153,
      "xla_temp_mb": 2.125244140625,
      "xla_arg_mb": 8.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.4679136440026923,
      "inner": 54,
      "median_ms": 0.1289613704742536,
      "q1_ms": 0.11151789813421029,
      "q3_ms": 0.13082629623197672,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 32,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007395902648568153,
      "xla_temp_mb": 2.125244140625,
      "xla_arg_mb": 8.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.4377975960014737,
      "inner": 200,
      "median_ms": 0.11385300502297468,
      "q1_ms": 0.1070407575025456,
      "q3_ms": 0.13394474250162602,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 32,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007412488339468837,
      "xla_temp_mb": 4.0,
      "xla_arg_mb": 8.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.7512319009983912,
      "inner": 199,
      "median_ms": 0.15313565326958717,
      "q1_ms": 0.13880998993949992,
      "q3_ms": 0.15385636683432363,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 32,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0004653937940020114,
      "xla_temp_mb": 4.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 1.0518487989975256,
      "inner": 177,
      "median_ms": 0.19845754237178617,
      "q1_ms": 0.16193067513394438,
      "q3_ms": 0.2356195932072995,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 32,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0004653937940020114,
      "xla_temp_mb": 4.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.9711685990041588,
      "inner": 95,
      "median_ms": 0.10178997893633909,
      "q1_ms": 0.08373147368394328,
      "q3_ms": 0.10718087371634809,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 32,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0004163338744547218,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.49651411199738504,
      "inner": 45,
      "median_ms": 0.08529348900386444,
      "q1_ms": 0.07717311115508588,
      "q3_ms": 0.12658388877753168,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 32,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0004163338744547218,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.48998266700073145,
      "inner": 43,
      "median_ms": 0.10505330232432286,
      "q1_ms": 0.08239797678597402,
      "q3_ms": 0.10819484880410654,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 32,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0004653937940020114,
      "xla_temp_mb": 4.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 1.1495645500035607,
      "inner": 10,
      "median_ms": 0.24779110026429407,
      "q1_ms": 0.21933025018370245,
      "q3_ms": 0.2945641997939674,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 32,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006948871887288988,
      "xla_temp_mb": 4.0,
      "xla_arg_mb": 8.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 2.663038457001676,
      "inner": 45,
      "median_ms": 0.11138933332404122,
      "q1_ms": 0.1106240999231684,
      "q3_ms": 0.13358579995434006,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 32,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006948871887288988,
      "xla_temp_mb": 4.0,
      "xla_arg_mb": 8.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 1.9783998549974058,
      "inner": 200,
      "median_ms": 0.17227090498636244,
      "q1_ms": 0.16146366997418227,
      "q3_ms": 0.1899375525135838,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 32,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007666696328669786,
      "xla_temp_mb": 2.125244140625,
      "xla_arg_mb": 8.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.8374787280044984,
      "inner": 16,
      "median_ms": 0.07514224989790819,
      "q1_ms": 0.0667801564304682,
      "q3_ms": 0.13705231231142534,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 32,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007666696328669786,
      "xla_temp_mb": 2.125244140625,
      "xla_arg_mb": 8.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.7592686360003427,
      "inner": 200,
      "median_ms": 0.06401850499969441,
      "q1_ms": 0.06114982501458144,
      "q3_ms": 0.06706655000016326,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 32,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006948871887288988,
      "xla_temp_mb": 4.0,
      "xla_arg_mb": 8.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 1.491698379999434,
      "inner": 109,
      "median_ms": 0.1674233119397578,
      "q1_ms": 0.1617394266033293,
      "q3_ms": 0.19300821561146303,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 64,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0005295496084727347,
      "xla_temp_mb": 16.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.5209942699948442,
      "inner": 149,
      "median_ms": 0.11408478523613948,
      "q1_ms": 0.1071147449668528,
      "q3_ms": 0.11768088927129189,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 64,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0005295496084727347,
      "xla_temp_mb": 16.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.2447844510024879,
      "inner": 25,
      "median_ms": 0.17496279993792996,
      "q1_ms": 0.10972753996611573,
      "q3_ms": 0.19741134005016647,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 64,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0005379546782933176,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.27401421099784784,
      "inner": 150,
      "median_ms": 0.07900371332652867,
      "q1_ms": 0.06904712667164858,
      "q3_ms": 0.09497498334288443,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 64,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0005379546782933176,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.2249002160024247,
      "inner": 196,
      "median_ms": 0.11333636226778736,
      "q1_ms": 0.10519712499704578,
      "q3_ms": 0.11867041329034173,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 64,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0005295496084727347,
      "xla_temp_mb": 16.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.45258903100329917,
      "inner": 197,
      "median_ms": 0.1520205583561656,
      "q1_ms": 0.1428236548217406,
      "q3_ms": 0.16751178682262913,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 64,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007830221438780427,
      "xla_temp_mb": 24.0,
      "xla_arg_mb": 16.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.9077144539987785,
      "inner": 138,
      "median_ms": 0.23421219563148782,
      "q1_ms": 0.23231171017750574,
      "q3_ms": 0.23909295289384952,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 64,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007830221438780427,
      "xla_temp_mb": 24.0,
      "xla_arg_mb": 16.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.394834617000015,
      "inner": 157,
      "median_ms": 0.23358414650510603,
      "q1_ms": 0.2305645477671805,
      "q3_ms": 0.2361588280121158,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 64,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007911636494100094,
      "xla_temp_mb": 4.250244140625,
      "xla_arg_mb": 16.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.709906090000004,
      "inner": 200,
      "median_ms": 0.10173575999942841,
      "q1_ms": 0.09760672000993509,
      "q3_ms": 0.10740995499872952,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 64,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007911636494100094,
      "xla_temp_mb": 4.250244140625,
      "xla_arg_mb": 16.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.5663534569976036,
      "inner": 200,
      "median_ms": 0.09784384499653243,
      "q1_ms": 0.09644917749028536,
      "q3_ms": 0.13402416749158874,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 64,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007830221438780427,
      "xla_temp_mb": 24.0,
      "xla_arg_mb": 16.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.6966347039997345,
      "inner": 18,
      "median_ms": 0.48671761073314374,
      "q1_ms": 0.3668890834281532,
      "q3_ms": 0.5589721664970663,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 64,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00041235165554098785,
      "xla_temp_mb": 16.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.7831873309987714,
      "inner": 199,
      "median_ms": 0.1076354924384645,
      "q1_ms": 0.1075348718678017,
      "q3_ms": 0.10965207286574527,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 64,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00041235165554098785,
      "xla_temp_mb": 16.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.9616567400007625,
      "inner": 11,
      "median_ms": 0.27599518216447905,
      "q1_ms": 0.18144259097012266,
      "q3_ms": 0.3212108633967794,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 64,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0003806430322583765,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.5498176100009005,
      "inner": 46,
      "median_ms": 0.05525808702931618,
      "q1_ms": 0.04692808690141765,
      "q3_ms": 0.060565434741084063,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 64,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0003806430322583765,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.36302812800568063,
      "inner": 98,
      "median_ms": 0.05956091842321888,
      "q1_ms": 0.04047331124323192,
      "q3_ms": 0.07914177550279716,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 64,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00041235165554098785,
      "xla_temp_mb": 16.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.7233993260015268,
      "inner": 200,
      "median_ms": 0.10742237500380725,
      "q1_ms": 0.10735758251030347,
      "q3_ms": 0.10750220500995056,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 64,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006174938171170652,
      "xla_temp_mb": 24.0,
      "xla_arg_mb": 16.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 3.1110477370020817,
      "inner": 110,
      "median_ms": 0.3241659091145266,
      "q1_ms": 0.3206249408818125,
      "q3_ms": 0.3247216454672691,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 64,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006174938171170652,
      "xla_temp_mb": 24.0,
      "xla_arg_mb": 16.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 4.186580124005559,
      "inner": 10,
      "median_ms": 0.33779760051402263,
      "q1_ms": 0.3373567000380717,
      "q3_ms": 0.35461485022096895,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 64,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006908720242790878,
      "xla_temp_mb": 4.250244140625,
      "xla_arg_mb": 16.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.6813333940008306,
      "inner": 200,
      "median_ms": 0.11977511498116655,
      "q1_ms": 0.11836059000415844,
      "q3_ms": 0.14388230751137598,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 64,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006908720242790878,
      "xla_temp_mb": 4.250244140625,
      "xla_arg_mb": 16.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 1.1307113760049106,
      "inner": 139,
      "median_ms": 0.11634933810409748,
      "q1_ms": 0.11581166907215304,
      "q3_ms": 0.11662328416350123,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 64,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006174938171170652,
      "xla_temp_mb": 24.0,
      "xla_arg_mb": 16.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 2.829857520002406,
      "inner": 110,
      "median_ms": 0.3153146090467503,
      "q1_ms": 0.3139784863544099,
      "q3_ms": 0.3166617545502959,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 128,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007525204564444721,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.5005968929981464,
      "inner": 200,
      "median_ms": 0.10404042499430943,
      "q1_ms": 0.09618530251827906,
      "q3_ms": 0.11655231501208618,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 128,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.000711066008079797,
      "xla_temp_mb": 64.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.7704624890029663,
      "inner": 8,
      "median_ms": 0.3327150006953161,
      "q1_ms": 0.3165606249240227,
      "q3_ms": 0.6376678747983533,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 128,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007525204564444721,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.8864964730018983,
      "inner": 24,
      "median_ms": 0.22762554150782913,
      "q1_ms": 0.07865204179324792,
      "q3_ms": 0.3323147707305907,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 128,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007525204564444721,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 2.119298523000907,
      "inner": 19,
      "median_ms": 0.21500342087452545,
      "q1_ms": 0.21123228917100564,
      "q3_ms": 0.2249482368872697,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 128,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.000711066008079797,
      "xla_temp_mb": 64.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.7469784269997035,
      "inner": 92,
      "median_ms": 0.2948832934634993,
      "q1_ms": 0.29187445109172305,
      "q3_ms": 0.30320143480861117,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 128,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006482516182586551,
      "xla_temp_mb": 8.500244140625,
      "xla_arg_mb": 32.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 3.006318455998553,
      "inner": 9,
      "median_ms": 0.44463522256895277,
      "q1_ms": 0.4441157778071809,
      "q3_ms": 0.5106156111449106,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 128,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006482516182586551,
      "xla_temp_mb": 96.0,
      "xla_arg_mb": 32.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 2.0808347610000055,
      "inner": 53,
      "median_ms": 0.6960459057085725,
      "q1_ms": 0.6951908962603413,
      "q3_ms": 0.6969853396031056,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 128,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006482516182586551,
      "xla_temp_mb": 8.500244140625,
      "xla_arg_mb": 32.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.630634637993353,
      "inner": 89,
      "median_ms": 0.293138011206946,
      "q1_ms": 0.2917218033657305,
      "q3_ms": 0.3182220281222103,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 128,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006482516182586551,
      "xla_temp_mb": 8.500244140625,
      "xla_arg_mb": 32.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.714821417001076,
      "inner": 89,
      "median_ms": 0.33505051682920334,
      "q1_ms": 0.33362759552729904,
      "q3_ms": 0.3377105618056611,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 128,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006482516182586551,
      "xla_temp_mb": 96.0,
      "xla_arg_mb": 32.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.4428358649965958,
      "inner": 61,
      "median_ms": 0.6985809672127676,
      "q1_ms": 0.6941550327908652,
      "q3_ms": 0.7063358033087379,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 128,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0005480899708345532,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.37085065800056327,
      "inner": 200,
      "median_ms": 0.057625000008556526,
      "q1_ms": 0.05710854749850114,
      "q3_ms": 0.06012856500092312,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 128,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007409274694509804,
      "xla_temp_mb": 64.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 1.5907766110030934,
      "inner": 109,
      "median_ms": 0.3580203026996652,
      "q1_ms": 0.3486842339177009,
      "q3_ms": 0.36433272019942897,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 128,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0005480899708345532,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.3662782310057082,
      "inner": 200,
      "median_ms": 0.0565812200147775,
      "q1_ms": 0.05500138249772135,
      "q3_ms": 0.057665499989525415,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 128,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0005480899708345532,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.3799640430006548,
      "inner": 200,
      "median_ms": 0.0550267300059204,
      "q1_ms": 0.054367332504625665,
      "q3_ms": 0.05635299001369276,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 128,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007409274694509804,
      "xla_temp_mb": 64.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 1.5768059300025925,
      "inner": 107,
      "median_ms": 0.35919406538052384,
      "q1_ms": 0.3522229719883654,
      "q3_ms": 0.387681088736444,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 128,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007646068697795272,
      "xla_temp_mb": 8.500244140625,
      "xla_arg_mb": 32.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.6529399079954601,
      "inner": 142,
      "median_ms": 0.2596161760099586,
      "q1_ms": 0.2590997711182277,
      "q3_ms": 0.2638052640795137,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 128,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007646068697795272,
      "xla_temp_mb": 96.0,
      "xla_arg_mb": 32.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 3.5306668670018553,
      "inner": 46,
      "median_ms": 0.9630907609433685,
      "q1_ms": 0.9561248804422339,
      "q3_ms": 1.023639326060378,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 128,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007646068697795272,
      "xla_temp_mb": 8.500244140625,
      "xla_arg_mb": 32.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.674736175002181,
      "inner": 136,
      "median_ms": 0.25967671321910424,
      "q1_ms": 0.2591664080736176,
      "q3_ms": 0.26008848162211295,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 128,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007646068697795272,
      "xla_temp_mb": 8.500244140625,
      "xla_arg_mb": 32.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.7161508730059722,
      "inner": 128,
      "median_ms": 0.26417021871338875,
      "q1_ms": 0.2636708750003436,
      "q3_ms": 0.26531426169640326,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 128,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007646068697795272,
      "xla_temp_mb": 96.0,
      "xla_arg_mb": 32.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 3.8363267319946317,
      "inner": 44,
      "median_ms": 1.1055223636834524,
      "q1_ms": 1.095140999992013,
      "q3_ms": 1.13239551129399,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 256,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006812745705246925,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.4950409029988805,
      "inner": 166,
      "median_ms": 0.17197765060174502,
      "q1_ms": 0.17193029818745015,
      "q3_ms": 0.17446832529217263,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 256,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007446179515682161,
      "xla_temp_mb": 256.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.8635916380007984,
      "inner": 50,
      "median_ms": 0.8893794199684635,
      "q1_ms": 0.8887858100206358,
      "q3_ms": 0.9014326500619063,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 256,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006812745705246925,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.3749678930034861,
      "inner": 185,
      "median_ms": 0.16978236758527723,
      "q1_ms": 0.1695236297343178,
      "q3_ms": 0.1731123945737144,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 256,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006812745705246925,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.45659930299734697,
      "inner": 168,
      "median_ms": 0.17199658928542963,
      "q1_ms": 0.17012989881203838,
      "q3_ms": 0.17224261907921617,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 256,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007446179515682161,
      "xla_temp_mb": 256.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.33920073599438183,
      "inner": 37,
      "median_ms": 0.9037079730045956,
      "q1_ms": 0.8956112973022312,
      "q3_ms": 0.9253018377765674,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 256,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0011262360494583845,
      "xla_temp_mb": 17.000244140625,
      "xla_arg_mb": 64.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.7171982159998151,
      "inner": 45,
      "median_ms": 0.9784297778322878,
      "q1_ms": 0.9777462444617413,
      "q3_ms": 0.9825334221305739,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 256,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0011259721359238029,
      "xla_temp_mb": 384.0,
      "xla_arg_mb": 64.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 3.4838512419955805,
      "inner": 21,
      "median_ms": 2.518001047497992,
      "q1_ms": 2.5008378570852243,
      "q3_ms": 2.5248180238469065,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 256,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0011262360494583845,
      "xla_temp_mb": 17.000244140625,
      "xla_arg_mb": 64.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.7827193240009365,
      "inner": 46,
      "median_ms": 0.9763020868337227,
      "q1_ms": 0.9753541087221769,
      "q3_ms": 0.9765192391331885,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 256,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0011262360494583845,
      "xla_temp_mb": 17.000244140625,
      "xla_arg_mb": 64.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.7030736540036742,
      "inner": 45,
      "median_ms": 0.9754962000038682,
      "q1_ms": 0.9738431889369774,
      "q3_ms": 0.9884944778200911,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 256,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0011259721359238029,
      "xla_temp_mb": 384.0,
      "xla_arg_mb": 64.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.5822168710001279,
      "inner": 21,
      "median_ms": 2.1614679523433247,
      "q1_ms": 2.1591158808968474,
      "q3_ms": 2.1643075954108055,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 256,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0004299705324228853,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.37870820500393165,
      "inner": 179,
      "median_ms": 0.15621334637089906,
      "q1_ms": 0.15034856146764708,
      "q3_ms": 0.16349705584723162,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 256,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0004971984890289605,
      "xla_temp_mb": 256.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 1.8743787849962246,
      "inner": 42,
      "median_ms": 1.0675385000511248,
      "q1_ms": 1.0653861072136177,
      "q3_ms": 1.1050043333289934,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 256,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0004299705324228853,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.37769526700139977,
      "inner": 200,
      "median_ms": 0.14301867002359359,
      "q1_ms": 0.1415388475106738,
      "q3_ms": 0.14355223500388092,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 256,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0004299705324228853,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.3687306540014106,
      "inner": 193,
      "median_ms": 0.14290542488343472,
      "q1_ms": 0.1422877538749596,
      "q3_ms": 0.14357135234246216,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 256,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0004971984890289605,
      "xla_temp_mb": 256.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 1.7584224289967096,
      "inner": 41,
      "median_ms": 1.0647213902102965,
      "q1_ms": 1.048254146309507,
      "q3_ms": 1.069106780583175,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 256,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007133711478672922,
      "xla_temp_mb": 17.000244140625,
      "xla_arg_mb": 64.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 1.1948360610040254,
      "inner": 34,
      "median_ms": 0.7000712941910879,
      "q1_ms": 0.6762016911285601,
      "q3_ms": 0.7224834410889718,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 256,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006073800614103675,
      "xla_temp_mb": 384.0,
      "xla_arg_mb": 64.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 5.785485875996528,
      "inner": 15,
      "median_ms": 3.1296699996649595,
      "q1_ms": 3.1215258667846983,
      "q3_ms": 3.3507231666590087,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 256,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007133711478672922,
      "xla_temp_mb": 17.000244140625,
      "xla_arg_mb": 64.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.7245757840064471,
      "inner": 66,
      "median_ms": 0.7255188939972979,
      "q1_ms": 0.6731903636550256,
      "q3_ms": 0.735775143932022,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 256,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007133711478672922,
      "xla_temp_mb": 17.000244140625,
      "xla_arg_mb": 64.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.8067051559992251,
      "inner": 61,
      "median_ms": 0.6728085574024112,
      "q1_ms": 0.6726865245641942,
      "q3_ms": 0.6748603196594231,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 256,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006073800614103675,
      "xla_temp_mb": 384.0,
      "xla_arg_mb": 64.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 4.452719968001475,
      "inner": 15,
      "median_ms": 3.132261533270745,
      "q1_ms": 3.1138142667866004,
      "q3_ms": 3.2239040667870236,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 512,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0005241255275905132,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.42977653999696486,
      "inner": 71,
      "median_ms": 0.6042825070195045,
      "q1_ms": 0.6035227676145207,
      "q3_ms": 0.6050867887097411,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 512,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0005897894152440131,
      "xla_temp_mb": 1024.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.8860350590039161,
      "inner": 15,
      "median_ms": 3.0472406001839167,
      "q1_ms": 3.0466040334431455,
      "q3_ms": 3.047358133335365,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 512,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0005241255275905132,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.36076580500230193,
      "inner": 72,
      "median_ms": 0.6111477500275618,
      "q1_ms": 0.6096949444124928,
      "q3_ms": 0.6130733472168179,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 512,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0005241255275905132,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.5332220999989659,
      "inner": 63,
      "median_ms": 0.6105984127590625,
      "q1_ms": 0.6082717778294214,
      "q3_ms": 0.6111589841205545,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 512,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0005897894152440131,
      "xla_temp_mb": 1024.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.28664284000115003,
      "inner": 11,
      "median_ms": 3.051687635888811,
      "q1_ms": 3.0503910909994207,
      "q3_ms": 3.0535006361870645,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 512,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0008133161463774741,
      "xla_temp_mb": 34.000244140625,
      "xla_arg_mb": 128.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 1.0219374079970294,
      "inner": 15,
      "median_ms": 3.442449466577576,
      "q1_ms": 3.4375435333155715,
      "q3_ms": 3.444779900019057,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 512,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0011549380142241716,
      "xla_temp_mb": 1537.0,
      "xla_arg_mb": 128.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 2.333153906998632,
      "inner": 5,
      "median_ms": 8.149214599689003,
      "q1_ms": 8.139772700087633,
      "q3_ms": 8.15045299968915,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 512,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0008133161463774741,
      "xla_temp_mb": 34.000244140625,
      "xla_arg_mb": 128.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 0.6563491429988062,
      "inner": 15,
      "median_ms": 3.5061012667332156,
      "q1_ms": 3.4748451335569066,
      "q3_ms": 3.5146762330744723,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 512,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0008133161463774741,
      "xla_temp_mb": 34.000244140625,
      "xla_arg_mb": 128.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 0.9173224149999442,
      "inner": 15,
      "median_ms": 3.444174066923248,
      "q1_ms": 3.419785900041461,
      "q3_ms": 3.5653268001624383,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 512,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0011549380142241716,
      "xla_temp_mb": 1537.0,
      "xla_arg_mb": 128.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 0.8487061449995963,
      "inner": 5,
      "median_ms": 8.129856000596192,
      "q1_ms": 8.112873800564557,
      "q3_ms": 8.132436799496645,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 512,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00037968141259625554,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.39253015400026925,
      "inner": 97,
      "median_ms": 0.4131993092536523,
      "q1_ms": 0.4123394845141165,
      "q3_ms": 0.4136350257748248,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 512,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006253995816223323,
      "xla_temp_mb": 1024.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 1.4606141190015478,
      "inner": 15,
      "median_ms": 3.113383799791336,
      "q1_ms": 3.1120185332838446,
      "q3_ms": 3.1154536998656113,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 512,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00037968141259625554,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.5700661759983632,
      "inner": 89,
      "median_ms": 0.4190161124099402,
      "q1_ms": 0.4106416124026448,
      "q3_ms": 0.42014263486594294,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 512,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00037968141259625554,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.5759660659969086,
      "inner": 95,
      "median_ms": 0.4232135788984842,
      "q1_ms": 0.4170971947084916,
      "q3_ms": 0.4275146105149026,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 512,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006253995816223323,
      "xla_temp_mb": 1024.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 1.4181273610010976,
      "inner": 15,
      "median_ms": 3.1578625998614975,
      "q1_ms": 3.14555566656054,
      "q3_ms": 3.19367979997575,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 512,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006316628423519433,
      "xla_temp_mb": 34.000244140625,
      "xla_arg_mb": 128.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 0.7692183339968324,
      "inner": 24,
      "median_ms": 2.075544541791411,
      "q1_ms": 2.059039979182368,
      "q3_ms": 2.081384166710147,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 512,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0005864849081262946,
      "xla_temp_mb": 1537.0,
      "xla_arg_mb": 128.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 1.9618294120009523,
      "inner": 5,
      "median_ms": 8.293607999803498,
      "q1_ms": 8.282656299707014,
      "q3_ms": 8.304912399762543,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 512,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006316628423519433,
      "xla_temp_mb": 34.000244140625,
      "xla_arg_mb": 128.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 1.208228206000058,
      "inner": 24,
      "median_ms": 2.075114999873525,
      "q1_ms": 2.072741312379852,
      "q3_ms": 2.0780993750122434,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 512,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006316628423519433,
      "xla_temp_mb": 34.000244140625,
      "xla_arg_mb": 128.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 0.8405786880030064,
      "inner": 23,
      "median_ms": 2.07685856509235,
      "q1_ms": 2.0736061085936734,
      "q3_ms": 2.080716521880316,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 512,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0005864849081262946,
      "xla_temp_mb": 1537.0,
      "xla_arg_mb": 128.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 1.8843149069944047,
      "inner": 5,
      "median_ms": 8.469981999951415,
      "q1_ms": 8.383086600224487,
      "q3_ms": 8.78114290026133,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 16,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006341451080515981,
      "xla_temp_mb": 3.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.2773891699980595,
      "inner": 200,
      "median_ms": 0.06314811002084753,
      "q1_ms": 0.0617635849994258,
      "q3_ms": 0.06589567999981227,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 16,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006341451080515981,
      "xla_temp_mb": 3.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.15377031800016994,
      "inner": 200,
      "median_ms": 0.07741005501884501,
      "q1_ms": 0.06606443250348093,
      "q3_ms": 0.09146010748736444,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 16,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0004391638212837279,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.19336949900025502,
      "inner": 200,
      "median_ms": 0.05072628500784049,
      "q1_ms": 0.04858260748733301,
      "q3_ms": 0.052200529989931965,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 16,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0004391638212837279,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.13282969199644867,
      "inner": 200,
      "median_ms": 0.053100629993423354,
      "q1_ms": 0.04687577500590123,
      "q3_ms": 0.05674898749930435,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 16,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006341451080515981,
      "xla_temp_mb": 3.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.23167255699809175,
      "inner": 200,
      "median_ms": 0.07633176497620298,
      "q1_ms": 0.06818288749855128,
      "q3_ms": 0.09058472749529756,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 16,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006048346986062825,
      "xla_temp_mb": 3.0,
      "xla_arg_mb": 8.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.43510949699702905,
      "inner": 188,
      "median_ms": 0.08231819149241733,
      "q1_ms": 0.07773285904748148,
      "q3_ms": 0.08469596277210896,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 16,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006048346986062825,
      "xla_temp_mb": 3.0,
      "xla_arg_mb": 8.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.36496026600070763,
      "inner": 129,
      "median_ms": 0.07864906201716658,
      "q1_ms": 0.07363650777472848,
      "q3_ms": 0.09354260852655158,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 16,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.000719865201972425,
      "xla_temp_mb": 2.062744140625,
      "xla_arg_mb": 8.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.306524232000811,
      "inner": 200,
      "median_ms": 0.0773149299857323,
      "q1_ms": 0.07455900749846478,
      "q3_ms": 0.09237491500243777,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 16,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.000719865201972425,
      "xla_temp_mb": 2.062744140625,
      "xla_arg_mb": 8.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.2914885950012831,
      "inner": 200,
      "median_ms": 0.07816994999302551,
      "q1_ms": 0.07425411000440363,
      "q3_ms": 0.08222294251027051,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 16,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006048346986062825,
      "xla_temp_mb": 3.0,
      "xla_arg_mb": 8.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.28179424100380857,
      "inner": 200,
      "median_ms": 0.0734585800091736,
      "q1_ms": 0.07282167000084883,
      "q3_ms": 0.0764436024837778,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 16,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0005328193074092269,
      "xla_temp_mb": 3.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.33577803700609365,
      "inner": 200,
      "median_ms": 0.06194112000230235,
      "q1_ms": 0.06094160249631386,
      "q3_ms": 0.0654341074914555,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 16,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0005328193074092269,
      "xla_temp_mb": 3.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.22655304499494378,
      "inner": 200,
      "median_ms": 0.0550592950094142,
      "q1_ms": 0.05232882000200334,
      "q3_ms": 0.057173269997292664,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 16,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0005072162603028119,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.21736134299862897,
      "inner": 200,
      "median_ms": 0.04888175000814954,
      "q1_ms": 0.04860802999246516,
      "q3_ms": 0.05324761248630239,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 16,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0005072162603028119,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.1849253120017238,
      "inner": 200,
      "median_ms": 0.061821545023121864,
      "q1_ms": 0.06032218250766164,
      "q3_ms": 0.09130193499004235,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 16,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0005328193074092269,
      "xla_temp_mb": 3.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.35931629100377904,
      "inner": 200,
      "median_ms": 0.0677633449959103,
      "q1_ms": 0.06135879248176934,
      "q3_ms": 0.06813298999986728,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 16,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00047061077202670276,
      "xla_temp_mb": 3.0,
      "xla_arg_mb": 8.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.779512913999497,
      "inner": 200,
      "median_ms": 0.07430588000715943,
      "q1_ms": 0.07416574000671972,
      "q3_ms": 0.07896828250522958,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 16,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00047061077202670276,
      "xla_temp_mb": 3.0,
      "xla_arg_mb": 8.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.8075170840020292,
      "inner": 200,
      "median_ms": 0.07043279500067001,
      "q1_ms": 0.070362415008276,
      "q3_ms": 0.07990042249730323,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 16,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0008460411918349564,
      "xla_temp_mb": 2.062744140625,
      "xla_arg_mb": 8.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.35527745399303967,
      "inner": 23,
      "median_ms": 0.07701100010941128,
      "q1_ms": 0.07402319565501428,
      "q3_ms": 0.10096460881987182,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 16,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0008460411918349564,
      "xla_temp_mb": 2.062744140625,
      "xla_arg_mb": 8.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.3561339060033788,
      "inner": 200,
      "median_ms": 0.06959040500078117,
      "q1_ms": 0.06742808498529484,
      "q3_ms": 0.07637496999450377,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 16,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00047061077202670276,
      "xla_temp_mb": 3.0,
      "xla_arg_mb": 8.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.720433241003775,
      "inner": 200,
      "median_ms": 0.08474294001644012,
      "q1_ms": 0.07894571499491576,
      "q3_ms": 0.08553143499739235,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 32,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00044052154407836497,
      "xla_temp_mb": 8.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.3926310850001755,
      "inner": 200,
      "median_ms": 0.06595301500055939,
      "q1_ms": 0.06493241000498529,
      "q3_ms": 0.06668658001217409,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 32,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00044052154407836497,
      "xla_temp_mb": 8.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.3704572529968573,
      "inner": 26,
      "median_ms": 0.10488280755593084,
      "q1_ms": 0.07931034619664737,
      "q3_ms": 0.14724842311317424,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 32,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00041910100844688714,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.2979269249990466,
      "inner": 88,
      "median_ms": 0.07879110225158828,
      "q1_ms": 0.06870632949523332,
      "q3_ms": 0.08160335797087595,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 32,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00041910100844688714,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.231372039001144,
      "inner": 200,
      "median_ms": 0.049474215011287015,
      "q1_ms": 0.0441669425163127,
      "q3_ms": 0.05261163498289534,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 32,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00044052154407836497,
      "xla_temp_mb": 8.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.19770771299954504,
      "inner": 200,
      "median_ms": 0.06744324000464985,
      "q1_ms": 0.06593505499040475,
      "q3_ms": 0.06892121000419138,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 32,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007663984433747828,
      "xla_temp_mb": 8.0,
      "xla_arg_mb": 16.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.682435009002802,
      "inner": 155,
      "median_ms": 0.14877683870802302,
      "q1_ms": 0.14827840646067936,
      "q3_ms": 0.15020075807115063,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 32,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007663984433747828,
      "xla_temp_mb": 8.0,
      "xla_arg_mb": 16.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.4401502709952183,
      "inner": 42,
      "median_ms": 0.15233319053160294,
      "q1_ms": 0.15111277380561833,
      "q3_ms": 0.15643211909697302,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 32,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007477875333279371,
      "xla_temp_mb": 4.125244140625,
      "xla_arg_mb": 16.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.5295159029992647,
      "inner": 200,
      "median_ms": 0.09149136996711604,
      "q1_ms": 0.09084814999368973,
      "q3_ms": 0.09303484001065954,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 32,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007477875333279371,
      "xla_temp_mb": 4.125244140625,
      "xla_arg_mb": 16.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.5054920300026424,
      "inner": 57,
      "median_ms": 0.14750242093884336,
      "q1_ms": 0.14099808769267902,
      "q3_ms": 0.1900417193067832,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 32,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007663984433747828,
      "xla_temp_mb": 8.0,
      "xla_arg_mb": 16.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.4519934130003094,
      "inner": 72,
      "median_ms": 0.15105251390600138,
      "q1_ms": 0.1503319791077552,
      "q3_ms": 0.15148060416928882,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 32,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006300578825175762,
      "xla_temp_mb": 8.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.4414701089990558,
      "inner": 200,
      "median_ms": 0.07635611000296194,
      "q1_ms": 0.075724939997599,
      "q3_ms": 0.07844528998248279,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 32,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006300578825175762,
      "xla_temp_mb": 8.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.42253834700386506,
      "inner": 200,
      "median_ms": 0.07400566501019057,
      "q1_ms": 0.06722592499500024,
      "q3_ms": 0.07555357002274832,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 32,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0004133675538469106,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.21546793799643638,
      "inner": 200,
      "median_ms": 0.05252514500170946,
      "q1_ms": 0.04766091000419692,
      "q3_ms": 0.05533393750738469,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 32,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0004133675538469106,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.24445274400204653,
      "inner": 200,
      "median_ms": 0.053745555014756974,
      "q1_ms": 0.04650660001061624,
      "q3_ms": 0.06015087999912794,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 32,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006300578825175762,
      "xla_temp_mb": 8.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.6200863769990974,
      "inner": 192,
      "median_ms": 0.07239101561632803,
      "q1_ms": 0.0676954062631315,
      "q3_ms": 0.07625481509876408,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 32,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007315461989492178,
      "xla_temp_mb": 8.0,
      "xla_arg_mb": 16.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 1.5682339109989698,
      "inner": 181,
      "median_ms": 0.15658850830501123,
      "q1_ms": 0.1564246325950442,
      "q3_ms": 0.15872080388040932,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 32,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007315461989492178,
      "xla_temp_mb": 8.0,
      "xla_arg_mb": 16.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 1.5313090839990764,
      "inner": 178,
      "median_ms": 0.1593838932513903,
      "q1_ms": 0.15850399438605062,
      "q3_ms": 0.16656721067540295,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 32,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.000715727626811713,
      "xla_temp_mb": 4.125244140625,
      "xla_arg_mb": 16.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.558229656999174,
      "inner": 102,
      "median_ms": 0.09198396076624027,
      "q1_ms": 0.09172139702879539,
      "q3_ms": 0.09234424510841742,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 32,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.000715727626811713,
      "xla_temp_mb": 4.125244140625,
      "xla_arg_mb": 16.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.7743091819938854,
      "inner": 26,
      "median_ms": 0.16822619223743318,
      "q1_ms": 0.14506117314830097,
      "q3_ms": 0.18570138462774383,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 32,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007315461989492178,
      "xla_temp_mb": 8.0,
      "xla_arg_mb": 16.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 1.7692233979978482,
      "inner": 145,
      "median_ms": 0.1608470827642154,
      "q1_ms": 0.1606649413839352,
      "q3_ms": 0.16108107241773817,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 64,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0005603734171018004,
      "xla_temp_mb": 16.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.6644174980028765,
      "inner": 163,
      "median_ms": 0.16988101226071958,
      "q1_ms": 0.16942168098957955,
      "q3_ms": 0.18548061351643225,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 64,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0005603734171018004,
      "xla_temp_mb": 16.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.3786319699938758,
      "inner": 113,
      "median_ms": 0.17445666368574778,
      "q1_ms": 0.1686113938110761,
      "q3_ms": 0.17874726546702668,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 64,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007154354825615883,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.4360244999988936,
      "inner": 200,
      "median_ms": 0.05615121001028456,
      "q1_ms": 0.04722488500192412,
      "q3_ms": 0.07150702998842462,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 64,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007154354825615883,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.4443042549974052,
      "inner": 200,
      "median_ms": 0.055429239982913714,
      "q1_ms": 0.05246415248620906,
      "q3_ms": 0.057157812516379636,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 64,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0005603734171018004,
      "xla_temp_mb": 16.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.4852770680008689,
      "inner": 179,
      "median_ms": 0.16905503910406905,
      "q1_ms": 0.16886714245771534,
      "q3_ms": 0.1703707709652892,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 64,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006689945585094392,
      "xla_temp_mb": 16.0,
      "xla_arg_mb": 32.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 1.0064671190048102,
      "inner": 93,
      "median_ms": 0.3921561075658638,
      "q1_ms": 0.39166620431338955,
      "q3_ms": 0.3926904516242024,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 64,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006689945585094392,
      "xla_temp_mb": 16.0,
      "xla_arg_mb": 32.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.8538601830005064,
      "inner": 75,
      "median_ms": 0.3926976133758823,
      "q1_ms": 0.39225424664133846,
      "q3_ms": 0.3935076799825765,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 64,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006689945585094392,
      "xla_temp_mb": 8.250244140625,
      "xla_arg_mb": 32.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 1.6069132119955611,
      "inner": 89,
      "median_ms": 0.3473785730553634,
      "q1_ms": 0.34456955619860086,
      "q3_ms": 0.3477838427090646,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 64,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006689945585094392,
      "xla_temp_mb": 8.250244140625,
      "xla_arg_mb": 32.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 1.4020215480049956,
      "inner": 112,
      "median_ms": 0.3448170268062053,
      "q1_ms": 0.34345516074349036,
      "q3_ms": 0.34491448660186763,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 64,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006689945585094392,
      "xla_temp_mb": 16.0,
      "xla_arg_mb": 32.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.4943244560054154,
      "inner": 94,
      "median_ms": 0.39755736167748223,
      "q1_ms": 0.39582296275066836,
      "q3_ms": 0.4002829999545282,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 64,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0005070326151326299,
      "xla_temp_mb": 16.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.8521672380011296,
      "inner": 150,
      "median_ms": 0.1798848332934237,
      "q1_ms": 0.1796373666487246,
      "q3_ms": 0.18303059332538396,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 64,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0005070326151326299,
      "xla_temp_mb": 16.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.8131996119991527,
      "inner": 163,
      "median_ms": 0.1776712822322191,
      "q1_ms": 0.17751586811236347,
      "q3_ms": 0.17800332823563997,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 64,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00039393448969349265,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.42157695700007025,
      "inner": 200,
      "median_ms": 0.04924593500618357,
      "q1_ms": 0.04767739001181326,
      "q3_ms": 0.057241922495450126,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 64,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00039393448969349265,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.4649726260031457,
      "inner": 200,
      "median_ms": 0.07416336000460433,
      "q1_ms": 0.06685144502625917,
      "q3_ms": 0.08032933499634964,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 64,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0005070326151326299,
      "xla_temp_mb": 16.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.9462245030008489,
      "inner": 162,
      "median_ms": 0.18009448764065314,
      "q1_ms": 0.17915656790586426,
      "q3_ms": 0.18035624693036456,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 64,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006008640630170703,
      "xla_temp_mb": 16.0,
      "xla_arg_mb": 32.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 3.788075830001617,
      "inner": 74,
      "median_ms": 0.49096655404994366,
      "q1_ms": 0.4887076148726329,
      "q3_ms": 0.5073567703024113,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 64,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006008640630170703,
      "xla_temp_mb": 16.0,
      "xla_arg_mb": 32.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 4.275364746994455,
      "inner": 85,
      "median_ms": 0.49788485887213885,
      "q1_ms": 0.49158758826225124,
      "q3_ms": 0.507963588247281,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 64,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007147961296141148,
      "xla_temp_mb": 8.250244140625,
      "xla_arg_mb": 32.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 1.8219877949959482,
      "inner": 15,
      "median_ms": 1.7748308668766792,
      "q1_ms": 1.7533158665173687,
      "q3_ms": 1.777158666664036,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 64,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007147961296141148,
      "xla_temp_mb": 8.250244140625,
      "xla_arg_mb": 32.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 2.205003975999716,
      "inner": 25,
      "median_ms": 1.7347263198462315,
      "q1_ms": 1.7083902400918305,
      "q3_ms": 1.7397987200820353,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 64,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006008640630170703,
      "xla_temp_mb": 16.0,
      "xla_arg_mb": 32.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 3.3842096439984743,
      "inner": 83,
      "median_ms": 0.48986842169591893,
      "q1_ms": 0.48206756019429026,
      "q3_ms": 0.4999655964156663,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 128,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00046745306462980807,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.4711865110002691,
      "inner": 200,
      "median_ms": 0.10379982497397577,
      "q1_ms": 0.10290833250110154,
      "q3_ms": 0.10407626998130581,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 128,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006089538801461458,
      "xla_temp_mb": 64.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.9829242259947932,
      "inner": 96,
      "median_ms": 0.40934927081555844,
      "q1_ms": 0.40864795312245406,
      "q3_ms": 0.4143088176912594,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 128,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00046745306462980807,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.39544508200197015,
      "inner": 16,
      "median_ms": 0.11535156227182597,
      "q1_ms": 0.10210490631834546,
      "q3_ms": 0.15140603113650286,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 128,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00046745306462980807,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.4588150389972725,
      "inner": 200,
      "median_ms": 0.10437060002004728,
      "q1_ms": 0.09919693000483676,
      "q3_ms": 0.10439697751280619,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 128,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006089538801461458,
      "xla_temp_mb": 64.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.24822332299663685,
      "inner": 75,
      "median_ms": 0.4082383733475581,
      "q1_ms": 0.40806777993566357,
      "q3_ms": 0.4092252400005236,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 128,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.000824376184027642,
      "xla_temp_mb": 16.500244140625,
      "xla_arg_mb": 64.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 1.2739684999978635,
      "inner": 35,
      "median_ms": 1.2284303142223507,
      "q1_ms": 1.2277968287110395,
      "q3_ms": 1.2325292143421913,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 128,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007913450244814157,
      "xla_temp_mb": 96.0,
      "xla_arg_mb": 64.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 1.5299878420046298,
      "inner": 40,
      "median_ms": 0.9818894750424079,
      "q1_ms": 0.9798468624467205,
      "q3_ms": 0.9826064500884968,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 128,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.000824376184027642,
      "xla_temp_mb": 16.500244140625,
      "xla_arg_mb": 64.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 1.3668647070007864,
      "inner": 36,
      "median_ms": 1.3944342222708253,
      "q1_ms": 1.254967388983156,
      "q3_ms": 1.4034071526920244,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 128,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.000824376184027642,
      "xla_temp_mb": 16.500244140625,
      "xla_arg_mb": 64.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 1.2241939529994852,
      "inner": 37,
      "median_ms": 1.2377122432037533,
      "q1_ms": 1.2290302566307711,
      "q3_ms": 1.2564848242421962,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 128,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007913450244814157,
      "xla_temp_mb": 96.0,
      "xla_arg_mb": 64.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.5422572160023265,
      "inner": 45,
      "median_ms": 0.9965743111226605,
      "q1_ms": 0.9944682221329357,
      "q3_ms": 1.005821844430304,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 128,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0003490043163765222,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.43386565199762117,
      "inner": 123,
      "median_ms": 0.0936276666494465,
      "q1_ms": 0.09360078455926661,
      "q3_ms": 0.09402391059836183,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 128,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0005110679776407778,
      "xla_temp_mb": 64.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 2.1666581370009226,
      "inner": 84,
      "median_ms": 0.47810851191614,
      "q1_ms": 0.4773396250086426,
      "q3_ms": 0.4783533868682016,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 128,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0003490043163765222,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.43862113599607255,
      "inner": 200,
      "median_ms": 0.0937389699902269,
      "q1_ms": 0.09192772000460536,
      "q3_ms": 0.09386651501699816,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 128,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0003490043163765222,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.4306580490010674,
      "inner": 200,
      "median_ms": 0.09372550000989577,
      "q1_ms": 0.09261131499442854,
      "q3_ms": 0.0938245074939914,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 128,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0005110679776407778,
      "xla_temp_mb": 64.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 1.7967323499979102,
      "inner": 81,
      "median_ms": 0.4774769629785837,
      "q1_ms": 0.4761894506597777,
      "q3_ms": 0.47826300614821027,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 128,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006633914308622479,
      "xla_temp_mb": 16.500244140625,
      "xla_arg_mb": 64.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 1.582500952994451,
      "inner": 9,
      "median_ms": 4.7444581117209355,
      "q1_ms": 4.706165000243345,
      "q3_ms": 4.754970722084787,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 128,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006111591355875134,
      "xla_temp_mb": 96.0,
      "xla_arg_mb": 64.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 4.714283986999362,
      "inner": 36,
      "median_ms": 1.2389704166303597,
      "q1_ms": 1.2381815555878952,
      "q3_ms": 1.2582306250502975,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 128,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006633914308622479,
      "xla_temp_mb": 16.500244140625,
      "xla_arg_mb": 64.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 1.7980439429957187,
      "inner": 10,
      "median_ms": 4.675815499649616,
      "q1_ms": 4.665422399557428,
      "q3_ms": 4.695081499812659,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 128,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006633914308622479,
      "xla_temp_mb": 16.500244140625,
      "xla_arg_mb": 64.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 1.9000445870042313,
      "inner": 10,
      "median_ms": 4.691758700209903,
      "q1_ms": 4.6876374501152895,
      "q3_ms": 4.751384100018186,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 128,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006111591355875134,
      "xla_temp_mb": 96.0,
      "xla_arg_mb": 64.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 6.619392020002124,
      "inner": 32,
      "median_ms": 1.4632531563165685,
      "q1_ms": 1.4589085000125124,
      "q3_ms": 1.4667414532141265,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 256,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0003583835787139833,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 1.2578443919992424,
      "inner": 98,
      "median_ms": 0.32926307139117555,
      "q1_ms": 0.3276141326249653,
      "q3_ms": 0.34575818363566435,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 256,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0004445998929440975,
      "xla_temp_mb": 256.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 2.4188444429964875,
      "inner": 36,
      "median_ms": 1.1630408332873936,
      "q1_ms": 1.1512387500260957,
      "q3_ms": 1.1743469583759887,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 256,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0003583835787139833,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.653794415993616,
      "inner": 78,
      "median_ms": 0.3408316923200917,
      "q1_ms": 0.3289476538693401,
      "q3_ms": 0.34605516665555436,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 256,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0003583835787139833,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.5927285430007032,
      "inner": 93,
      "median_ms": 0.3801058709961913,
      "q1_ms": 0.3428527634241618,
      "q3_ms": 0.3863429570159379,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 256,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0004445998929440975,
      "xla_temp_mb": 256.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.6561099789978471,
      "inner": 34,
      "median_ms": 1.1674897941702511,
      "q1_ms": 1.1491509411665475,
      "q3_ms": 1.193906852944672,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 256,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.000901060935575515,
      "xla_temp_mb": 33.000244140625,
      "xla_arg_mb": 128.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 2.818186402997526,
      "inner": 9,
      "median_ms": 5.330715777947464,
      "q1_ms": 4.90439472196158,
      "q3_ms": 5.384647944285663,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 256,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.000895481207408011,
      "xla_temp_mb": 384.0,
      "xla_arg_mb": 128.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 2.842835750001541,
      "inner": 7,
      "median_ms": 2.846859857007595,
      "q1_ms": 2.803467000506186,
      "q3_ms": 2.940576642556282,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 256,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.000901060935575515,
      "xla_temp_mb": 33.000244140625,
      "xla_arg_mb": 128.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 2.7667059770028573,
      "inner": 9,
      "median_ms": 4.893965555109187,
      "q1_ms": 4.834339944257711,
      "q3_ms": 4.972930611175898,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 256,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.000901060935575515,
      "xla_temp_mb": 33.000244140625,
      "xla_arg_mb": 128.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 2.177869182996801,
      "inner": 7,
      "median_ms": 4.842260571812013,
      "q1_ms": 4.824385285506391,
      "q3_ms": 4.904727356915828,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 256,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.000895481207408011,
      "xla_temp_mb": 384.0,
      "xla_arg_mb": 128.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 0.9193500839974149,
      "inner": 16,
      "median_ms": 2.7823828127111483,
      "q1_ms": 2.7612459064130235,
      "q3_ms": 2.812003218650716,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 256,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0003968327073380351,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.4776372740016086,
      "inner": 139,
      "median_ms": 0.24144856114545168,
      "q1_ms": 0.23965185611565737,
      "q3_ms": 0.24262326978910848,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 256,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0005120915011502802,
      "xla_temp_mb": 256.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 2.7253188799950294,
      "inner": 32,
      "median_ms": 1.3008012185764528,
      "q1_ms": 1.2928293906497856,
      "q3_ms": 1.3101202499683495,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 256,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0003968327073380351,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.49959794099413557,
      "inner": 145,
      "median_ms": 0.23865120687144648,
      "q1_ms": 0.2362880827698471,
      "q3_ms": 0.24065828276497056,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 256,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0003968327073380351,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.5933273260015994,
      "inner": 68,
      "median_ms": 0.23508905884655715,
      "q1_ms": 0.23421358090512084,
      "q3_ms": 0.2691305588300833,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 256,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0005120915011502802,
      "xla_temp_mb": 256.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 2.787751201001811,
      "inner": 33,
      "median_ms": 1.3242653636332375,
      "q1_ms": 1.316462575756084,
      "q3_ms": 1.3316648333839662,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 256,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.000795340514741838,
      "xla_temp_mb": 33.000244140625,
      "xla_arg_mb": 128.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 1.776009144996351,
      "inner": 3,
      "median_ms": 15.142742665678574,
      "q1_ms": 14.996263499294098,
      "q3_ms": 15.310660500593562,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 256,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0005537637043744326,
      "xla_temp_mb": 384.0,
      "xla_arg_mb": 128.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 7.845223955002439,
      "inner": 12,
      "median_ms": 3.7506084166428386,
      "q1_ms": 3.7055029582309844,
      "q3_ms": 3.8151960834511556,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 256,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.000795340514741838,
      "xla_temp_mb": 33.000244140625,
      "xla_arg_mb": 128.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 1.9118788330015377,
      "inner": 3,
      "median_ms": 15.56146499933675,
      "q1_ms": 14.988485500604536,
      "q3_ms": 15.791357167472597,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 256,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.000795340514741838,
      "xla_temp_mb": 33.000244140625,
      "xla_arg_mb": 128.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 2.2316370690023177,
      "inner": 2,
      "median_ms": 17.382464502588846,
      "q1_ms": 17.151294749055523,
      "q3_ms": 17.845942249550717,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 256,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0005537637043744326,
      "xla_temp_mb": 384.0,
      "xla_arg_mb": 128.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 7.574196869005391,
      "inner": 12,
      "median_ms": 3.7264125833947523,
      "q1_ms": 3.691865500210649,
      "q3_ms": 3.9485310417148867,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 512,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006677013006992638,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 192.0,
      "xla_out_mb": 64.0,
      "compile_s": 0.3365042009972967,
      "inner": 38,
      "median_ms": 1.2456159209633417,
      "q1_ms": 1.241144684076732,
      "q3_ms": 1.2637266316137117,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 512,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0008860177476890385,
      "xla_temp_mb": 1024.0,
      "xla_arg_mb": 192.0,
      "xla_out_mb": 64.0,
      "compile_s": 1.1153312270034803,
      "inner": 13,
      "median_ms": 3.5480443078155917,
      "q1_ms": 3.5461592306091916,
      "q3_ms": 3.6032739229155526,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 512,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006677013006992638,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 192.0,
      "xla_out_mb": 64.0,
      "compile_s": 0.40053462800278794,
      "inner": 39,
      "median_ms": 1.2537315898300267,
      "q1_ms": 1.2529494743853902,
      "q3_ms": 1.4283951025460668,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 512,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006677013006992638,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 192.0,
      "xla_out_mb": 64.0,
      "compile_s": 0.39488640499621397,
      "inner": 40,
      "median_ms": 1.2618269749509636,
      "q1_ms": 1.2423329250850657,
      "q3_ms": 1.3361999000153446,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 512,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0008860177476890385,
      "xla_temp_mb": 1024.0,
      "xla_arg_mb": 192.0,
      "xla_out_mb": 64.0,
      "compile_s": 0.38082623400259763,
      "inner": 13,
      "median_ms": 3.5304001536972534,
      "q1_ms": 3.5289327693485224,
      "q3_ms": 3.576825384908051,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 512,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.001239205477759242,
      "xla_temp_mb": 66.000244140625,
      "xla_arg_mb": 256.0,
      "xla_out_mb": 192.0000228881836,
      "compile_s": 1.3941961119999178,
      "inner": 2,
      "median_ms": 19.771247501921607,
      "q1_ms": 19.708235498910653,
      "q3_ms": 19.798584002273856,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 512,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0012390939518809319,
      "xla_temp_mb": 1537.0,
      "xla_arg_mb": 256.0,
      "xla_out_mb": 192.0000228881836,
      "compile_s": 2.669990319001954,
      "inner": 5,
      "median_ms": 9.626222999941092,
      "q1_ms": 9.248062100232346,
      "q3_ms": 11.329863400169415,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 512,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.001239205477759242,
      "xla_temp_mb": 66.000244140625,
      "xla_arg_mb": 256.0,
      "xla_out_mb": 192.0000228881836,
      "compile_s": 1.1760434719981276,
      "inner": 2,
      "median_ms": 17.3236580012599,
      "q1_ms": 17.283962999499636,
      "q3_ms": 17.357058999550645,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 512,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.001239205477759242,
      "xla_temp_mb": 66.000244140625,
      "xla_arg_mb": 256.0,
      "xla_out_mb": 192.0000228881836,
      "compile_s": 1.2085555150042637,
      "inner": 2,
      "median_ms": 19.62096249917522,
      "q1_ms": 19.593646000430454,
      "q3_ms": 19.683260750753107,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 512,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0012390939518809319,
      "xla_temp_mb": 1537.0,
      "xla_arg_mb": 256.0,
      "xla_out_mb": 192.0000228881836,
      "compile_s": 0.7747064940049313,
      "inner": 4,
      "median_ms": 9.598212000128115,
      "q1_ms": 9.461710999858042,
      "q3_ms": 9.887390749099723,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 512,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0004773139953613281,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 192.0,
      "xla_out_mb": 64.0,
      "compile_s": 0.3683336030007922,
      "inner": 64,
      "median_ms": 0.7708877657250923,
      "q1_ms": 0.7589942422328022,
      "q3_ms": 0.7731348281936334,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 512,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0005617492133751512,
      "xla_temp_mb": 1024.0,
      "xla_arg_mb": 192.0,
      "xla_out_mb": 64.0,
      "compile_s": 1.0439839679966099,
      "inner": 13,
      "median_ms": 3.611713384788555,
      "q1_ms": 3.5948009230756166,
      "q3_ms": 3.633654461457179,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 512,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0004773139953613281,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 192.0,
      "xla_out_mb": 64.0,
      "compile_s": 0.38002774299820885,
      "inner": 63,
      "median_ms": 0.7684809681940972,
      "q1_ms": 0.7614401746693711,
      "q3_ms": 0.770367000003471,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 512,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0004773139953613281,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 192.0,
      "xla_out_mb": 64.0,
      "compile_s": 0.4311806389960111,
      "inner": 62,
      "median_ms": 0.7878574192985862,
      "q1_ms": 0.7778500160883037,
      "q3_ms": 0.8300269434966463,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 512,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0005617492133751512,
      "xla_temp_mb": 1024.0,
      "xla_arg_mb": 192.0,
      "xla_out_mb": 64.0,
      "compile_s": 1.2334547560021747,
      "inner": 13,
      "median_ms": 3.605567769017608,
      "q1_ms": 3.596714884714367,
      "q3_ms": 3.6483133844740223,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 512,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007104802643880248,
      "xla_temp_mb": 66.000244140625,
      "xla_arg_mb": 256.0,
      "xla_out_mb": 192.0000228881836,
      "compile_s": 1.4031557249982143,
      "inner": 1,
      "median_ms": 51.90352799399989,
      "q1_ms": 51.52686399742379,
      "q3_ms": 52.443820000917185,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 512,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007121484959498048,
      "xla_temp_mb": 1537.0,
      "xla_arg_mb": 256.0,
      "xla_out_mb": 192.0000228881836,
      "compile_s": 1.4348004869971192,
      "inner": 5,
      "median_ms": 9.506875000079162,
      "q1_ms": 9.5046780006669,
      "q3_ms": 9.849194100388559,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 512,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007104802643880248,
      "xla_temp_mb": 66.000244140625,
      "xla_arg_mb": 256.0,
      "xla_out_mb": 192.0000228881836,
      "compile_s": 1.510731112997746,
      "inner": 1,
      "median_ms": 51.82351600524271,
      "q1_ms": 51.46414050250314,
      "q3_ms": 52.1251165009744,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 512,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007104802643880248,
      "xla_temp_mb": 66.000244140625,
      "xla_arg_mb": 256.0,
      "xla_out_mb": 192.0000228881836,
      "compile_s": 1.4518520400015404,
      "inner": 1,
      "median_ms": 52.584025004762225,
      "q1_ms": 52.48378100077389,
      "q3_ms": 52.826253999228356,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 512,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007121484959498048,
      "xla_temp_mb": 1537.0,
      "xla_arg_mb": 256.0,
      "xla_out_mb": 192.0000228881836,
      "compile_s": 1.7771444839963806,
      "inner": 2,
      "median_ms": 13.790919998427853,
      "q1_ms": 10.933171999568003,
      "q3_ms": 20.644309999624966,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 16,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004193125292658806,
      "xla_temp_mb": 1.0,
      "xla_arg_mb": 1.5,
      "xla_out_mb": 0.5,
      "compile_s": 0.37611156800267054,
      "inner": 200,
      "median_ms": 0.039961615002539475,
      "q1_ms": 0.0397805674947449,
      "q3_ms": 0.04097593498954666,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 16,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004193125292658806,
      "xla_temp_mb": 1.0,
      "xla_arg_mb": 1.5,
      "xla_out_mb": 0.5,
      "compile_s": 0.12871147099940572,
      "inner": 200,
      "median_ms": 0.043615145004878286,
      "q1_ms": 0.04230753500451101,
      "q3_ms": 0.04656749249988934,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 16,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0033779738005250692,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 1.5,
      "xla_out_mb": 0.5,
      "compile_s": 0.10146517099929042,
      "inner": 200,
      "median_ms": 0.039994375001697335,
      "q1_ms": 0.037234010014799424,
      "q3_ms": 0.04112872751647956,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 16,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0033779738005250692,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 1.5,
      "xla_out_mb": 0.5,
      "compile_s": 0.09751081000285922,
      "inner": 200,
      "median_ms": 0.0359228000161238,
      "q1_ms": 0.03543264749168884,
      "q3_ms": 0.03663900748506421,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 16,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004193125292658806,
      "xla_temp_mb": 1.0,
      "xla_arg_mb": 1.5,
      "xla_out_mb": 0.5,
      "compile_s": 0.1318154999971739,
      "inner": 200,
      "median_ms": 0.051133255001332145,
      "q1_ms": 0.04853790749621112,
      "q3_ms": 0.05528242500076885,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 16,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0033779738005250692,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 1.5,
      "xla_out_mb": 0.5,
      "compile_s": 0.24430164899968076,
      "inner": 200,
      "median_ms": 0.06086997502279701,
      "q1_ms": 0.06006660500133876,
      "q3_ms": 0.06099957499827724,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 16,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.006359456572681665,
      "xla_temp_mb": 1.0,
      "xla_arg_mb": 2.5,
      "xla_out_mb": 1.5000228881835938,
      "compile_s": 0.37146185700112255,
      "inner": 200,
      "median_ms": 0.05477145499753533,
      "q1_ms": 0.0536535374885716,
      "q3_ms": 0.05589804499322781,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 16,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.006359456572681665,
      "xla_temp_mb": 1.0,
      "xla_arg_mb": 2.5,
      "xla_out_mb": 1.5000228881835938,
      "compile_s": 0.21941811199940275,
      "inner": 200,
      "median_ms": 0.05367778499930864,
      "q1_ms": 0.05313154500981909,
      "q3_ms": 0.05615230251351022,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 16,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0043426030315458775,
      "xla_temp_mb": 1.062744140625,
      "xla_arg_mb": 2.5,
      "xla_out_mb": 1.5000228881835938,
      "compile_s": 0.18508606000250438,
      "inner": 200,
      "median_ms": 0.049141725030494854,
      "q1_ms": 0.04883048748524743,
      "q3_ms": 0.054868390016054036,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 16,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0043426030315458775,
      "xla_temp_mb": 1.062744140625,
      "xla_arg_mb": 2.5,
      "xla_out_mb": 1.5000228881835938,
      "compile_s": 0.1783898869980476,
      "inner": 200,
      "median_ms": 0.05007224499422591,
      "q1_ms": 0.04901638749288395,
      "q3_ms": 0.05018869000195991,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 16,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.006359456572681665,
      "xla_temp_mb": 1.0,
      "xla_arg_mb": 2.5,
      "xla_out_mb": 1.5000228881835938,
      "compile_s": 0.2140672910027206,
      "inner": 200,
      "median_ms": 0.058025380021717865,
      "q1_ms": 0.0545039025018923,
      "q3_ms": 0.061685927503276616,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 16,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0043426030315458775,
      "xla_temp_mb": 3.063232421875,
      "xla_arg_mb": 2.5,
      "xla_out_mb": 1.5000228881835938,
      "compile_s": 0.5553503259943682,
      "inner": 200,
      "median_ms": 0.10998726000252645,
      "q1_ms": 0.1096725849856739,
      "q3_ms": 0.11253246249907534,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 16,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00428376067429781,
      "xla_temp_mb": 1.0,
      "xla_arg_mb": 1.5,
      "xla_out_mb": 0.5,
      "compile_s": 0.22188099400227657,
      "inner": 200,
      "median_ms": 0.04342698499385733,
      "q1_ms": 0.04300000251532765,
      "q3_ms": 0.043928762497671414,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 16,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00428376067429781,
      "xla_temp_mb": 1.0,
      "xla_arg_mb": 1.5,
      "xla_out_mb": 0.5,
      "compile_s": 0.22285494200332323,
      "inner": 200,
      "median_ms": 0.04444305999641074,
      "q1_ms": 0.04225686750942259,
      "q3_ms": 0.04586796749208588,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 16,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0031301137059926987,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 1.5,
      "xla_out_mb": 0.5,
      "compile_s": 0.12914969199482584,
      "inner": 200,
      "median_ms": 0.03587766499549616,
      "q1_ms": 0.035823984999296954,
      "q3_ms": 0.03599178249714896,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 16,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0031301137059926987,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 1.5,
      "xla_out_mb": 0.5,
      "compile_s": 0.11831440700188978,
      "inner": 200,
      "median_ms": 0.038266695009951945,
      "q1_ms": 0.037302617529348936,
      "q3_ms": 0.040449162497679936,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 16,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00428376067429781,
      "xla_temp_mb": 1.0,
      "xla_arg_mb": 1.5,
      "xla_out_mb": 0.5,
      "compile_s": 0.22259971799940104,
      "inner": 200,
      "median_ms": 0.05095436998090008,
      "q1_ms": 0.04672827248214162,
      "q3_ms": 0.0544924525092938,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 16,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0031301137059926987,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 1.5,
      "xla_out_mb": 0.5,
      "compile_s": 0.23902049000025727,
      "inner": 200,
      "median_ms": 0.05919364499277435,
      "q1_ms": 0.058634062497731065,
      "q3_ms": 0.06252701999983401,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 16,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00491351168602705,
      "xla_temp_mb": 1.0,
      "xla_arg_mb": 2.5,
      "xla_out_mb": 1.5000228881835938,
      "compile_s": 0.46715618999587605,
      "inner": 200,
      "median_ms": 0.05912005999562098,
      "q1_ms": 0.05753913248554454,
      "q3_ms": 0.06589603999600513,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 16,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00491351168602705,
      "xla_temp_mb": 1.0,
      "xla_arg_mb": 2.5,
      "xla_out_mb": 1.5000228881835938,
      "compile_s": 0.4725989060025313,
      "inner": 200,
      "median_ms": 0.054967779979051556,
      "q1_ms": 0.054853314995853,
      "q3_ms": 0.06385677001162549,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 16,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004201397765427828,
      "xla_temp_mb": 1.062744140625,
      "xla_arg_mb": 2.5,
      "xla_out_mb": 1.5000228881835938,
      "compile_s": 0.24484393800230464,
      "inner": 200,
      "median_ms": 0.0521516899971175,
      "q1_ms": 0.051479562498570886,
      "q3_ms": 0.05459852749481797,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 16,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004201397765427828,
      "xla_temp_mb": 1.062744140625,
      "xla_arg_mb": 2.5,
      "xla_out_mb": 1.5000228881835938,
      "compile_s": 0.23688799500087043,
      "inner": 200,
      "median_ms": 0.04917504000331974,
      "q1_ms": 0.04903815748548368,
      "q3_ms": 0.04960640750141465,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 16,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00491351168602705,
      "xla_temp_mb": 1.0,
      "xla_arg_mb": 2.5,
      "xla_out_mb": 1.5000228881835938,
      "compile_s": 0.49292818299727514,
      "inner": 200,
      "median_ms": 0.060702404989569914,
      "q1_ms": 0.05789976999949431,
      "q3_ms": 0.07049947249470279,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 16,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004201397765427828,
      "xla_temp_mb": 3.063232421875,
      "xla_arg_mb": 2.5,
      "xla_out_mb": 1.5000228881835938,
      "compile_s": 0.5538295789956464,
      "inner": 200,
      "median_ms": 0.11551128500286723,
      "q1_ms": 0.11506073500640923,
      "q3_ms": 0.11621737750829197,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 32,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004176608286798,
      "xla_temp_mb": 3.0,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.28474260000075446,
      "inner": 200,
      "median_ms": 0.041489210016152356,
      "q1_ms": 0.040992570011439966,
      "q3_ms": 0.04201235251457547,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 32,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004176608286798,
      "xla_temp_mb": 3.0,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.15701668000110658,
      "inner": 200,
      "median_ms": 0.046540740004274994,
      "q1_ms": 0.04518461750194547,
      "q3_ms": 0.04944698250255897,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 32,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0022217121440917253,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.10635681099665817,
      "inner": 200,
      "median_ms": 0.0365201849854202,
      "q1_ms": 0.03610161498727393,
      "q3_ms": 0.037424672500492306,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 32,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0022217121440917253,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.10599797999748262,
      "inner": 200,
      "median_ms": 0.03519892499753041,
      "q1_ms": 0.034827394993044436,
      "q3_ms": 0.03555399749529897,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 32,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004176608286798,
      "xla_temp_mb": 3.0,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.1517363219973049,
      "inner": 200,
      "median_ms": 0.04352296997240046,
      "q1_ms": 0.04244211748300586,
      "q3_ms": 0.048338257474824786,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 32,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0022217121440917253,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.050084797003364656,
      "inner": 200,
      "median_ms": 0.06712759997753892,
      "q1_ms": 0.06634982501054765,
      "q3_ms": 0.0673294575062755,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 32,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.006654238328337669,
      "xla_temp_mb": 5.0625,
      "xla_arg_mb": 5.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.5550410290015861,
      "inner": 200,
      "median_ms": 0.0624007849910413,
      "q1_ms": 0.0623560725034622,
      "q3_ms": 0.06411936750737368,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 32,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.006654238328337669,
      "xla_temp_mb": 5.0625,
      "xla_arg_mb": 5.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.318913426999643,
      "inner": 200,
      "median_ms": 0.06298889500612859,
      "q1_ms": 0.060146077503304696,
      "q3_ms": 0.06548030001795269,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 32,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.006654238328337669,
      "xla_temp_mb": 2.125244140625,
      "xla_arg_mb": 5.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.22255356000096072,
      "inner": 200,
      "median_ms": 0.05669931499141967,
      "q1_ms": 0.050811319997592364,
      "q3_ms": 0.060191115026100306,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 32,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.006654238328337669,
      "xla_temp_mb": 2.125244140625,
      "xla_arg_mb": 5.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.2206821400031913,
      "inner": 200,
      "median_ms": 0.05236684501142008,
      "q1_ms": 0.050454265019652667,
      "q3_ms": 0.05427699248684803,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 32,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.006654238328337669,
      "xla_temp_mb": 5.0625,
      "xla_arg_mb": 5.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.2998298120000982,
      "inner": 200,
      "median_ms": 0.059013505015172996,
      "q1_ms": 0.058699735018308274,
      "q3_ms": 0.06255828000576003,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 32,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.006654238328337669,
      "xla_temp_mb": 6.125732421875,
      "xla_arg_mb": 5.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.11170874500385253,
      "inner": 200,
      "median_ms": 0.11799024501669919,
      "q1_ms": 0.11362077750163735,
      "q3_ms": 0.13106857502862113,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 32,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004325577989220619,
      "xla_temp_mb": 3.0,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.3727008170026238,
      "inner": 200,
      "median_ms": 0.04504149499553023,
      "q1_ms": 0.04389239251395338,
      "q3_ms": 0.04740956501336768,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 32,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004325577989220619,
      "xla_temp_mb": 3.0,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.3599412410039804,
      "inner": 200,
      "median_ms": 0.03845655999612063,
      "q1_ms": 0.038156062491907505,
      "q3_ms": 0.03854074000628316,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 32,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0032848461996763945,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.13606683599937242,
      "inner": 200,
      "median_ms": 0.036104745013290085,
      "q1_ms": 0.03525615749822464,
      "q3_ms": 0.0378101049864199,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 32,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0032848461996763945,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.1290764030054561,
      "inner": 200,
      "median_ms": 0.04279438500816468,
      "q1_ms": 0.03886302751197945,
      "q3_ms": 0.043603720023384085,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 32,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004325577989220619,
      "xla_temp_mb": 3.0,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.3559495089939446,
      "inner": 200,
      "median_ms": 0.041712134989211336,
      "q1_ms": 0.03932162500859704,
      "q3_ms": 0.04462057500859373,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 32,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0032848461996763945,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.045902812998974696,
      "inner": 200,
      "median_ms": 0.061712025017186534,
      "q1_ms": 0.06076770749132265,
      "q3_ms": 0.06324282749119448,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 32,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004901744890958071,
      "xla_temp_mb": 5.0625,
      "xla_arg_mb": 5.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.6950385999953141,
      "inner": 200,
      "median_ms": 0.06417625001631677,
      "q1_ms": 0.06415576248400612,
      "q3_ms": 0.0644604650005931,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 32,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004901744890958071,
      "xla_temp_mb": 5.0625,
      "xla_arg_mb": 5.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.7203562299982877,
      "inner": 200,
      "median_ms": 0.06104029998823534,
      "q1_ms": 0.060454137510532746,
      "q3_ms": 0.061882722493464826,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 32,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.005748792551457882,
      "xla_temp_mb": 2.125244140625,
      "xla_arg_mb": 5.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.29720209800143493,
      "inner": 200,
      "median_ms": 0.05054113000369398,
      "q1_ms": 0.050356782503513386,
      "q3_ms": 0.05438397500256542,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 32,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.005748792551457882,
      "xla_temp_mb": 2.125244140625,
      "xla_arg_mb": 5.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.30440132600051584,
      "inner": 200,
      "median_ms": 0.05091513998195296,
      "q1_ms": 0.05062914749942138,
      "q3_ms": 0.05474022500493447,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 32,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004901744890958071,
      "xla_temp_mb": 5.0625,
      "xla_arg_mb": 5.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.715414446000068,
      "inner": 200,
      "median_ms": 0.06514739499834832,
      "q1_ms": 0.06429144499634276,
      "q3_ms": 0.06640124500336242,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 32,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.005748792551457882,
      "xla_temp_mb": 6.125732421875,
      "xla_arg_mb": 5.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.1207195690003573,
      "inner": 179,
      "median_ms": 0.11843865918283034,
      "q1_ms": 0.11635699998632895,
      "q3_ms": 0.12754798044421967,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 64,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004283299203962088,
      "xla_temp_mb": 12.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.358806418000313,
      "inner": 200,
      "median_ms": 0.059076075012853835,
      "q1_ms": 0.059050477502751164,
      "q3_ms": 0.059282517486280994,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 64,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004283299203962088,
      "xla_temp_mb": 12.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.199551022000378,
      "inner": 200,
      "median_ms": 0.05902021497604437,
      "q1_ms": 0.059001407516916515,
      "q3_ms": 0.05908864748562337,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 64,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0027185636572539806,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.13205183800164377,
      "inner": 200,
      "median_ms": 0.03657323501101928,
      "q1_ms": 0.03628403499533306,
      "q3_ms": 0.04176591250143247,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 64,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0027185636572539806,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.12831046999781393,
      "inner": 200,
      "median_ms": 0.0374508499953663,
      "q1_ms": 0.036964025002816925,
      "q3_ms": 0.038395940009650076,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 64,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004283299203962088,
      "xla_temp_mb": 12.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.18845042300381465,
      "inner": 200,
      "median_ms": 0.05906079499254702,
      "q1_ms": 0.058994470000470756,
      "q3_ms": 0.05910948750170064,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 64,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0027185636572539806,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.2022887150014867,
      "inner": 200,
      "median_ms": 0.05668899499141844,
      "q1_ms": 0.05626697000479908,
      "q3_ms": 0.05752125001890818,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 64,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004668043926358223,
      "xla_temp_mb": 20.125,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.6844836759992177,
      "inner": 194,
      "median_ms": 0.15530834537354712,
      "q1_ms": 0.15525605670415277,
      "q3_ms": 0.15533910309667842,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 64,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004668043926358223,
      "xla_temp_mb": 20.125,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.3414190420007799,
      "inner": 200,
      "median_ms": 0.15523670499533182,
      "q1_ms": 0.15519690248765983,
      "q3_ms": 0.15536632499788539,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 64,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004016230348497629,
      "xla_temp_mb": 4.250244140625,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.2758369350049179,
      "inner": 200,
      "median_ms": 0.05781011997896712,
      "q1_ms": 0.05755248252171441,
      "q3_ms": 0.05787629750557244,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 64,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004016230348497629,
      "xla_temp_mb": 4.250244140625,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.26794437199714594,
      "inner": 200,
      "median_ms": 0.05811445000290405,
      "q1_ms": 0.057482124993839534,
      "q3_ms": 0.0583206900046207,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 64,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004668043926358223,
      "xla_temp_mb": 20.125,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.3502316609956324,
      "inner": 200,
      "median_ms": 0.15531901000940707,
      "q1_ms": 0.15522517751378473,
      "q3_ms": 0.15538833999016788,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 64,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004016230348497629,
      "xla_temp_mb": 12.250732421875,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.5276581099969917,
      "inner": 164,
      "median_ms": 0.13459271341583565,
      "q1_ms": 0.1343151463471458,
      "q3_ms": 0.13645054267970674,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 64,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004315928090363741,
      "xla_temp_mb": 12.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 1.2901892910012975,
      "inner": 200,
      "median_ms": 0.08696693999809213,
      "q1_ms": 0.0860305099922698,
      "q3_ms": 0.08821903749776538,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 64,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004315928090363741,
      "xla_temp_mb": 12.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 1.3189787180017447,
      "inner": 200,
      "median_ms": 0.08275474501715507,
      "q1_ms": 0.08132826500514057,
      "q3_ms": 0.08441764748567948,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 64,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0032096244394779205,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.1678374379989691,
      "inner": 200,
      "median_ms": 0.03769992501474917,
      "q1_ms": 0.037545827490248485,
      "q3_ms": 0.04232692250297987,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 64,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0032096244394779205,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.17947812099737348,
      "inner": 200,
      "median_ms": 0.037383519993454684,
      "q1_ms": 0.03725572749317507,
      "q3_ms": 0.037651855000149226,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 64,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004315928090363741,
      "xla_temp_mb": 12.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 1.3371033929943223,
      "inner": 200,
      "median_ms": 0.08814556997094769,
      "q1_ms": 0.08729157751076855,
      "q3_ms": 0.08852509250573348,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 64,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0032096244394779205,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.04846235900186002,
      "inner": 200,
      "median_ms": 0.05992528502247296,
      "q1_ms": 0.05980743748295936,
      "q3_ms": 0.06154518498078687,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 64,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004813428036868572,
      "xla_temp_mb": 20.125,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.8628899230025127,
      "inner": 192,
      "median_ms": 0.158410104139269,
      "q1_ms": 0.1582579140517737,
      "q3_ms": 0.16099923436740937,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 64,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004813428036868572,
      "xla_temp_mb": 20.125,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.8614498020033352,
      "inner": 200,
      "median_ms": 0.16012497500923928,
      "q1_ms": 0.15927676498904475,
      "q3_ms": 0.16079644999990705,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 64,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0061174426227808,
      "xla_temp_mb": 4.250244140625,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.4504081280028913,
      "inner": 200,
      "median_ms": 0.060178970015840605,
      "q1_ms": 0.05986048749036854,
      "q3_ms": 0.06205219749972457,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 64,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0061174426227808,
      "xla_temp_mb": 4.250244140625,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.37452449800184695,
      "inner": 200,
      "median_ms": 0.06094957501773024,
      "q1_ms": 0.06024807002177113,
      "q3_ms": 0.062103190011839615,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 64,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004813428036868572,
      "xla_temp_mb": 20.125,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.924103772995295,
      "inner": 181,
      "median_ms": 0.1627765966946573,
      "q1_ms": 0.1602052541392423,
      "q3_ms": 0.1675107955960917,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 64,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0061174426227808,
      "xla_temp_mb": 12.250732421875,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.3642144570039818,
      "inner": 174,
      "median_ms": 0.13828101727158773,
      "q1_ms": 0.13724925576815591,
      "q3_ms": 0.13877557182564781,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 128,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0025094568263739347,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.33409959599521244,
      "inner": 200,
      "median_ms": 0.04410967998410342,
      "q1_ms": 0.041981494996434776,
      "q3_ms": 0.047492849989794195,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 128,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004196838941425085,
      "xla_temp_mb": 48.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.5360445599944796,
      "inner": 138,
      "median_ms": 0.18849055798626438,
      "q1_ms": 0.18754002171656303,
      "q3_ms": 0.1903259963791757,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 128,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0025094568263739347,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.2723486880058772,
      "inner": 200,
      "median_ms": 0.03395285999431508,
      "q1_ms": 0.0337487475189846,
      "q3_ms": 0.034630047503014794,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 128,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0025094568263739347,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.27042957700177794,
      "inner": 200,
      "median_ms": 0.03513330000714632,
      "q1_ms": 0.03312268749141367,
      "q3_ms": 0.037134572503418894,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 128,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004196838941425085,
      "xla_temp_mb": 48.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.2632459579981514,
      "inner": 164,
      "median_ms": 0.18462983537507927,
      "q1_ms": 0.1834125914683682,
      "q3_ms": 0.19312986280899014,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 128,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.002333720214664936,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.05093724399921484,
      "inner": 200,
      "median_ms": 0.07103250496584224,
      "q1_ms": 0.07024153001111699,
      "q3_ms": 0.0718836200030637,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 128,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.005233302712440491,
      "xla_temp_mb": 8.500244140625,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.3635326840012567,
      "inner": 188,
      "median_ms": 0.1526363244688292,
      "q1_ms": 0.15253274999308136,
      "q3_ms": 0.15389037765844282,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 128,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.005233302712440491,
      "xla_temp_mb": 80.25,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.9989973249976174,
      "inner": 80,
      "median_ms": 0.5124826125211257,
      "q1_ms": 0.5120043999795598,
      "q3_ms": 0.5186491312542785,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 128,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.005233302712440491,
      "xla_temp_mb": 8.500244140625,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.36432480700023007,
      "inner": 162,
      "median_ms": 0.15254896914731586,
      "q1_ms": 0.15243891667287918,
      "q3_ms": 0.15266635183465554,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 128,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.005233302712440491,
      "xla_temp_mb": 8.500244140625,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.3605436739962897,
      "inner": 197,
      "median_ms": 0.15493463961966572,
      "q1_ms": 0.15413688323939181,
      "q3_ms": 0.1556614010281882,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 128,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.005233302712440491,
      "xla_temp_mb": 80.25,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.4876833429952967,
      "inner": 80,
      "median_ms": 0.5227770125202369,
      "q1_ms": 0.5205047125400597,
      "q3_ms": 0.5265338937078923,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 128,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.005233302712440491,
      "xla_temp_mb": 24.500732421875,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.12386923500162084,
      "inner": 93,
      "median_ms": 0.340454376293897,
      "q1_ms": 0.3378003118565977,
      "q3_ms": 0.34233784945856943,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 128,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.003435371909290552,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.3403102379961638,
      "inner": 200,
      "median_ms": 0.03501858500385424,
      "q1_ms": 0.03496515750157414,
      "q3_ms": 0.035193337498640176,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 128,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004306195303797722,
      "xla_temp_mb": 48.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 1.4509770139993634,
      "inner": 131,
      "median_ms": 0.26298532828302945,
      "q1_ms": 0.26287959159537616,
      "q3_ms": 0.26315653056594124,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 128,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.003435371909290552,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.3402767239967943,
      "inner": 191,
      "median_ms": 0.03446758641246999,
      "q1_ms": 0.03410108117078899,
      "q3_ms": 0.036550413600640874,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 128,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.003435371909290552,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.34163229600380873,
      "inner": 200,
      "median_ms": 0.03391288501006784,
      "q1_ms": 0.033766460001061205,
      "q3_ms": 0.03430988997934037,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 128,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004306195303797722,
      "xla_temp_mb": 48.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 1.4581409750026069,
      "inner": 126,
      "median_ms": 0.25752491268259103,
      "q1_ms": 0.2520115079003893,
      "q3_ms": 0.26316246427115175,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 128,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.003435371909290552,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.053930602996842936,
      "inner": 200,
      "median_ms": 0.05741126500652172,
      "q1_ms": 0.05726701749154017,
      "q3_ms": 0.05807943251056713,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 128,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00502068130299449,
      "xla_temp_mb": 8.500244140625,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.41971113099862123,
      "inner": 197,
      "median_ms": 0.1349376700257145,
      "q1_ms": 0.1348672055894613,
      "q3_ms": 0.1350217614257011,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 128,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00502068130299449,
      "xla_temp_mb": 80.25,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 1.020126411000092,
      "inner": 80,
      "median_ms": 0.5422374375484651,
      "q1_ms": 0.5385139313148102,
      "q3_ms": 0.5476627937241574,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 128,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00502068130299449,
      "xla_temp_mb": 8.500244140625,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.4562116969973431,
      "inner": 200,
      "median_ms": 0.13690351999684935,
      "q1_ms": 0.13484235749274376,
      "q3_ms": 0.1370231499822694,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 128,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00502068130299449,
      "xla_temp_mb": 8.500244140625,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.43799177199980477,
      "inner": 200,
      "median_ms": 0.1356959800250479,
      "q1_ms": 0.13459023501127376,
      "q3_ms": 0.13816306251101196,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 128,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00502068130299449,
      "xla_temp_mb": 80.25,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 1.0233567989998846,
      "inner": 78,
      "median_ms": 0.5452228077233602,
      "q1_ms": 0.5392803910254793,
      "q3_ms": 0.5522618717711735,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 128,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00502068130299449,
      "xla_temp_mb": 24.500732421875,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.11486126499949023,
      "inner": 126,
      "median_ms": 0.27424567457591553,
      "q1_ms": 0.2735992896752635,
      "q3_ms": 0.2742981944302012,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 256,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.002204917138442397,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.32919394000055036,
      "inner": 200,
      "median_ms": 0.08912302499084035,
      "q1_ms": 0.08670153250932344,
      "q3_ms": 0.08966983999926015,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 256,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0032324623316526413,
      "xla_temp_mb": 192.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.5591898290003883,
      "inner": 70,
      "median_ms": 0.6233042000011275,
      "q1_ms": 0.6187384571343759,
      "q3_ms": 0.6259414214582648,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 256,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.002204917138442397,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.33092971199948806,
      "inner": 200,
      "median_ms": 0.09799966999707976,
      "q1_ms": 0.09245973000361118,
      "q3_ms": 0.09934912750395597,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 256,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.002204917138442397,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.34039320599549683,
      "inner": 200,
      "median_ms": 0.086280985015037,
      "q1_ms": 0.08394568249059375,
      "q3_ms": 0.08651567999550025,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 256,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0032324623316526413,
      "xla_temp_mb": 192.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.24776743399706902,
      "inner": 70,
      "median_ms": 0.6217110285693448,
      "q1_ms": 0.6196313713839378,
      "q3_ms": 0.6287013643096933,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 256,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.002204917138442397,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.20610574699821882,
      "inner": 155,
      "median_ms": 0.1874048838808742,
      "q1_ms": 0.1766127000014568,
      "q3_ms": 0.19313772579827795,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 256,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0038350424729287624,
      "xla_temp_mb": 17.000244140625,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.39963025599718094,
      "inner": 101,
      "median_ms": 0.43714833664846153,
      "q1_ms": 0.43701238120216784,
      "q3_ms": 0.44030650001737326,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 256,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004831663798540831,
      "xla_temp_mb": 320.5,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 1.0790447749968735,
      "inner": 27,
      "median_ms": 1.8032884814871337,
      "q1_ms": 1.7887308519372405,
      "q3_ms": 1.8493779259895544,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 256,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0038350424729287624,
      "xla_temp_mb": 17.000244140625,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.4084481239988236,
      "inner": 99,
      "median_ms": 0.44591515148622973,
      "q1_ms": 0.4428627424093782,
      "q3_ms": 0.46315537878803703,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 256,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0038350424729287624,
      "xla_temp_mb": 17.000244140625,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.4013573219999671,
      "inner": 90,
      "median_ms": 0.46583051104486606,
      "q1_ms": 0.4427506722600406,
      "q3_ms": 0.49381867219280035,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 256,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004831663798540831,
      "xla_temp_mb": 320.5,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.4386480070024845,
      "inner": 27,
      "median_ms": 1.7947860369868405,
      "q1_ms": 1.7823116295740302,
      "q3_ms": 1.8041920369863942,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 256,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0038350424729287624,
      "xla_temp_mb": 49.000732421875,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.27550850200350396,
      "inner": 50,
      "median_ms": 0.8687142599956132,
      "q1_ms": 0.8632818500336725,
      "q3_ms": 0.8857420600543264,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 256,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.003190818941220641,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.35849946600501426,
      "inner": 200,
      "median_ms": 0.07435375500790542,
      "q1_ms": 0.07433978249537176,
      "q3_ms": 0.07442391499353107,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 256,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004415563773363829,
      "xla_temp_mb": 192.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 1.7005532789989957,
      "inner": 51,
      "median_ms": 0.8568883529412287,
      "q1_ms": 0.843340107844686,
      "q3_ms": 0.8649192842473413,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 256,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.003190818941220641,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.36135778799507534,
      "inner": 200,
      "median_ms": 0.07442434500262607,
      "q1_ms": 0.07160847500927048,
      "q3_ms": 0.0758321875036927,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 256,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.003190818941220641,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.3647763640037738,
      "inner": 200,
      "median_ms": 0.07362262000242481,
      "q1_ms": 0.07105863249307731,
      "q3_ms": 0.07400908250929206,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 256,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004415563773363829,
      "xla_temp_mb": 192.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 1.735719844000414,
      "inner": 45,
      "median_ms": 0.8435276000657015,
      "q1_ms": 0.8377280332853034,
      "q3_ms": 0.8499383332996837,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 256,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.003190818941220641,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.21828215700224973,
      "inner": 135,
      "median_ms": 0.1439586667132702,
      "q1_ms": 0.14291271110306736,
      "q3_ms": 0.14421455923217796,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 256,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.005156615749001503,
      "xla_temp_mb": 17.000244140625,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.4378215690012439,
      "inner": 117,
      "median_ms": 0.35230674356536573,
      "q1_ms": 0.3429172692348301,
      "q3_ms": 0.36372909828339917,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 256,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00421260017901659,
      "xla_temp_mb": 320.5,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 1.182301766006276,
      "inner": 25,
      "median_ms": 1.86967695975909,
      "q1_ms": 1.8192287000420038,
      "q3_ms": 1.9974006999109406,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 256,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.005156615749001503,
      "xla_temp_mb": 17.000244140625,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.44653866700537037,
      "inner": 122,
      "median_ms": 0.34402256558857264,
      "q1_ms": 0.34095691802069833,
      "q3_ms": 0.3447647991453748,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 256,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.005156615749001503,
      "xla_temp_mb": 17.000244140625,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.4443777359992964,
      "inner": 101,
      "median_ms": 0.34590652476163986,
      "q1_ms": 0.3445335692814828,
      "q3_ms": 0.35418435144681937,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 256,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00421260017901659,
      "xla_temp_mb": 320.5,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 1.1199674700037576,
      "inner": 26,
      "median_ms": 1.8477138076447703,
      "q1_ms": 1.8261714039209227,
      "q3_ms": 1.850521807612565,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 256,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.005156615749001503,
      "xla_temp_mb": 49.000732421875,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.3012264119970496,
      "inner": 60,
      "median_ms": 0.6936118666393061,
      "q1_ms": 0.6879703333349122,
      "q3_ms": 0.7055014332460511,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 512,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.003060600720345974,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.3297638209987781,
      "inner": 130,
      "median_ms": 0.304405569253364,
      "q1_ms": 0.30343970001222964,
      "q3_ms": 0.31872415383986663,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 512,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00421844283118844,
      "xla_temp_mb": 768.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.6989217169975745,
      "inner": 21,
      "median_ms": 2.2083261428633705,
      "q1_ms": 2.188608976125362,
      "q3_ms": 2.220224642986154,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 512,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.003060600720345974,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.32861473099910654,
      "inner": 136,
      "median_ms": 0.3058690146965947,
      "q1_ms": 0.30514485291967186,
      "q3_ms": 0.3090076580680462,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 512,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.003060600720345974,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.3206275080010528,
      "inner": 137,
      "median_ms": 0.3056506569268705,
      "q1_ms": 0.304351587591601,
      "q3_ms": 0.3165463284639485,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 512,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00421844283118844,
      "xla_temp_mb": 768.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.2843513379993965,
      "inner": 21,
      "median_ms": 2.1948972855517197,
      "q1_ms": 2.189709618895514,
      "q3_ms": 2.2042681666789576,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 512,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.003060600720345974,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.04604606499924557,
      "inner": 59,
      "median_ms": 0.5923481526147274,
      "q1_ms": 0.5852938474918582,
      "q3_ms": 0.5954820762627159,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 512,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.005082274321466684,
      "xla_temp_mb": 34.000244140625,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.40388838299986674,
      "inner": 33,
      "median_ms": 1.5351997274344533,
      "q1_ms": 1.5330444394115177,
      "q3_ms": 1.5405463938915283,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 512,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0050018467009067535,
      "xla_temp_mb": 1281.0,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 1.4331575279938988,
      "inner": 7,
      "median_ms": 6.463911285598962,
      "q1_ms": 6.434660857069373,
      "q3_ms": 6.499819500147298,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 512,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.005082274321466684,
      "xla_temp_mb": 34.000244140625,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.4169226720041479,
      "inner": 32,
      "median_ms": 1.5335051873535122,
      "q1_ms": 1.5192393593679299,
      "q3_ms": 1.5532609376123219,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 512,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.005082274321466684,
      "xla_temp_mb": 34.000244140625,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.3944098220017622,
      "inner": 33,
      "median_ms": 1.5903006060398184,
      "q1_ms": 1.5266281667264234,
      "q3_ms": 1.6658201060305475,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 512,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0050018467009067535,
      "xla_temp_mb": 1281.0,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.6319602619987563,
      "inner": 7,
      "median_ms": 6.466519142642418,
      "q1_ms": 6.4029986434824036,
      "q3_ms": 6.484294285785706,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 512,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.005082274321466684,
      "xla_temp_mb": 98.000732421875,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.1163426659986726,
      "inner": 17,
      "median_ms": 2.7627889413612565,
      "q1_ms": 2.7469834118057728,
      "q3_ms": 2.7771500296212785,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 512,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0030704804230481386,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.3603106609953102,
      "inner": 154,
      "median_ms": 0.21117464934064487,
      "q1_ms": 0.21111564283628115,
      "q3_ms": 0.21289993508782407,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 512,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004382818937301636,
      "xla_temp_mb": 768.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 1.0362232580009731,
      "inner": 18,
      "median_ms": 2.472485611279909,
      "q1_ms": 2.446638805647834,
      "q3_ms": 2.660053944459327,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 512,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0030704804230481386,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.3607022739961394,
      "inner": 162,
      "median_ms": 0.20936768518884402,
      "q1_ms": 0.2081950617453621,
      "q3_ms": 0.21105935494559186,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 512,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0030704804230481386,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.36632916999951703,
      "inner": 176,
      "median_ms": 0.21253069884220496,
      "q1_ms": 0.21223934657593252,
      "q3_ms": 0.21331194885864202,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 512,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004382818937301636,
      "xla_temp_mb": 768.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 1.0330476129965973,
      "inner": 20,
      "median_ms": 2.471612449880922,
      "q1_ms": 2.444954625207174,
      "q3_ms": 2.4910763750085607,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 512,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0030704804230481386,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.04753902700031176,
      "inner": 100,
      "median_ms": 0.41101264003373217,
      "q1_ms": 0.4081821699946886,
      "q3_ms": 0.41211027000827016,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 512,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004873078316450119,
      "xla_temp_mb": 34.000244140625,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.4450344830038375,
      "inner": 49,
      "median_ms": 1.0849073877004547,
      "q1_ms": 0.9919668775591144,
      "q3_ms": 1.1278621326598615,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 512,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00808752328157425,
      "xla_temp_mb": 1281.0,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 2.2829222720029065,
      "inner": 6,
      "median_ms": 8.841934333759127,
      "q1_ms": 7.879996000459263,
      "q3_ms": 8.947770915862444,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 512,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004873078316450119,
      "xla_temp_mb": 34.000244140625,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.4478700419931556,
      "inner": 47,
      "median_ms": 0.998907383039663,
      "q1_ms": 0.9980989255362804,
      "q3_ms": 1.0066331062958813,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 512,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004873078316450119,
      "xla_temp_mb": 34.000244140625,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.4509345559999929,
      "inner": 47,
      "median_ms": 1.0035257234468442,
      "q1_ms": 0.9972823829733162,
      "q3_ms": 1.0153531595975398,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 512,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00808752328157425,
      "xla_temp_mb": 1281.0,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 2.4685774219979066,
      "inner": 6,
      "median_ms": 8.098518833624743,
      "q1_ms": 7.934677167213522,
      "q3_ms": 8.232764333418649,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 512,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004873078316450119,
      "xla_temp_mb": 98.000732421875,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.1231382379992283,
      "inner": 24,
      "median_ms": 1.8613165833206342,
      "q1_ms": 1.8610095209320814,
      "q3_ms": 1.87579129169535,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 16,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004541471600532532,
      "xla_temp_mb": 1.75,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.23452620399621082,
      "inner": 200,
      "median_ms": 0.051688210005522706,
      "q1_ms": 0.0505872424946574,
      "q3_ms": 0.05299150749124237,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 16,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004541471600532532,
      "xla_temp_mb": 1.75,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.21175925200077472,
      "inner": 200,
      "median_ms": 0.04584653001074912,
      "q1_ms": 0.04511190998528036,
      "q3_ms": 0.04889686000751681,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 16,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0027530461084097624,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.10931787899608025,
      "inner": 200,
      "median_ms": 0.07999701501830714,
      "q1_ms": 0.06445955998060526,
      "q3_ms": 0.09304173749114852,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 16,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0027530461084097624,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.13939826499699848,
      "inner": 200,
      "median_ms": 0.0497771850132267,
      "q1_ms": 0.04389489749883069,
      "q3_ms": 0.05967147997580469,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 16,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004541471600532532,
      "xla_temp_mb": 1.75,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.14287140400119824,
      "inner": 200,
      "median_ms": 0.05423309001344023,
      "q1_ms": 0.049682877524901414,
      "q3_ms": 0.055764267508493504,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 16,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0027530461084097624,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.36509763300273335,
      "inner": 200,
      "median_ms": 0.05892583998502232,
      "q1_ms": 0.05826673999763443,
      "q3_ms": 0.059229964990663575,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 16,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.005383178126066923,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 5.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.3768500189980841,
      "inner": 200,
      "median_ms": 0.06583680002222536,
      "q1_ms": 0.06429865748941666,
      "q3_ms": 0.06724697499521426,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 16,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.005383178126066923,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 5.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.2703738109994447,
      "inner": 200,
      "median_ms": 0.05855844498000806,
      "q1_ms": 0.05684534249667195,
      "q3_ms": 0.05907908001972828,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 16,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.005383178126066923,
      "xla_temp_mb": 2.062744140625,
      "xla_arg_mb": 5.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.24546262100193417,
      "inner": 200,
      "median_ms": 0.06561072001204593,
      "q1_ms": 0.05796173250928405,
      "q3_ms": 0.06678607749563525,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 16,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.005383178126066923,
      "xla_temp_mb": 2.062744140625,
      "xla_arg_mb": 5.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.21825718900072388,
      "inner": 200,
      "median_ms": 0.05005236002034508,
      "q1_ms": 0.049713737516867695,
      "q3_ms": 0.05478623499584501,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 16,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.005383178126066923,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 5.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.2332557639965671,
      "inner": 200,
      "median_ms": 0.0705060100153787,
      "q1_ms": 0.06610335751247476,
      "q3_ms": 0.08017555001060828,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 16,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.005383178126066923,
      "xla_temp_mb": 6.063232421875,
      "xla_arg_mb": 5.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.778218792998814,
      "inner": 153,
      "median_ms": 0.1453336862694271,
      "q1_ms": 0.1420704673340748,
      "q3_ms": 0.20277639216111115,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 16,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004810147453099489,
      "xla_temp_mb": 1.75,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.23172373699344462,
      "inner": 200,
      "median_ms": 0.051006664980377536,
      "q1_ms": 0.05063972001153161,
      "q3_ms": 0.05261106000034488,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 16,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004810147453099489,
      "xla_temp_mb": 1.75,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.22867943099845434,
      "inner": 200,
      "median_ms": 0.04825970499950927,
      "q1_ms": 0.04460898251636536,
      "q3_ms": 0.05027750748922699,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 16,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0030694082379341125,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.1280754449981032,
      "inner": 200,
      "median_ms": 0.05951946499408223,
      "q1_ms": 0.05635944751702482,
      "q3_ms": 0.06330572499791742,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 16,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0030694082379341125,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.13236916000460042,
      "inner": 200,
      "median_ms": 0.041767190014070366,
      "q1_ms": 0.03909634500814718,
      "q3_ms": 0.04326410251451307,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 16,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004810147453099489,
      "xla_temp_mb": 1.75,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.2664830290013924,
      "inner": 200,
      "median_ms": 0.06144987499283161,
      "q1_ms": 0.058430247499927646,
      "q3_ms": 0.06182195997098461,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 16,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0030694082379341125,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.05095996699674288,
      "inner": 200,
      "median_ms": 0.05984944000374526,
      "q1_ms": 0.059111767513968516,
      "q3_ms": 0.06039830999725382,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 16,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00432160496711731,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 5.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.49675766799919074,
      "inner": 200,
      "median_ms": 0.0533248600186198,
      "q1_ms": 0.051614320000226144,
      "q3_ms": 0.05829116000313661,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 16,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00432160496711731,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 5.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.47378263199789217,
      "inner": 200,
      "median_ms": 0.07423874500091188,
      "q1_ms": 0.062013392489461694,
      "q3_ms": 0.08359664001545752,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 16,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004320277366787195,
      "xla_temp_mb": 2.062744140625,
      "xla_arg_mb": 5.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.29871642000216525,
      "inner": 200,
      "median_ms": 0.060517510028148536,
      "q1_ms": 0.06011138999383547,
      "q3_ms": 0.06139146749774227,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 16,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004320277366787195,
      "xla_temp_mb": 2.062744140625,
      "xla_arg_mb": 5.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.2596677510009613,
      "inner": 200,
      "median_ms": 0.053693859990744386,
      "q1_ms": 0.05191681500946288,
      "q3_ms": 0.059330082513042726,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 16,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00432160496711731,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 5.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.5840555120012141,
      "inner": 200,
      "median_ms": 0.05389620000642026,
      "q1_ms": 0.05281376750644995,
      "q3_ms": 0.054633522504445864,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 16,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004320277366787195,
      "xla_temp_mb": 6.063232421875,
      "xla_arg_mb": 5.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.11914158600120572,
      "inner": 176,
      "median_ms": 0.14464281817212893,
      "q1_ms": 0.14411521874535538,
      "q3_ms": 0.14611190623658,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 32,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004443137440830469,
      "xla_temp_mb": 4.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.3533161130035296,
      "inner": 200,
      "median_ms": 0.03828212498774519,
      "q1_ms": 0.0373839224994299,
      "q3_ms": 0.03886402750140405,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 32,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004443137440830469,
      "xla_temp_mb": 4.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.1570605929955491,
      "inner": 200,
      "median_ms": 0.045069299994793255,
      "q1_ms": 0.0412803900144354,
      "q3_ms": 0.050207274998683715,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 32,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0027243313379585743,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.1381294379971223,
      "inner": 200,
      "median_ms": 0.04580389999318868,
      "q1_ms": 0.04548151002381928,
      "q3_ms": 0.046624025017081294,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 32,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0027243313379585743,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.12572276099672308,
      "inner": 200,
      "median_ms": 0.04646870496799238,
      "q1_ms": 0.046081045002210885,
      "q3_ms": 0.0515577599981043,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 32,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004443137440830469,
      "xla_temp_mb": 4.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.1619454589963425,
      "inner": 200,
      "median_ms": 0.04106650500034448,
      "q1_ms": 0.04012301749753533,
      "q3_ms": 0.041657102501631016,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 32,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0027243313379585743,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.047706614997878205,
      "inner": 200,
      "median_ms": 0.05694417999620782,
      "q1_ms": 0.056711912493483396,
      "q3_ms": 0.058051754986081505,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 32,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0059395041316747665,
      "xla_temp_mb": 4.0625,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.5864202180018765,
      "inner": 200,
      "median_ms": 0.08779691001109313,
      "q1_ms": 0.08777819999522762,
      "q3_ms": 0.08783164999840665,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 32,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0059395041316747665,
      "xla_temp_mb": 4.0625,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.31590124200010905,
      "inner": 200,
      "median_ms": 0.09000214500701986,
      "q1_ms": 0.08793842998784385,
      "q3_ms": 0.09027180749399122,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 32,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0059395041316747665,
      "xla_temp_mb": 4.125244140625,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.3450912280022749,
      "inner": 200,
      "median_ms": 0.07355239002208691,
      "q1_ms": 0.06725975001245388,
      "q3_ms": 0.0961345774885558,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 32,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0059395041316747665,
      "xla_temp_mb": 4.125244140625,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.22980528700281866,
      "inner": 200,
      "median_ms": 0.06416933501895983,
      "q1_ms": 0.06309457499810378,
      "q3_ms": 0.06595725251827389,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 32,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0059395041316747665,
      "xla_temp_mb": 4.0625,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.3031518350035185,
      "inner": 200,
      "median_ms": 0.08715599498827942,
      "q1_ms": 0.08492135499182041,
      "q3_ms": 0.0909990374930203,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 32,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0059395041316747665,
      "xla_temp_mb": 12.125732421875,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.17194378799467813,
      "inner": 103,
      "median_ms": 0.19654373784416165,
      "q1_ms": 0.18517543206859416,
      "q3_ms": 0.22857936894017028,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 32,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00452336436137557,
      "xla_temp_mb": 4.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.36209469599998556,
      "inner": 200,
      "median_ms": 0.03837385000224458,
      "q1_ms": 0.03787968749747961,
      "q3_ms": 0.03916762749213376,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 32,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00452336436137557,
      "xla_temp_mb": 4.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.3713408179974067,
      "inner": 200,
      "median_ms": 0.04212330997688696,
      "q1_ms": 0.03693652250149171,
      "q3_ms": 0.04607117747582379,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 32,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0031470218673348427,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.15225002599618165,
      "inner": 200,
      "median_ms": 0.04918250000628177,
      "q1_ms": 0.04435474500496639,
      "q3_ms": 0.05041053500463022,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 32,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0031470218673348427,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.15888176699809264,
      "inner": 200,
      "median_ms": 0.04060538500198163,
      "q1_ms": 0.039079362504708115,
      "q3_ms": 0.04177613000138081,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 32,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00452336436137557,
      "xla_temp_mb": 4.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.3628073990039411,
      "inner": 200,
      "median_ms": 0.0386654549947707,
      "q1_ms": 0.038298282506730175,
      "q3_ms": 0.03892429249390261,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 32,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0031470218673348427,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.049959548996412195,
      "inner": 200,
      "median_ms": 0.057454650013823994,
      "q1_ms": 0.057371115017303964,
      "q3_ms": 0.05772102751507191,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 32,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.005252570379525423,
      "xla_temp_mb": 4.0625,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.7333298029989237,
      "inner": 200,
      "median_ms": 0.09632824498112313,
      "q1_ms": 0.08983778499896289,
      "q3_ms": 0.09811533498577774,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 32,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.005252570379525423,
      "xla_temp_mb": 4.0625,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.7286670379980933,
      "inner": 200,
      "median_ms": 0.08792404998530401,
      "q1_ms": 0.08647485499750474,
      "q3_ms": 0.08877305999703822,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 32,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004545743577182293,
      "xla_temp_mb": 4.125244140625,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.3052220820027287,
      "inner": 200,
      "median_ms": 0.05943911000940716,
      "q1_ms": 0.059163414989598095,
      "q3_ms": 0.06034342748534982,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 32,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004545743577182293,
      "xla_temp_mb": 4.125244140625,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.31055043100059265,
      "inner": 200,
      "median_ms": 0.056238990000565536,
      "q1_ms": 0.05530739999812795,
      "q3_ms": 0.057411165016674204,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 32,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.005252570379525423,
      "xla_temp_mb": 4.0625,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.7257534589953138,
      "inner": 200,
      "median_ms": 0.0914809499954572,
      "q1_ms": 0.08970237500761868,
      "q3_ms": 0.09205729247696581,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 32,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004545743577182293,
      "xla_temp_mb": 12.125732421875,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.11997865299781552,
      "inner": 136,
      "median_ms": 0.18460720591118682,
      "q1_ms": 0.17331066914373558,
      "q3_ms": 0.18515390438016552,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 64,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004548877011984587,
      "xla_temp_mb": 12.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.4085004640000989,
      "inner": 200,
      "median_ms": 0.08532431998901302,
      "q1_ms": 0.08529060249202303,
      "q3_ms": 0.08537025751138572,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 64,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004548877011984587,
      "xla_temp_mb": 12.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.2125654109986499,
      "inner": 200,
      "median_ms": 0.0841962299818988,
      "q1_ms": 0.08396321500185877,
      "q3_ms": 0.08425962250839802,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 64,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0023481957614421844,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.1546979500053567,
      "inner": 200,
      "median_ms": 0.031548849983664695,
      "q1_ms": 0.03110227999059134,
      "q3_ms": 0.039381887509080116,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 64,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0023481957614421844,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.14942267299920786,
      "inner": 200,
      "median_ms": 0.03892223998263944,
      "q1_ms": 0.031685637495684205,
      "q3_ms": 0.04563094249533606,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 64,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004548877011984587,
      "xla_temp_mb": 12.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.2107053720028489,
      "inner": 200,
      "median_ms": 0.08421361497312319,
      "q1_ms": 0.0838643524912186,
      "q3_ms": 0.08424109000770841,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 64,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0023481957614421844,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.2856220879984903,
      "inner": 200,
      "median_ms": 0.05584168997302186,
      "q1_ms": 0.05578002252150327,
      "q3_ms": 0.056173645025410224,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 64,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.006896344944834709,
      "xla_temp_mb": 20.125,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.6984317770038615,
      "inner": 151,
      "median_ms": 0.2261237549712721,
      "q1_ms": 0.22470430133803376,
      "q3_ms": 0.2263221059600508,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 64,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.006896344944834709,
      "xla_temp_mb": 20.125,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.3769672409980558,
      "inner": 111,
      "median_ms": 0.22434813516090735,
      "q1_ms": 0.22424108556970987,
      "q3_ms": 0.2244191711544592,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 64,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004327160306274891,
      "xla_temp_mb": 8.250244140625,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.36654995899880305,
      "inner": 200,
      "median_ms": 0.1084017549874261,
      "q1_ms": 0.10815969249961199,
      "q3_ms": 0.10863471750781173,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 64,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004327160306274891,
      "xla_temp_mb": 8.250244140625,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.3575799110039952,
      "inner": 199,
      "median_ms": 0.10854703515777411,
      "q1_ms": 0.10813204775258417,
      "q3_ms": 0.10856807536250168,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 64,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.006896344944834709,
      "xla_temp_mb": 20.125,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.3608777679983177,
      "inner": 152,
      "median_ms": 0.22425373026635498,
      "q1_ms": 0.22419117108017272,
      "q3_ms": 0.22431729935066724,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 64,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004327160306274891,
      "xla_temp_mb": 24.250732421875,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.35942076300125336,
      "inner": 123,
      "median_ms": 0.23478142277345182,
      "q1_ms": 0.23473454877706737,
      "q3_ms": 0.23493615041909988,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 64,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004492104984819889,
      "xla_temp_mb": 12.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 1.2928466020021006,
      "inner": 200,
      "median_ms": 0.11362686502252473,
      "q1_ms": 0.11189149498022744,
      "q3_ms": 0.114417165004852,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 64,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004492104984819889,
      "xla_temp_mb": 12.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 1.2905907919994206,
      "inner": 200,
      "median_ms": 0.10779389998788247,
      "q1_ms": 0.10755261499070912,
      "q3_ms": 0.10852498498934438,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 64,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0031647554133087397,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.20403543900465593,
      "inner": 200,
      "median_ms": 0.04152437501034001,
      "q1_ms": 0.03678308252347051,
      "q3_ms": 0.04234660998918116,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 64,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0031647554133087397,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.2005544950006879,
      "inner": 200,
      "median_ms": 0.03832279999187449,
      "q1_ms": 0.037365260013757506,
      "q3_ms": 0.04186081501757144,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 64,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004492104984819889,
      "xla_temp_mb": 12.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 1.3164246720043593,
      "inner": 200,
      "median_ms": 0.1289434250065824,
      "q1_ms": 0.12774109000019962,
      "q3_ms": 0.12905079749543802,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 64,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0031647554133087397,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.05255154699989362,
      "inner": 200,
      "median_ms": 0.061855995008954785,
      "q1_ms": 0.06115627251347178,
      "q3_ms": 0.06223764748938265,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 64,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00452549010515213,
      "xla_temp_mb": 20.125,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.9359819380042609,
      "inner": 154,
      "median_ms": 0.2554373896163707,
      "q1_ms": 0.2447959642657744,
      "q3_ms": 0.2570495357109134,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 64,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00452549010515213,
      "xla_temp_mb": 20.125,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.8721042939941981,
      "inner": 158,
      "median_ms": 0.2267282341770486,
      "q1_ms": 0.22669503798883678,
      "q3_ms": 0.22693536392799166,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 64,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0038649465423077345,
      "xla_temp_mb": 8.250244140625,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.47411131500120973,
      "inner": 199,
      "median_ms": 0.11508416081135205,
      "q1_ms": 0.11458706029400641,
      "q3_ms": 0.11509568594649419,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 64,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0038649465423077345,
      "xla_temp_mb": 8.250244140625,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.46891889400285436,
      "inner": 168,
      "median_ms": 0.11507283332840805,
      "q1_ms": 0.1145598035648055,
      "q3_ms": 0.11519263989612227,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 64,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00452549010515213,
      "xla_temp_mb": 20.125,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.8336737869976787,
      "inner": 148,
      "median_ms": 0.2267878851460016,
      "q1_ms": 0.22658346959950587,
      "q3_ms": 0.22682524659840791,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 64,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0038649465423077345,
      "xla_temp_mb": 24.250732421875,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.11427151599491481,
      "inner": 133,
      "median_ms": 0.2336300676688552,
      "q1_ms": 0.23353642483611234,
      "q3_ms": 0.23369680827744374,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00237796432338655,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.3967501510051079,
      "inner": 200,
      "median_ms": 0.06121903999883215,
      "q1_ms": 0.061169672499090666,
      "q3_ms": 0.06604850248550065,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0028126207180321217,
      "xla_temp_mb": 48.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.6160274959984235,
      "inner": 139,
      "median_ms": 0.25471185612292935,
      "q1_ms": 0.2536241762551962,
      "q3_ms": 0.25500661869706726,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00237796432338655,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.39403001200116705,
      "inner": 200,
      "median_ms": 0.06009255499520805,
      "q1_ms": 0.05880285749299219,
      "q3_ms": 0.06299002749074134,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00237796432338655,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.38713829700282076,
      "inner": 200,
      "median_ms": 0.05955959000857547,
      "q1_ms": 0.058714989991131006,
      "q3_ms": 0.06038007250026567,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0028126207180321217,
      "xla_temp_mb": 48.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.22484634900320088,
      "inner": 145,
      "median_ms": 0.2517460758601123,
      "q1_ms": 0.24962263449761568,
      "q3_ms": 0.251806600035393,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00237796432338655,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.047817885999393184,
      "inner": 200,
      "median_ms": 0.061743935002596124,
      "q1_ms": 0.061523267522716196,
      "q3_ms": 0.06281454001509701,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.005381712689995766,
      "xla_temp_mb": 16.500244140625,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.5590445530033321,
      "inner": 129,
      "median_ms": 0.29493689148094826,
      "q1_ms": 0.29248945736967014,
      "q3_ms": 0.2990803798518953,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.005381712689995766,
      "xla_temp_mb": 80.25,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 1.203650215997186,
      "inner": 64,
      "median_ms": 0.6814579531919662,
      "q1_ms": 0.6813045469016288,
      "q3_ms": 0.6815493125031935,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.005381712689995766,
      "xla_temp_mb": 16.500244140625,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.5254620330015314,
      "inner": 109,
      "median_ms": 0.2914226788561791,
      "q1_ms": 0.2907946605490063,
      "q3_ms": 0.29150563305021704,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.005381712689995766,
      "xla_temp_mb": 16.500244140625,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.5175008299993351,
      "inner": 112,
      "median_ms": 0.2904396517676235,
      "q1_ms": 0.29003513840767225,
      "q3_ms": 0.29053242409255675,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.005381712689995766,
      "xla_temp_mb": 80.25,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.4786877870064927,
      "inner": 64,
      "median_ms": 0.681493812521694,
      "q1_ms": 0.6813475858962192,
      "q3_ms": 0.6815918906113438,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.005381712689995766,
      "xla_temp_mb": 48.500732421875,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.39639542800432537,
      "inner": 86,
      "median_ms": 0.41806766281762636,
      "q1_ms": 0.41803608719208063,
      "q3_ms": 0.4182339127806318,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0028695594519376755,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.47887907000404084,
      "inner": 200,
      "median_ms": 0.11305046999041224,
      "q1_ms": 0.11042426251151483,
      "q3_ms": 0.11887604750882019,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0046579716727137566,
      "xla_temp_mb": 48.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 1.496421226998791,
      "inner": 114,
      "median_ms": 0.32771162283065597,
      "q1_ms": 0.32395106140220853,
      "q3_ms": 0.35248805263683225,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0028695594519376755,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.49242121100542136,
      "inner": 200,
      "median_ms": 0.11076317998231389,
      "q1_ms": 0.10853718249563826,
      "q3_ms": 0.11141486749693286,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0028695594519376755,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.48784428799990565,
      "inner": 200,
      "median_ms": 0.10775242000818253,
      "q1_ms": 0.10772163750516484,
      "q3_ms": 0.11004028248862596,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0046579716727137566,
      "xla_temp_mb": 48.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 1.5194062619993929,
      "inner": 112,
      "median_ms": 0.3229810356840192,
      "q1_ms": 0.3198115892765989,
      "q3_ms": 0.35138434376286953,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0028695594519376755,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.0520327720005298,
      "inner": 200,
      "median_ms": 0.07274655501532834,
      "q1_ms": 0.07248792000609683,
      "q3_ms": 0.07455977500285371,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004662639927119017,
      "xla_temp_mb": 16.500244140625,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.58198309999716,
      "inner": 124,
      "median_ms": 0.3049686370838073,
      "q1_ms": 0.29942764514646936,
      "q3_ms": 0.3086858145134982,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004662639927119017,
      "xla_temp_mb": 80.25,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 1.0139362960035214,
      "inner": 62,
      "median_ms": 0.7241006452175263,
      "q1_ms": 0.7104475564354347,
      "q3_ms": 0.7754125806733951,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004662639927119017,
      "xla_temp_mb": 16.500244140625,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.6090734839963261,
      "inner": 122,
      "median_ms": 0.3039658360768178,
      "q1_ms": 0.3009340614910604,
      "q3_ms": 0.32119073360069605,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004662639927119017,
      "xla_temp_mb": 16.500244140625,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.5898680379978032,
      "inner": 130,
      "median_ms": 0.3025109769413785,
      "q1_ms": 0.3020400461537065,
      "q3_ms": 0.30285874997319373,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004662639927119017,
      "xla_temp_mb": 80.25,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 1.0691822939988924,
      "inner": 62,
      "median_ms": 0.7169785806787156,
      "q1_ms": 0.710814701667386,
      "q3_ms": 0.7198144435425355,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004662639927119017,
      "xla_temp_mb": 48.500732421875,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.12014512900350383,
      "inner": 87,
      "median_ms": 0.4685824137648043,
      "q1_ms": 0.4388132816195306,
      "q3_ms": 0.4792475804282438,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 256,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0034007239155471325,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.43757279699639184,
      "inner": 136,
      "median_ms": 0.24178170588129433,
      "q1_ms": 0.24163532353828915,
      "q3_ms": 0.24180317645646832,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 256,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0037950000260025263,
      "xla_temp_mb": 192.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.7572839870044845,
      "inner": 58,
      "median_ms": 0.7266355861592556,
      "q1_ms": 0.7263707672697411,
      "q3_ms": 0.726868000003441,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 256,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0034007239155471325,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.4343574369995622,
      "inner": 125,
      "median_ms": 0.24044271203456447,
      "q1_ms": 0.2396645919943694,
      "q3_ms": 0.24064477198407985,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 256,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0034007239155471325,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.43231087499589194,
      "inner": 129,
      "median_ms": 0.24066541083120965,
      "q1_ms": 0.24035117056946337,
      "q3_ms": 0.24078920927591796,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 256,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0037950000260025263,
      "xla_temp_mb": 192.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.2595722499972908,
      "inner": 53,
      "median_ms": 0.7270865282647856,
      "q1_ms": 0.7265338018390499,
      "q3_ms": 0.7281414622892318,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 256,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0034007239155471325,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.049323603001539595,
      "inner": 173,
      "median_ms": 0.1670905722660481,
      "q1_ms": 0.16546752311829666,
      "q3_ms": 0.16715752314661786,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 256,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004902470391243696,
      "xla_temp_mb": 33.000244140625,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.5092934499989497,
      "inner": 50,
      "median_ms": 0.8949935001146514,
      "q1_ms": 0.8927090700308327,
      "q3_ms": 0.8960867799760308,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 256,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.006926695816218853,
      "xla_temp_mb": 320.5,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 1.217389560995798,
      "inner": 23,
      "median_ms": 2.0432318696807865,
      "q1_ms": 2.042959434734455,
      "q3_ms": 2.043459065095268,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 256,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004902470391243696,
      "xla_temp_mb": 33.000244140625,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.486582374003774,
      "inner": 50,
      "median_ms": 0.8936611800163519,
      "q1_ms": 0.8935763999761548,
      "q3_ms": 0.8939090499916347,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 256,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004902470391243696,
      "xla_temp_mb": 33.000244140625,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.48016193300281884,
      "inner": 49,
      "median_ms": 0.8936161020941729,
      "q1_ms": 0.8933136938436774,
      "q3_ms": 0.8941538571332561,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 256,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.006926695816218853,
      "xla_temp_mb": 320.5,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.41469557600066764,
      "inner": 23,
      "median_ms": 2.0430873043116424,
      "q1_ms": 2.042807825921994,
      "q3_ms": 2.0433604345864453,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 256,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004902470391243696,
      "xla_temp_mb": 97.000732421875,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.11633601300127339,
      "inner": 44,
      "median_ms": 1.0011556590722068,
      "q1_ms": 1.0006723750848323,
      "q3_ms": 1.001349795511818,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 256,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0035469671711325645,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.5983545589988353,
      "inner": 132,
      "median_ms": 0.2718941515481812,
      "q1_ms": 0.2717984015169239,
      "q3_ms": 0.2719764621140326,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 256,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004809652455151081,
      "xla_temp_mb": 192.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 1.9200620530027663,
      "inner": 47,
      "median_ms": 0.9673782128919947,
      "q1_ms": 0.9664111594828044,
      "q3_ms": 0.968475606412617,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 256,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0035469671711325645,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.5623718950009788,
      "inner": 133,
      "median_ms": 0.270116654160359,
      "q1_ms": 0.2698180413650895,
      "q3_ms": 0.27138989474436204,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 256,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0035469671711325645,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.5169710539994412,
      "inner": 129,
      "median_ms": 0.27127386820493554,
      "q1_ms": 0.269972325588961,
      "q3_ms": 0.2775727131897169,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 256,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004809652455151081,
      "xla_temp_mb": 192.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 1.888266403999296,
      "inner": 46,
      "median_ms": 0.9459506521955557,
      "q1_ms": 0.9446146522468948,
      "q3_ms": 0.9568803261177179,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 256,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0035469671711325645,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.052055084001040086,
      "inner": 76,
      "median_ms": 0.13874386843925,
      "q1_ms": 0.13358467104377555,
      "q3_ms": 0.17349513157234087,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 256,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004253170918673277,
      "xla_temp_mb": 33.000244140625,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.6100235369958682,
      "inner": 58,
      "median_ms": 0.7637675690007042,
      "q1_ms": 0.7634300345044125,
      "q3_ms": 0.7669825516696337,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 256,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004358398262411356,
      "xla_temp_mb": 320.5,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 1.2283354469982442,
      "inner": 22,
      "median_ms": 2.096726818060481,
      "q1_ms": 2.094952091118268,
      "q3_ms": 2.0984423636863623,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 256,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004253170918673277,
      "xla_temp_mb": 33.000244140625,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.6047255459998269,
      "inner": 55,
      "median_ms": 0.7623132908272303,
      "q1_ms": 0.7605283636869651,
      "q3_ms": 0.7626105726971715,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 256,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004253170918673277,
      "xla_temp_mb": 33.000244140625,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.5855004240002017,
      "inner": 55,
      "median_ms": 0.7639148726329124,
      "q1_ms": 0.7616864726349541,
      "q3_ms": 0.7790929454627489,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 256,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004358398262411356,
      "xla_temp_mb": 320.5,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 1.4669934990015463,
      "inner": 22,
      "median_ms": 2.135747227426195,
      "q1_ms": 2.1143992273044367,
      "q3_ms": 2.1753543865501834,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 256,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004253170918673277,
      "xla_temp_mb": 97.000732421875,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.3185836510019726,
      "inner": 48,
      "median_ms": 0.9102631249030916,
      "q1_ms": 0.9074715729487555,
      "q3_ms": 0.914124520780509,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0028048453386873007,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.4338374409999233,
      "inner": 60,
      "median_ms": 0.6881910833423414,
      "q1_ms": 0.6876910333327639,
      "q3_ms": 0.6884947333067732,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004870464093983173,
      "xla_temp_mb": 768.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.9424347590029356,
      "inner": 19,
      "median_ms": 2.4150215789344847,
      "q1_ms": 2.4137981578716885,
      "q3_ms": 2.415365394737915,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0028048453386873007,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.4580070039955899,
      "inner": 56,
      "median_ms": 0.6938896072564863,
      "q1_ms": 0.6937684821813102,
      "q3_ms": 0.6965611874615466,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0028048453386873007,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.4492157149943523,
      "inner": 62,
      "median_ms": 0.688698161308909,
      "q1_ms": 0.6885278306717444,
      "q3_ms": 0.6894940242120591,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004870464093983173,
      "xla_temp_mb": 768.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.28080950799630955,
      "inner": 19,
      "median_ms": 2.40848757905616,
      "q1_ms": 2.4067796579114846,
      "q3_ms": 2.408836158122684,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0028048453386873007,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.3212864180022734,
      "inner": 70,
      "median_ms": 0.6017945285878211,
      "q1_ms": 0.601431421403374,
      "q3_ms": 0.6035253142597088,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004695911426097155,
      "xla_temp_mb": 66.000244140625,
      "xla_arg_mb": 160.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 0.5886697659952915,
      "inner": 17,
      "median_ms": 2.98332029428018,
      "q1_ms": 2.9808152355733055,
      "q3_ms": 2.9926989704750557,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004695911426097155,
      "xla_temp_mb": 1281.0,
      "xla_arg_mb": 160.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 1.7226877819994115,
      "inner": 7,
      "median_ms": 6.920072572289167,
      "q1_ms": 6.907652500168687,
      "q3_ms": 6.926519642417718,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004695911426097155,
      "xla_temp_mb": 66.000244140625,
      "xla_arg_mb": 160.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 0.5971133150014794,
      "inner": 17,
      "median_ms": 2.980876882194185,
      "q1_ms": 2.9787681470385423,
      "q3_ms": 2.990003558807075,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004695911426097155,
      "xla_temp_mb": 66.000244140625,
      "xla_arg_mb": 160.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 0.5958940310010803,
      "inner": 17,
      "median_ms": 2.9850081764914864,
      "q1_ms": 2.98380799979307,
      "q3_ms": 2.9903489412390627,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004695911426097155,
      "xla_temp_mb": 1281.0,
      "xla_arg_mb": 160.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 0.7070563190063694,
      "inner": 7,
      "median_ms": 6.91037971409969,
      "q1_ms": 6.900082499799153,
      "q3_ms": 6.917718714248622,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004695911426097155,
      "xla_temp_mb": 194.000732421875,
      "xla_arg_mb": 160.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 0.4205655120022129,
      "inner": 15,
      "median_ms": 2.9067183335428126,
      "q1_ms": 2.9060387668626695,
      "q3_ms": 2.9169190997587675,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0036932043731212616,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.5397186400005012,
      "inner": 58,
      "median_ms": 0.7044966895873899,
      "q1_ms": 0.7028800776205226,
      "q3_ms": 0.7049850172389315,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004430409986525774,
      "xla_temp_mb": 768.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 1.1485090170026524,
      "inner": 18,
      "median_ms": 2.7118051665537575,
      "q1_ms": 2.7096741667062614,
      "q3_ms": 2.7140385556978996,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0036932043731212616,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.5675823900019168,
      "inner": 53,
      "median_ms": 0.7023711508702276,
      "q1_ms": 0.7019894999860488,
      "q3_ms": 0.7029972547793735,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0036932043731212616,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.5819475730022532,
      "inner": 60,
      "median_ms": 0.7021515667778052,
      "q1_ms": 0.7017028833312604,
      "q3_ms": 0.7023857583286978,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004430409986525774,
      "xla_temp_mb": 768.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 1.1869571269999142,
      "inner": 18,
      "median_ms": 2.715912388844622,
      "q1_ms": 2.71451305545472,
      "q3_ms": 2.7190049444875006,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0036932043731212616,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.3814982769981725,
      "inner": 83,
      "median_ms": 0.49682266267538294,
      "q1_ms": 0.496586301219888,
      "q3_ms": 0.49796922889598416,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004600926768034697,
      "xla_temp_mb": 66.000244140625,
      "xla_arg_mb": 160.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 0.678197089000605,
      "inner": 22,
      "median_ms": 2.141059545630758,
      "q1_ms": 2.140002795468635,
      "q3_ms": 2.14334640903954,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flax",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.005244561471045017,
      "xla_temp_mb": 1281.0,
      "xla_arg_mb": 160.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 2.6329301579971798,
      "inner": 5,
      "median_ms": 8.41255160048604,
      "q1_ms": 8.401871400565142,
      "q3_ms": 8.442263000324601,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004600926768034697,
      "xla_temp_mb": 66.000244140625,
      "xla_arg_mb": 160.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 0.7200900349998847,
      "inner": 22,
      "median_ms": 2.1410877727083224,
      "q1_ms": 2.136571750022333,
      "q3_ms": 2.146534749871617,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004600926768034697,
      "xla_temp_mb": 66.000244140625,
      "xla_arg_mb": 160.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 0.6967520309990505,
      "inner": 23,
      "median_ms": 2.142342956453238,
      "q1_ms": 2.138334478292128,
      "q3_ms": 2.142456217489266,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.005244561471045017,
      "xla_temp_mb": 1281.0,
      "xla_arg_mb": 160.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 2.5773133670008974,
      "inner": 5,
      "median_ms": 8.42040459974669,
      "q1_ms": 8.391085899347672,
      "q3_ms": 8.424756699969294,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004600926768034697,
      "xla_temp_mb": 194.000732421875,
      "xla_arg_mb": 160.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 0.4432049820024986,
      "inner": 18,
      "median_ms": 2.3712346664979123,
      "q1_ms": 2.368442916728479,
      "q3_ms": 2.377532388992323,
      "tag": null
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 16,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0005622011376544833,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.20398781699623214,
      "inner": 200,
      "median_ms": 0.05232526502368273,
      "q1_ms": 0.048265179993904894,
      "q3_ms": 0.05245410500720027,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 16,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0005374412285163999,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.107580894997227,
      "inner": 200,
      "median_ms": 0.04362866500741802,
      "q1_ms": 0.03699657001561718,
      "q3_ms": 0.05071596749985474,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 16,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006108639063313603,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 4.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.3223988640020252,
      "inner": 200,
      "median_ms": 0.05013101497752359,
      "q1_ms": 0.047261734998755855,
      "q3_ms": 0.05506193501787493,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 16,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006692991591989994,
      "xla_temp_mb": 1.062744140625,
      "xla_arg_mb": 4.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.21167462499579415,
      "inner": 200,
      "median_ms": 0.04650801500247326,
      "q1_ms": 0.04625625751941698,
      "q3_ms": 0.04659327250919887,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 16,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0004711545188911259,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.19379895600286545,
      "inner": 200,
      "median_ms": 0.04103450999537017,
      "q1_ms": 0.04101647500647232,
      "q3_ms": 0.04114012999707484,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 16,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0005859605153091252,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.11899912000080803,
      "inner": 200,
      "median_ms": 0.03698169002745999,
      "q1_ms": 0.03566074250556994,
      "q3_ms": 0.044683535015792586,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 16,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006589461117982864,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 4.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.58598292300303,
      "inner": 200,
      "median_ms": 0.047568914997100364,
      "q1_ms": 0.047551460011163726,
      "q3_ms": 0.04765816000144696,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 16,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007603430422022939,
      "xla_temp_mb": 1.062744140625,
      "xla_arg_mb": 4.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.24981009800103493,
      "inner": 200,
      "median_ms": 0.04844567498366814,
      "q1_ms": 0.04685544250605744,
      "q3_ms": 0.05202156247833045,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 32,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.000846682523842901,
      "xla_temp_mb": 4.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.21491274899744894,
      "inner": 200,
      "median_ms": 0.03776808000111487,
      "q1_ms": 0.037697375009884126,
      "q3_ms": 0.037842169986106455,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 32,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006542886258102953,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.13441842199972598,
      "inner": 200,
      "median_ms": 0.03615442499722121,
      "q1_ms": 0.03607868749895715,
      "q3_ms": 0.036204602492944105,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 32,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007412488339468837,
      "xla_temp_mb": 4.0,
      "xla_arg_mb": 8.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.3988392340033897,
      "inner": 200,
      "median_ms": 0.08650069499708479,
      "q1_ms": 0.08546989498427138,
      "q3_ms": 0.08651672500491259,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 32,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007395902648568153,
      "xla_temp_mb": 2.125244140625,
      "xla_arg_mb": 8.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.3180054049953469,
      "inner": 200,
      "median_ms": 0.04750509498990141,
      "q1_ms": 0.047021877489896724,
      "q3_ms": 0.047561025003233226,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 32,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0004653937940020114,
      "xla_temp_mb": 4.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.3146383759958553,
      "inner": 200,
      "median_ms": 0.03930112998205004,
      "q1_ms": 0.03927821249817498,
      "q3_ms": 0.03932244500902016,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 32,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0004163338744547218,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.1655720620037755,
      "inner": 200,
      "median_ms": 0.03726312999788206,
      "q1_ms": 0.037166192505537765,
      "q3_ms": 0.0374091499907081,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 32,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006948871887288988,
      "xla_temp_mb": 4.0,
      "xla_arg_mb": 8.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 1.0844252330061863,
      "inner": 200,
      "median_ms": 0.09807288497540867,
      "q1_ms": 0.09797810500458581,
      "q3_ms": 0.09809159251744859,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 32,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007666696328669786,
      "xla_temp_mb": 2.125244140625,
      "xla_arg_mb": 8.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.39276006600266555,
      "inner": 200,
      "median_ms": 0.050314010004512966,
      "q1_ms": 0.04923772499751067,
      "q3_ms": 0.05054927751189098,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 64,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0005379546782933176,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.1923738709956524,
      "inner": 200,
      "median_ms": 0.03166268001223216,
      "q1_ms": 0.03143845749946195,
      "q3_ms": 0.03190266250385321,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 64,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0005379546782933176,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.1862402509941603,
      "inner": 200,
      "median_ms": 0.033374470003764145,
      "q1_ms": 0.03244701498260838,
      "q3_ms": 0.033880249993671896,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 64,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007911636494100094,
      "xla_temp_mb": 4.250244140625,
      "xla_arg_mb": 16.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.5154227449966129,
      "inner": 200,
      "median_ms": 0.09873628998320783,
      "q1_ms": 0.0987113425071584,
      "q3_ms": 0.09876100748442695,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 64,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007911636494100094,
      "xla_temp_mb": 4.250244140625,
      "xla_arg_mb": 16.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.5118885499978205,
      "inner": 200,
      "median_ms": 0.09746909003297333,
      "q1_ms": 0.09669600249253563,
      "q3_ms": 0.09758365749803488,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 64,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0003806430322583765,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.25016915800370043,
      "inner": 200,
      "median_ms": 0.03451417502219556,
      "q1_ms": 0.034437832509865984,
      "q3_ms": 0.03675964500871487,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 64,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0003806430322583765,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.24781788200198207,
      "inner": 200,
      "median_ms": 0.03262359998188913,
      "q1_ms": 0.03228308249163092,
      "q3_ms": 0.032697490005375585,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 64,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006908720242790878,
      "xla_temp_mb": 4.250244140625,
      "xla_arg_mb": 16.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.5978433740019682,
      "inner": 161,
      "median_ms": 0.12000155280328281,
      "q1_ms": 0.11998782919481923,
      "q3_ms": 0.12007016149079804,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 64,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006908720242790878,
      "xla_temp_mb": 4.250244140625,
      "xla_arg_mb": 16.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.5998103840029216,
      "inner": 196,
      "median_ms": 0.11826760712619015,
      "q1_ms": 0.11785442346044132,
      "q3_ms": 0.11829319132351773,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 128,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007525204564444721,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.31749940699955914,
      "inner": 200,
      "median_ms": 0.057427559986535925,
      "q1_ms": 0.05741892000514781,
      "q3_ms": 0.05745578750065761,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 128,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007525204564444721,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.3149000120029086,
      "inner": 200,
      "median_ms": 0.05712323498300975,
      "q1_ms": 0.05711717749363743,
      "q3_ms": 0.05712543748813914,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 128,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006482516182586551,
      "xla_temp_mb": 8.500244140625,
      "xla_arg_mb": 32.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.5438713580006151,
      "inner": 112,
      "median_ms": 0.31046405352494083,
      "q1_ms": 0.3095514285372961,
      "q3_ms": 0.31063758930421115,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 128,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006482516182586551,
      "xla_temp_mb": 8.500244140625,
      "xla_arg_mb": 32.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.5395689920042059,
      "inner": 107,
      "median_ms": 0.3108637289510753,
      "q1_ms": 0.3095272710265605,
      "q3_ms": 0.31114589252534197,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 128,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0005480899708345532,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.3477119879971724,
      "inner": 200,
      "median_ms": 0.057655345008242875,
      "q1_ms": 0.05760268999438267,
      "q3_ms": 0.057715155016921926,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 128,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0005480899708345532,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.35244232499826467,
      "inner": 200,
      "median_ms": 0.057308784998895135,
      "q1_ms": 0.05728794751121313,
      "q3_ms": 0.05733816749852849,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 128,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007646068697795272,
      "xla_temp_mb": 8.500244140625,
      "xla_arg_mb": 32.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.6220950299975812,
      "inner": 112,
      "median_ms": 0.2681526875026715,
      "q1_ms": 0.26804520088522366,
      "q3_ms": 0.26820157144616785,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 128,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007646068697795272,
      "xla_temp_mb": 8.500244140625,
      "xla_arg_mb": 32.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.618937769002514,
      "inner": 115,
      "median_ms": 0.2672105825774412,
      "q1_ms": 0.26678992172768945,
      "q3_ms": 0.26757418694931245,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 256,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006812745705246925,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.3153391219966579,
      "inner": 143,
      "median_ms": 0.1660857972008391,
      "q1_ms": 0.16577958043459542,
      "q3_ms": 0.16613998601454882,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 256,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006812745705246925,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.3131847279946669,
      "inner": 160,
      "median_ms": 0.16850086249178275,
      "q1_ms": 0.16309204377193964,
      "q3_ms": 0.16852885937623796,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 256,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0011262360494583845,
      "xla_temp_mb": 17.000244140625,
      "xla_arg_mb": 64.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.5432818330009468,
      "inner": 46,
      "median_ms": 0.9991688259659618,
      "q1_ms": 0.9972000000783263,
      "q3_ms": 1.0004450435624912,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 256,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0011262360494583845,
      "xla_temp_mb": 17.000244140625,
      "xla_arg_mb": 64.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.5431391940001049,
      "inner": 47,
      "median_ms": 0.9983402766359157,
      "q1_ms": 0.9971306383104479,
      "q3_ms": 0.9992333297633833,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 256,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0004299705324228853,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.3542157080009929,
      "inner": 200,
      "median_ms": 0.13984172997879796,
      "q1_ms": 0.1397803624968219,
      "q3_ms": 0.13990927998747793,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 256,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0004299705324228853,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.3492904529994121,
      "inner": 200,
      "median_ms": 0.14156203000311507,
      "q1_ms": 0.1383546900069632,
      "q3_ms": 0.14182690751113114,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 256,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007133711478672922,
      "xla_temp_mb": 17.000244140625,
      "xla_arg_mb": 64.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.6148431350011379,
      "inner": 61,
      "median_ms": 0.6881176393152764,
      "q1_ms": 0.6875035409781257,
      "q3_ms": 0.689910139342709,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 256,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007133711478672922,
      "xla_temp_mb": 17.000244140625,
      "xla_arg_mb": 64.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.6136532249947777,
      "inner": 63,
      "median_ms": 0.6893512063903646,
      "q1_ms": 0.6883748413071902,
      "q3_ms": 0.6895445872537999,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 512,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0005241255275905132,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.3194499969977187,
      "inner": 78,
      "median_ms": 0.5912437179806427,
      "q1_ms": 0.5907771282158464,
      "q3_ms": 0.59314037823172,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 512,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0005241255275905132,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.3152933050005231,
      "inner": 75,
      "median_ms": 0.5970804400082367,
      "q1_ms": 0.5948319266705464,
      "q3_ms": 0.5971930599965466,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 512,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0008133161463774741,
      "xla_temp_mb": 34.000244140625,
      "xla_arg_mb": 128.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 0.542187703998934,
      "inner": 14,
      "median_ms": 3.5142535000756783,
      "q1_ms": 3.5137365357513772,
      "q3_ms": 3.519966857181446,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 512,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0008133161463774741,
      "xla_temp_mb": 34.000244140625,
      "xla_arg_mb": 128.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 0.5387151070026448,
      "inner": 14,
      "median_ms": 3.5156211425990995,
      "q1_ms": 3.509106357211879,
      "q3_ms": 3.519971178580558,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 512,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00037968141259625554,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.35748874799901387,
      "inner": 86,
      "median_ms": 0.41053876745092227,
      "q1_ms": 0.41045965702133524,
      "q3_ms": 0.4109755814275812,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 512,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00037968141259625554,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.35587778099579737,
      "inner": 111,
      "median_ms": 0.41588356760110257,
      "q1_ms": 0.41484602702480555,
      "q3_ms": 0.41618161262007775,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 512,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006316628423519433,
      "xla_temp_mb": 34.000244140625,
      "xla_arg_mb": 128.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 0.6203048959941952,
      "inner": 24,
      "median_ms": 2.087189500040646,
      "q1_ms": 2.086023791586437,
      "q3_ms": 2.0944221666165204,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 32,
      "seq": 512,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006316628423519433,
      "xla_temp_mb": 34.000244140625,
      "xla_arg_mb": 128.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 0.618342414003564,
      "inner": 24,
      "median_ms": 2.0943422083898136,
      "q1_ms": 2.087619312684789,
      "q3_ms": 2.0957213124953946,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 16,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006341451080515981,
      "xla_temp_mb": 3.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.18642382800317137,
      "inner": 200,
      "median_ms": 0.04903203000139911,
      "q1_ms": 0.04883393499767408,
      "q3_ms": 0.04913595250400249,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 16,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0004391638212837279,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.10923207800078671,
      "inner": 200,
      "median_ms": 0.04432097000972135,
      "q1_ms": 0.04354815999249695,
      "q3_ms": 0.04445873499207664,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 16,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006048346986062825,
      "xla_temp_mb": 3.0,
      "xla_arg_mb": 8.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.3140958310032147,
      "inner": 200,
      "median_ms": 0.07202581498859217,
      "q1_ms": 0.07193695249952725,
      "q3_ms": 0.07205718002296635,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 16,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.000719865201972425,
      "xla_temp_mb": 2.062744140625,
      "xla_arg_mb": 8.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.22740114699990954,
      "inner": 200,
      "median_ms": 0.04271378002158599,
      "q1_ms": 0.04259884002749459,
      "q3_ms": 0.04437582751052105,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 16,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0005328193074092269,
      "xla_temp_mb": 3.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.19375483300245833,
      "inner": 200,
      "median_ms": 0.03687634500238346,
      "q1_ms": 0.03683100751004531,
      "q3_ms": 0.036949304994777776,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 16,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0005072162603028119,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.12469789700116962,
      "inner": 200,
      "median_ms": 0.036286479989939835,
      "q1_ms": 0.036134189995209454,
      "q3_ms": 0.0365420599882782,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 16,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00047061077202670276,
      "xla_temp_mb": 3.0,
      "xla_arg_mb": 8.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.5756289099954301,
      "inner": 200,
      "median_ms": 0.0744231000135187,
      "q1_ms": 0.07437841748469509,
      "q3_ms": 0.07447573001627461,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 16,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0008460411918349564,
      "xla_temp_mb": 2.062744140625,
      "xla_arg_mb": 8.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.258858370994858,
      "inner": 200,
      "median_ms": 0.04540075999102555,
      "q1_ms": 0.044766992486984236,
      "q3_ms": 0.0459803874946374,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 32,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00044052154407836497,
      "xla_temp_mb": 8.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.27242842599662254,
      "inner": 200,
      "median_ms": 0.06483977998868795,
      "q1_ms": 0.06477057500887895,
      "q3_ms": 0.06492611750218202,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 32,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00041910100844688714,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.14537430200289236,
      "inner": 200,
      "median_ms": 0.04769921499246266,
      "q1_ms": 0.04218265252347919,
      "q3_ms": 0.04818287499801954,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 32,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007663984433747828,
      "xla_temp_mb": 8.0,
      "xla_arg_mb": 16.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.4897191920026671,
      "inner": 197,
      "median_ms": 0.15108991370135727,
      "q1_ms": 0.1507054289435344,
      "q3_ms": 0.1512646370701525,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 32,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007477875333279371,
      "xla_temp_mb": 4.125244140625,
      "xla_arg_mb": 16.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.3316489479984739,
      "inner": 200,
      "median_ms": 0.08876212003087858,
      "q1_ms": 0.0885645500056853,
      "q3_ms": 0.08884157499778667,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 32,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006300578825175762,
      "xla_temp_mb": 8.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.3220219909999287,
      "inner": 200,
      "median_ms": 0.06678896497760434,
      "q1_ms": 0.0667119450008613,
      "q3_ms": 0.06679099498796859,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 32,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0004133675538469106,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.16457146900211228,
      "inner": 200,
      "median_ms": 0.03512412997224601,
      "q1_ms": 0.03484765751636587,
      "q3_ms": 0.0357483350126131,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 32,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007315461989492178,
      "xla_temp_mb": 8.0,
      "xla_arg_mb": 16.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 1.0801946810024674,
      "inner": 177,
      "median_ms": 0.15868706782256375,
      "q1_ms": 0.15858733049599785,
      "q3_ms": 0.15872513840753072,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 32,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.000715727626811713,
      "xla_temp_mb": 4.125244140625,
      "xla_arg_mb": 16.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.3539193870019517,
      "inner": 200,
      "median_ms": 0.09196517999953358,
      "q1_ms": 0.09155330248177052,
      "q3_ms": 0.09197254248647369,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 64,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007218203973025084,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.20904994000011357,
      "inner": 200,
      "median_ms": 0.04563906499242876,
      "q1_ms": 0.04560056748232455,
      "q3_ms": 0.04566276251352974,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 64,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007218203973025084,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.19826478300092276,
      "inner": 200,
      "median_ms": 0.045631155007868074,
      "q1_ms": 0.045618475014634896,
      "q3_ms": 0.045702562492806464,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 64,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006689945585094392,
      "xla_temp_mb": 8.250244140625,
      "xla_arg_mb": 32.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.3333481519948691,
      "inner": 147,
      "median_ms": 0.21813915645288698,
      "q1_ms": 0.21803607145260204,
      "q3_ms": 0.21815626190751192,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 64,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006689945585094392,
      "xla_temp_mb": 8.250244140625,
      "xla_arg_mb": 32.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.33186068900249666,
      "inner": 165,
      "median_ms": 0.21786392123127976,
      "q1_ms": 0.21775000908140402,
      "q3_ms": 0.21788802425197157,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 64,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00039393448969349265,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.22327280000172323,
      "inner": 200,
      "median_ms": 0.047032805014168844,
      "q1_ms": 0.045972744992468506,
      "q3_ms": 0.0472858300054213,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 64,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00039393448969349265,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.21473631499975454,
      "inner": 200,
      "median_ms": 0.04592730998410843,
      "q1_ms": 0.04591140748743783,
      "q3_ms": 0.04600014750394621,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 64,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007203149725683033,
      "xla_temp_mb": 8.250244140625,
      "xla_arg_mb": 32.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.3601614020008128,
      "inner": 142,
      "median_ms": 0.19935703522592954,
      "q1_ms": 0.1989342288678171,
      "q3_ms": 0.19945821479499065,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 64,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007203149725683033,
      "xla_temp_mb": 8.250244140625,
      "xla_arg_mb": 32.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.36043635299574817,
      "inner": 149,
      "median_ms": 0.19896535572811538,
      "q1_ms": 0.19888675167798003,
      "q3_ms": 0.199012261727081,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 128,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0004716230323538184,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.2091249320001225,
      "inner": 200,
      "median_ms": 0.09941402000549715,
      "q1_ms": 0.09698835749077261,
      "q3_ms": 0.09947786500561051,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 128,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0004716230323538184,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.20445464899967192,
      "inner": 200,
      "median_ms": 0.0988481950116693,
      "q1_ms": 0.09612687750632176,
      "q3_ms": 0.09986239499994554,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 128,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0008243123884312809,
      "xla_temp_mb": 16.500244140625,
      "xla_arg_mb": 64.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.3340904050055542,
      "inner": 66,
      "median_ms": 0.5966977576353009,
      "q1_ms": 0.5963476212469465,
      "q3_ms": 0.5970324849234506,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 128,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0008243123884312809,
      "xla_temp_mb": 16.500244140625,
      "xla_arg_mb": 64.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.3367441840018728,
      "inner": 66,
      "median_ms": 0.5966028940306525,
      "q1_ms": 0.5962685454230241,
      "q3_ms": 0.5968156363616635,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 128,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0003490043163765222,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.22356615400349256,
      "inner": 200,
      "median_ms": 0.0920166300056735,
      "q1_ms": 0.09177905250908225,
      "q3_ms": 0.09208217748891911,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 128,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0003490043163765222,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.21868866000295384,
      "inner": 200,
      "median_ms": 0.09224280500347959,
      "q1_ms": 0.09035340750415344,
      "q3_ms": 0.09239096250894363,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 128,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006634605815634131,
      "xla_temp_mb": 16.500244140625,
      "xla_arg_mb": 64.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.3629180260031717,
      "inner": 84,
      "median_ms": 0.4706341428356522,
      "q1_ms": 0.4697066785974173,
      "q3_ms": 0.4707129821240198,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 128,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006634605815634131,
      "xla_temp_mb": 16.500244140625,
      "xla_arg_mb": 64.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.3599492220018874,
      "inner": 79,
      "median_ms": 0.47087549365308323,
      "q1_ms": 0.47020561391888543,
      "q3_ms": 0.47101974052589246,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 256,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00035923015093430877,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.20958732600411167,
      "inner": 123,
      "median_ms": 0.3185668861781194,
      "q1_ms": 0.31852208943805677,
      "q3_ms": 0.3186120487864572,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 256,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00035923015093430877,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.2028430920036044,
      "inner": 119,
      "median_ms": 0.32119293279697364,
      "q1_ms": 0.32015747897925656,
      "q3_ms": 0.3212581764617563,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 256,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0009009960340335965,
      "xla_temp_mb": 33.000244140625,
      "xla_arg_mb": 128.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 0.3335896290009259,
      "inner": 24,
      "median_ms": 1.9473867917743821,
      "q1_ms": 1.943628624909858,
      "q3_ms": 1.9556954166546348,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 256,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0009009960340335965,
      "xla_temp_mb": 33.000244140625,
      "xla_arg_mb": 128.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 0.33618720099912025,
      "inner": 24,
      "median_ms": 1.94521183342052,
      "q1_ms": 1.9432654374516762,
      "q3_ms": 1.945492999918012,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 256,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0003968327073380351,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.2257597439966048,
      "inner": 149,
      "median_ms": 0.22852138924567128,
      "q1_ms": 0.22847610400288815,
      "q3_ms": 0.22879503017401592,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 256,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0003968327073380351,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.2203819520000252,
      "inner": 145,
      "median_ms": 0.23062748965959803,
      "q1_ms": 0.23053421035252414,
      "q3_ms": 0.2309040517439857,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 256,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007929968414828181,
      "xla_temp_mb": 33.000244140625,
      "xla_arg_mb": 128.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 0.3624370760007878,
      "inner": 36,
      "median_ms": 1.2902814166106207,
      "q1_ms": 1.2891805416883368,
      "q3_ms": 1.2906808749701466,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 256,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007929968414828181,
      "xla_temp_mb": 33.000244140625,
      "xla_arg_mb": 128.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 0.3629551569974865,
      "inner": 36,
      "median_ms": 1.2907166387271394,
      "q1_ms": 1.2898468055734864,
      "q3_ms": 1.2908205971042561,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 512,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006678033387288451,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 192.0,
      "xla_out_mb": 64.0,
      "compile_s": 0.20801525300339563,
      "inner": 39,
      "median_ms": 1.180016359005374,
      "q1_ms": 1.1782777948657241,
      "q3_ms": 1.190064461507763,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 512,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0006678033387288451,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 192.0,
      "xla_out_mb": 64.0,
      "compile_s": 0.20302733999415068,
      "inner": 42,
      "median_ms": 1.1860918095321504,
      "q1_ms": 1.185805654781039,
      "q3_ms": 1.1902450000239437,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 512,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.001239205477759242,
      "xla_temp_mb": 66.000244140625,
      "xla_arg_mb": 256.0,
      "xla_out_mb": 192.0000228881836,
      "compile_s": 0.33799584199732635,
      "inner": 7,
      "median_ms": 7.042892857238518,
      "q1_ms": 7.030951071750938,
      "q3_ms": 7.058806500156477,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 512,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.001239205477759242,
      "xla_temp_mb": 66.000244140625,
      "xla_arg_mb": 256.0,
      "xla_out_mb": 192.0000228881836,
      "compile_s": 0.3394923119994928,
      "inner": 7,
      "median_ms": 7.041294143294051,
      "q1_ms": 7.029072428330048,
      "q3_ms": 7.0446289999901115,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 512,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0004773139953613281,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 192.0,
      "xla_out_mb": 64.0,
      "compile_s": 0.23114913600147702,
      "inner": 59,
      "median_ms": 0.7289133050184997,
      "q1_ms": 0.7286860254372202,
      "q3_ms": 0.7334743983184295,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 512,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0004773139953613281,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 192.0,
      "xla_out_mb": 64.0,
      "compile_s": 0.21859130499797175,
      "inner": 64,
      "median_ms": 0.7339692343748538,
      "q1_ms": 0.7339231249829936,
      "q3_ms": 0.7349344140834546,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 512,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007104802643880248,
      "xla_temp_mb": 66.000244140625,
      "xla_arg_mb": 256.0,
      "xla_out_mb": 192.0000228881836,
      "compile_s": 0.3607448859984288,
      "inner": 12,
      "median_ms": 4.136974999710219,
      "q1_ms": 4.1278259999065385,
      "q3_ms": 4.145845499806455,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 512,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0007104802643880248,
      "xla_temp_mb": 66.000244140625,
      "xla_arg_mb": 256.0,
      "xla_out_mb": 192.0000228881836,
      "compile_s": 0.36354649499844527,
      "inner": 12,
      "median_ms": 4.125223500523134,
      "q1_ms": 4.121894833284992,
      "q3_ms": 4.127421500015771,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 16,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004193125292658806,
      "xla_temp_mb": 1.0,
      "xla_arg_mb": 1.5,
      "xla_out_mb": 0.5,
      "compile_s": 0.35205756599316373,
      "inner": 200,
      "median_ms": 0.04461172498849919,
      "q1_ms": 0.04425136499776272,
      "q3_ms": 0.044937037491763476,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 16,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0033779738005250692,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 1.5,
      "xla_out_mb": 0.5,
      "compile_s": 0.09786658200027887,
      "inner": 200,
      "median_ms": 0.04137640498811379,
      "q1_ms": 0.04014863250631606,
      "q3_ms": 0.04241885000737966,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 16,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0033779738005250692,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 1.5,
      "xla_out_mb": 0.5,
      "compile_s": 0.0611848269982147,
      "inner": 200,
      "median_ms": 0.06592483001440996,
      "q1_ms": 0.06567460752194165,
      "q3_ms": 0.06632489250478102,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 16,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.006359456572681665,
      "xla_temp_mb": 1.0,
      "xla_arg_mb": 2.5,
      "xla_out_mb": 1.5000228881835938,
      "compile_s": 0.3504167930004769,
      "inner": 200,
      "median_ms": 0.05439026997919427,
      "q1_ms": 0.0525610324984882,
      "q3_ms": 0.059102527502545854,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 16,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0043426030315458775,
      "xla_temp_mb": 1.062744140625,
      "xla_arg_mb": 2.5,
      "xla_out_mb": 1.5000228881835938,
      "compile_s": 0.1748023589971126,
      "inner": 200,
      "median_ms": 0.04892197503068019,
      "q1_ms": 0.04823950999707449,
      "q3_ms": 0.049170719994435785,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 16,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0043426030315458775,
      "xla_temp_mb": 3.063232421875,
      "xla_arg_mb": 2.5,
      "xla_out_mb": 1.5000228881835938,
      "compile_s": 0.10750683399965055,
      "inner": 200,
      "median_ms": 0.1088836299823015,
      "q1_ms": 0.10780558748592739,
      "q3_ms": 0.10936543998468551,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 16,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00428376067429781,
      "xla_temp_mb": 1.0,
      "xla_arg_mb": 1.5,
      "xla_out_mb": 0.5,
      "compile_s": 0.2147465439993539,
      "inner": 200,
      "median_ms": 0.039257150019693654,
      "q1_ms": 0.03899154251485015,
      "q3_ms": 0.039422862500941847,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 16,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0031301137059926987,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 1.5,
      "xla_out_mb": 0.5,
      "compile_s": 0.11712076699768659,
      "inner": 200,
      "median_ms": 0.0359760399805964,
      "q1_ms": 0.03529795751092024,
      "q3_ms": 0.03635606000898406,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 16,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0031301137059926987,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 1.5,
      "xla_out_mb": 0.5,
      "compile_s": 0.04610003800189588,
      "inner": 177,
      "median_ms": 0.059025192073885735,
      "q1_ms": 0.05892474577307706,
      "q3_ms": 0.059296949172989694,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 16,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00491351168602705,
      "xla_temp_mb": 1.0,
      "xla_arg_mb": 2.5,
      "xla_out_mb": 1.5000228881835938,
      "compile_s": 0.4464600369974505,
      "inner": 200,
      "median_ms": 0.05314100500982022,
      "q1_ms": 0.053058322500874056,
      "q3_ms": 0.053223885006445926,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 16,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004201397765427828,
      "xla_temp_mb": 1.062744140625,
      "xla_arg_mb": 2.5,
      "xla_out_mb": 1.5000228881835938,
      "compile_s": 0.23519798400229774,
      "inner": 200,
      "median_ms": 0.04875300499406876,
      "q1_ms": 0.048640617496857885,
      "q3_ms": 0.048855047498364,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 16,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004201397765427828,
      "xla_temp_mb": 3.063232421875,
      "xla_arg_mb": 2.5,
      "xla_out_mb": 1.5000228881835938,
      "compile_s": 0.11044118399877334,
      "inner": 200,
      "median_ms": 0.1158396850223653,
      "q1_ms": 0.1155697524882271,
      "q3_ms": 0.1160491075279424,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 32,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004176608286798,
      "xla_temp_mb": 3.0,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.2798211179979262,
      "inner": 200,
      "median_ms": 0.04067863501404645,
      "q1_ms": 0.04027156750453287,
      "q3_ms": 0.042553357507131295,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 32,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0022217121440917253,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.10245448900241172,
      "inner": 200,
      "median_ms": 0.04540166501101339,
      "q1_ms": 0.038330620009219274,
      "q3_ms": 0.04558964752504835,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 32,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0022217121440917253,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.04208756899606669,
      "inner": 200,
      "median_ms": 0.05985634998069145,
      "q1_ms": 0.05951914250545087,
      "q3_ms": 0.060495840007206425,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 32,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.006654238328337669,
      "xla_temp_mb": 5.0625,
      "xla_arg_mb": 5.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.5412438290004502,
      "inner": 200,
      "median_ms": 0.06534027499583317,
      "q1_ms": 0.06525574000988854,
      "q3_ms": 0.06536845499795163,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 32,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.006654238328337669,
      "xla_temp_mb": 2.125244140625,
      "xla_arg_mb": 5.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.208715608001512,
      "inner": 200,
      "median_ms": 0.0481859299907228,
      "q1_ms": 0.04809498248505406,
      "q3_ms": 0.04839577748498414,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 32,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.006654238328337669,
      "xla_temp_mb": 6.125732421875,
      "xla_arg_mb": 5.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.10557324800174683,
      "inner": 170,
      "median_ms": 0.10915337646592355,
      "q1_ms": 0.1087838264699231,
      "q3_ms": 0.10942483234376757,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 32,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004325577989220619,
      "xla_temp_mb": 3.0,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.3502949030007585,
      "inner": 200,
      "median_ms": 0.03704957001900766,
      "q1_ms": 0.03699228998812032,
      "q3_ms": 0.03711945499162539,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 32,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0032848461996763945,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.13071456699981354,
      "inner": 200,
      "median_ms": 0.03446367998549249,
      "q1_ms": 0.03433488749578828,
      "q3_ms": 0.034474125004635425,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 32,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0032848461996763945,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.04637116999947466,
      "inner": 200,
      "median_ms": 0.06077949998143595,
      "q1_ms": 0.06041694248779095,
      "q3_ms": 0.06161620998682338,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 32,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004901744890958071,
      "xla_temp_mb": 5.0625,
      "xla_arg_mb": 5.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.6854354249953758,
      "inner": 200,
      "median_ms": 0.06608318002690794,
      "q1_ms": 0.06604445748962462,
      "q3_ms": 0.06615058249735739,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 32,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.005748792551457882,
      "xla_temp_mb": 2.125244140625,
      "xla_arg_mb": 5.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.2889189640045515,
      "inner": 200,
      "median_ms": 0.052045929987798445,
      "q1_ms": 0.051454665008350275,
      "q3_ms": 0.052110452506894944,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 32,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.005748792551457882,
      "xla_temp_mb": 6.125732421875,
      "xla_arg_mb": 5.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.105174023003201,
      "inner": 191,
      "median_ms": 0.11342247117296134,
      "q1_ms": 0.113071704189493,
      "q3_ms": 0.11366069634872883,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 64,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0027185636572539806,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.12156028400204377,
      "inner": 200,
      "median_ms": 0.03591001499444246,
      "q1_ms": 0.03579973250452895,
      "q3_ms": 0.036061754981346894,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 64,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0027185636572539806,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.11721133800165262,
      "inner": 159,
      "median_ms": 0.03610696853541201,
      "q1_ms": 0.03586213208120698,
      "q3_ms": 0.03615718865920878,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 64,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0027185636572539806,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.042665752000175416,
      "inner": 200,
      "median_ms": 0.055545899995195214,
      "q1_ms": 0.055418637493858114,
      "q3_ms": 0.055629359994782135,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 64,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004016230348497629,
      "xla_temp_mb": 4.250244140625,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.26255383000534493,
      "inner": 200,
      "median_ms": 0.05965552500128979,
      "q1_ms": 0.05951640998318908,
      "q3_ms": 0.059778682498290436,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 64,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004016230348497629,
      "xla_temp_mb": 4.250244140625,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.2638046460051555,
      "inner": 200,
      "median_ms": 0.05896682501770556,
      "q1_ms": 0.05862202249772963,
      "q3_ms": 0.05950261001999024,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 64,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004016230348497629,
      "xla_temp_mb": 12.250732421875,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.10329388699756237,
      "inner": 173,
      "median_ms": 0.14117277455125865,
      "q1_ms": 0.14106465319024852,
      "q3_ms": 0.1419174797690643,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 64,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0032096244394779205,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.16853815100330394,
      "inner": 200,
      "median_ms": 0.038265069997578394,
      "q1_ms": 0.03795756500039715,
      "q3_ms": 0.039556907504447736,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 64,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0032096244394779205,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.15732228900014888,
      "inner": 200,
      "median_ms": 0.03847276500891894,
      "q1_ms": 0.038242177506617736,
      "q3_ms": 0.039006697516015265,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 64,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0032096244394779205,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.044071530996006913,
      "inner": 176,
      "median_ms": 0.05717711366576789,
      "q1_ms": 0.05643648294360239,
      "q3_ms": 0.05819437784669969,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 64,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0061174426227808,
      "xla_temp_mb": 4.250244140625,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.3529332929974771,
      "inner": 200,
      "median_ms": 0.06216827499883947,
      "q1_ms": 0.06212779751876951,
      "q3_ms": 0.062387889993260615,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 64,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0061174426227808,
      "xla_temp_mb": 4.250244140625,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.3555352440016577,
      "inner": 200,
      "median_ms": 0.0595778299975791,
      "q1_ms": 0.059537177494348725,
      "q3_ms": 0.05985843499729526,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 64,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0061174426227808,
      "xla_temp_mb": 12.250732421875,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.1053570590011077,
      "inner": 152,
      "median_ms": 0.13856007894661926,
      "q1_ms": 0.13848029278195004,
      "q3_ms": 0.13857659209143572,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 128,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.002333720214664936,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.14225942100165412,
      "inner": 200,
      "median_ms": 0.031210025008476805,
      "q1_ms": 0.03043100248760311,
      "q3_ms": 0.036530262495944044,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 128,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.002333720214664936,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.13572402099816827,
      "inner": 200,
      "median_ms": 0.031180975020106416,
      "q1_ms": 0.031158179990598,
      "q3_ms": 0.03601734000767465,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 128,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.002333720214664936,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.04202742499910528,
      "inner": 200,
      "median_ms": 0.05677908498910256,
      "q1_ms": 0.05541407748751226,
      "q3_ms": 0.06168256000819383,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 128,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.005233302712440491,
      "xla_temp_mb": 8.500244140625,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.30486840900266543,
      "inner": 182,
      "median_ms": 0.14390678572433918,
      "q1_ms": 0.1435519917412331,
      "q3_ms": 0.1439734478081336,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 128,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.005233302712440491,
      "xla_temp_mb": 8.500244140625,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.3058436840001377,
      "inner": 184,
      "median_ms": 0.14345979893139724,
      "q1_ms": 0.1417104293503381,
      "q3_ms": 0.14362696196179098,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 128,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.005233302712440491,
      "xla_temp_mb": 24.500732421875,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.10247576900292188,
      "inner": 110,
      "median_ms": 0.30058512724777264,
      "q1_ms": 0.30030682725985325,
      "q3_ms": 0.30063161816955963,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 128,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.003435371909290552,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.16337217100226553,
      "inner": 200,
      "median_ms": 0.039534719981020316,
      "q1_ms": 0.0350149474797945,
      "q3_ms": 0.04278882499420433,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 128,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.003435371909290552,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.1598738750035409,
      "inner": 200,
      "median_ms": 0.04785115001141094,
      "q1_ms": 0.046336987488757586,
      "q3_ms": 0.04988126247553737,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 128,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.003435371909290552,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.04343768199760234,
      "inner": 200,
      "median_ms": 0.058971430007659364,
      "q1_ms": 0.05833146751683671,
      "q3_ms": 0.060503202494146535,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 128,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00502068130299449,
      "xla_temp_mb": 8.500244140625,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.3622708330003661,
      "inner": 200,
      "median_ms": 0.13104263500281377,
      "q1_ms": 0.13100774502163404,
      "q3_ms": 0.13105030251608696,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 128,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00502068130299449,
      "xla_temp_mb": 8.500244140625,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.3590101349982433,
      "inner": 200,
      "median_ms": 0.12965646499651484,
      "q1_ms": 0.1287704399874201,
      "q3_ms": 0.12967520000529476,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 128,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00502068130299449,
      "xla_temp_mb": 24.500732421875,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.10969447200477589,
      "inner": 128,
      "median_ms": 0.2701614999978119,
      "q1_ms": 0.270142749968727,
      "q3_ms": 0.2701931953197345,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 256,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.002204917138442397,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.14404360600019572,
      "inner": 200,
      "median_ms": 0.08663750501000322,
      "q1_ms": 0.08412137251070817,
      "q3_ms": 0.08699865999005851,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 256,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.002204917138442397,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.141081863999716,
      "inner": 200,
      "median_ms": 0.08648556999105494,
      "q1_ms": 0.08338838499184931,
      "q3_ms": 0.08675237000716152,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 256,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.002204917138442397,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.04248694400303066,
      "inner": 173,
      "median_ms": 0.17066769365987045,
      "q1_ms": 0.17027162715948774,
      "q3_ms": 0.17128223697977907,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 256,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0038350424729287624,
      "xla_temp_mb": 17.000244140625,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.31119995399785694,
      "inner": 102,
      "median_ms": 0.43493915683368456,
      "q1_ms": 0.4348093970845951,
      "q3_ms": 0.4353472304144673,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 256,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0038350424729287624,
      "xla_temp_mb": 17.000244140625,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.3115342960009002,
      "inner": 97,
      "median_ms": 0.43546864954927533,
      "q1_ms": 0.43479246390412185,
      "q3_ms": 0.4356655051956342,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 256,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0038350424729287624,
      "xla_temp_mb": 49.000732421875,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.10727828199742362,
      "inner": 51,
      "median_ms": 0.8503838039092336,
      "q1_ms": 0.8501912647368632,
      "q3_ms": 0.8507340686514462,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 256,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.003190818941220641,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.1620371060053003,
      "inner": 200,
      "median_ms": 0.06133840997790685,
      "q1_ms": 0.06021736249749665,
      "q3_ms": 0.06372452000505291,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 256,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.003190818941220641,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.16027720700367354,
      "inner": 200,
      "median_ms": 0.06386118002410512,
      "q1_ms": 0.06166677998407977,
      "q3_ms": 0.06397567000021809,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 256,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.003190818941220641,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.042969996997271664,
      "inner": 175,
      "median_ms": 0.1467857200103546,
      "q1_ms": 0.14306677715337301,
      "q3_ms": 0.1482033485496816,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 256,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.005156615749001503,
      "xla_temp_mb": 17.000244140625,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.36242756900173845,
      "inner": 107,
      "median_ms": 0.3282750654659913,
      "q1_ms": 0.3258009439088931,
      "q3_ms": 0.32835334113599934,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 256,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.005156615749001503,
      "xla_temp_mb": 17.000244140625,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.36391255800117506,
      "inner": 118,
      "median_ms": 0.328198169450319,
      "q1_ms": 0.3279169999956848,
      "q3_ms": 0.32828681355182443,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 256,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.005156615749001503,
      "xla_temp_mb": 49.000732421875,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.10600278999481816,
      "inner": 59,
      "median_ms": 0.680165525409393,
      "q1_ms": 0.6799242288190911,
      "q3_ms": 0.6802806609642684,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 512,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.003060600720345974,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.14256644200213486,
      "inner": 125,
      "median_ms": 0.3077566879801452,
      "q1_ms": 0.30759233998833224,
      "q3_ms": 0.30811785999685526,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 512,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.003060600720345974,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.13669304399809334,
      "inner": 127,
      "median_ms": 0.30757043307636106,
      "q1_ms": 0.3075300905267767,
      "q3_ms": 0.3078843031354032,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 512,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.003060600720345974,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.04250409700034652,
      "inner": 69,
      "median_ms": 0.5795112318583711,
      "q1_ms": 0.5791978768209683,
      "q3_ms": 0.5795989347751711,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 512,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.005082274321466684,
      "xla_temp_mb": 34.000244140625,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.31512675099656917,
      "inner": 33,
      "median_ms": 1.514980424183181,
      "q1_ms": 1.5135634698019589,
      "q3_ms": 1.5155341817731052,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 512,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.005082274321466684,
      "xla_temp_mb": 34.000244140625,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.3105780080004479,
      "inner": 32,
      "median_ms": 1.5151277186760126,
      "q1_ms": 1.5137034687313644,
      "q3_ms": 1.5156936093490003,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 512,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.005082274321466684,
      "xla_temp_mb": 98.000732421875,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.1061396090008202,
      "inner": 17,
      "median_ms": 2.718004588078817,
      "q1_ms": 2.713302147158709,
      "q3_ms": 2.7183590881416904,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 512,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0030704804230481386,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.1638938180039986,
      "inner": 151,
      "median_ms": 0.1941351126050669,
      "q1_ms": 0.1922786821355266,
      "q3_ms": 0.19422345694512227,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 512,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0030704804230481386,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.1607789670015336,
      "inner": 188,
      "median_ms": 0.19422413831133337,
      "q1_ms": 0.1942009813571163,
      "q3_ms": 0.1943017021481052,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 512,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0030704804230481386,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.04235306099872105,
      "inner": 93,
      "median_ms": 0.4083918494714414,
      "q1_ms": 0.40604298923700366,
      "q3_ms": 0.4089246774314507,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 512,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004873078316450119,
      "xla_temp_mb": 34.000244140625,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.35681081600341713,
      "inner": 48,
      "median_ms": 0.9796542500832099,
      "q1_ms": 0.9787241561601453,
      "q3_ms": 0.9798894791401835,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 512,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004873078316450119,
      "xla_temp_mb": 34.000244140625,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.3631597729981877,
      "inner": 48,
      "median_ms": 0.980115750129092,
      "q1_ms": 0.9785256667479795,
      "q3_ms": 0.9818592604157553,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 32,
      "seq": 512,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004873078316450119,
      "xla_temp_mb": 98.000732421875,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.10454394100088393,
      "inner": 23,
      "median_ms": 1.8641081304572846,
      "q1_ms": 1.855177565195608,
      "q3_ms": 1.8647656957593584,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 16,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004541471600532532,
      "xla_temp_mb": 1.75,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.21305602599750273,
      "inner": 200,
      "median_ms": 0.04078117501194356,
      "q1_ms": 0.04057508749610861,
      "q3_ms": 0.043823332507599844,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 16,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0027530461084097624,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.09983466499397764,
      "inner": 200,
      "median_ms": 0.034713804998318665,
      "q1_ms": 0.03465281250100816,
      "q3_ms": 0.03483432999928482,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 16,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0027530461084097624,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.04350196300219977,
      "inner": 200,
      "median_ms": 0.056734350000624545,
      "q1_ms": 0.056603555003675865,
      "q3_ms": 0.056909605027613,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 16,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.005383178126066923,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 5.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.3391737700003432,
      "inner": 200,
      "median_ms": 0.04943294497934403,
      "q1_ms": 0.049392410001019016,
      "q3_ms": 0.04949333000695333,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 16,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.005383178126066923,
      "xla_temp_mb": 2.062744140625,
      "xla_arg_mb": 5.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.18822601700230734,
      "inner": 200,
      "median_ms": 0.056599700001243036,
      "q1_ms": 0.05026383749282104,
      "q3_ms": 0.05794989752757829,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 16,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.005383178126066923,
      "xla_temp_mb": 6.063232421875,
      "xla_arg_mb": 5.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.11004197900183499,
      "inner": 162,
      "median_ms": 0.14001820369387896,
      "q1_ms": 0.13993746912923216,
      "q3_ms": 0.14040347532837014,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 16,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004810147453099489,
      "xla_temp_mb": 1.75,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.2161157739974442,
      "inner": 200,
      "median_ms": 0.04193776501779212,
      "q1_ms": 0.041665499993541744,
      "q3_ms": 0.042290307483199285,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 16,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0030694082379341125,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.12542191700049443,
      "inner": 200,
      "median_ms": 0.035273409994260874,
      "q1_ms": 0.03525144498780719,
      "q3_ms": 0.03540407999025774,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 16,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0030694082379341125,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.044771414999559056,
      "inner": 200,
      "median_ms": 0.05673204501363216,
      "q1_ms": 0.05672127999787335,
      "q3_ms": 0.05683153500285698,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 16,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00432160496711731,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 5.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.45376212800329085,
      "inner": 200,
      "median_ms": 0.04992316502466565,
      "q1_ms": 0.0499031824983831,
      "q3_ms": 0.05144239250512328,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 16,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004320277366787195,
      "xla_temp_mb": 2.062744140625,
      "xla_arg_mb": 5.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.24663543500355445,
      "inner": 200,
      "median_ms": 0.049503224981890526,
      "q1_ms": 0.04921981249935925,
      "q3_ms": 0.04961763999745017,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 16,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004320277366787195,
      "xla_temp_mb": 6.063232421875,
      "xla_arg_mb": 5.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.10790699200151721,
      "inner": 155,
      "median_ms": 0.14134680644811823,
      "q1_ms": 0.13340808708885415,
      "q3_ms": 0.14200121292942053,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 32,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004443137440830469,
      "xla_temp_mb": 4.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.3228757460019551,
      "inner": 200,
      "median_ms": 0.03665242500574095,
      "q1_ms": 0.03661647750050179,
      "q3_ms": 0.03822301750915358,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 32,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0027243313379585743,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.11882870699628256,
      "inner": 200,
      "median_ms": 0.03600319498218596,
      "q1_ms": 0.03593448751416872,
      "q3_ms": 0.03626826748586609,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 32,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0027243313379585743,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.04612608699972043,
      "inner": 200,
      "median_ms": 0.056413445017824415,
      "q1_ms": 0.056342235002375674,
      "q3_ms": 0.05901601751247654,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 32,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0059395041316747665,
      "xla_temp_mb": 4.0625,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.5277621519999229,
      "inner": 200,
      "median_ms": 0.0890897950012004,
      "q1_ms": 0.0890547749986581,
      "q3_ms": 0.0891033699917898,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 32,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0059395041316747665,
      "xla_temp_mb": 4.125244140625,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.21850839899707353,
      "inner": 200,
      "median_ms": 0.05513988497114042,
      "q1_ms": 0.0538219675036089,
      "q3_ms": 0.05527085999347037,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 32,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0059395041316747665,
      "xla_temp_mb": 12.125732421875,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.1057463760007522,
      "inner": 168,
      "median_ms": 0.1640905773976036,
      "q1_ms": 0.16403756251834456,
      "q3_ms": 0.16412114285101,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 32,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00452336436137557,
      "xla_temp_mb": 4.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.347788783001306,
      "inner": 200,
      "median_ms": 0.03697642998304218,
      "q1_ms": 0.036881200030620676,
      "q3_ms": 0.04500193750573089,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 32,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0031470218673348427,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.14792874300474068,
      "inner": 200,
      "median_ms": 0.03713072499522241,
      "q1_ms": 0.036957544998585945,
      "q3_ms": 0.03776459501750651,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 32,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0031470218673348427,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.0441424249947886,
      "inner": 161,
      "median_ms": 0.05679462109193472,
      "q1_ms": 0.056675447210537924,
      "q3_ms": 0.057034950310390395,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 32,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.005252570379525423,
      "xla_temp_mb": 4.0625,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.677135707999696,
      "inner": 200,
      "median_ms": 0.09047340499819256,
      "q1_ms": 0.09046073499121121,
      "q3_ms": 0.09049412250533351,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 32,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004545743577182293,
      "xla_temp_mb": 4.125244140625,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.2897917950031115,
      "inner": 200,
      "median_ms": 0.05525750999368029,
      "q1_ms": 0.05503303249497549,
      "q3_ms": 0.05580171002293355,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 32,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004545743577182293,
      "xla_temp_mb": 12.125732421875,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.1101135749995592,
      "inner": 143,
      "median_ms": 0.16437507689591838,
      "q1_ms": 0.1643536678231256,
      "q3_ms": 0.1645706888030599,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 64,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0023481957614421844,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.15513235700200312,
      "inner": 200,
      "median_ms": 0.0313540000206558,
      "q1_ms": 0.03109901748757693,
      "q3_ms": 0.031558185000903904,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 64,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0023481957614421844,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.1478992870033835,
      "inner": 200,
      "median_ms": 0.03159470499667805,
      "q1_ms": 0.031540024992864346,
      "q3_ms": 0.03201207002348383,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 64,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0023481957614421844,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.04350219399930211,
      "inner": 200,
      "median_ms": 0.056694875020184554,
      "q1_ms": 0.05656368000927614,
      "q3_ms": 0.0568653750087833,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 64,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004327160306274891,
      "xla_temp_mb": 8.250244140625,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.34862840099958703,
      "inner": 200,
      "median_ms": 0.11026092997781234,
      "q1_ms": 0.1102299000012863,
      "q3_ms": 0.11027854250642122,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 64,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004327160306274891,
      "xla_temp_mb": 8.250244140625,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.35292093900352484,
      "inner": 200,
      "median_ms": 0.10859383000934031,
      "q1_ms": 0.10800677000588621,
      "q3_ms": 0.10862724249818712,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 64,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004327160306274891,
      "xla_temp_mb": 24.250732421875,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.10433234099764377,
      "inner": 128,
      "median_ms": 0.23370192189986483,
      "q1_ms": 0.23343034766298842,
      "q3_ms": 0.2338184609129712,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 64,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0031647554133087397,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.19649610700435005,
      "inner": 200,
      "median_ms": 0.031501529992965516,
      "q1_ms": 0.031234172492986545,
      "q3_ms": 0.03165641999657964,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 64,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0031647554133087397,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.19455547699908493,
      "inner": 200,
      "median_ms": 0.031809459978831,
      "q1_ms": 0.03174055249473895,
      "q3_ms": 0.03194108499883441,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 64,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0031647554133087397,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.04347894999955315,
      "inner": 200,
      "median_ms": 0.061107575020287186,
      "q1_ms": 0.059848867495020386,
      "q3_ms": 0.06237867499294225,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 64,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0038649465423077345,
      "xla_temp_mb": 8.250244140625,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.4695809660042869,
      "inner": 190,
      "median_ms": 0.11781993686638184,
      "q1_ms": 0.11752888682454604,
      "q3_ms": 0.11786444213756017,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 64,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0038649465423077345,
      "xla_temp_mb": 8.250244140625,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.4644909179987735,
      "inner": 190,
      "median_ms": 0.11519644735687364,
      "q1_ms": 0.11463195788612786,
      "q3_ms": 0.11528279474556544,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 64,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0038649465423077345,
      "xla_temp_mb": 24.250732421875,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.10548639299668139,
      "inner": 126,
      "median_ms": 0.23335848412003427,
      "q1_ms": 0.2329956984058732,
      "q3_ms": 0.23340480552250814,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00237796432338655,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.17697288399358513,
      "inner": 200,
      "median_ms": 0.048626515017531347,
      "q1_ms": 0.048214137514150934,
      "q3_ms": 0.05115554748044815,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00237796432338655,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.17579003600258147,
      "inner": 200,
      "median_ms": 0.050417449965607375,
      "q1_ms": 0.050387195024086395,
      "q3_ms": 0.05096992248581955,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.00237796432338655,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.04315295400010655,
      "inner": 200,
      "median_ms": 0.0633165450199158,
      "q1_ms": 0.06274311248489539,
      "q3_ms": 0.06575864248588914,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.005381712689995766,
      "xla_temp_mb": 16.500244140625,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.41769760999886785,
      "inner": 124,
      "median_ms": 0.2660768548066112,
      "q1_ms": 0.26541576208992707,
      "q3_ms": 0.2661193507998588,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.005381712689995766,
      "xla_temp_mb": 16.500244140625,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.42202384200209053,
      "inner": 130,
      "median_ms": 0.2665212923042978,
      "q1_ms": 0.2657871346365517,
      "q3_ms": 0.26655335386749357,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.005381712689995766,
      "xla_temp_mb": 48.500732421875,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.1064447660028236,
      "inner": 85,
      "median_ms": 0.4190626823829542,
      "q1_ms": 0.4189201765043168,
      "q3_ms": 0.41913110587355573,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0028695594519376755,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.20029738699668087,
      "inner": 200,
      "median_ms": 0.050284749995626044,
      "q1_ms": 0.050253967492608353,
      "q3_ms": 0.05033476998505648,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0028695594519376755,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.1949564550013747,
      "inner": 200,
      "median_ms": 0.05037838000134798,
      "q1_ms": 0.05031706499721622,
      "q3_ms": 0.050387370010867016,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0028695594519376755,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.046154351002769545,
      "inner": 200,
      "median_ms": 0.0653430799866328,
      "q1_ms": 0.06528414249260095,
      "q3_ms": 0.06541100749018369,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004662639927119017,
      "xla_temp_mb": 16.500244140625,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.4817675840022275,
      "inner": 122,
      "median_ms": 0.2465194344547287,
      "q1_ms": 0.24651310654842776,
      "q3_ms": 0.2465977786861734,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004662639927119017,
      "xla_temp_mb": 16.500244140625,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.48010324600181775,
      "inner": 129,
      "median_ms": 0.2452640930373879,
      "q1_ms": 0.2451240232389406,
      "q3_ms": 0.24536307753368766,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004662639927119017,
      "xla_temp_mb": 48.500732421875,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.10943145499913953,
      "inner": 86,
      "median_ms": 0.4253908139849785,
      "q1_ms": 0.42515757559158374,
      "q3_ms": 0.4254007732480775,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 256,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0034007239155471325,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.18045955799607327,
      "inner": 176,
      "median_ms": 0.15531948863347433,
      "q1_ms": 0.15510374150422152,
      "q3_ms": 0.15533881818638207,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 256,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0034007239155471325,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.17925585200282512,
      "inner": 178,
      "median_ms": 0.1559846685349737,
      "q1_ms": 0.15436778931268308,
      "q3_ms": 0.15599860393542278,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 256,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0034007239155471325,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.042494248002185486,
      "inner": 170,
      "median_ms": 0.16714069411939228,
      "q1_ms": 0.1646706117753638,
      "q3_ms": 0.1677121941183406,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 256,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004902470391243696,
      "xla_temp_mb": 33.000244140625,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.42371435899985954,
      "inner": 56,
      "median_ms": 0.7843491251021207,
      "q1_ms": 0.7820879018254345,
      "q3_ms": 0.785903758891696,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 256,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004902470391243696,
      "xla_temp_mb": 33.000244140625,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.43150563400558895,
      "inner": 58,
      "median_ms": 0.7838010000153285,
      "q1_ms": 0.781907862017698,
      "q3_ms": 0.7847872586416056,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 256,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004902470391243696,
      "xla_temp_mb": 97.000732421875,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.1057194340028218,
      "inner": 43,
      "median_ms": 1.0008717674577914,
      "q1_ms": 1.0001679767683949,
      "q3_ms": 1.0008901745307008,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 256,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0035469671711325645,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.20389825699385256,
      "inner": 200,
      "median_ms": 0.11322458998620277,
      "q1_ms": 0.11066522249166155,
      "q3_ms": 0.11326458748953883,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 256,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0035469671711325645,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.19506962999730604,
      "inner": 200,
      "median_ms": 0.11398026501410641,
      "q1_ms": 0.1101245500103687,
      "q3_ms": 0.11420967499361723,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 256,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0035469671711325645,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.04340067200246267,
      "inner": 200,
      "median_ms": 0.1416705900192028,
      "q1_ms": 0.14078180749493185,
      "q3_ms": 0.14179645000695018,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 256,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004253170918673277,
      "xla_temp_mb": 33.000244140625,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.4773168559986516,
      "inner": 69,
      "median_ms": 0.5899669565184487,
      "q1_ms": 0.5895474565424619,
      "q3_ms": 0.5903974637529219,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 256,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004253170918673277,
      "xla_temp_mb": 33.000244140625,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.478221637000388,
      "inner": 68,
      "median_ms": 0.5893979558844448,
      "q1_ms": 0.5889557867382844,
      "q3_ms": 0.590574654405132,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 256,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004253170918673277,
      "xla_temp_mb": 97.000732421875,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.10969530299917096,
      "inner": 47,
      "median_ms": 0.9051764042096887,
      "q1_ms": 0.9049846489010924,
      "q3_ms": 0.9070589042861916,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0028048453386873007,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.177762627004995,
      "inner": 71,
      "median_ms": 0.5588508873085924,
      "q1_ms": 0.5587136478373997,
      "q3_ms": 0.559103830953614,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0028048453386873007,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.17540101100166794,
      "inner": 79,
      "median_ms": 0.5640368734573344,
      "q1_ms": 0.5635000948880393,
      "q3_ms": 0.5703760189710839,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": false,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0028048453386873007,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.046803537996311206,
      "inner": 69,
      "median_ms": 0.6005476811813002,
      "q1_ms": 0.6004154782441101,
      "q3_ms": 0.6006682028321624,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004695911426097155,
      "xla_temp_mb": 66.000244140625,
      "xla_arg_mb": 160.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 0.42198679099965375,
      "inner": 19,
      "median_ms": 2.720806841879446,
      "q1_ms": 2.71936226310668,
      "q3_ms": 2.740672579085712,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004695911426097155,
      "xla_temp_mb": 66.000244140625,
      "xla_arg_mb": 160.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 0.42997577100322815,
      "inner": 19,
      "median_ms": 2.731566157691352,
      "q1_ms": 2.723388315644115,
      "q3_ms": 2.740750052716497,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004695911426097155,
      "xla_temp_mb": 194.000732421875,
      "xla_arg_mb": 160.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 0.1076859019958647,
      "inner": 16,
      "median_ms": 2.906348000124126,
      "q1_ms": 2.9016287187459966,
      "q3_ms": 2.907097499928568,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0036932043731212616,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.2028677280031843,
      "inner": 103,
      "median_ms": 0.3513130679352489,
      "q1_ms": 0.35006547089156387,
      "q3_ms": 0.3513985242693486,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0036932043731212616,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.19883472000219626,
      "inner": 110,
      "median_ms": 0.35204071821433236,
      "q1_ms": 0.3520103818366558,
      "q3_ms": 0.3532226727243555,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": true,
      "mode": "fwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.0036932043731212616,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.04510714999923948,
      "inner": 84,
      "median_ms": 0.49380916667966324,
      "q1_ms": 0.4934799702217182,
      "q3_ms": 0.4940870714909536,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004600926768034697,
      "xla_temp_mb": 66.000244140625,
      "xla_arg_mb": 160.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 0.4814242239954183,
      "inner": 29,
      "median_ms": 1.7622535863132536,
      "q1_ms": 1.7604140862360884,
      "q3_ms": 1.7652634137605139,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004600926768034697,
      "xla_temp_mb": 66.000244140625,
      "xla_arg_mb": 160.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 0.48132217099919217,
      "inner": 28,
      "median_ms": 1.766444785711688,
      "q1_ms": 1.7622092321159601,
      "q3_ms": 1.7687945534947045,
      "tag": "after"
    },
    {
      "suite": "dispatch",
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 64,
      "heads": 8,
      "rel_err": 0.004600926768034697,
      "xla_temp_mb": 194.000732421875,
      "xla_arg_mb": 160.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 0.1140623930041329,
      "inner": 19,
      "median_ms": 2.3694966316964527,
      "q1_ms": 2.368286710598592,
      "q3_ms": 2.3754700525483043,
      "tag": "after"
    }
  ]
}