{
  "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": {
    "output": "docs/benchmarks/gpu_attention.json",
    "batch": 4,
    "heads": 8,
    "head_dims": [
      64,
      128
    ],
    "seq_lens": [
      128,
      512,
      1024,
      2048,
      4096,
      8192
    ],
    "ragged_seq_lens": [
      600,
      1500,
      3000
    ],
    "dtypes": [
      "bfloat16",
      "float32"
    ],
    "modes": [
      "fwd",
      "fwd_bwd"
    ],
    "causal": [
      0,
      1
    ],
    "impls": [
      "probjax_mha",
      "probjax_flex",
      "probjax_default",
      "jax_xla",
      "jax_cudnn",
      "torch_flash",
      "torch_efficient",
      "torch_math"
    ],
    "append": false,
    "tag": null
  },
  "results": [
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0017409920692443848,
      "rel_err": 0.0018631259445101023,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 1.5,
      "xla_out_mb": 0.5,
      "compile_s": 0.1784586110006785,
      "inner": 200,
      "median_ms": 0.04134164504648652,
      "q1_ms": 0.04126564752368722,
      "q3_ms": 0.04152240497205639,
      "tflops": 3.2465502485225053
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0017409920692443848,
      "rel_err": 0.0018631259445101023,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 1.5,
      "xla_out_mb": 0.5,
      "compile_s": 0.17938975199649576,
      "inner": 200,
      "median_ms": 0.04282880501705222,
      "q1_ms": 0.04259110748535022,
      "q3_ms": 0.042965659995388705,
      "tflops": 3.1338191188514704
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0017409920692443848,
      "rel_err": 0.0018631259445101023,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 1.5,
      "xla_out_mb": 0.5,
      "compile_s": 0.1779300109919859,
      "inner": 200,
      "median_ms": 0.041823999999905936,
      "q1_ms": 0.04157345254498068,
      "q3_ms": 0.04217561498080613,
      "tflops": 3.2091078806499107
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0027137398719787598,
      "rel_err": 0.0029041138477623463,
      "xla_temp_mb": 3.0,
      "xla_arg_mb": 1.5,
      "xla_out_mb": 0.5,
      "compile_s": 0.6111467899900163,
      "inner": 200,
      "median_ms": 0.041703269962454215,
      "q1_ms": 0.041472462471574545,
      "q3_ms": 0.04807617751794169,
      "tflops": 3.2183981764700293
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0017409920692443848,
      "rel_err": 0.0018631259445101023,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 1.5,
      "xla_out_mb": 0.5,
      "compile_s": 0.07084445899818093,
      "inner": 200,
      "median_ms": 0.05560025500017218,
      "q1_ms": 0.055420642529497854,
      "q3_ms": 0.05586503502854612,
      "tflops": 2.41397684236492
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0027071237564086914,
      "rel_err": 0.002897033467888832,
      "peak_mb": 3.5166015625,
      "compile_s": 0.00010461799683980644,
      "inner": 200,
      "median_ms": 0.026083429984282702,
      "q1_ms": 0.026073889966937713,
      "q3_ms": 0.026224850007565692,
      "tflops": 5.145708523797546
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0017409920692443848,
      "rel_err": 0.0018631259445101023,
      "peak_mb": 3.5009765625,
      "compile_s": 0.00010530899453442544,
      "inner": 200,
      "median_ms": 0.03194666998751927,
      "q1_ms": 0.03176505251758499,
      "q3_ms": 0.032147450001502875,
      "tflops": 4.201305740236318
    },
    {
      "impl": "torch_math",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0017409920692443848,
      "rel_err": 0.0018631259445101023,
      "peak_mb": 19.62939453125,
      "compile_s": 0.00028309600020293146,
      "inner": 200,
      "median_ms": 0.17104239501350094,
      "q1_ms": 0.16624228748696623,
      "q3_ms": 0.17136473252321593,
      "tflops": 0.7847044470431191
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0032582879066467285,
      "rel_err": 0.003487034235149622,
      "xla_temp_mb": 1.031494140625,
      "xla_arg_mb": 2.5,
      "xla_out_mb": 1.5000228881835938,
      "compile_s": 0.42165933999058325,
      "inner": 200,
      "median_ms": 0.049919009979930706,
      "q1_ms": 0.04967201999534154,
      "q3_ms": 0.050547527462185826,
      "tflops": 9.410484065867127
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0032582879066467285,
      "rel_err": 0.003487034235149622,
      "xla_temp_mb": 1.031494140625,
      "xla_arg_mb": 2.5,
      "xla_out_mb": 1.5000228881835938,
      "compile_s": 0.42328347101283725,
      "inner": 200,
      "median_ms": 0.04892323006060906,
      "q1_ms": 0.04859287750150543,
      "q3_ms": 0.04904596251435578,
      "tflops": 9.60202438428596
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0032582879066467285,
      "rel_err": 0.003487034235149622,
      "xla_temp_mb": 1.031494140625,
      "xla_arg_mb": 2.5,
      "xla_out_mb": 1.5000228881835938,
      "compile_s": 0.4214675969997188,
      "inner": 191,
      "median_ms": 0.0590049476254273,
      "q1_ms": 0.05066784816283333,
      "q3_ms": 0.059813884821664345,
      "tflops": 7.96140098254342
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.003368556499481201,
      "rel_err": 0.0037138124462217093,
      "xla_temp_mb": 5.015625,
      "xla_arg_mb": 2.5,
      "xla_out_mb": 1.5000228881835938,
      "compile_s": 0.9139604350057198,
      "inner": 200,
      "median_ms": 0.05782872001873329,
      "q1_ms": 0.05429872249806067,
      "q3_ms": 0.05891707998671336,
      "tflops": 8.123334700263523
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0032582879066467285,
      "rel_err": 0.003487034235149622,
      "xla_temp_mb": 3.031982421875,
      "xla_arg_mb": 2.5,
      "xla_out_mb": 1.5000228881835938,
      "compile_s": 0.10760049600503407,
      "inner": 193,
      "median_ms": 0.11085369948411139,
      "q1_ms": 0.110713253860842,
      "q3_ms": 0.111540829008228,
      "tflops": 4.237675875376002
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0032582879066467285,
      "rel_err": 0.003487034235149622,
      "peak_mb": 16.658203125,
      "compile_s": 0.0005313550063874573,
      "inner": 144,
      "median_ms": 0.24206657639701412,
      "q1_ms": 0.24019950003801366,
      "q3_ms": 0.2428392534865351,
      "tflops": 1.9406316022313705
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0035500526428222656,
      "rel_err": 0.0036161637399345636,
      "peak_mb": 17.158203125,
      "compile_s": 0.0003606120008043945,
      "inner": 175,
      "median_ms": 0.24838542856741697,
      "q1_ms": 0.24786578857206873,
      "q3_ms": 0.24870303426203982,
      "tflops": 1.891262505652568
    },
    {
      "impl": "torch_math",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.003372371196746826,
      "rel_err": 0.0038937218487262726,
      "peak_mb": 31.7509765625,
      "compile_s": 0.0007293800008483231,
      "inner": 75,
      "median_ms": 0.5958589601020018,
      "q1_ms": 0.5925411466159858,
      "q3_ms": 0.5993545667539972,
      "tflops": 0.7883779207072492
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.005925178527832031,
      "rel_err": 0.00204979139380157,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 1.5,
      "xla_out_mb": 0.5,
      "compile_s": 0.2058926970057655,
      "inner": 200,
      "median_ms": 0.03708353506226558,
      "q1_ms": 0.03591533000872005,
      "q3_ms": 0.039293390000239015,
      "tflops": 1.8096673870848614
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.005925178527832031,
      "rel_err": 0.00204979139380157,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 1.5,
      "xla_out_mb": 0.5,
      "compile_s": 0.19832208899606485,
      "inner": 200,
      "median_ms": 0.03537545999279246,
      "q1_ms": 0.03521878501487663,
      "q3_ms": 0.03591835502447793,
      "tflops": 1.8970456925132009
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.005925178527832031,
      "rel_err": 0.00204979139380157,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 1.5,
      "xla_out_mb": 0.5,
      "compile_s": 0.20421429300040472,
      "inner": 200,
      "median_ms": 0.03570377499272581,
      "q1_ms": 0.0355282474993146,
      "q3_ms": 0.0357573050132487,
      "tflops": 1.8796013590628045
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0073577165603637695,
      "rel_err": 0.0025453721173107624,
      "xla_temp_mb": 3.0,
      "xla_arg_mb": 1.5,
      "xla_out_mb": 0.5,
      "compile_s": 0.3603209599968977,
      "inner": 200,
      "median_ms": 0.03957474502385594,
      "q1_ms": 0.039315432477451395,
      "q3_ms": 0.03962631999456789,
      "tflops": 1.6957497504922974
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.005925178527832031,
      "rel_err": 0.00204979139380157,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 1.5,
      "xla_out_mb": 0.5,
      "compile_s": 0.051109524007188156,
      "inner": 200,
      "median_ms": 0.05838366996613331,
      "q1_ms": 0.0579874725008267,
      "q3_ms": 0.058736892460728996,
      "tflops": 1.1494457960407065
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.005925178527832031,
      "rel_err": 0.00204979139380157,
      "peak_mb": 19.7666015625,
      "compile_s": 9.518000297248363e-05,
      "inner": 200,
      "median_ms": 0.027485434984555468,
      "q1_ms": 0.026605317507346626,
      "q3_ms": 0.02882950251660077,
      "tflops": 2.4416154969972133
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.005925178527832031,
      "rel_err": 0.00204979139380157,
      "peak_mb": 19.7509765625,
      "compile_s": 0.00010721299622673541,
      "inner": 200,
      "median_ms": 0.03187342997989617,
      "q1_ms": 0.031786817489773966,
      "q3_ms": 0.03207193251000717,
      "tflops": 2.1054798320208463
    },
    {
      "impl": "torch_math",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.005925178527832031,
      "rel_err": 0.00204979139380157,
      "peak_mb": 27.81689453125,
      "compile_s": 0.00030832299671601504,
      "inner": 200,
      "median_ms": 0.2006534000247484,
      "q1_ms": 0.20027125003252877,
      "q3_ms": 0.2013270500174258,
      "tflops": 0.33445166636460105
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.014573812484741211,
      "rel_err": 0.004595478065311909,
      "xla_temp_mb": 1.031494140625,
      "xla_arg_mb": 2.5,
      "xla_out_mb": 1.5000228881835938,
      "compile_s": 0.4816983720083954,
      "inner": 200,
      "median_ms": 0.057964029983850196,
      "q1_ms": 0.050780642486643046,
      "q3_ms": 0.058504252519924194,
      "tflops": 4.052185882614476
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.014573812484741211,
      "rel_err": 0.004595478065311909,
      "xla_temp_mb": 1.031494140625,
      "xla_arg_mb": 2.5,
      "xla_out_mb": 1.5000228881835938,
      "compile_s": 0.48328126600245014,
      "inner": 200,
      "median_ms": 0.05281674501020461,
      "q1_ms": 0.05265886753477389,
      "q3_ms": 0.053258480002114084,
      "tflops": 4.447093889534827
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.014573812484741211,
      "rel_err": 0.004595478065311909,
      "xla_temp_mb": 1.031494140625,
      "xla_arg_mb": 2.5,
      "xla_out_mb": 1.5000228881835938,
      "compile_s": 0.47491438199358527,
      "inner": 200,
      "median_ms": 0.053348750006989576,
      "q1_ms": 0.05309143747581402,
      "q3_ms": 0.05358709997381084,
      "tflops": 4.402746530504026
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.014573812484741211,
      "rel_err": 0.004490324761718512,
      "xla_temp_mb": 5.015625,
      "xla_arg_mb": 2.5,
      "xla_out_mb": 1.5000228881835938,
      "compile_s": 0.5805231909907889,
      "inner": 200,
      "median_ms": 0.05608983003185131,
      "q1_ms": 0.05599270250968402,
      "q3_ms": 0.05731282002670923,
      "tflops": 4.18758665994566
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.014573812484741211,
      "rel_err": 0.004595478065311909,
      "xla_temp_mb": 3.031982421875,
      "xla_arg_mb": 2.5,
      "xla_out_mb": 1.5000228881835938,
      "compile_s": 0.1157392089953646,
      "inner": 186,
      "median_ms": 0.11307245695597982,
      "q1_ms": 0.1129879166656274,
      "q3_ms": 0.11382627417333424,
      "tflops": 2.0772611679557067
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.014573812484741211,
      "rel_err": 0.004595478065311909,
      "peak_mb": 24.783203125,
      "compile_s": 0.00044951000018045306,
      "inner": 165,
      "median_ms": 0.23472132725697575,
      "q1_ms": 0.22911477876319125,
      "q3_ms": 0.24088346359679813,
      "tflops": 1.0006803674165041
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.014573812484741211,
      "rel_err": 0.006634362041950226,
      "peak_mb": 25.283203125,
      "compile_s": 0.0003647900011856109,
      "inner": 162,
      "median_ms": 0.24802946912435192,
      "q1_ms": 0.2444469876390491,
      "q3_ms": 0.2517898241040027,
      "tflops": 0.9469883753298692
    },
    {
      "impl": "torch_math",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 128,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.011228084564208984,
      "rel_err": 0.0034594754688441753,
      "peak_mb": 31.7509765625,
      "compile_s": 0.001005340993287973,
      "inner": 54,
      "median_ms": 0.6855430555232386,
      "q1_ms": 0.6538151851768985,
      "q3_ms": 0.7000091203803163,
      "tflops": 0.34262038264063194
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0010340511798858643,
      "rel_err": 0.002302922774106264,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.1791836430056719,
      "inner": 200,
      "median_ms": 0.054021320029278286,
      "q1_ms": 0.0540128300053766,
      "q3_ms": 0.0540526824988774,
      "tflops": 39.752520798012974
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0010340511798858643,
      "rel_err": 0.002302922774106264,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.17745843000011519,
      "inner": 200,
      "median_ms": 0.05402858500019647,
      "q1_ms": 0.05401924252510071,
      "q3_ms": 0.05408058499597246,
      "tflops": 39.74717546262207
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0010340511798858643,
      "rel_err": 0.002302922774106264,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.17660666900337674,
      "inner": 200,
      "median_ms": 0.05405403506301809,
      "q1_ms": 0.05405032748967642,
      "q3_ms": 0.05406274998676963,
      "tflops": 39.72846144596584
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0012623369693756104,
      "rel_err": 0.0028113352600485086,
      "xla_temp_mb": 48.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.7029671010095626,
      "inner": 183,
      "median_ms": 0.16027936609169274,
      "q1_ms": 0.1600220764616932,
      "q3_ms": 0.16033176229019838,
      "tflops": 13.398378720636229
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0010340511798858643,
      "rel_err": 0.002302922774106264,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.045587008993607014,
      "inner": 200,
      "median_ms": 0.06974160503887106,
      "q1_ms": 0.06926577749254648,
      "q3_ms": 0.07019535998551873,
      "tflops": 30.79200208832421
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0010340511798858643,
      "rel_err": 0.002302922774106264,
      "peak_mb": 30.3134765625,
      "compile_s": 9.555100405123085e-05,
      "inner": 200,
      "median_ms": 0.04600866501277778,
      "q1_ms": 0.045593635040859226,
      "q3_ms": 0.046614212478743866,
      "tflops": 46.67563484842668
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0010340511798858643,
      "rel_err": 0.002302922774106264,
      "peak_mb": 30.2509765625,
      "compile_s": 0.0001102579990401864,
      "inner": 200,
      "median_ms": 0.0617554999917047,
      "q1_ms": 0.05792314997961512,
      "q3_ms": 0.06187118997331709,
      "tflops": 34.77396585386663
    },
    {
      "impl": "torch_math",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0009598731994628906,
      "rel_err": 0.0021377219818532467,
      "peak_mb": 116.26611328125,
      "compile_s": 0.0006047839997336268,
      "inner": 83,
      "median_ms": 0.5666934699240893,
      "q1_ms": 0.566405825283945,
      "q3_ms": 0.5672146506251274,
      "tflops": 3.7894977831448515
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0026992112398147583,
      "rel_err": 0.005563188809901476,
      "xla_temp_mb": 4.125244140625,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.42887019799673,
      "inner": 133,
      "median_ms": 0.20704663155558414,
      "q1_ms": 0.2059263195378594,
      "q3_ms": 0.20707875939437953,
      "tflops": 36.301932137360986
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0026992112398147583,
      "rel_err": 0.005563188809901476,
      "xla_temp_mb": 4.125244140625,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.431159828993259,
      "inner": 136,
      "median_ms": 0.207385683846523,
      "q1_ms": 0.20674241177938016,
      "q3_ms": 0.20786202939608997,
      "tflops": 36.242582557253094
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0026992112398147583,
      "rel_err": 0.005563188809901476,
      "xla_temp_mb": 4.125244140625,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.42916090799553785,
      "inner": 138,
      "median_ms": 0.20683466668149617,
      "q1_ms": 0.2050829493057698,
      "q3_ms": 0.20770809783573274,
      "tflops": 36.339134481620306
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0029108524322509766,
      "rel_err": 0.007554607931524515,
      "xla_temp_mb": 80.0625,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 1.2230064570030663,
      "inner": 90,
      "median_ms": 0.45813402215976057,
      "q1_ms": 0.4581188332748651,
      "q3_ms": 0.4582599888736796,
      "tflops": 16.406100408275183
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0026992112398147583,
      "rel_err": 0.005563188809901476,
      "xla_temp_mb": 12.125732421875,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.10420399199938402,
      "inner": 115,
      "median_ms": 0.23544521739139504,
      "q1_ms": 0.23539115221786272,
      "q3_ms": 0.23551774350420127,
      "tflops": 31.923318941345798
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0026992112398147583,
      "rel_err": 0.005563188809901476,
      "peak_mb": 50.376953125,
      "compile_s": 0.0004301029985072091,
      "inner": 120,
      "median_ms": 0.2460848417589053,
      "q1_ms": 0.24063039581960766,
      "q3_ms": 0.2555695291448501,
      "tflops": 30.54309527672484
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0029108524322509766,
      "rel_err": 0.007554607931524515,
      "peak_mb": 48.380859375,
      "compile_s": 0.0006121280021034181,
      "inner": 92,
      "median_ms": 0.3238048152327436,
      "q1_ms": 0.32280497828690824,
      "q3_ms": 0.32408185325039085,
      "tflops": 23.21210931529085
    },
    {
      "impl": "torch_math",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0016099810600280762,
      "rel_err": 0.0033246709499508142,
      "peak_mb": 174.2509765625,
      "compile_s": 0.0014088549942243844,
      "inner": 36,
      "median_ms": 1.3676363608586446,
      "q1_ms": 1.36457819462521,
      "q3_ms": 1.3699198332485847,
      "tflops": 5.495753829827324
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.007252931594848633,
      "rel_err": 0.002509122248739004,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.19967591000022367,
      "inner": 200,
      "median_ms": 0.041950689992518164,
      "q1_ms": 0.041924720026145224,
      "q3_ms": 0.04200401748676086,
      "tflops": 25.59533166657092
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.007252931594848633,
      "rel_err": 0.002509122248739004,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.20002412899339106,
      "inner": 200,
      "median_ms": 0.041984150011558086,
      "q1_ms": 0.04196194000542164,
      "q3_ms": 0.04204687251331052,
      "tflops": 25.574933009347642
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.007252931594848633,
      "rel_err": 0.002509122248739004,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.1995495809969725,
      "inner": 139,
      "median_ms": 0.042054496379736676,
      "q1_ms": 0.04202826979792871,
      "q3_ms": 0.04217209710685952,
      "tflops": 25.532152716905824
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.008372068405151367,
      "rel_err": 0.0028962830547243357,
      "xla_temp_mb": 48.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.865534436001326,
      "inner": 147,
      "median_ms": 0.2050698911224739,
      "q1_ms": 0.2050418469585244,
      "q3_ms": 0.20528665981746475,
      "tflops": 5.235979880433686
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.007252931594848633,
      "rel_err": 0.002509122248739004,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.05059954000171274,
      "inner": 200,
      "median_ms": 0.05524332998902537,
      "q1_ms": 0.054867249964445364,
      "q3_ms": 0.055677877498965245,
      "tflops": 19.436587624484428
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.007252931594848633,
      "rel_err": 0.002509122248739004,
      "peak_mb": 30.3134765625,
      "compile_s": 9.626199607737362e-05,
      "inner": 200,
      "median_ms": 0.03932797997549642,
      "q1_ms": 0.03932357249141205,
      "q3_ms": 0.039334295033768285,
      "tflops": 27.30223684687092
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.007252931594848633,
      "rel_err": 0.002509122248739004,
      "peak_mb": 30.2509765625,
      "compile_s": 0.0001063310046447441,
      "inner": 200,
      "median_ms": 0.04734208501758985,
      "q1_ms": 0.04708340002252953,
      "q3_ms": 0.04788543999893591,
      "tflops": 22.68049291874351
    },
    {
      "impl": "torch_math",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.007252931594848633,
      "rel_err": 0.002509122248739004,
      "peak_mb": 117.26611328125,
      "compile_s": 0.0007382059993688017,
      "inner": 71,
      "median_ms": 0.661317239602385,
      "q1_ms": 0.6597410845668645,
      "q3_ms": 0.6616080703884429,
      "tflops": 1.6236410601447258
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.01679849624633789,
      "rel_err": 0.004870529752224684,
      "xla_temp_mb": 4.125244140625,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.48390325299988035,
      "inner": 161,
      "median_ms": 0.13877881370175352,
      "q1_ms": 0.13734404651501636,
      "q3_ms": 0.13882302174233305,
      "tflops": 27.07975579094113
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.01679849624633789,
      "rel_err": 0.004870529752224684,
      "xla_temp_mb": 4.125244140625,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.4813658830098575,
      "inner": 171,
      "median_ms": 0.13821069590273524,
      "q1_ms": 0.1374940789231148,
      "q3_ms": 0.13822505260104237,
      "tflops": 27.191067662699073
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.01679849624633789,
      "rel_err": 0.004870529752224684,
      "xla_temp_mb": 4.125244140625,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.47948387399083003,
      "inner": 162,
      "median_ms": 0.13877975926547667,
      "q1_ms": 0.13740135186652214,
      "q3_ms": 0.13910540431503632,
      "tflops": 27.079571285399087
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.01679849624633789,
      "rel_err": 0.005369177088141441,
      "xla_temp_mb": 80.0625,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 1.6260668210015865,
      "inner": 75,
      "median_ms": 0.5289564268120254,
      "q1_ms": 0.5249590600336281,
      "q3_ms": 0.5291726333477225,
      "tflops": 7.104737164551949
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.01679849624633789,
      "rel_err": 0.004870529752224684,
      "xla_temp_mb": 12.125732421875,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.10555946599924937,
      "inner": 120,
      "median_ms": 0.21709542500805884,
      "q1_ms": 0.2170702166646758,
      "q3_ms": 0.2171362999433768,
      "tflops": 17.31080414919151
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.01679849624633789,
      "rel_err": 0.004870529752224684,
      "peak_mb": 50.376953125,
      "compile_s": 0.0003857489937217906,
      "inner": 152,
      "median_ms": 0.24387736840780186,
      "q1_ms": 0.2413612565808427,
      "q3_ms": 0.24426718415557652,
      "tflops": 15.409779138324403
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.01679849624633789,
      "rel_err": 0.005369177088141441,
      "peak_mb": 48.380859375,
      "compile_s": 0.0004616939986590296,
      "inner": 125,
      "median_ms": 0.2563265518983826,
      "q1_ms": 0.25571671599755064,
      "q3_ms": 0.2567185800289735,
      "tflops": 14.661362064004395
    },
    {
      "impl": "torch_math",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 512,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.01679849624633789,
      "rel_err": 0.0038863555528223515,
      "peak_mb": 174.2509765625,
      "compile_s": 0.001537007003207691,
      "inner": 33,
      "median_ms": 1.4441926667297427,
      "q1_ms": 1.4427439697707691,
      "q3_ms": 1.445762424158271,
      "tflops": 2.6022126206400875
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 600,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.001130521297454834,
      "rel_err": 0.0033199135214090347,
      "xla_temp_mb": 10.0,
      "xla_arg_mb": 7.03125,
      "xla_out_mb": 2.34375,
      "compile_s": 0.21827725999173708,
      "inner": 200,
      "median_ms": 0.09193282494379673,
      "q1_ms": 0.09191566750814673,
      "q3_ms": 0.09194449750793865,
      "tflops": 32.07907514865282
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 600,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.001130521297454834,
      "rel_err": 0.0033199135214090347,
      "xla_temp_mb": 10.0,
      "xla_arg_mb": 7.03125,
      "xla_out_mb": 2.34375,
      "compile_s": 0.21494667099614162,
      "inner": 200,
      "median_ms": 0.08748220505367499,
      "q1_ms": 0.08584566247009207,
      "q3_ms": 0.0875923149578739,
      "tflops": 33.711084422146854
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 600,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.001130521297454834,
      "rel_err": 0.0033199135214090347,
      "xla_temp_mb": 10.0,
      "xla_arg_mb": 7.03125,
      "xla_out_mb": 2.34375,
      "compile_s": 0.21747885100194253,
      "inner": 200,
      "median_ms": 0.08741327503230423,
      "q1_ms": 0.08648599752632435,
      "q3_ms": 0.08743118251004489,
      "tflops": 33.73766740704007
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 600,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0012695491313934326,
      "rel_err": 0.0037281857803463936,
      "xla_temp_mb": 65.91796875,
      "xla_arg_mb": 7.03125,
      "xla_out_mb": 2.34375,
      "compile_s": 0.6951585019996855,
      "inner": 163,
      "median_ms": 0.21577970552589584,
      "q1_ms": 0.215380113504747,
      "q3_ms": 0.21581728834668862,
      "tflops": 13.66727233598006
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 600,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.001130521297454834,
      "rel_err": 0.0033199135214090347,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 7.03125,
      "xla_out_mb": 2.34375,
      "compile_s": 0.05002240800240543,
      "inner": 200,
      "median_ms": 0.10842676500033122,
      "q1_ms": 0.10839330501767108,
      "q3_ms": 0.10843147749255877,
      "tflops": 27.199188318410048
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 600,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0008531510829925537,
      "rel_err": 0.0025053820572793484,
      "peak_mb": 33.23046875,
      "compile_s": 0.00011829299910459667,
      "inner": 200,
      "median_ms": 0.06185969497892074,
      "q1_ms": 0.06060162002540892,
      "q3_ms": 0.06186327500472544,
      "tflops": 47.67433788680884
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 600,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.001130521297454834,
      "rel_err": 0.0033199135214090347,
      "peak_mb": 33.1572265625,
      "compile_s": 0.0001232229988090694,
      "inner": 200,
      "median_ms": 0.07728291500825435,
      "q1_ms": 0.07363066000834806,
      "q3_ms": 0.07736574996670242,
      "tflops": 38.160051282809576
    },
    {
      "impl": "torch_math",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 600,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0008531510829925537,
      "rel_err": 0.0025053820572793484,
      "peak_mb": 148.72412109375,
      "compile_s": 0.0007717489934293553,
      "inner": 67,
      "median_ms": 0.7334214627808794,
      "q1_ms": 0.7330219702466753,
      "q3_ms": 0.7338962463814932,
      "tflops": 4.02104403765055
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 600,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0017500221729278564,
      "rel_err": 0.004071441944688559,
      "xla_temp_mb": 20.00048828125,
      "xla_arg_mb": 11.71875,
      "xla_out_mb": 7.031272888183594,
      "compile_s": 0.4797334260074422,
      "inner": 119,
      "median_ms": 0.3101269663419152,
      "q1_ms": 0.3096280000006248,
      "q3_ms": 0.31065312605507733,
      "tflops": 33.2828844964745
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 600,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0017500221729278564,
      "rel_err": 0.004071441944688559,
      "xla_temp_mb": 20.00048828125,
      "xla_arg_mb": 11.71875,
      "xla_out_mb": 7.031272888183594,
      "compile_s": 0.48080686599132605,
      "inner": 120,
      "median_ms": 0.3139176333206706,
      "q1_ms": 0.3136058291905405,
      "q3_ms": 0.31432014584424905,
      "tflops": 32.88098183849404
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 600,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0017500221729278564,
      "rel_err": 0.004071441944688559,
      "xla_temp_mb": 20.00048828125,
      "xla_arg_mb": 11.71875,
      "xla_out_mb": 7.031272888183594,
      "compile_s": 0.4842866370017873,
      "inner": 112,
      "median_ms": 0.313566089321934,
      "q1_ms": 0.31319136152368243,
      "q3_ms": 0.3142783258616484,
      "tflops": 32.917845237412216
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 600,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.002234041690826416,
      "rel_err": 0.005197517573833466,
      "xla_temp_mb": 87.96388244628906,
      "xla_arg_mb": 11.71875,
      "xla_out_mb": 7.031272888183594,
      "compile_s": 1.031083091002074,
      "inner": 59,
      "median_ms": 0.7414489321536937,
      "q1_ms": 0.741305779784097,
      "q3_ms": 0.7416588050546125,
      "tflops": 13.921282440879404
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 600,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0017500221729278564,
      "rel_err": 0.004071441944688559,
      "xla_temp_mb": 14.209716796875,
      "xla_arg_mb": 11.71875,
      "xla_out_mb": 7.031272888183594,
      "compile_s": 0.10458858001220506,
      "inner": 99,
      "median_ms": 0.33991389897313307,
      "q1_ms": 0.3396169697152065,
      "q3_ms": 0.3400043636734475,
      "tflops": 30.36627814038239
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 600,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0017500221729278564,
      "rel_err": 0.004071441944688559,
      "peak_mb": 58.1533203125,
      "compile_s": 0.0004525760014075786,
      "inner": 120,
      "median_ms": 0.2967249500215985,
      "q1_ms": 0.29666329167715355,
      "q3_ms": 0.2967638125483063,
      "tflops": 34.786154650118476
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 600,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0023759305477142334,
      "rel_err": 0.005527622997760773,
      "peak_mb": 54.873046875,
      "compile_s": 0.0007274860108736902,
      "inner": 90,
      "median_ms": 0.44612265547685737,
      "q1_ms": 0.4451150777135303,
      "q3_ms": 0.4463329944378024,
      "tflops": 23.136955438783918
    },
    {
      "impl": "torch_math",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 600,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0017395317554473877,
      "rel_err": 0.003820978570729494,
      "peak_mb": 227.9072265625,
      "compile_s": 0.001735532991006039,
      "inner": 29,
      "median_ms": 1.746384689889463,
      "q1_ms": 1.7449217414431244,
      "q3_ms": 1.7506803966986402,
      "tflops": 5.910450349088506
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 600,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00457155704498291,
      "rel_err": 0.001817264943383634,
      "xla_temp_mb": 10.0,
      "xla_arg_mb": 7.03125,
      "xla_out_mb": 2.34375,
      "compile_s": 0.21628028400300536,
      "inner": 200,
      "median_ms": 0.07429926998156589,
      "q1_ms": 0.07380088252830319,
      "q3_ms": 0.07452820500475354,
      "tflops": 19.84622460443888
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 600,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00457155704498291,
      "rel_err": 0.001817264943383634,
      "xla_temp_mb": 10.0,
      "xla_arg_mb": 7.03125,
      "xla_out_mb": 2.34375,
      "compile_s": 0.2188260389957577,
      "inner": 200,
      "median_ms": 0.07463646499672905,
      "q1_ms": 0.07351897249463946,
      "q3_ms": 0.07475660997442901,
      "tflops": 19.756562694449997
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 600,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00457155704498291,
      "rel_err": 0.001817264943383634,
      "xla_temp_mb": 10.0,
      "xla_arg_mb": 7.03125,
      "xla_out_mb": 2.34375,
      "compile_s": 0.21626390299934428,
      "inner": 200,
      "median_ms": 0.07434396000462584,
      "q1_ms": 0.07383872747595888,
      "q3_ms": 0.07454013000824489,
      "tflops": 19.834294539976746
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 600,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.006919384002685547,
      "rel_err": 0.0027505625039339066,
      "xla_temp_mb": 65.91796875,
      "xla_arg_mb": 7.03125,
      "xla_out_mb": 2.34375,
      "compile_s": 0.6660494780080626,
      "inner": 122,
      "median_ms": 0.2649389507943673,
      "q1_ms": 0.26486048358731107,
      "q3_ms": 0.26506065979283155,
      "tflops": 5.565659543750822
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 600,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00457155704498291,
      "rel_err": 0.001817264943383634,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 7.03125,
      "xla_out_mb": 2.34375,
      "compile_s": 0.04948018299182877,
      "inner": 183,
      "median_ms": 0.08179250814509197,
      "q1_ms": 0.08178443171286223,
      "q3_ms": 0.0818203196780059,
      "tflops": 18.028057012071002
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 600,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00457155704498291,
      "rel_err": 0.001817264943383634,
      "peak_mb": 33.23046875,
      "compile_s": 0.0001094170002033934,
      "inner": 200,
      "median_ms": 0.054360260037356056,
      "q1_ms": 0.05395071246312,
      "q3_ms": 0.054435102465504315,
      "tflops": 27.12569805565115
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 600,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00457155704498291,
      "rel_err": 0.001817264943383634,
      "peak_mb": 33.1572265625,
      "compile_s": 0.00013311200018506497,
      "inner": 200,
      "median_ms": 0.06390736998582724,
      "q1_ms": 0.06163840251247166,
      "q3_ms": 0.06396785749529954,
      "tflops": 23.073395139355814
    },
    {
      "impl": "torch_math",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 600,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00390017032623291,
      "rel_err": 0.0015503782778978348,
      "peak_mb": 150.72412109375,
      "compile_s": 0.0009211220021825284,
      "inner": 55,
      "median_ms": 0.8637030000417408,
      "q1_ms": 0.863496772812637,
      "q3_ms": 0.8647949634574946,
      "tflops": 1.707253534986839
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 600,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0126495361328125,
      "rel_err": 0.0030066915787756443,
      "xla_temp_mb": 20.00048828125,
      "xla_arg_mb": 11.71875,
      "xla_out_mb": 7.031272888183594,
      "compile_s": 0.5105029709957307,
      "inner": 153,
      "median_ms": 0.21397881694914153,
      "q1_ms": 0.2130835621048593,
      "q3_ms": 0.2143033202360932,
      "tflops": 24.119022964907114
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 600,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0126495361328125,
      "rel_err": 0.0030066915787756443,
      "xla_temp_mb": 20.00048828125,
      "xla_arg_mb": 11.71875,
      "xla_out_mb": 7.031272888183594,
      "compile_s": 0.5124870630097575,
      "inner": 169,
      "median_ms": 0.2150924319398628,
      "q1_ms": 0.2149405414596773,
      "q3_ms": 0.21548604144240846,
      "tflops": 23.99414964745454
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 600,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0126495361328125,
      "rel_err": 0.0030066915787756443,
      "xla_temp_mb": 20.00048828125,
      "xla_arg_mb": 11.71875,
      "xla_out_mb": 7.031272888183594,
      "compile_s": 0.511372915992979,
      "inner": 142,
      "median_ms": 0.21509046487602024,
      "q1_ms": 0.21444424649890723,
      "q3_ms": 0.21521337322143075,
      "tflops": 23.99436908081823
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 600,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0126495361328125,
      "rel_err": 0.0036837407387793064,
      "xla_temp_mb": 87.96388244628906,
      "xla_arg_mb": 11.71875,
      "xla_out_mb": 7.031272888183594,
      "compile_s": 0.5126504319923697,
      "inner": 58,
      "median_ms": 0.7520649827522997,
      "q1_ms": 0.751644086292423,
      "q3_ms": 0.7522656981971774,
      "tflops": 6.862385722457995
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 600,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0126495361328125,
      "rel_err": 0.0030066915787756443,
      "xla_temp_mb": 14.209716796875,
      "xla_arg_mb": 11.71875,
      "xla_out_mb": 7.031272888183594,
      "compile_s": 0.10730121099913958,
      "inner": 101,
      "median_ms": 0.2983492177549178,
      "q1_ms": 0.29827299511110833,
      "q3_ms": 0.2984431584361056,
      "tflops": 17.298386229520894
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 600,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.011621952056884766,
      "rel_err": 0.002953536109998822,
      "peak_mb": 58.1533203125,
      "compile_s": 0.0004007680108770728,
      "inner": 131,
      "median_ms": 0.2676094121269469,
      "q1_ms": 0.2674234923035731,
      "q3_ms": 0.2684147213363508,
      "tflops": 19.285420340715728
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 600,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.01287078857421875,
      "rel_err": 0.005205949302762747,
      "peak_mb": 54.873046875,
      "compile_s": 0.000501558999530971,
      "inner": 105,
      "median_ms": 0.3547284190842349,
      "q1_ms": 0.35453420478872777,
      "q3_ms": 0.35558795719961855,
      "tflops": 14.549045755407779
    },
    {
      "impl": "torch_math",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 600,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.011621952056884766,
      "rel_err": 0.0027624431531876326,
      "peak_mb": 227.9072265625,
      "compile_s": 0.0018880010029533878,
      "inner": 26,
      "median_ms": 1.8642501152219038,
      "q1_ms": 1.8615931535000878,
      "q3_ms": 1.864642192511103,
      "tflops": 2.7683838975574826
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 1024,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0008517801761627197,
      "rel_err": 0.0030269771814346313,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.1768312030035304,
      "inner": 148,
      "median_ms": 0.16256245942566092,
      "q1_ms": 0.1625454290144655,
      "q3_ms": 0.1625888547317339,
      "tflops": 52.84082574998281
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 1024,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0008517801761627197,
      "rel_err": 0.0030269771814346313,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.1743090630043298,
      "inner": 173,
      "median_ms": 0.15715902891263064,
      "q1_ms": 0.15476890177650354,
      "q3_ms": 0.15722947973910062,
      "tflops": 54.65759524879349
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 1024,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0008517801761627197,
      "rel_err": 0.0030269771814346313,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.1785176230041543,
      "inner": 176,
      "median_ms": 0.15716969879907663,
      "q1_ms": 0.15517092044154362,
      "q3_ms": 0.15717491189868515,
      "tflops": 54.653884671378314
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 1024,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0008517801761627197,
      "rel_err": 0.0030269771814346313,
      "xla_temp_mb": 192.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.7453525450109737,
      "inner": 75,
      "median_ms": 0.5594422132708132,
      "q1_ms": 0.5592962133232504,
      "q3_ms": 0.5595605666167103,
      "tflops": 15.354462692006061
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 1024,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0008517801761627197,
      "rel_err": 0.0030269771814346313,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.04839091200847179,
      "inner": 146,
      "median_ms": 0.16866186302799527,
      "q1_ms": 0.16863102403115354,
      "q3_ms": 0.16895262673991568,
      "tflops": 50.92991644811965
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 1024,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0008517801761627197,
      "rel_err": 0.0030269771814346313,
      "peak_mb": 44.3759765625,
      "compile_s": 0.00020189100177958608,
      "inner": 200,
      "median_ms": 0.16220989004068542,
      "q1_ms": 0.161590370007616,
      "q3_ms": 0.16239062999375165,
      "tflops": 52.95567730084446
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 1024,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0008517801761627197,
      "rel_err": 0.0030269771814346313,
      "peak_mb": 44.2509765625,
      "compile_s": 0.00022557600459549576,
      "inner": 200,
      "median_ms": 0.20290645996283274,
      "q1_ms": 0.2023115575138945,
      "q3_ms": 0.20307883751229383,
      "tflops": 42.3344559536126
    },
    {
      "impl": "torch_math",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 1024,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0008517801761627197,
      "rel_err": 0.0030269771814346313,
      "peak_mb": 360.28173828125,
      "compile_s": 0.001969424993149005,
      "inner": 25,
      "median_ms": 1.9915486400714144,
      "q1_ms": 1.9904662200133316,
      "q3_ms": 1.9919986996683292,
      "tflops": 4.31319347123351
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 1024,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0014092624187469482,
      "rel_err": 0.004609355237334967,
      "xla_temp_mb": 8.250244140625,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.42536358498909976,
      "inner": 62,
      "median_ms": 0.6822212742389961,
      "q1_ms": 0.6820178145238559,
      "q3_ms": 0.6824080645005339,
      "tflops": 44.06894391491475
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 1024,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0014092624187469482,
      "rel_err": 0.004609355237334967,
      "xla_temp_mb": 8.250244140625,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.44135905800794717,
      "inner": 62,
      "median_ms": 0.6816889677617339,
      "q1_ms": 0.6813186691649589,
      "q3_ms": 0.6817670967598866,
      "tflops": 44.103355773403585
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 1024,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0014092624187469482,
      "rel_err": 0.004609355237334967,
      "xla_temp_mb": 8.250244140625,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.43045969400554895,
      "inner": 61,
      "median_ms": 0.6822703933900558,
      "q1_ms": 0.6820794590224023,
      "q3_ms": 0.6824594263472121,
      "tflops": 44.06577122981194
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 1024,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0016392767429351807,
      "rel_err": 0.004950827918946743,
      "xla_temp_mb": 320.125,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 1.2666406329954043,
      "inner": 28,
      "median_ms": 1.6186678211462484,
      "q1_ms": 1.617928053747164,
      "q3_ms": 1.6189760894381573,
      "tflops": 18.57377448246907
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 1024,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0014092624187469482,
      "rel_err": 0.004609355237334967,
      "xla_temp_mb": 24.250732421875,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.10725140600698069,
      "inner": 57,
      "median_ms": 0.7276586841550868,
      "q1_ms": 0.7276135087554474,
      "q3_ms": 0.7284743509157316,
      "tflops": 41.31713360489801
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 1024,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0014092624187469482,
      "rel_err": 0.004609355237334967,
      "peak_mb": 84.501953125,
      "compile_s": 0.0009087289945455268,
      "inner": 64,
      "median_ms": 0.7302528749733028,
      "q1_ms": 0.7301673045958523,
      "q3_ms": 0.7305453672188378,
      "tflops": 41.170356327729806
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 1024,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0014559626579284668,
      "rel_err": 0.004609355237334967,
      "peak_mb": 80.509765625,
      "compile_s": 0.001055396001902409,
      "inner": 47,
      "median_ms": 1.0546757873677788,
      "q1_ms": 1.0543071063906648,
      "q3_ms": 1.0587219574595208,
      "tflops": 28.506173586325097
    },
    {
      "impl": "torch_math",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 1024,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0009740293025970459,
      "rel_err": 0.00294169457629323,
      "peak_mb": 588.2509765625,
      "compile_s": 0.004869348995271139,
      "inner": 10,
      "median_ms": 4.853680900123436,
      "q1_ms": 4.852516600658419,
      "q3_ms": 4.85510649959906,
      "tflops": 6.194220776078504
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 1024,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0050013065338134766,
      "rel_err": 0.002148212166503072,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.2043885729945032,
      "inner": 200,
      "median_ms": 0.10898848493525293,
      "q1_ms": 0.10018558252340881,
      "q3_ms": 0.10929448500974104,
      "tflops": 39.40753281001679
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 1024,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0050013065338134766,
      "rel_err": 0.002148212166503072,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.19932475499808788,
      "inner": 199,
      "median_ms": 0.1025036633004149,
      "q1_ms": 0.10159818339161575,
      "q3_ms": 0.10256667334929193,
      "tflops": 41.90062245299887
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 1024,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0050013065338134766,
      "rel_err": 0.002148212166503072,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.1990126960008638,
      "inner": 200,
      "median_ms": 0.10240713505481835,
      "q1_ms": 0.10052760502730962,
      "q3_ms": 0.10296506254235283,
      "tflops": 41.94011768516824
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 1024,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00864255428314209,
      "rel_err": 0.003712238045409322,
      "xla_temp_mb": 192.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.637670315001742,
      "inner": 70,
      "median_ms": 0.586634071286036,
      "q1_ms": 0.5861767143609801,
      "q3_ms": 0.586838535777393,
      "tflops": 7.321373759598467
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 1024,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0050013065338134766,
      "rel_err": 0.002148212166503072,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.04823611999745481,
      "inner": 162,
      "median_ms": 0.13814943209851022,
      "q1_ms": 0.13813845375639006,
      "q3_ms": 0.13817887039845062,
      "tflops": 31.08928665691067
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 1024,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0050013065338134766,
      "rel_err": 0.002148212166503072,
      "peak_mb": 44.3759765625,
      "compile_s": 0.00015020399587228894,
      "inner": 200,
      "median_ms": 0.11406548495870084,
      "q1_ms": 0.11374500252713915,
      "q3_ms": 0.11440617501648376,
      "tflops": 37.653522426657446
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 1024,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0050013065338134766,
      "rel_err": 0.002148212166503072,
      "peak_mb": 44.2509765625,
      "compile_s": 0.00016509200213477015,
      "inner": 200,
      "median_ms": 0.12819916504668072,
      "q1_ms": 0.1272736149985576,
      "q3_ms": 0.12880631249572616,
      "tflops": 33.50230318922973
    },
    {
      "impl": "torch_math",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 1024,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0050013065338134766,
      "rel_err": 0.002148212166503072,
      "peak_mb": 364.28173828125,
      "compile_s": 0.002459252005792223,
      "inner": 20,
      "median_ms": 2.465637749992311,
      "q1_ms": 2.465476724682958,
      "q3_ms": 2.470524324962753,
      "tflops": 1.7419295660984235
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 1024,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.01181650161743164,
      "rel_err": 0.004610817413777113,
      "xla_temp_mb": 8.250244140625,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.4854635929950746,
      "inner": 94,
      "median_ms": 0.41193494689650834,
      "q1_ms": 0.41145558505841684,
      "q3_ms": 0.41223386173480825,
      "tflops": 36.49213461798528
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 1024,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.01181650161743164,
      "rel_err": 0.004610817413777113,
      "xla_temp_mb": 8.250244140625,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.48530817999562714,
      "inner": 96,
      "median_ms": 0.4130642499452127,
      "q1_ms": 0.41276029696746264,
      "q3_ms": 0.4131386614668978,
      "tflops": 36.39236641271628
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 1024,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.01181650161743164,
      "rel_err": 0.004610817413777113,
      "xla_temp_mb": 8.250244140625,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.48731962400779594,
      "inner": 90,
      "median_ms": 0.4122477999771945,
      "q1_ms": 0.41036501109677676,
      "q3_ms": 0.41251073893767576,
      "tflops": 36.464440894121424
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 1024,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.01181650161743164,
      "rel_err": 0.005088115110993385,
      "xla_temp_mb": 320.125,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 1.2890617890079739,
      "inner": 22,
      "median_ms": 2.0252599543627268,
      "q1_ms": 2.024822977546137,
      "q3_ms": 2.0267680450739465,
      "tflops": 7.422447426375014
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 1024,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.01181650161743164,
      "rel_err": 0.004610817413777113,
      "xla_temp_mb": 24.250732421875,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.10922726300486829,
      "inner": 74,
      "median_ms": 0.5403750000295002,
      "q1_ms": 0.5402741284145836,
      "q3_ms": 0.5403887973099355,
      "tflops": 27.818432635076288
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 1024,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.01181650161743164,
      "rel_err": 0.004610817413777113,
      "peak_mb": 84.501953125,
      "compile_s": 0.0006175779999466613,
      "inner": 86,
      "median_ms": 0.526831058080059,
      "q1_ms": 0.5265878023428107,
      "q3_ms": 0.5281271046125005,
      "tflops": 28.5335978307407
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 1024,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.01181650161743164,
      "rel_err": 0.004610817413777113,
      "peak_mb": 80.509765625,
      "compile_s": 0.0008401790109928697,
      "inner": 68,
      "median_ms": 0.6940180441340679,
      "q1_ms": 0.6931906028672956,
      "q3_ms": 0.6943016029682065,
      "tflops": 21.659934728002693
    },
    {
      "impl": "torch_math",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 1024,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.01178121566772461,
      "rel_err": 0.0036301971413195133,
      "peak_mb": 588.2509765625,
      "compile_s": 0.005198061000555754,
      "inner": 9,
      "median_ms": 5.293382556475181,
      "q1_ms": 5.2912957778769645,
      "q3_ms": 5.293958111400975,
      "tflops": 2.839844914970578
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 1500,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0007491111755371094,
      "rel_err": 0.0027269322890788317,
      "xla_temp_mb": 24.0,
      "xla_arg_mb": 17.578125,
      "xla_out_mb": 5.859375,
      "compile_s": 0.21513231299468316,
      "inner": 100,
      "median_ms": 0.39651023005717434,
      "q1_ms": 0.38827544507512357,
      "q3_ms": 0.3966974300419679,
      "tflops": 46.48555977317967
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 1500,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0007491111755371094,
      "rel_err": 0.0027269322890788317,
      "xla_temp_mb": 24.0,
      "xla_arg_mb": 17.578125,
      "xla_out_mb": 5.859375,
      "compile_s": 0.21084067499032244,
      "inner": 99,
      "median_ms": 0.38632864648161097,
      "q1_ms": 0.38597155549043716,
      "q3_ms": 0.38637763131503955,
      "tflops": 47.71067371748047
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 1500,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0007491111755371094,
      "rel_err": 0.0027269322890788317,
      "xla_temp_mb": 24.0,
      "xla_arg_mb": 17.578125,
      "xla_out_mb": 5.859375,
      "compile_s": 0.21452731899626087,
      "inner": 99,
      "median_ms": 0.3864091010811278,
      "q1_ms": 0.38636300502335563,
      "q3_ms": 0.38644851516400736,
      "tflops": 47.70073983358416
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 1500,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.001270890235900879,
      "rel_err": 0.004626324865967035,
      "xla_temp_mb": 411.98731994628906,
      "xla_arg_mb": 17.578125,
      "xla_out_mb": 5.859375,
      "compile_s": 0.7635749000037322,
      "inner": 40,
      "median_ms": 1.155215374819818,
      "q1_ms": 1.1549135499080876,
      "q3_ms": 1.1556681127331103,
      "tflops": 15.95546631542615
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 1500,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0007491111755371094,
      "rel_err": 0.0027269322890788317,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 17.578125,
      "xla_out_mb": 5.859375,
      "compile_s": 0.0471200489992043,
      "inner": 91,
      "median_ms": 0.3654805824715787,
      "q1_ms": 0.36510646704310595,
      "q3_ms": 0.36564991207522146,
      "tflops": 50.432227822755394
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 1500,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0007491111755371094,
      "rel_err": 0.0027269322890788317,
      "peak_mb": 57.60595703125,
      "compile_s": 0.0003492010000627488,
      "inner": 150,
      "median_ms": 0.31307933337908855,
      "q1_ms": 0.31303565003327094,
      "q3_ms": 0.3131980866601225,
      "tflops": 58.87325682299771
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 1500,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0007491111755371094,
      "rel_err": 0.0027269322890788317,
      "peak_mb": 57.4228515625,
      "compile_s": 0.00043446199561003596,
      "inner": 120,
      "median_ms": 0.41472779169756296,
      "q1_ms": 0.413699979192946,
      "q3_ms": 0.41477071257152903,
      "tflops": 44.443609444532704
    },
    {
      "impl": "torch_math",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 1500,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0007491111755371094,
      "rel_err": 0.0027269322890788317,
      "peak_mb": 717.58984375,
      "compile_s": 0.003991888996097259,
      "inner": 12,
      "median_ms": 4.040296416254326,
      "q1_ms": 4.0358130008826265,
      "q3_ms": 4.04043425017638,
      "tflops": 4.562041518995263
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 1500,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0013424158096313477,
      "rel_err": 0.0044233533553779125,
      "xla_temp_mb": 48.00048828125,
      "xla_arg_mb": 29.296875,
      "xla_out_mb": 17.578147888183594,
      "compile_s": 0.4913856439961819,
      "inner": 32,
      "median_ms": 1.537245218969474,
      "q1_ms": 1.5370118123883003,
      "q3_ms": 1.5374391875866422,
      "tflops": 41.965978624573005
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 1500,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0013424158096313477,
      "rel_err": 0.0044233533553779125,
      "xla_temp_mb": 48.00048828125,
      "xla_arg_mb": 29.296875,
      "xla_out_mb": 17.578147888183594,
      "compile_s": 0.48754606199509,
      "inner": 32,
      "median_ms": 1.5496580626859213,
      "q1_ms": 1.5494753752136603,
      "q3_ms": 1.5498266720896936,
      "tflops": 41.62982889798641
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 1500,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0013424158096313477,
      "rel_err": 0.0044233533553779125,
      "xla_temp_mb": 48.00048828125,
      "xla_arg_mb": 29.296875,
      "xla_out_mb": 17.578147888183594,
      "compile_s": 0.4903800920001231,
      "inner": 31,
      "median_ms": 1.5621567741149074,
      "q1_ms": 1.5590741935621706,
      "q3_ms": 1.5624628224853248,
      "tflops": 41.296751432999706
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 1500,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0013798773288726807,
      "rel_err": 0.0044233533553779125,
      "xla_temp_mb": 686.8286285400391,
      "xla_arg_mb": 29.296875,
      "xla_out_mb": 17.578147888183594,
      "compile_s": 1.2416311830020277,
      "inner": 14,
      "median_ms": 3.370687928898925,
      "q1_ms": 3.370228785830217,
      "q3_ms": 3.3719852851520824,
      "tflops": 19.139119776381555
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 1500,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0013424158096313477,
      "rel_err": 0.0044233533553779125,
      "xla_temp_mb": 35.554443359375,
      "xla_arg_mb": 29.296875,
      "xla_out_mb": 17.578147888183594,
      "compile_s": 0.10406849600258283,
      "inner": 32,
      "median_ms": 1.413593562119786,
      "q1_ms": 1.4123665469014668,
      "q3_ms": 1.4180111093082814,
      "tflops": 45.636880167492826
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 1500,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0013424158096313477,
      "rel_err": 0.0044233533553779125,
      "peak_mb": 116.95068359375,
      "compile_s": 0.0014663129986729473,
      "inner": 35,
      "median_ms": 1.327720285709282,
      "q1_ms": 1.3269495431034426,
      "q3_ms": 1.327903914248704,
      "tflops": 48.5885473727902
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 1500,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0015041232109069824,
      "rel_err": 0.0044233533553779125,
      "peak_mb": 111.09912109375,
      "compile_s": 0.002119879995007068,
      "inner": 23,
      "median_ms": 2.207672348089071,
      "q1_ms": 2.203289717726875,
      "q3_ms": 2.207974238958939,
      "tflops": 29.221727606381737
    },
    {
      "impl": "torch_math",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 1500,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0011610090732574463,
      "rel_err": 0.004025899805128574,
      "peak_mb": 1203.4931640625,
      "compile_s": 0.009371334002935328,
      "inner": 5,
      "median_ms": 9.669491997919977,
      "q1_ms": 9.648008499061689,
      "q3_ms": 9.671094999066554,
      "tflops": 6.671705195461902
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 1500,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.006905078887939453,
      "rel_err": 0.0023887839633971453,
      "xla_temp_mb": 24.0,
      "xla_arg_mb": 17.578125,
      "xla_out_mb": 5.859375,
      "compile_s": 0.22206739899411332,
      "inner": 117,
      "median_ms": 0.26063465811292297,
      "q1_ms": 0.258494790583256,
      "q3_ms": 0.26097726492255,
      "tflops": 35.359840731569406
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 1500,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.006905078887939453,
      "rel_err": 0.0023887839633971453,
      "xla_temp_mb": 24.0,
      "xla_arg_mb": 17.578125,
      "xla_out_mb": 5.859375,
      "compile_s": 0.21862934599630535,
      "inner": 142,
      "median_ms": 0.2604252958020956,
      "q1_ms": 0.259793820410957,
      "q3_ms": 0.26062073237893935,
      "tflops": 35.38826737861707
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 1500,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.006905078887939453,
      "rel_err": 0.0023887839633971453,
      "xla_temp_mb": 24.0,
      "xla_arg_mb": 17.578125,
      "xla_out_mb": 5.859375,
      "compile_s": 0.218038909995812,
      "inner": 128,
      "median_ms": 0.2608076406431792,
      "q1_ms": 0.26037104299803104,
      "q3_ms": 0.26096559366806105,
      "tflops": 35.336388064676214
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 1500,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.01177072525024414,
      "rel_err": 0.004072034731507301,
      "xla_temp_mb": 411.98731994628906,
      "xla_arg_mb": 17.578125,
      "xla_out_mb": 5.859375,
      "compile_s": 0.5492628479987616,
      "inner": 38,
      "median_ms": 1.202719447304078,
      "q1_ms": 1.2021100130943139,
      "q3_ms": 1.2031646447691235,
      "tflops": 7.662634890171489
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 1500,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.006905078887939453,
      "rel_err": 0.0023887839633971453,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 17.578125,
      "xla_out_mb": 5.859375,
      "compile_s": 0.04814724098832812,
      "inner": 124,
      "median_ms": 0.255756451576505,
      "q1_ms": 0.25567327417863245,
      "q3_ms": 0.2558202056294953,
      "tflops": 36.03428161124294
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 1500,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.006905078887939453,
      "rel_err": 0.0023887839633971453,
      "peak_mb": 57.88720703125,
      "compile_s": 0.0002652710099937394,
      "inner": 197,
      "median_ms": 0.2246363299602151,
      "q1_ms": 0.22446585279596298,
      "q3_ms": 0.22512987818330205,
      "tflops": 41.02631129004033
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 1500,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.006905078887939453,
      "rel_err": 0.0023887839633971453,
      "peak_mb": 57.7041015625,
      "compile_s": 0.00029697200807277113,
      "inner": 177,
      "median_ms": 0.25731083620629025,
      "q1_ms": 0.2565538417942569,
      "q3_ms": 0.2575015932345226,
      "tflops": 35.81660273573315
    },
    {
      "impl": "torch_math",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 1500,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.006905078887939453,
      "rel_err": 0.0023887839633971453,
      "peak_mb": 726.17333984375,
      "compile_s": 0.005017950999899767,
      "inner": 9,
      "median_ms": 5.050623445034337,
      "q1_ms": 5.048311221697885,
      "q3_ms": 5.052574333805953,
      "tflops": 1.8247252245781598
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 1500,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.014543533325195312,
      "rel_err": 0.004484972450882196,
      "xla_temp_mb": 48.00048828125,
      "xla_arg_mb": 29.296875,
      "xla_out_mb": 17.578147888183594,
      "compile_s": 0.5113424779992783,
      "inner": 50,
      "median_ms": 0.916871480003465,
      "q1_ms": 0.916321050026454,
      "q3_ms": 0.9193265299836639,
      "tflops": 35.180503160462685
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 1500,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.014543533325195312,
      "rel_err": 0.004484972450882196,
      "xla_temp_mb": 48.00048828125,
      "xla_arg_mb": 29.296875,
      "xla_out_mb": 17.578147888183594,
      "compile_s": 0.5095993919967441,
      "inner": 49,
      "median_ms": 0.9107345101014921,
      "q1_ms": 0.9102601632511491,
      "q3_ms": 0.9113749183597974,
      "tflops": 35.41756641724864
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 1500,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.014543533325195312,
      "rel_err": 0.004484972450882196,
      "xla_temp_mb": 48.00048828125,
      "xla_arg_mb": 29.296875,
      "xla_out_mb": 17.578147888183594,
      "compile_s": 0.5066961910051759,
      "inner": 49,
      "median_ms": 0.910767449104056,
      "q1_ms": 0.9099673570373703,
      "q3_ms": 0.9113415000349645,
      "tflops": 35.41628549826962
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 1500,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.014543533325195312,
      "rel_err": 0.005273605231195688,
      "xla_temp_mb": 686.8286285400391,
      "xla_arg_mb": 29.296875,
      "xla_out_mb": 17.578147888183594,
      "compile_s": 0.9600043880054727,
      "inner": 13,
      "median_ms": 3.5720951541756785,
      "q1_ms": 3.5700651927842184,
      "q3_ms": 3.578434076469141,
      "tflops": 9.029994613187625
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 1500,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.014543533325195312,
      "rel_err": 0.004484972450882196,
      "xla_temp_mb": 35.554443359375,
      "xla_arg_mb": 29.296875,
      "xla_out_mb": 17.578147888183594,
      "compile_s": 0.10660905099939555,
      "inner": 42,
      "median_ms": 0.9853815239953941,
      "q1_ms": 0.9851646786306185,
      "q3_ms": 0.9881590477085584,
      "tflops": 32.73452892562127
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 1500,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.014543533325195312,
      "rel_err": 0.004484972450882196,
      "peak_mb": 117.79443359375,
      "compile_s": 0.0009865059982985258,
      "inner": 51,
      "median_ms": 0.9493631762726342,
      "q1_ms": 0.9479837254712395,
      "q3_ms": 0.949705313627526,
      "tflops": 33.97646001674796
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 1500,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.014543533325195312,
      "rel_err": 0.004484972450882196,
      "peak_mb": 111.94287109375,
      "compile_s": 0.0013992860040161759,
      "inner": 36,
      "median_ms": 1.3842435833390079,
      "q1_ms": 1.3819566389025162,
      "q3_ms": 1.3846805278363816,
      "tflops": 23.30225719536556
    },
    {
      "impl": "torch_math",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 1500,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.011868476867675781,
      "rel_err": 0.003660031361505389,
      "peak_mb": 1204.0556640625,
      "compile_s": 0.010342420995584689,
      "inner": 4,
      "median_ms": 10.578304500086233,
      "q1_ms": 10.559899499639869,
      "q3_ms": 10.581296250165906,
      "tflops": 3.0492599262705147
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 2048,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0007493197917938232,
      "rel_err": 0.0029883224051445723,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.18444056500447914,
      "inner": 70,
      "median_ms": 0.5682391428438547,
      "q1_ms": 0.5682048714204159,
      "q3_ms": 0.5683049142784772,
      "tflops": 60.46703892315571
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 2048,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0007493197917938232,
      "rel_err": 0.0029883224051445723,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.17769749200670049,
      "inner": 76,
      "median_ms": 0.5586033422380415,
      "q1_ms": 0.5584571314868395,
      "q3_ms": 0.5600840328330186,
      "tflops": 61.51008375699629
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 2048,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0007493197917938232,
      "rel_err": 0.0029883224051445723,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.17666970800200943,
      "inner": 71,
      "median_ms": 0.5587470421770876,
      "q1_ms": 0.5585256408331063,
      "q3_ms": 0.5620504789092106,
      "tflops": 61.494264442316506
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 2048,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0007819384336471558,
      "rel_err": 0.003118406981229782,
      "xla_temp_mb": 768.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.8573716670071008,
      "inner": 22,
      "median_ms": 2.0516256818981757,
      "q1_ms": 2.0502412953646854,
      "q3_ms": 2.051964999819492,
      "tflops": 16.747566903242397
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 2048,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0007493197917938232,
      "rel_err": 0.0029883224051445723,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.04783562201191671,
      "inner": 62,
      "median_ms": 0.6316174837974682,
      "q1_ms": 0.6313382581274416,
      "q3_ms": 0.6317787579797779,
      "tflops": 54.399599835994486
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 2048,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0007493197917938232,
      "rel_err": 0.0029883224051445723,
      "peak_mb": 72.5009765625,
      "compile_s": 0.0005791750008938834,
      "inner": 88,
      "median_ms": 0.5524068181495585,
      "q1_ms": 0.5511405056908685,
      "q3_ms": 0.5525657044761993,
      "tflops": 62.20006205408104
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 2048,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0007493197917938232,
      "rel_err": 0.0029883224051445723,
      "peak_mb": 72.2509765625,
      "compile_s": 0.0007063060038490221,
      "inner": 73,
      "median_ms": 0.7105656849277805,
      "q1_ms": 0.7105340411900924,
      "q3_ms": 0.7107288630285351,
      "tflops": 48.355470995608535
    },
    {
      "impl": "torch_math",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 2048,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0007493197917938232,
      "rel_err": 0.0029883224051445723,
      "peak_mb": 1280.31298828125,
      "compile_s": 0.007265700987773016,
      "inner": 6,
      "median_ms": 7.335837833428134,
      "q1_ms": 7.33441166812554,
      "q3_ms": 7.341573583592739,
      "tflops": 4.683819237583014
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 2048,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0009163618087768555,
      "rel_err": 0.003617694368585944,
      "xla_temp_mb": 16.50048828125,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.2676046940032393,
      "inner": 20,
      "median_ms": 2.387044150236761,
      "q1_ms": 2.3863107497163583,
      "q3_ms": 2.387733424984617,
      "tflops": 50.379916213980366
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 2048,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0009163618087768555,
      "rel_err": 0.003617694368585944,
      "xla_temp_mb": 16.50048828125,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.2628140439919662,
      "inner": 20,
      "median_ms": 2.3879879496234935,
      "q1_ms": 2.3860319997766055,
      "q3_ms": 2.3976967997441534,
      "tflops": 50.360004667092596
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 2048,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0009163618087768555,
      "rel_err": 0.003617694368585944,
      "xla_temp_mb": 16.50048828125,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.26407227299932856,
      "inner": 20,
      "median_ms": 2.3874774502473883,
      "q1_ms": 2.386246125024627,
      "q3_ms": 2.393666674834094,
      "tflops": 50.37077283202773
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 2048,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0012067556381225586,
      "rel_err": 0.003785555250942707,
      "xla_temp_mb": 1280.25,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 1.4035320669936482,
      "inner": 8,
      "median_ms": 6.112728624430019,
      "q1_ms": 6.103204937062401,
      "q3_ms": 6.11962049970316,
      "tflops": 19.673551972743358
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 2048,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0009163618087768555,
      "rel_err": 0.003617694368585944,
      "xla_temp_mb": 48.500732421875,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.1058063529926585,
      "inner": 19,
      "median_ms": 2.3730545263357556,
      "q1_ms": 2.3721812892625547,
      "q3_ms": 2.3800648686071666,
      "tflops": 50.67691574440668
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 2048,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0009163618087768555,
      "rel_err": 0.003617694368585944,
      "peak_mb": 152.751953125,
      "compile_s": 0.0022522600047523156,
      "inner": 22,
      "median_ms": 2.243770499427972,
      "q1_ms": 2.2433672271621288,
      "q3_ms": 2.2529811136667957,
      "tflops": 53.59687379732415
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 2048,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.001443326473236084,
      "rel_err": 0.004759587347507477,
      "peak_mb": 144.767578125,
      "compile_s": 0.003335058005177416,
      "inner": 15,
      "median_ms": 3.554215467496154,
      "q1_ms": 3.538459000022461,
      "q3_ms": 3.555144199344795,
      "tflops": 33.83562009331955
    },
    {
      "impl": "torch_math",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 2048,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0009163618087768555,
      "rel_err": 0.0030359115917235613,
      "peak_mb": 2184.2509765625,
      "compile_s": 0.017277227001613937,
      "inner": 2,
      "median_ms": 17.613768002775032,
      "q1_ms": 17.57822350191418,
      "q3_ms": 17.632506249356084,
      "tflops": 6.827561500131789
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 2048,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.006710529327392578,
      "rel_err": 0.002321480307728052,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.20723564700165298,
      "inner": 115,
      "median_ms": 0.3424371390239052,
      "q1_ms": 0.33747143910580274,
      "q3_ms": 0.34309716087907716,
      "tflops": 50.16940987467102
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 2048,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.006710529327392578,
      "rel_err": 0.002321480307728052,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.19706522100022994,
      "inner": 108,
      "median_ms": 0.332012750058547,
      "q1_ms": 0.3312693148068825,
      "q3_ms": 0.3327598610835762,
      "tflops": 51.7446067386584
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 2048,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.006710529327392578,
      "rel_err": 0.002321480307728052,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.2048493239999516,
      "inner": 119,
      "median_ms": 0.33195749582613215,
      "q1_ms": 0.3314176176311722,
      "q3_ms": 0.3324114621427114,
      "tflops": 51.75321961399004
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 2048,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00992274284362793,
      "rel_err": 0.0034327327739447355,
      "xla_temp_mb": 768.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.5200527330016484,
      "inner": 19,
      "median_ms": 2.4314844738814587,
      "q1_ms": 2.4297823160611047,
      "q3_ms": 2.4324898155387737,
      "tflops": 7.065588684008831
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 2048,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.006710529327392578,
      "rel_err": 0.002321480307728052,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.048963255001581274,
      "inner": 87,
      "median_ms": 0.4015222759502565,
      "q1_ms": 0.40147005171908445,
      "q3_ms": 0.4016327182963041,
      "tflops": 42.78683951803552
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 2048,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.006710529327392578,
      "rel_err": 0.002321480307728052,
      "peak_mb": 72.5009765625,
      "compile_s": 0.000364930005162023,
      "inner": 140,
      "median_ms": 0.3423464071862067,
      "q1_ms": 0.3414438142499421,
      "q3_ms": 0.34260966073946164,
      "tflops": 50.18270623957694
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 2048,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.006710529327392578,
      "rel_err": 0.002321480307728052,
      "peak_mb": 72.2509765625,
      "compile_s": 0.00044622500718105584,
      "inner": 113,
      "median_ms": 0.42366606195712775,
      "q1_ms": 0.4231672478011455,
      "q3_ms": 0.42381998220050654,
      "tflops": 40.550496550602844
    },
    {
      "impl": "torch_math",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 2048,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.006710529327392578,
      "rel_err": 0.002321480307728052,
      "peak_mb": 1296.31298828125,
      "compile_s": 0.00917718700657133,
      "inner": 5,
      "median_ms": 9.231066799839027,
      "q1_ms": 9.230098999978509,
      "q3_ms": 9.23701709980378,
      "tflops": 1.8610924995472446
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 2048,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.013099074363708496,
      "rel_err": 0.006403794977813959,
      "xla_temp_mb": 16.50048828125,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.29694258199015167,
      "inner": 34,
      "median_ms": 1.3641848234007792,
      "q1_ms": 1.3630016619572416,
      "q3_ms": 1.3657681028074449,
      "tflops": 44.077269525769196
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 2048,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.013099074363708496,
      "rel_err": 0.006403794977813959,
      "xla_temp_mb": 16.50048828125,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.29389097999955993,
      "inner": 35,
      "median_ms": 1.371499085715706,
      "q1_ms": 1.3665356858317472,
      "q3_ms": 1.3726496856959005,
      "tflops": 43.84220359331984
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 2048,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.013099074363708496,
      "rel_err": 0.006403794977813959,
      "xla_temp_mb": 16.50048828125,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.29759869299596176,
      "inner": 34,
      "median_ms": 1.3626527939135593,
      "q1_ms": 1.3620977793727755,
      "q3_ms": 1.365101691429599,
      "tflops": 44.126825565966115
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 2048,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.012175559997558594,
      "rel_err": 0.004521289374679327,
      "xla_temp_mb": 1280.25,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 1.5766525840008399,
      "inner": 6,
      "median_ms": 7.046049666920832,
      "q1_ms": 7.041901832659884,
      "q3_ms": 7.053471167940492,
      "tflops": 8.53379481928588
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 2048,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.013099074363708496,
      "rel_err": 0.006403794977813959,
      "xla_temp_mb": 48.500732421875,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.109038745998987,
      "inner": 29,
      "median_ms": 1.535499310381871,
      "q1_ms": 1.5346663278941033,
      "q3_ms": 1.536224620739913,
      "tflops": 39.15960218116027
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 2048,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.013099074363708496,
      "rel_err": 0.006403794977813959,
      "peak_mb": 152.751953125,
      "compile_s": 0.0014811509900027886,
      "inner": 34,
      "median_ms": 1.4579987058395465,
      "q1_ms": 1.4565658530045766,
      "q3_ms": 1.4597944412825574,
      "tflops": 41.24114918838431
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 2048,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.013099074363708496,
      "rel_err": 0.006403794977813959,
      "peak_mb": 144.767578125,
      "compile_s": 0.0020718390005640686,
      "inner": 24,
      "median_ms": 2.153232417185791,
      "q1_ms": 2.1492826666265805,
      "q3_ms": 2.154902666613149,
      "tflops": 27.92524469912425
    },
    {
      "impl": "torch_math",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 2048,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.012173175811767578,
      "rel_err": 0.0032567521557211876,
      "peak_mb": 2184.2509765625,
      "compile_s": 0.019027225993340835,
      "inner": 2,
      "median_ms": 19.390889501664788,
      "q1_ms": 19.331394247274147,
      "q3_ms": 19.417827996221604,
      "tflops": 3.1009171672520557
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 3000,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0005662888288497925,
      "rel_err": 0.003144341055303812,
      "xla_temp_mb": 48.0,
      "xla_arg_mb": 35.15625,
      "xla_out_mb": 11.71875,
      "compile_s": 0.21819964399037417,
      "inner": 34,
      "median_ms": 1.3224544116597186,
      "q1_ms": 1.3182240588374792,
      "q3_ms": 1.3228497058769468,
      "tflops": 55.75088210978043
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 3000,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0005662888288497925,
      "rel_err": 0.003144341055303812,
      "xla_temp_mb": 48.0,
      "xla_arg_mb": 35.15625,
      "xla_out_mb": 11.71875,
      "compile_s": 0.21452741899702232,
      "inner": 36,
      "median_ms": 1.300818860828359,
      "q1_ms": 1.3005073055359795,
      "q3_ms": 1.3010370556811623,
      "tflops": 56.67814499018729
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 3000,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0005662888288497925,
      "rel_err": 0.003144341055303812,
      "xla_temp_mb": 48.0,
      "xla_arg_mb": 35.15625,
      "xla_out_mb": 11.71875,
      "compile_s": 0.21834278499591164,
      "inner": 36,
      "median_ms": 1.3006760834994362,
      "q1_ms": 1.3004114443901926,
      "q3_ms": 1.3012975693224385,
      "tflops": 56.68436664233625
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 3000,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0005277097225189209,
      "rel_err": 0.0029301291797310114,
      "xla_temp_mb": 1647.949234008789,
      "xla_arg_mb": 35.15625,
      "xla_out_mb": 11.71875,
      "compile_s": 0.8599983059975784,
      "inner": 10,
      "median_ms": 4.5019443990895525,
      "q1_ms": 4.488652500731405,
      "q3_ms": 4.507680499955313,
      "tflops": 16.376923716541306
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 3000,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0005662888288497925,
      "rel_err": 0.003144341055303812,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 35.15625,
      "xla_out_mb": 11.71875,
      "compile_s": 0.048185574007220566,
      "inner": 39,
      "median_ms": 1.184315281884315,
      "q1_ms": 1.1766268974334264,
      "q3_ms": 1.1850083974422887,
      "tflops": 62.25369302226214
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 3000,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0005277097225189209,
      "rel_err": 0.0029301291797310114,
      "peak_mb": 99.9296875,
      "compile_s": 0.0012135349970776588,
      "inner": 42,
      "median_ms": 1.1996458096767864,
      "q1_ms": 1.198073309621707,
      "q3_ms": 1.2000227023755934,
      "tflops": 61.458139898695684
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 3000,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0005662888288497925,
      "rel_err": 0.003144341055303812,
      "peak_mb": 99.5634765625,
      "compile_s": 0.0013979140057926998,
      "inner": 36,
      "median_ms": 1.4834250557921929,
      "q1_ms": 1.4822375414951239,
      "q3_ms": 1.4916115140092896,
      "tflops": 49.70119637127678
    },
    {
      "impl": "torch_math",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 3000,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0004800260066986084,
      "rel_err": 0.0026653632521629333,
      "peak_mb": 2655.85986328125,
      "compile_s": 0.01563984900712967,
      "inner": 3,
      "median_ms": 16.12977966336378,
      "q1_ms": 16.103750499799695,
      "q3_ms": 16.144886668674495,
      "tflops": 4.57092418735647
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 3000,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0010789334774017334,
      "rel_err": 0.00497435312718153,
      "xla_temp_mb": 96.00048828125,
      "xla_arg_mb": 58.59375,
      "xla_out_mb": 35.156272888183594,
      "compile_s": 0.3224735179974232,
      "inner": 9,
      "median_ms": 5.588811333937984,
      "q1_ms": 5.588144499698602,
      "q3_ms": 5.603476777549885,
      "tflops": 46.17225105327977
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 3000,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0010789334774017334,
      "rel_err": 0.00497435312718153,
      "xla_temp_mb": 96.00048828125,
      "xla_arg_mb": 58.59375,
      "xla_out_mb": 35.156272888183594,
      "compile_s": 0.31909918700694107,
      "inner": 9,
      "median_ms": 5.679714333382435,
      "q1_ms": 5.675131277207079,
      "q3_ms": 5.688147444743663,
      "tflops": 45.43327091000454
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 3000,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0010789334774017334,
      "rel_err": 0.00497435312718153,
      "xla_temp_mb": 96.00048828125,
      "xla_arg_mb": 58.59375,
      "xla_out_mb": 35.156272888183594,
      "compile_s": 0.3121126530022593,
      "inner": 8,
      "median_ms": 5.665642000167281,
      "q1_ms": 5.642504813295091,
      "q3_ms": 5.682910000359698,
      "tflops": 45.546118161433604
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 3000,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0010789334774017334,
      "rel_err": 0.00497435312718153,
      "xla_temp_mb": 2746.9482421875,
      "xla_arg_mb": 58.59375,
      "xla_out_mb": 35.156272888183594,
      "compile_s": 1.5085847760055913,
      "inner": 3,
      "median_ms": 12.71427666749029,
      "q1_ms": 12.707178168056998,
      "q3_ms": 12.743109500055045,
      "tflops": 20.295924553837548
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 3000,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0010789334774017334,
      "rel_err": 0.00497435312718153,
      "xla_temp_mb": 71.045654296875,
      "xla_arg_mb": 58.59375,
      "xla_out_mb": 35.156272888183594,
      "compile_s": 0.10971776000224054,
      "inner": 10,
      "median_ms": 4.89430470042862,
      "q1_ms": 4.8780007993627805,
      "q3_ms": 4.941394150228007,
      "tflops": 52.72413872748899
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 3000,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0010789334774017334,
      "rel_err": 0.00497435312718153,
      "peak_mb": 220.5869140625,
      "compile_s": 0.004503157004364766,
      "inner": 11,
      "median_ms": 4.607848545261236,
      "q1_ms": 4.60775877251184,
      "q3_ms": 4.607882272632031,
      "tflops": 56.001840656281885
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 3000,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0010789334774017334,
      "rel_err": 0.00497435312718153,
      "peak_mb": 208.5791015625,
      "compile_s": 0.0072993759968085214,
      "inner": 6,
      "median_ms": 7.811412166726465,
      "q1_ms": 7.809457582576821,
      "q3_ms": 7.818917166635705,
      "tflops": 33.034743845572855
    },
    {
      "impl": "torch_math",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 3000,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0007157474756240845,
      "rel_err": 0.0031312471255660057,
      "peak_mb": 4589.2509765625,
      "compile_s": 0.037430954005685635,
      "inner": 1,
      "median_ms": 37.81480999896303,
      "q1_ms": 37.746996502391994,
      "q3_ms": 37.92647099180613,
      "tflops": 6.823993033604461
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 3000,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.007532358169555664,
      "rel_err": 0.0026057888753712177,
      "xla_temp_mb": 48.0,
      "xla_arg_mb": 35.15625,
      "xla_out_mb": 11.71875,
      "compile_s": 0.2192406420072075,
      "inner": 56,
      "median_ms": 0.8224066784480654,
      "q1_ms": 0.8155633838344198,
      "q3_ms": 0.8321925001187733,
      "tflops": 44.824538717954916
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 3000,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.007532358169555664,
      "rel_err": 0.0026057888753712177,
      "xla_temp_mb": 48.0,
      "xla_arg_mb": 35.15625,
      "xla_out_mb": 11.71875,
      "compile_s": 0.21762137999758124,
      "inner": 57,
      "median_ms": 0.7894130177558971,
      "q1_ms": 0.7874687280650905,
      "q3_ms": 0.7904054299088832,
      "tflops": 46.69798846843835
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 3000,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.007532358169555664,
      "rel_err": 0.0026057888753712177,
      "xla_temp_mb": 48.0,
      "xla_arg_mb": 35.15625,
      "xla_out_mb": 11.71875,
      "compile_s": 0.2193085710023297,
      "inner": 53,
      "median_ms": 0.7853498679634078,
      "q1_ms": 0.7835249056850718,
      "q3_ms": 0.787739867793686,
      "tflops": 46.93958897019592
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 3000,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.008092641830444336,
      "rel_err": 0.0027996166609227657,
      "xla_temp_mb": 1647.949234008789,
      "xla_arg_mb": 35.15625,
      "xla_out_mb": 11.71875,
      "compile_s": 0.5034067099913955,
      "inner": 10,
      "median_ms": 4.49499139940599,
      "q1_ms": 4.483910551061854,
      "q3_ms": 4.498512599093374,
      "tflops": 8.201128038837082
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 3000,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.007532358169555664,
      "rel_err": 0.0026057888753712177,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 35.15625,
      "xla_out_mb": 11.71875,
      "compile_s": 0.04668372399464715,
      "inner": 53,
      "median_ms": 0.7825623395993322,
      "q1_ms": 0.7821866227390755,
      "q3_ms": 0.7833247263816994,
      "tflops": 47.10679026398609
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 3000,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.007532358169555664,
      "rel_err": 0.0026057888753712177,
      "peak_mb": 100.2109375,
      "compile_s": 0.0007381460018223152,
      "inner": 70,
      "median_ms": 0.7261163143474342,
      "q1_ms": 0.7230728357431612,
      "q3_ms": 0.726797021447315,
      "tflops": 50.768725714598396
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 3000,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.007532358169555664,
      "rel_err": 0.0026057888753712177,
      "peak_mb": 99.8447265625,
      "compile_s": 0.000840370004880242,
      "inner": 62,
      "median_ms": 0.8363694193371151,
      "q1_ms": 0.8362617983009606,
      "q3_ms": 0.8371178467519536,
      "tflops": 44.07621697744217
    },
    {
      "impl": "torch_math",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 3000,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.007532358169555664,
      "rel_err": 0.0026057888753712177,
      "peak_mb": 2690.4736328125,
      "compile_s": 0.02009916400129441,
      "inner": 2,
      "median_ms": 20.383211005537305,
      "q1_ms": 20.370649748656433,
      "q3_ms": 20.39769800103386,
      "tflops": 1.8085472396859121
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 3000,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.012875914573669434,
      "rel_err": 0.003490941599011421,
      "xla_temp_mb": 96.00048828125,
      "xla_arg_mb": 58.59375,
      "xla_out_mb": 35.156272888183594,
      "compile_s": 0.3166914120083675,
      "inner": 15,
      "median_ms": 3.094087934005074,
      "q1_ms": 3.090985000017099,
      "q3_ms": 3.094916499685496,
      "tflops": 41.700172313133876
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 3000,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.012875914573669434,
      "rel_err": 0.003490941599011421,
      "xla_temp_mb": 96.00048828125,
      "xla_arg_mb": 58.59375,
      "xla_out_mb": 35.156272888183594,
      "compile_s": 0.3176818549982272,
      "inner": 16,
      "median_ms": 2.9988186242917436,
      "q1_ms": 2.9930767500445654,
      "q3_ms": 3.0000149058651004,
      "tflops": 43.0249428741202
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 3000,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.012875914573669434,
      "rel_err": 0.003490941599011421,
      "xla_temp_mb": 96.00048828125,
      "xla_arg_mb": 58.59375,
      "xla_out_mb": 35.156272888183594,
      "compile_s": 0.3118229550018441,
      "inner": 16,
      "median_ms": 2.9872341874579433,
      "q1_ms": 2.9868337187508587,
      "q3_ms": 2.993511718159425,
      "tflops": 43.19179277664735
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 3000,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.012875914573669434,
      "rel_err": 0.0037529929541051388,
      "xla_temp_mb": 2746.9482421875,
      "xla_arg_mb": 58.59375,
      "xla_out_mb": 35.156272888183594,
      "compile_s": 0.9560593770002015,
      "inner": 3,
      "median_ms": 12.79665233353929,
      "q1_ms": 12.787009167368524,
      "q3_ms": 12.798459167242981,
      "tflops": 10.082636977003393
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 3000,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.012875914573669434,
      "rel_err": 0.003490941599011421,
      "xla_temp_mb": 71.045654296875,
      "xla_arg_mb": 58.59375,
      "xla_out_mb": 35.156272888183594,
      "compile_s": 0.11019416199997067,
      "inner": 16,
      "median_ms": 2.951692562419339,
      "q1_ms": 2.9502821253117872,
      "q3_ms": 2.9607496562675806,
      "tflops": 43.71186946862994
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 3000,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.012875914573669434,
      "rel_err": 0.003490941599011421,
      "peak_mb": 220.8681640625,
      "compile_s": 0.0027489689964568242,
      "inner": 18,
      "median_ms": 2.7662050552963695,
      "q1_ms": 2.763523666443588,
      "q3_ms": 2.767530416652638,
      "tflops": 46.64296298387628
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 3000,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.016532182693481445,
      "rel_err": 0.006153587251901627,
      "peak_mb": 208.8603515625,
      "compile_s": 0.004237344008288346,
      "inner": 11,
      "median_ms": 4.455856544981626,
      "q1_ms": 4.45042900008213,
      "q3_ms": 4.459698363709983,
      "tflops": 28.95604889823311
    },
    {
      "impl": "torch_math",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 3000,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.012837886810302734,
      "rel_err": 0.003093926003202796,
      "peak_mb": 4589.2509765625,
      "compile_s": 0.04125188999751117,
      "inner": 1,
      "median_ms": 41.69686199747957,
      "q1_ms": 41.67060749750817,
      "q3_ms": 41.8060979936854,
      "tflops": 3.094333573778263
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 4096,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0006000101566314697,
      "rel_err": 0.0036908327601850033,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.18793096599983983,
      "inner": 22,
      "median_ms": 2.133927727382715,
      "q1_ms": 2.129034386598505,
      "q3_ms": 2.1344701364631127,
      "tflops": 64.4065643406632
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 4096,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0006000101566314697,
      "rel_err": 0.0036908327601850033,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.18044397598714568,
      "inner": 24,
      "median_ms": 2.11597945841883,
      "q1_ms": 2.1153345208707224,
      "q3_ms": 2.1177225832313225,
      "tflops": 64.95287698808832
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 4096,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0006000101566314697,
      "rel_err": 0.0036908327601850033,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.17946954299986828,
      "inner": 22,
      "median_ms": 2.1180467725075274,
      "q1_ms": 2.1163187272825534,
      "q3_ms": 2.12107936386019,
      "tflops": 64.88947990005332
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 4096,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0006000101566314697,
      "rel_err": 0.0036908327601850033,
      "xla_temp_mb": 3072.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 1.016846844999236,
      "inner": 6,
      "median_ms": 7.8457153334359955,
      "q1_ms": 7.835024332356019,
      "q3_ms": 7.859785166753379,
      "tflops": 17.51770840910809
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 4096,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0006000101566314697,
      "rel_err": 0.0036908327601850033,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.04924023899366148,
      "inner": 20,
      "median_ms": 2.187969099759357,
      "q1_ms": 2.185506200476084,
      "q3_ms": 2.1925688750343397,
      "tflops": 62.815765308164615
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 4096,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0005109310150146484,
      "rel_err": 0.0031428816728293896,
      "peak_mb": 128.7509765625,
      "compile_s": 0.001973392005311325,
      "inner": 25,
      "median_ms": 2.0358039595885202,
      "q1_ms": 2.0346621799399145,
      "q3_ms": 2.0434233397827484,
      "tflops": 67.51089800403933
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 4096,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0006000101566314697,
      "rel_err": 0.0036908327601850033,
      "peak_mb": 128.2509765625,
      "compile_s": 0.002458178991219029,
      "inner": 20,
      "median_ms": 2.6718403001723345,
      "q1_ms": 2.6699426747654798,
      "q3_ms": 2.6726798747404246,
      "tflops": 51.43980853314292
    },
    {
      "impl": "torch_math",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 4096,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00046722590923309326,
      "rel_err": 0.002874039113521576,
      "peak_mb": 4848.37548828125,
      "compile_s": 0.028199844993650913,
      "inner": 1,
      "median_ms": 28.642952995141968,
      "q1_ms": 28.588941000634804,
      "q3_ms": 28.723179493681528,
      "tflops": 4.798351395378491
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 4096,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0008118301630020142,
      "rel_err": 0.004207376390695572,
      "xla_temp_mb": 33.00048828125,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.27266000500821974,
      "inner": 6,
      "median_ms": 8.828691332989061,
      "q1_ms": 8.826855166262249,
      "q3_ms": 8.849354499640564,
      "tflops": 54.48557651512541
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 4096,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0008118301630020142,
      "rel_err": 0.004207376390695572,
      "xla_temp_mb": 33.00048828125,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.2767119280033512,
      "inner": 5,
      "median_ms": 8.813557200483046,
      "q1_ms": 8.781192899914458,
      "q3_ms": 8.82214449957246,
      "tflops": 54.57913600715449
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 4096,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0008118301630020142,
      "rel_err": 0.004207376390695572,
      "xla_temp_mb": 33.00048828125,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.26908333000028506,
      "inner": 5,
      "median_ms": 8.872352601611055,
      "q1_ms": 8.80453919962747,
      "q3_ms": 8.87793599977158,
      "tflops": 54.21745040482866
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 4096,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0006997585296630859,
      "rel_err": 0.003550565103068948,
      "xla_temp_mb": 5120.5,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 1.7796943580033258,
      "inner": 2,
      "median_ms": 23.282003494387027,
      "q1_ms": 23.250706752151018,
      "q3_ms": 23.296363004192244,
      "tflops": 20.66129477508116
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 4096,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0008118301630020142,
      "rel_err": 0.004207376390695572,
      "xla_temp_mb": 97.000732421875,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.11262927700590808,
      "inner": 6,
      "median_ms": 8.231348166494476,
      "q1_ms": 8.222771084547276,
      "q3_ms": 8.35862791670176,
      "tflops": 58.439556609942464
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 4096,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0008118301630020142,
      "rel_err": 0.004207376390695572,
      "peak_mb": 289.251953125,
      "compile_s": 0.007504102002712898,
      "inner": 6,
      "median_ms": 7.770219332693766,
      "q1_ms": 7.769527916631584,
      "q3_ms": 7.771520834163918,
      "tflops": 61.90769096156713
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 4096,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.001096099615097046,
      "rel_err": 0.005173790268599987,
      "peak_mb": 273.283203125,
      "compile_s": 0.011971612009801902,
      "inner": 4,
      "median_ms": 13.018441997701302,
      "q1_ms": 13.016866625548573,
      "q3_ms": 13.01893300114898,
      "tflops": 36.950376799077624
    },
    {
      "impl": "torch_math",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 4096,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "error": "OutOfMemoryError: CUDA out of memory. Tried to allocate 2.00 GiB. GPU 0 has a total capacity of 23.61 GiB of which 1.29 GiB is free. Including non-PyTorch memory, this process has 22.28 GiB memory in use. Of the allocated memory 4.20 GiB is allocated by PyTorch, and 33.75 MiB is reserved by PyTorch but unallocated. I"
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 4096,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.004708766937255859,
      "rel_err": 0.0018718079663813114,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.2034883700107457,
      "inner": 36,
      "median_ms": 1.1916930274714508,
      "q1_ms": 1.188263347204257,
      "q3_ms": 1.2029249860966553,
      "tflops": 57.66541815035189
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 4096,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.004708766937255859,
      "rel_err": 0.0018718079663813114,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.20054477300436702,
      "inner": 42,
      "median_ms": 1.1906344762315886,
      "q1_ms": 1.1902593690874275,
      "q3_ms": 1.1921171667165167,
      "tflops": 57.716686445616986
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 4096,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.004708766937255859,
      "rel_err": 0.0018718079663813114,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.2009215360012604,
      "inner": 40,
      "median_ms": 1.1900397497811355,
      "q1_ms": 1.1890852998476475,
      "q3_ms": 1.1940774998947745,
      "tflops": 57.74553055781409
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 4096,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0070722103118896484,
      "rel_err": 0.002811313373968005,
      "xla_temp_mb": 3072.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.4792784059973201,
      "inner": 6,
      "median_ms": 7.863269832644923,
      "q1_ms": 7.851197417039657,
      "q3_ms": 7.880845250231991,
      "tflops": 8.739300341787358
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 4096,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.004708766937255859,
      "rel_err": 0.0018718079663813114,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.047575689997756854,
      "inner": 35,
      "median_ms": 1.3008305428749216,
      "q1_ms": 1.3001028855796903,
      "q3_ms": 1.3023379996801459,
      "tflops": 52.82738563635307
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 4096,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.004708766937255859,
      "rel_err": 0.0018718079663813114,
      "peak_mb": 128.7509765625,
      "compile_s": 0.001155965990619734,
      "inner": 44,
      "median_ms": 1.1672353182141457,
      "q1_ms": 1.1658598864001264,
      "q3_ms": 1.1688283408478708,
      "tflops": 58.87371266416088
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 4096,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.004708766937255859,
      "rel_err": 0.0018718079663813114,
      "peak_mb": 128.2509765625,
      "compile_s": 0.001388035001582466,
      "inner": 36,
      "median_ms": 1.4903336947706218,
      "q1_ms": 1.4776587358937832,
      "q3_ms": 1.4922023474355228,
      "tflops": 46.11012753528105
    },
    {
      "impl": "torch_math",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 4096,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0038379430770874023,
      "rel_err": 0.0015256420010700822,
      "peak_mb": 4912.37548828125,
      "compile_s": 0.03616822499316186,
      "inner": 1,
      "median_ms": 36.635369004216045,
      "q1_ms": 36.56795150163816,
      "q3_ms": 36.675755502074026,
      "tflops": 1.8757686520938732
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 4096,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.01760578155517578,
      "rel_err": 0.0043429541401565075,
      "xla_temp_mb": 33.00048828125,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.3020527670014417,
      "inner": 10,
      "median_ms": 4.937743800110184,
      "q1_ms": 4.928549950273009,
      "q3_ms": 4.9514752499817405,
      "tflops": 48.71013529916902
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 4096,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.01760578155517578,
      "rel_err": 0.0043429541401565075,
      "xla_temp_mb": 33.00048828125,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.2998228590004146,
      "inner": 10,
      "median_ms": 4.933647099824157,
      "q1_ms": 4.923854050139198,
      "q3_ms": 4.939446550270077,
      "tflops": 48.75058221828887
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 4096,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.01760578155517578,
      "rel_err": 0.0043429541401565075,
      "xla_temp_mb": 33.00048828125,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.3011981590098003,
      "inner": 10,
      "median_ms": 4.917496499547269,
      "q1_ms": 4.915916549362009,
      "q3_ms": 4.923017450346379,
      "tflops": 48.91069441496163
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 4096,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.01760578155517578,
      "rel_err": 0.004249929450452328,
      "xla_temp_mb": 5120.5,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 1.1335033700015629,
      "inner": 1,
      "median_ms": 28.079786992748268,
      "q1_ms": 27.954871504334733,
      "q3_ms": 28.15121250023367,
      "tflops": 8.56552681963417
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 4096,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.01760578155517578,
      "rel_err": 0.0043429541401565075,
      "xla_temp_mb": 97.000732421875,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.11105785499967169,
      "inner": 9,
      "median_ms": 4.9632982215068,
      "q1_ms": 4.961025167075503,
      "q3_ms": 4.969128667047093,
      "tflops": 48.459342526265026
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 4096,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.01760578155517578,
      "rel_err": 0.0043429541401565075,
      "peak_mb": 289.251953125,
      "compile_s": 0.00446332099090796,
      "inner": 11,
      "median_ms": 4.579127817785114,
      "q1_ms": 4.571311226754915,
      "q3_ms": 4.5811152270868085,
      "tflops": 52.524886429646905
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 4096,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.01760578155517578,
      "rel_err": 0.006654990371316671,
      "peak_mb": 273.283203125,
      "compile_s": 0.006921259991941042,
      "inner": 7,
      "median_ms": 7.43187457140136,
      "q1_ms": 7.427825499949645,
      "q3_ms": 7.435887929334838,
      "tflops": 32.363055412902064
    },
    {
      "impl": "torch_math",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 4096,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "error": "OutOfMemoryError: CUDA out of memory. Tried to allocate 2.00 GiB. GPU 0 has a total capacity of 23.61 GiB of which 1.29 GiB is free. Including non-PyTorch memory, this process has 22.28 GiB memory in use. Of the allocated memory 4.20 GiB is allocated by PyTorch, and 33.75 MiB is reserved by PyTorch but unallocated. I"
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 8192,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00034452974796295166,
      "rel_err": 0.0028339633718132973,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.18880291499954183,
      "inner": 5,
      "median_ms": 8.458541601430625,
      "q1_ms": 8.455956898978911,
      "q3_ms": 8.46396600099979,
      "tflops": 64.99416090771696
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 8192,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00034452974796295166,
      "rel_err": 0.0028339633718132973,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.18136204199981876,
      "inner": 6,
      "median_ms": 8.424007498736804,
      "q1_ms": 8.410912583106741,
      "q3_ms": 8.461615166500755,
      "tflops": 65.26060357501306
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 8192,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00034452974796295166,
      "rel_err": 0.0028339633718132973,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.18586039000365417,
      "inner": 6,
      "median_ms": 8.43195916725866,
      "q1_ms": 8.416050666710362,
      "q3_ms": 8.447381666580139,
      "tflops": 65.1990602638002
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 8192,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00032208114862442017,
      "rel_err": 0.0026493100449442863,
      "xla_temp_mb": 12288.000015258789,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 2.358572266995907,
      "inner": 1,
      "median_ms": 31.107264003367163,
      "q1_ms": 30.98630100430455,
      "q3_ms": 31.202207996102516,
      "tflops": 17.67290796864978
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 8192,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00034452974796295166,
      "rel_err": 0.0028339633718132973,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.05397378100315109,
      "inner": 6,
      "median_ms": 8.172032500927648,
      "q1_ms": 8.159973998166,
      "q3_ms": 8.181957916046183,
      "tflops": 67.27283742760378
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 8192,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0003210902214050293,
      "rel_err": 0.0026411591097712517,
      "peak_mb": 241.2509765625,
      "compile_s": 0.007400336995488033,
      "inner": 6,
      "median_ms": 7.834757833431164,
      "q1_ms": 7.833990749835114,
      "q3_ms": 7.844438834581524,
      "tflops": 70.16883298449561
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 8192,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00034452974796295166,
      "rel_err": 0.0028339633718132973,
      "peak_mb": 240.2509765625,
      "compile_s": 0.009459570006583817,
      "inner": 5,
      "median_ms": 10.475570400012657,
      "q1_ms": 10.475144699739758,
      "q3_ms": 10.477003299456555,
      "tflops": 52.47979755711782
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 8192,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0006198734045028687,
      "rel_err": 0.004101343918591738,
      "xla_temp_mb": 66.00048828125,
      "xla_arg_mb": 160.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 0.2960554229939589,
      "inner": 1,
      "median_ms": 33.92856000573374,
      "q1_ms": 33.77991299930727,
      "q3_ms": 33.97752749879146,
      "tflops": 56.71167147332011
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 8192,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0006198734045028687,
      "rel_err": 0.004101343918591738,
      "xla_temp_mb": 66.00048828125,
      "xla_arg_mb": 160.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 0.29872138499922585,
      "inner": 1,
      "median_ms": 33.970338001381606,
      "q1_ms": 33.76386800664477,
      "q3_ms": 34.095850998710375,
      "tflops": 56.64192533291082
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 8192,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0006198734045028687,
      "rel_err": 0.004101343918591738,
      "xla_temp_mb": 66.00048828125,
      "xla_arg_mb": 160.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 0.2949238040018827,
      "inner": 1,
      "median_ms": 33.80051800922956,
      "q1_ms": 33.789927503676154,
      "q3_ms": 34.24800349603174,
      "tflops": 56.926504738258544
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 8192,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 20.00GiB with allocator GPU_0_bfc on device 0. [tf-allocator-allocation-error=''] [executable_name='jit_loss']"
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 8192,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0006198734045028687,
      "rel_err": 0.004101343918591738,
      "xla_temp_mb": 194.000732421875,
      "xla_arg_mb": 160.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 0.14337550800701138,
      "inner": 1,
      "median_ms": 32.00658300193027,
      "q1_ms": 31.03256800386589,
      "q3_ms": 33.47725100320531,
      "tflops": 60.1171749102976
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 8192,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0006198734045028687,
      "rel_err": 0.004101343918591738,
      "peak_mb": 562.251953125,
      "compile_s": 0.02841970999725163,
      "inner": 1,
      "median_ms": 28.60044300905429,
      "q1_ms": 28.591776004759595,
      "q3_ms": 28.604179504327476,
      "tflops": 67.27676728639682
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 8192,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0006056204438209534,
      "rel_err": 0.00489960378035903,
      "peak_mb": 530.314453125,
      "compile_s": 0.04968036100035533,
      "inner": 1,
      "median_ms": 50.48983199230861,
      "q1_ms": 50.434723001671955,
      "q3_ms": 50.8673565054778,
      "tflops": 38.10956132516178
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 8192,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.007262229919433594,
      "rel_err": 0.002886849222704768,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.2107865730067715,
      "inner": 10,
      "median_ms": 4.52553430077387,
      "q1_ms": 4.511659399577184,
      "q3_ms": 4.533449199516326,
      "tflops": 60.73932682313238
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 8192,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.007262229919433594,
      "rel_err": 0.002886849222704768,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.21558746400114615,
      "inner": 11,
      "median_ms": 4.511802273389714,
      "q1_ms": 4.509771226862954,
      "q3_ms": 4.519918544577773,
      "tflops": 60.92419177258945
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 8192,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.007262229919433594,
      "rel_err": 0.002886849222704768,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.2096244360000128,
      "inner": 11,
      "median_ms": 4.514184091187252,
      "q1_ms": 4.505702545099087,
      "q3_ms": 4.524413045807953,
      "tflops": 60.89204635686575
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 8192,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.008362770080566406,
      "rel_err": 0.0033243310172110796,
      "xla_temp_mb": 12288.000015258789,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.4729520289984066,
      "inner": 1,
      "median_ms": 31.092235003598034,
      "q1_ms": 31.057549502293114,
      "q3_ms": 31.10955800366355,
      "tflops": 8.840725245778914
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 8192,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.007262229919433594,
      "rel_err": 0.002886849222704768,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.04928664700128138,
      "inner": 10,
      "median_ms": 4.619807600101922,
      "q1_ms": 4.617631850123871,
      "q3_ms": 4.636151900194818,
      "tflops": 59.49986032706982
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 8192,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.007262229919433594,
      "rel_err": 0.002886849222704768,
      "peak_mb": 241.2509765625,
      "compile_s": 0.004130572997382842,
      "inner": 12,
      "median_ms": 4.295040750245486,
      "q1_ms": 4.291414750696276,
      "q3_ms": 4.316695083010321,
      "tflops": 63.99890546516681
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 8192,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.007262229919433594,
      "rel_err": 0.002886849222704768,
      "peak_mb": 240.2509765625,
      "compile_s": 0.005071241001132876,
      "inner": 9,
      "median_ms": 5.553937555911641,
      "q1_ms": 5.544058333422678,
      "q3_ms": 5.566716166665881,
      "tflops": 49.4924374242232
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 8192,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.014099597930908203,
      "rel_err": 0.004575144965201616,
      "xla_temp_mb": 66.00048828125,
      "xla_arg_mb": 160.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 0.31856476499524433,
      "inner": 2,
      "median_ms": 18.76278149575228,
      "q1_ms": 18.74172874886426,
      "q3_ms": 18.87288199941395,
      "tflops": 51.27558909758686
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 8192,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.014099597930908203,
      "rel_err": 0.004575144965201616,
      "xla_temp_mb": 66.00048828125,
      "xla_arg_mb": 160.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 0.31806246600172017,
      "inner": 2,
      "median_ms": 18.72116250160616,
      "q1_ms": 18.676625750231324,
      "q3_ms": 18.779801001073793,
      "tflops": 51.38957979887521
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 8192,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.014099597930908203,
      "rel_err": 0.004575144965201616,
      "xla_temp_mb": 66.00048828125,
      "xla_arg_mb": 160.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 0.3156805409962544,
      "inner": 2,
      "median_ms": 18.65548299974762,
      "q1_ms": 18.650706500920933,
      "q3_ms": 18.66937724480522,
      "tflops": 51.57050473134442
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 8192,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 20.00GiB with allocator GPU_0_bfc on device 0. [tf-allocator-allocation-error=''] [executable_name='jit_loss']"
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 8192,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.014099597930908203,
      "rel_err": 0.004575144965201616,
      "xla_temp_mb": 194.000732421875,
      "xla_arg_mb": 160.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 0.13183684399700724,
      "inner": 2,
      "median_ms": 19.208629499189556,
      "q1_ms": 18.623998497787397,
      "q3_ms": 19.215791253373027,
      "tflops": 50.08544073092729
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 8192,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.014099597930908203,
      "rel_err": 0.004575144965201616,
      "peak_mb": 562.251953125,
      "compile_s": 0.015500175999477506,
      "inner": 3,
      "median_ms": 16.01983866324493,
      "q1_ms": 15.998628332454246,
      "q3_ms": 16.03599566684958,
      "tflops": 60.05507886364228
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 64,
      "seq": 8192,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.014099597930908203,
      "rel_err": 0.0040384577587246895,
      "peak_mb": 530.314453125,
      "compile_s": 0.026479862004634924,
      "inner": 1,
      "median_ms": 27.740585996070877,
      "q1_ms": 27.720928999769967,
      "q3_ms": 27.753219997975975,
      "tflops": 34.68105087759379
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 128,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.002167940139770508,
      "rel_err": 0.002893048105761409,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.24491043199668638,
      "inner": 200,
      "median_ms": 0.04651136994652916,
      "q1_ms": 0.042619080013537314,
      "q3_ms": 0.04674628500652034,
      "tflops": 5.771394313016395
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 128,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.002167940139770508,
      "rel_err": 0.002893048105761409,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.23690555299981497,
      "inner": 200,
      "median_ms": 0.040786140016280115,
      "q1_ms": 0.040593874982732814,
      "q3_ms": 0.0432628724956885,
      "tflops": 6.581536176084616
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 128,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.002167940139770508,
      "rel_err": 0.002893048105761409,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.24017176000052132,
      "inner": 200,
      "median_ms": 0.036008199967909604,
      "q1_ms": 0.03588389252399793,
      "q3_ms": 0.03660718499304494,
      "tflops": 7.454842403653302
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 128,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0023674368858337402,
      "rel_err": 0.003159270389005542,
      "xla_temp_mb": 3.0,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.5799718979978934,
      "inner": 200,
      "median_ms": 0.03830834997643251,
      "q1_ms": 0.038179035000212025,
      "q3_ms": 0.03841838501102757,
      "tflops": 7.007230960486234
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 128,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.002167940139770508,
      "rel_err": 0.002893048105761409,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.04369149400736205,
      "inner": 200,
      "median_ms": 0.05714036997233052,
      "q1_ms": 0.05651561499689706,
      "q3_ms": 0.05718330248782877,
      "tflops": 4.697824955105933
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 128,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.002339482307434082,
      "rel_err": 0.0031219657976180315,
      "peak_mb": 23.2666015625,
      "compile_s": 8.199499279726297e-05,
      "inner": 200,
      "median_ms": 0.02664459498191718,
      "q1_ms": 0.02655567499459721,
      "q3_ms": 0.029073012483422644,
      "tflops": 10.07466828383687
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 128,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.002167940139770508,
      "rel_err": 0.002893048105761409,
      "peak_mb": 23.2509765625,
      "compile_s": 0.00010271499922964722,
      "inner": 200,
      "median_ms": 0.0360923599509988,
      "q1_ms": 0.03598270497604972,
      "q3_ms": 0.03632918502262328,
      "tflops": 7.437459239696281
    },
    {
      "impl": "torch_math",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 128,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0018723011016845703,
      "rel_err": 0.002498527057468891,
      "peak_mb": 36.25,
      "compile_s": 0.00025227699370589107,
      "inner": 200,
      "median_ms": 0.16190997506782878,
      "q1_ms": 0.16160068749741185,
      "q3_ms": 0.1647308549945592,
      "tflops": 1.6579303152109348
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 128,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.006493330001831055,
      "rel_err": 0.0053722672164440155,
      "xla_temp_mb": 2.031494140625,
      "xla_arg_mb": 5.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.4817590060119983,
      "inner": 200,
      "median_ms": 0.055612180003663525,
      "q1_ms": 0.05560464000154752,
      "q3_ms": 0.05579011503868969,
      "tflops": 16.894214467731846
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 128,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.006493330001831055,
      "rel_err": 0.0053722672164440155,
      "xla_temp_mb": 2.031494140625,
      "xla_arg_mb": 5.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.48326360300416127,
      "inner": 200,
      "median_ms": 0.055121549958130345,
      "q1_ms": 0.0551027150140726,
      "q3_ms": 0.05593058249360183,
      "tflops": 17.04458776492408
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 128,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.006493330001831055,
      "rel_err": 0.0053722672164440155,
      "xla_temp_mb": 2.031494140625,
      "xla_arg_mb": 5.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.4799695120018441,
      "inner": 200,
      "median_ms": 0.05841733996930998,
      "q1_ms": 0.058281699966755696,
      "q3_ms": 0.05857085248862859,
      "tflops": 16.082966059282853
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 128,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.006493330001831055,
      "rel_err": 0.0053722672164440155,
      "xla_temp_mb": 5.015625,
      "xla_arg_mb": 5.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.9419509540020954,
      "inner": 200,
      "median_ms": 0.07405099997413345,
      "q1_ms": 0.07395269247354008,
      "q3_ms": 0.07407888249872485,
      "tflops": 12.687527465235886
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 128,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.006493330001831055,
      "rel_err": 0.0053722672164440155,
      "xla_temp_mb": 6.031982421875,
      "xla_arg_mb": 5.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.10448707800242119,
      "inner": 182,
      "median_ms": 0.11410988460468886,
      "q1_ms": 0.11282748628118597,
      "q3_ms": 0.11709396426590252,
      "tflops": 8.233503164557526
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 128,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0037539005279541016,
      "rel_err": 0.00389949232339859,
      "peak_mb": 33.283203125,
      "compile_s": 0.0004903779918095097,
      "inner": 149,
      "median_ms": 0.24312156375228808,
      "q1_ms": 0.24254039595420818,
      "q3_ms": 0.24399065764952177,
      "tflops": 3.8644210801361214
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 128,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.006493330001831055,
      "rel_err": 0.0053722672164440155,
      "peak_mb": 36.28515625,
      "compile_s": 0.0005489780014613643,
      "inner": 166,
      "median_ms": 0.2505187469770212,
      "q1_ms": 0.2500119156478371,
      "q3_ms": 0.2511385331650584,
      "tflops": 3.7503145267055706
    },
    {
      "impl": "torch_math",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 128,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.004971742630004883,
      "rel_err": 0.004337208345532417,
      "peak_mb": 40.25,
      "compile_s": 0.00102754301042296,
      "inner": 71,
      "median_ms": 0.6197885494157803,
      "q1_ms": 0.6165009295907278,
      "q3_ms": 0.6216458731755601,
      "tflops": 1.5158784344848026
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 128,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0054264068603515625,
      "rel_err": 0.001877243397757411,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.26528663000499364,
      "inner": 200,
      "median_ms": 0.03750677999050822,
      "q1_ms": 0.037252904985507485,
      "q3_ms": 0.037729834984929767,
      "tflops": 3.5784924228090564
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 128,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0054264068603515625,
      "rel_err": 0.001877243397757411,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.2633176560048014,
      "inner": 200,
      "median_ms": 0.03668983998068143,
      "q1_ms": 0.03661171747808112,
      "q3_ms": 0.036773302490473725,
      "tflops": 3.6581715284305045
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 128,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0054264068603515625,
      "rel_err": 0.001877243397757411,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.26566646801074967,
      "inner": 200,
      "median_ms": 0.0367940900468966,
      "q1_ms": 0.036469724982453045,
      "q3_ms": 0.03730688255018322,
      "tflops": 3.6478066947417442
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 128,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0069255828857421875,
      "rel_err": 0.002395877381786704,
      "xla_temp_mb": 3.0,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.3648681699996814,
      "inner": 200,
      "median_ms": 0.03685184994537849,
      "q1_ms": 0.03681750746181933,
      "q3_ms": 0.037041555042378604,
      "tflops": 3.6420892899253747
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 128,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0054264068603515625,
      "rel_err": 0.001877243397757411,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.044077924991142936,
      "inner": 200,
      "median_ms": 0.05826690001413226,
      "q1_ms": 0.05804608499602182,
      "q3_ms": 0.05833535749843577,
      "tflops": 2.3034986925243386
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 128,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0054264068603515625,
      "rel_err": 0.001877243397757411,
      "peak_mb": 23.2666015625,
      "compile_s": 0.0001022240030579269,
      "inner": 200,
      "median_ms": 0.027083770037279464,
      "q1_ms": 0.027014492479793262,
      "q3_ms": 0.027531119994819164,
      "tflops": 4.955651588211536
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 128,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0054264068603515625,
      "rel_err": 0.001877243397757411,
      "peak_mb": 23.2509765625,
      "compile_s": 9.131299157161266e-05,
      "inner": 200,
      "median_ms": 0.03291645996796433,
      "q1_ms": 0.032791265002742875,
      "q3_ms": 0.03315229998406721,
      "tflops": 4.0775262021075855
    },
    {
      "impl": "torch_math",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 128,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0054264068603515625,
      "rel_err": 0.001877243397757411,
      "peak_mb": 36.3125,
      "compile_s": 0.0005872699985047802,
      "inner": 152,
      "median_ms": 0.21559334872060726,
      "q1_ms": 0.2154453321916088,
      "q3_ms": 0.2162652434842465,
      "tflops": 0.6225504116731173
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 128,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.019839763641357422,
      "rel_err": 0.006080516614019871,
      "xla_temp_mb": 2.031494140625,
      "xla_arg_mb": 5.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.507357401002082,
      "inner": 200,
      "median_ms": 0.0565260599978501,
      "q1_ms": 0.055673169990768656,
      "q3_ms": 0.06760330998076824,
      "tflops": 8.310539386928204
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 128,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.019839763641357422,
      "rel_err": 0.006080516614019871,
      "xla_temp_mb": 2.031494140625,
      "xla_arg_mb": 5.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.5071429450035794,
      "inner": 200,
      "median_ms": 0.05674747495504562,
      "q1_ms": 0.05609722251392668,
      "q3_ms": 0.05688746747182449,
      "tflops": 8.278113667121533
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 128,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.019839763641357422,
      "rel_err": 0.006080516614019871,
      "xla_temp_mb": 2.031494140625,
      "xla_arg_mb": 5.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.5035435589961708,
      "inner": 200,
      "median_ms": 0.055517705040983856,
      "q1_ms": 0.05542768001760123,
      "q3_ms": 0.05593696747382637,
      "tflops": 8.461481749888902
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 128,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.019839763641357422,
      "rel_err": 0.004750313237309456,
      "xla_temp_mb": 5.015625,
      "xla_arg_mb": 5.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.5821025169861969,
      "inner": 200,
      "median_ms": 0.07583938997413497,
      "q1_ms": 0.07581063004181487,
      "q3_ms": 0.07588965749164345,
      "tflops": 6.194169654584674
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 128,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.019839763641357422,
      "rel_err": 0.006080516614019871,
      "xla_temp_mb": 6.031982421875,
      "xla_arg_mb": 5.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.10552618298970629,
      "inner": 200,
      "median_ms": 0.12096573504095431,
      "q1_ms": 0.11703908003255492,
      "q3_ms": 0.12278988251637202,
      "tflops": 3.883430690856024
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 128,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.019839763641357422,
      "rel_err": 0.006080516614019871,
      "peak_mb": 33.283203125,
      "compile_s": 0.00045220600441098213,
      "inner": 171,
      "median_ms": 0.2450036140783739,
      "q1_ms": 0.23777523395102962,
      "q3_ms": 0.24655130412701887,
      "tflops": 1.9173678305403625
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 128,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.019839763641357422,
      "rel_err": 0.006080516614019871,
      "peak_mb": 36.28515625,
      "compile_s": 0.0005688560049748048,
      "inner": 159,
      "median_ms": 0.24897725156247733,
      "q1_ms": 0.24724270436040025,
      "q3_ms": 0.24921427045783334,
      "tflops": 1.8867669437748606
    },
    {
      "impl": "torch_math",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 128,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.019839763641357422,
      "rel_err": 0.004750313237309456,
      "peak_mb": 40.3125,
      "compile_s": 0.0011492839985294268,
      "inner": 65,
      "median_ms": 0.6558575078648013,
      "q1_ms": 0.6545187384919979,
      "q3_ms": 0.667331446087561,
      "tflops": 0.7162562635432037
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 512,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.001112908124923706,
      "rel_err": 0.0026897622738033533,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.24846781000087503,
      "inner": 200,
      "median_ms": 0.10790422995341942,
      "q1_ms": 0.10784501995658502,
      "q3_ms": 0.1080023175018141,
      "tflops": 39.803511853558206
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 512,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.001112908124923706,
      "rel_err": 0.0026897622738033533,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.23941669199848548,
      "inner": 200,
      "median_ms": 0.09688664002169389,
      "q1_ms": 0.09593203747499501,
      "q3_ms": 0.09705618747830158,
      "tflops": 44.329819829011655
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 512,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.001112908124923706,
      "rel_err": 0.0026897622738033533,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.23959060000197496,
      "inner": 174,
      "median_ms": 0.09702798848886085,
      "q1_ms": 0.09694931033935837,
      "q3_ms": 0.09728923279239432,
      "tflops": 44.2652410185034
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 512,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0012347698211669922,
      "rel_err": 0.0029842869844287634,
      "xla_temp_mb": 48.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.6841787559969816,
      "inner": 178,
      "median_ms": 0.19802384829316078,
      "q1_ms": 0.1978448623519265,
      "q3_ms": 0.19813360954736872,
      "tflops": 21.689141651472173
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 512,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.001112908124923706,
      "rel_err": 0.0026897622738033533,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.0434978880075505,
      "inner": 200,
      "median_ms": 0.0941897950542625,
      "q1_ms": 0.09381833751831437,
      "q3_ms": 0.09421899998415029,
      "tflops": 45.59907252718493
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 512,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0010547637939453125,
      "rel_err": 0.0025492345448583364,
      "peak_mb": 44.3134765625,
      "compile_s": 0.00013021699851378798,
      "inner": 200,
      "median_ms": 0.0875595050456468,
      "q1_ms": 0.0867021825251868,
      "q3_ms": 0.0876247024643817,
      "tflops": 49.05198234915711
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 512,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0010657310485839844,
      "rel_err": 0.0025757411494851112,
      "peak_mb": 44.2509765625,
      "compile_s": 0.00017351801216136664,
      "inner": 200,
      "median_ms": 0.1257933150191093,
      "q1_ms": 0.1250108375097625,
      "q3_ms": 0.12584872001752956,
      "tflops": 34.14304882057962
    },
    {
      "impl": "torch_math",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 512,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0009698271751403809,
      "rel_err": 0.0023439531214535236,
      "peak_mb": 144.26611328125,
      "compile_s": 0.0006825810123700649,
      "inner": 73,
      "median_ms": 0.6749517533909019,
      "q1_ms": 0.674814575376493,
      "q3_ms": 0.6754846096418364,
      "tflops": 6.363369343693736
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 512,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0019362568855285645,
      "rel_err": 0.004246557131409645,
      "xla_temp_mb": 8.125244140625,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.4806370349979261,
      "inner": 91,
      "median_ms": 0.44495793418308904,
      "q1_ms": 0.44441057692453323,
      "q3_ms": 0.44505387911787975,
      "tflops": 33.783835237365494
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 512,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0019362568855285645,
      "rel_err": 0.004246557131409645,
      "xla_temp_mb": 8.125244140625,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.4766963120055152,
      "inner": 86,
      "median_ms": 0.4456407557374278,
      "q1_ms": 0.4453009884834268,
      "q3_ms": 0.44605263373828064,
      "tflops": 33.73207082715097
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 512,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0019362568855285645,
      "rel_err": 0.004246557131409645,
      "xla_temp_mb": 8.125244140625,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.48332947700691875,
      "inner": 90,
      "median_ms": 0.4449562222968476,
      "q1_ms": 0.44451661112058594,
      "q3_ms": 0.4449701889017079,
      "tflops": 33.783965214383066
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 512,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0019362568855285645,
      "rel_err": 0.004246557131409645,
      "xla_temp_mb": 80.0625,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 1.4150422969978536,
      "inner": 77,
      "median_ms": 0.5449661559821408,
      "q1_ms": 0.544605993434436,
      "q3_ms": 0.5454189610303162,
      "tflops": 27.58407172810311
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 512,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0019362568855285645,
      "rel_err": 0.004246557131409645,
      "xla_temp_mb": 24.125732421875,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.10369666300539393,
      "inner": 83,
      "median_ms": 0.44506307234089953,
      "q1_ms": 0.4439562890841345,
      "q3_ms": 0.4451998494739408,
      "tflops": 33.775854413025364
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 512,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.001931697130203247,
      "rel_err": 0.0043192761950194836,
      "peak_mb": 84.376953125,
      "compile_s": 0.0005714810104109347,
      "inner": 97,
      "median_ms": 0.46312034018481746,
      "q1_ms": 0.4624124072112873,
      "q3_ms": 0.4631690361438626,
      "tflops": 32.45891884170111
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 512,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0019917190074920654,
      "rel_err": 0.004453485365957022,
      "peak_mb": 92.384765625,
      "compile_s": 0.0011569880007300526,
      "inner": 43,
      "median_ms": 1.108589000068605,
      "q1_ms": 1.1077547558586027,
      "q3_ms": 1.1117258836748078,
      "tflops": 13.559926659086209
    },
    {
      "impl": "torch_math",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 512,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0015252530574798584,
      "rel_err": 0.003345152596011758,
      "peak_mb": 204.2509765625,
      "compile_s": 0.0016076010069809854,
      "inner": 31,
      "median_ms": 1.645553870701171,
      "q1_ms": 1.6449677582826436,
      "q3_ms": 1.6465287739148124,
      "tflops": 9.135152487955132
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 512,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.005219221115112305,
      "rel_err": 0.0018961260793730617,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.2861385180003708,
      "inner": 200,
      "median_ms": 0.07841254999220837,
      "q1_ms": 0.07839619498554384,
      "q3_ms": 0.07843193750886712,
      "tflops": 27.38698904975529
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 512,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.005219221115112305,
      "rel_err": 0.0018961260793730617,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.26790892900316976,
      "inner": 200,
      "median_ms": 0.07841781502065714,
      "q1_ms": 0.07839093501388561,
      "q3_ms": 0.07844756997656077,
      "tflops": 27.385150267631165
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 512,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.005219221115112305,
      "rel_err": 0.0018961260793730617,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.2668941400042968,
      "inner": 200,
      "median_ms": 0.07841250000637956,
      "q1_ms": 0.07840564001526218,
      "q3_ms": 0.07842277005693177,
      "tflops": 27.3870065082134
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 512,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.01041102409362793,
      "rel_err": 0.003782291430979967,
      "xla_temp_mb": 48.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.8636951880034758,
      "inner": 144,
      "median_ms": 0.24009068056329852,
      "q1_ms": 0.23979599654921913,
      "q3_ms": 0.24014014929384253,
      "tflops": 8.944468993805149
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 512,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.005219221115112305,
      "rel_err": 0.0018961260793730617,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.04347719899669755,
      "inner": 200,
      "median_ms": 0.07929998006147798,
      "q1_ms": 0.07929537252493901,
      "q3_ms": 0.0793122524919454,
      "tflops": 27.08050678367315
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 512,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.005219221115112305,
      "rel_err": 0.0018961260793730617,
      "peak_mb": 44.3134765625,
      "compile_s": 0.00011151100625284016,
      "inner": 200,
      "median_ms": 0.062243319989647716,
      "q1_ms": 0.06125700499978847,
      "q3_ms": 0.062349147483473644,
      "tflops": 34.50143161317823
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 512,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.005219221115112305,
      "rel_err": 0.0018961260793730617,
      "peak_mb": 44.2509765625,
      "compile_s": 0.00013124800170771778,
      "inner": 200,
      "median_ms": 0.08236743502493482,
      "q1_ms": 0.08215272500819992,
      "q3_ms": 0.08237449499574723,
      "tflops": 26.071998567757987
    },
    {
      "impl": "torch_math",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 512,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.005219221115112305,
      "rel_err": 0.0018961260793730617,
      "peak_mb": 145.26611328125,
      "compile_s": 0.0007812770054442808,
      "inner": 64,
      "median_ms": 0.7670976876852365,
      "q1_ms": 0.7669768436926461,
      "q3_ms": 0.7691446875242036,
      "tflops": 2.7994917498450054
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 512,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.01401376724243164,
      "rel_err": 0.005147535353899002,
      "xla_temp_mb": 8.125244140625,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.5130365529912524,
      "inner": 105,
      "median_ms": 0.3110659427918671,
      "q1_ms": 0.31099547625109114,
      "q3_ms": 0.31276338572968687,
      "tflops": 24.162699074482262
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 512,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.01401376724243164,
      "rel_err": 0.005147535353899002,
      "xla_temp_mb": 8.125244140625,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.5056766030029394,
      "inner": 105,
      "median_ms": 0.31184255231415764,
      "q1_ms": 0.31174155239049073,
      "q3_ms": 0.3127810952719301,
      "tflops": 24.102524534329774
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 512,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.01401376724243164,
      "rel_err": 0.005147535353899002,
      "xla_temp_mb": 8.125244140625,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.5159933250106405,
      "inner": 111,
      "median_ms": 0.3110975675234521,
      "q1_ms": 0.311021477525198,
      "q3_ms": 0.3128663242400229,
      "tflops": 24.160242806891738
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 512,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.01401376724243164,
      "rel_err": 0.006148071028292179,
      "xla_temp_mb": 80.0625,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 1.6348487000068417,
      "inner": 67,
      "median_ms": 0.6104956866231567,
      "q1_ms": 0.6102574776350729,
      "q3_ms": 0.6108226417694171,
      "tflops": 12.311623051056136
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 512,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.01401376724243164,
      "rel_err": 0.005147535353899002,
      "xla_temp_mb": 24.125732421875,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.10556986599112861,
      "inner": 94,
      "median_ms": 0.3409593296732674,
      "q1_ms": 0.34076273399516643,
      "q3_ms": 0.3418296435754994,
      "tflops": 22.044250190198856
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 512,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.01401376724243164,
      "rel_err": 0.0040749916806817055,
      "peak_mb": 84.376953125,
      "compile_s": 0.0005256049917079508,
      "inner": 122,
      "median_ms": 0.35665347539060216,
      "q1_ms": 0.35648808190045633,
      "q3_ms": 0.35688202461031204,
      "tflops": 21.074217094809924
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 512,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.01401376724243164,
      "rel_err": 0.005914298817515373,
      "peak_mb": 92.384765625,
      "compile_s": 0.000729410006897524,
      "inner": 68,
      "median_ms": 0.6821548822164699,
      "q1_ms": 0.6818836251006681,
      "q3_ms": 0.6823109118105255,
      "tflops": 11.01830825219377
    },
    {
      "impl": "torch_math",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 512,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.01401376724243164,
      "rel_err": 0.003250203561037779,
      "peak_mb": 204.2509765625,
      "compile_s": 0.0017175180109916255,
      "inner": 29,
      "median_ms": 1.7284960340475664,
      "q1_ms": 1.7281415689624202,
      "q3_ms": 1.7300059312376483,
      "tflops": 4.348400354960353
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 600,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0011281371116638184,
      "rel_err": 0.0022587981075048447,
      "xla_temp_mb": 20.0,
      "xla_arg_mb": 14.0625,
      "xla_out_mb": 4.6875,
      "compile_s": 0.28105377999600023,
      "inner": 149,
      "median_ms": 0.18503122819250062,
      "q1_ms": 0.18447648989538282,
      "q3_ms": 0.1850774193896706,
      "tflops": 31.87699750802961
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 600,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0011281371116638184,
      "rel_err": 0.0022587981075048447,
      "xla_temp_mb": 20.0,
      "xla_arg_mb": 14.0625,
      "xla_out_mb": 4.6875,
      "compile_s": 0.27700014399306383,
      "inner": 172,
      "median_ms": 0.1728738778944907,
      "q1_ms": 0.17190177615787455,
      "q3_ms": 0.17298626741204887,
      "tflops": 34.11874640539876
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 600,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0011281371116638184,
      "rel_err": 0.0022587981075048447,
      "xla_temp_mb": 20.0,
      "xla_arg_mb": 14.0625,
      "xla_out_mb": 4.6875,
      "compile_s": 0.2821170509996591,
      "inner": 181,
      "median_ms": 0.17279675693472446,
      "q1_ms": 0.1726277927978695,
      "q3_ms": 0.17282022925438514,
      "tflops": 34.1339739508428
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 600,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0014718472957611084,
      "rel_err": 0.0029469875153154135,
      "xla_temp_mb": 65.91798400878906,
      "xla_arg_mb": 14.0625,
      "xla_out_mb": 4.6875,
      "compile_s": 0.7446916549961315,
      "inner": 136,
      "median_ms": 0.27060164703631445,
      "q1_ms": 0.2700308235516251,
      "q3_ms": 0.27073568745596704,
      "tflops": 21.796763118771636
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 600,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0011281371116638184,
      "rel_err": 0.0022587981075048447,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 14.0625,
      "xla_out_mb": 4.6875,
      "compile_s": 0.04395454199402593,
      "inner": 178,
      "median_ms": 0.11914657864437093,
      "q1_ms": 0.11819542695969615,
      "q3_ms": 0.1191735112554985,
      "tflops": 49.504065220412954
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 600,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0011044740676879883,
      "rel_err": 0.002211419166997075,
      "peak_mb": 49.63671875,
      "compile_s": 0.00015493300452362746,
      "inner": 200,
      "median_ms": 0.10997750003298279,
      "q1_ms": 0.10942515749775339,
      "q3_ms": 0.11002699498931179,
      "tflops": 53.63133366580518
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 600,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0011281371116638184,
      "rel_err": 0.0022587981075048447,
      "peak_mb": 49.5634765625,
      "compile_s": 0.00020565900194924325,
      "inner": 200,
      "median_ms": 0.1586855750065297,
      "q1_ms": 0.15834144749533152,
      "q3_ms": 0.15870939248998184,
      "tflops": 37.16935203314666
    },
    {
      "impl": "torch_math",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 600,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0009754598140716553,
      "rel_err": 0.0019531019497662783,
      "peak_mb": 181.38037109375,
      "compile_s": 0.0008778709889156744,
      "inner": 57,
      "median_ms": 0.8792861578209946,
      "q1_ms": 0.8783507981110903,
      "q3_ms": 0.8793692192595794,
      "tflops": 6.707986868139423
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 600,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0020381510257720947,
      "rel_err": 0.0037066892255097628,
      "xla_temp_mb": 40.000526428222656,
      "xla_arg_mb": 23.4375,
      "xla_out_mb": 14.062522888183594,
      "compile_s": 0.5275755309994565,
      "inner": 62,
      "median_ms": 0.6750286129611214,
      "q1_ms": 0.6749302096244308,
      "q3_ms": 0.6756594113774612,
      "tflops": 30.582170301555784
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 600,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0020381510257720947,
      "rel_err": 0.0037066892255097628,
      "xla_temp_mb": 40.00048828125,
      "xla_arg_mb": 23.4375,
      "xla_out_mb": 14.062522888183594,
      "compile_s": 0.5221658890077379,
      "inner": 64,
      "median_ms": 0.6622107187013171,
      "q1_ms": 0.6621665859256609,
      "q3_ms": 0.6623415234798813,
      "tflops": 31.17412542111264
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 600,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0020381510257720947,
      "rel_err": 0.0037066892255097628,
      "xla_temp_mb": 40.00048828125,
      "xla_arg_mb": 23.4375,
      "xla_out_mb": 14.062522888183594,
      "compile_s": 0.5221819590078667,
      "inner": 64,
      "median_ms": 0.6622379687541979,
      "q1_ms": 0.6620368125140885,
      "q3_ms": 0.662481796780412,
      "tflops": 31.172842654786454
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 600,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0020381510257720947,
      "rel_err": 0.003886493155732751,
      "xla_temp_mb": 87.96388244628906,
      "xla_arg_mb": 23.4375,
      "xla_out_mb": 14.062522888183594,
      "compile_s": 1.175565611993079,
      "inner": 52,
      "median_ms": 0.8638746154387122,
      "q1_ms": 0.8634602020790155,
      "q3_ms": 0.8645325960969785,
      "tflops": 23.896801261507356
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 600,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0020381510257720947,
      "rel_err": 0.0037066892255097628,
      "xla_temp_mb": 28.272216796875,
      "xla_arg_mb": 23.4375,
      "xla_out_mb": 14.062522888183594,
      "compile_s": 0.1025956500088796,
      "inner": 66,
      "median_ms": 0.5878929545305612,
      "q1_ms": 0.5876949999770713,
      "q3_ms": 0.5883631667088731,
      "tflops": 35.11496411193484
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 600,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0020381510257720947,
      "rel_err": 0.0037066892255097628,
      "peak_mb": 97.8408203125,
      "compile_s": 0.0006869490025565028,
      "inner": 80,
      "median_ms": 0.5753148625444737,
      "q1_ms": 0.5752226124059234,
      "q3_ms": 0.5754712000452855,
      "tflops": 35.88268154362893
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 600,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.002734154462814331,
      "rel_err": 0.005480120424181223,
      "peak_mb": 104.5341796875,
      "compile_s": 0.0016008579987101257,
      "inner": 31,
      "median_ms": 1.5645028390516076,
      "q1_ms": 1.562300790195173,
      "q3_ms": 1.5647559031945533,
      "tflops": 13.195143840399915
    },
    {
      "impl": "torch_math",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 600,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0020381510257720947,
      "rel_err": 0.0037066892255097628,
      "peak_mb": 263.0634765625,
      "compile_s": 0.0020411099976627156,
      "inner": 24,
      "median_ms": 2.1086180834875754,
      "q1_ms": 2.108183291663105,
      "q3_ms": 2.1130412290707077,
      "tflops": 9.790222402843034
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 600,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.007790327072143555,
      "rel_err": 0.002695032162591815,
      "xla_temp_mb": 20.0,
      "xla_arg_mb": 14.0625,
      "xla_out_mb": 4.6875,
      "compile_s": 0.2874426940106787,
      "inner": 173,
      "median_ms": 0.17329903468866484,
      "q1_ms": 0.16538817629808927,
      "q3_ms": 0.1733842254276495,
      "tflops": 17.017521218731268
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 600,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.007790327072143555,
      "rel_err": 0.002695032162591815,
      "xla_temp_mb": 20.0,
      "xla_arg_mb": 14.0625,
      "xla_out_mb": 4.6875,
      "compile_s": 0.29158640099922195,
      "inner": 187,
      "median_ms": 0.16228725665552732,
      "q1_ms": 0.16226529147330282,
      "q3_ms": 0.16231739304942344,
      "tflops": 18.172221656687647
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 600,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.007790327072143555,
      "rel_err": 0.002695032162591815,
      "xla_temp_mb": 20.0,
      "xla_arg_mb": 14.0625,
      "xla_out_mb": 4.6875,
      "compile_s": 0.28546399100741837,
      "inner": 173,
      "median_ms": 0.16222856070770952,
      "q1_ms": 0.1622109826670346,
      "q3_ms": 0.1622526214163309,
      "tflops": 18.178796551819804
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 600,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.007834672927856445,
      "rel_err": 0.0027103733737021685,
      "xla_temp_mb": 65.91798400878906,
      "xla_arg_mb": 14.0625,
      "xla_out_mb": 4.6875,
      "compile_s": 0.6613317250012187,
      "inner": 122,
      "median_ms": 0.31789697536847505,
      "q1_ms": 0.3174925287403013,
      "q3_ms": 0.31819845485113574,
      "tflops": 9.276967786754398
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 600,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.007790327072143555,
      "rel_err": 0.002695032162591815,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 14.0625,
      "xla_out_mb": 4.6875,
      "compile_s": 0.044290867997915484,
      "inner": 145,
      "median_ms": 0.10680193111605557,
      "q1_ms": 0.1067763240606492,
      "q3_ms": 0.1068450068926503,
      "tflops": 27.612983858834532
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 600,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.007790327072143555,
      "rel_err": 0.002695032162591815,
      "peak_mb": 49.63671875,
      "compile_s": 0.0001272410008823499,
      "inner": 200,
      "median_ms": 0.0839509250363335,
      "q1_ms": 0.0821474900294561,
      "q3_ms": 0.08413056752033299,
      "tflops": 35.12909475057764
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 600,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.007790327072143555,
      "rel_err": 0.002695032162591815,
      "peak_mb": 49.5634765625,
      "compile_s": 0.000168108002981171,
      "inner": 200,
      "median_ms": 0.11587540502659976,
      "q1_ms": 0.11531184252817184,
      "q3_ms": 0.11587751003389712,
      "tflops": 25.450784826366
    },
    {
      "impl": "torch_math",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 600,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.007790327072143555,
      "rel_err": 0.002695032162591815,
      "peak_mb": 183.38037109375,
      "compile_s": 0.001010170002700761,
      "inner": 49,
      "median_ms": 1.0064238776234264,
      "q1_ms": 1.0060274081687177,
      "q3_ms": 1.0066156632537782,
      "tflops": 2.93029613622052
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 600,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.014764547348022461,
      "rel_err": 0.00416430551558733,
      "xla_temp_mb": 40.00048828125,
      "xla_arg_mb": 23.4375,
      "xla_out_mb": 14.062522888183594,
      "compile_s": 0.524848914006725,
      "inner": 86,
      "median_ms": 0.4759456278407015,
      "q1_ms": 0.47589984303137267,
      "q3_ms": 0.47595162788300427,
      "tflops": 21.687183149111174
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 600,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.014764547348022461,
      "rel_err": 0.00416430551558733,
      "xla_temp_mb": 40.00048828125,
      "xla_arg_mb": 23.4375,
      "xla_out_mb": 14.062522888183594,
      "compile_s": 0.5342263990023639,
      "inner": 84,
      "median_ms": 0.46692011909770026,
      "q1_ms": 0.46674228577189963,
      "q3_ms": 0.4669904286683526,
      "tflops": 22.10639374449444
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 600,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.014764547348022461,
      "rel_err": 0.00416430551558733,
      "xla_temp_mb": 40.00048828125,
      "xla_arg_mb": 23.4375,
      "xla_out_mb": 14.062522888183594,
      "compile_s": 0.5338866170059191,
      "inner": 83,
      "median_ms": 0.46690826511840866,
      "q1_ms": 0.46688664456652706,
      "q3_ms": 0.46694018062434145,
      "tflops": 22.106954986933772
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 600,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.014764547348022461,
      "rel_err": 0.004817801062017679,
      "xla_temp_mb": 87.96388244628906,
      "xla_arg_mb": 23.4375,
      "xla_out_mb": 14.062522888183594,
      "compile_s": 0.520300912001403,
      "inner": 51,
      "median_ms": 0.8737240588707922,
      "q1_ms": 0.8733924608118376,
      "q3_ms": 0.8741319901424041,
      "tflops": 11.8137069652633
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 600,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.014764547348022461,
      "rel_err": 0.00416430551558733,
      "xla_temp_mb": 28.272216796875,
      "xla_arg_mb": 23.4375,
      "xla_out_mb": 14.062522888183594,
      "compile_s": 0.10778270100126974,
      "inner": 81,
      "median_ms": 0.46814429623299036,
      "q1_ms": 0.4678468148433891,
      "q3_ms": 0.4681961172199981,
      "tflops": 22.048586478693082
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 600,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.014764547348022461,
      "rel_err": 0.0061126695945858955,
      "peak_mb": 97.8408203125,
      "compile_s": 0.000570158998016268,
      "inner": 99,
      "median_ms": 0.4614305960583604,
      "q1_ms": 0.4610897827807654,
      "q3_ms": 0.46173646967303045,
      "tflops": 22.369387916995677
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 600,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.014764547348022461,
      "rel_err": 0.0061126695945858955,
      "peak_mb": 104.5341796875,
      "compile_s": 0.0010323629976483062,
      "inner": 49,
      "median_ms": 0.9585899388184771,
      "q1_ms": 0.9581018674332287,
      "q3_ms": 0.9592312552350839,
      "tflops": 10.767815915867446
    },
    {
      "impl": "torch_math",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 600,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.011423587799072266,
      "rel_err": 0.003221995895728469,
      "peak_mb": 263.0634765625,
      "compile_s": 0.0021712460002163425,
      "inner": 23,
      "median_ms": 2.2236303478206065,
      "q1_ms": 2.223230478272814,
      "q3_ms": 2.224079913107942,
      "tflops": 4.641922615472745
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 1024,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.000802159309387207,
      "rel_err": 0.0029085648711770773,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.24939857999561355,
      "inner": 106,
      "median_ms": 0.37881126421402683,
      "q1_ms": 0.3664732594531741,
      "q3_ms": 0.37899666028163165,
      "tflops": 45.35205472214639
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 1024,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.000802159309387207,
      "rel_err": 0.0029085648711770773,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.24126485800661612,
      "inner": 116,
      "median_ms": 0.339813905156916,
      "q1_ms": 0.3396899568981788,
      "q3_ms": 0.3417321810863133,
      "tflops": 50.55669860262735
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 1024,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.000802159309387207,
      "rel_err": 0.0029085648711770773,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.24235447800310794,
      "inner": 107,
      "median_ms": 0.3402136448742474,
      "q1_ms": 0.34018559347365973,
      "q3_ms": 0.3422416775720194,
      "tflops": 50.49729616326872
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 1024,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.000802159309387207,
      "rel_err": 0.0029085648711770773,
      "xla_temp_mb": 192.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.7740963280084543,
      "inner": 69,
      "median_ms": 0.6513196666660868,
      "q1_ms": 0.6507355289574226,
      "q3_ms": 0.65138994925318,
      "tflops": 26.37701586985494
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 1024,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.000802159309387207,
      "rel_err": 0.0029085648711770773,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.04243915600818582,
      "inner": 109,
      "median_ms": 0.33031684403102674,
      "q1_ms": 0.32972751375941406,
      "q3_ms": 0.33060702751070226,
      "tflops": 52.01027284695869
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 1024,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.000802159309387207,
      "rel_err": 0.0029085648711770773,
      "peak_mb": 72.3759765625,
      "compile_s": 0.0003570060071069747,
      "inner": 145,
      "median_ms": 0.3164233516395927,
      "q1_ms": 0.31615657243147044,
      "q3_ms": 0.3164420724276389,
      "tflops": 54.293935940505214
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 1024,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.000802159309387207,
      "rel_err": 0.0029085648711770773,
      "peak_mb": 72.2509765625,
      "compile_s": 0.00046628199925180525,
      "inner": 112,
      "median_ms": 0.42244558036535246,
      "q1_ms": 0.42239767857868826,
      "q3_ms": 0.42255538846802665,
      "tflops": 40.66765042053931
    },
    {
      "impl": "torch_math",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 1024,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.000802159309387207,
      "rel_err": 0.0029085648711770773,
      "peak_mb": 416.28173828125,
      "compile_s": 0.0022755340032745153,
      "inner": 22,
      "median_ms": 2.31063790918349,
      "q1_ms": 2.308767726803622,
      "q3_ms": 2.311044545587554,
      "tflops": 7.435119590014364
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 1024,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0025502443313598633,
      "rel_err": 0.005480848718434572,
      "xla_temp_mb": 16.250244140625,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.4831420929986052,
      "inner": 31,
      "median_ms": 1.4988312257393714,
      "q1_ms": 1.496796225593801,
      "q3_ms": 1.5068256775961228,
      "tflops": 40.11762039074024
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 1024,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0025502443313598633,
      "rel_err": 0.005480848718434572,
      "xla_temp_mb": 16.250244140625,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.4795370850042673,
      "inner": 31,
      "median_ms": 1.4959848711207029,
      "q1_ms": 1.4958025644920134,
      "q3_ms": 1.497766290530337,
      "tflops": 40.19395069079444
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 1024,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0025502443313598633,
      "rel_err": 0.005480848718434572,
      "xla_temp_mb": 16.250244140625,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.48025025401148014,
      "inner": 31,
      "median_ms": 1.495348838826401,
      "q1_ms": 1.4950133387878117,
      "q3_ms": 1.4960406131708184,
      "tflops": 40.211046802424804
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 1024,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0024344325065612793,
      "rel_err": 0.005231952294707298,
      "xla_temp_mb": 320.125,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 1.4322241339978063,
      "inner": 25,
      "median_ms": 1.8681348400423303,
      "q1_ms": 1.8677210400346667,
      "q3_ms": 1.8684899000800215,
      "tflops": 32.186939002024886
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 1024,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0025502443313598633,
      "rel_err": 0.005480848718434572,
      "xla_temp_mb": 48.250732421875,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.10341902800428215,
      "inner": 34,
      "median_ms": 1.3166617055129095,
      "q1_ms": 1.3163584705633933,
      "q3_ms": 1.3182449705816586,
      "tflops": 45.66817876773924
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 1024,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0015751421451568604,
      "rel_err": 0.0041829561814665794,
      "peak_mb": 152.501953125,
      "compile_s": 0.0013605539861600846,
      "inner": 37,
      "median_ms": 1.3429948920856003,
      "q1_ms": 1.3428465134819743,
      "q3_ms": 1.345444256691514,
      "tflops": 44.77272586690333
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 1024,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0018747150897979736,
      "rel_err": 0.0041829561814665794,
      "peak_mb": 156.517578125,
      "compile_s": 0.003872114000841975,
      "inner": 12,
      "median_ms": 3.925011750349465,
      "q1_ms": 3.911702458329576,
      "q3_ms": 3.940359917275297,
      "tflops": 15.31958270918459
    },
    {
      "impl": "torch_math",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 1024,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0015667378902435303,
      "rel_err": 0.0041829561814665794,
      "peak_mb": 648.2509765625,
      "compile_s": 0.005438495994894765,
      "inner": 9,
      "median_ms": 5.612139889207255,
      "q1_ms": 5.611326666742874,
      "q3_ms": 5.636335610890657,
      "tflops": 10.714191615151208
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 1024,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.007625579833984375,
      "rel_err": 0.002638038480654359,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.27162501799466554,
      "inner": 134,
      "median_ms": 0.24348797768340516,
      "q1_ms": 0.2248056269091825,
      "q3_ms": 0.24352659330809306,
      "tflops": 35.27868058918723
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 1024,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.007625579833984375,
      "rel_err": 0.002638038480654359,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.26368666300550103,
      "inner": 155,
      "median_ms": 0.21818478708155453,
      "q1_ms": 0.2180721193942572,
      "q3_ms": 0.2182304838976462,
      "tflops": 39.36999782110931
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 1024,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.007625579833984375,
      "rel_err": 0.002638038480654359,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.2681307799939532,
      "inner": 140,
      "median_ms": 0.21827167854228588,
      "q1_ms": 0.2179766892887918,
      "q3_ms": 0.21828870719348614,
      "tflops": 39.35432507491286
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 1024,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.007625579833984375,
      "rel_err": 0.002638038480654359,
      "xla_temp_mb": 192.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.6427441600098973,
      "inner": 66,
      "median_ms": 0.687756530224402,
      "q1_ms": 0.6871772425118693,
      "q3_ms": 0.6878493485132684,
      "tflops": 12.489789939467773
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 1024,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.007625579833984375,
      "rel_err": 0.002638038480654359,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.04532377200666815,
      "inner": 133,
      "median_ms": 0.2106634737605131,
      "q1_ms": 0.21011333078228261,
      "q3_ms": 0.21075596994610113,
      "tflops": 40.775623978199604
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 1024,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.007625579833984375,
      "rel_err": 0.002638038480654359,
      "peak_mb": 72.3759765625,
      "compile_s": 0.00021545701019931585,
      "inner": 200,
      "median_ms": 0.1807830650068354,
      "q1_ms": 0.18029239003226394,
      "q3_ms": 0.1810239249607548,
      "tflops": 47.5151507785047
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 1024,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.007625579833984375,
      "rel_err": 0.002638038480654359,
      "peak_mb": 72.2509765625,
      "compile_s": 0.0003010890068253502,
      "inner": 178,
      "median_ms": 0.25347101685572726,
      "q1_ms": 0.2522464578198217,
      "q3_ms": 0.2534913117905786,
      "tflops": 33.88921817790825
    },
    {
      "impl": "torch_math",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 1024,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.007625579833984375,
      "rel_err": 0.002638038480654359,
      "peak_mb": 420.28173828125,
      "compile_s": 0.002758707996690646,
      "inner": 18,
      "median_ms": 2.7868792223873444,
      "q1_ms": 2.784720639384533,
      "q3_ms": 2.787247111377332,
      "tflops": 3.0822773096860443
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 1024,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.011097908020019531,
      "rel_err": 0.003424442373216152,
      "xla_temp_mb": 16.250244140625,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.5106813890015474,
      "inner": 46,
      "median_ms": 0.9371942825331959,
      "q1_ms": 0.9371128152175199,
      "q3_ms": 0.9379794564270212,
      "tflops": 32.079550240891585
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 1024,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.011097908020019531,
      "rel_err": 0.003424442373216152,
      "xla_temp_mb": 16.250244140625,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.5116869500052417,
      "inner": 49,
      "median_ms": 0.9368024083396078,
      "q1_ms": 0.9362787449000669,
      "q3_ms": 0.9374325714556366,
      "tflops": 32.09296945050229
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 1024,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.011097908020019531,
      "rel_err": 0.003424442373216152,
      "xla_temp_mb": 16.250244140625,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.5113417870015837,
      "inner": 46,
      "median_ms": 0.9369773261309804,
      "q1_ms": 0.9368384783991369,
      "q3_ms": 0.9374969128617997,
      "tflops": 32.08697823686422
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 1024,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.011097908020019531,
      "rel_err": 0.003372065955772996,
      "xla_temp_mb": 320.125,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 1.287208923997241,
      "inner": 21,
      "median_ms": 2.268240428834577,
      "q1_ms": 2.2675421905227093,
      "q3_ms": 2.26897371439485,
      "tflops": 13.25466678479375
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 1024,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.011097908020019531,
      "rel_err": 0.003424442373216152,
      "xla_temp_mb": 48.250732421875,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.10371411600499414,
      "inner": 49,
      "median_ms": 0.8699103672537307,
      "q1_ms": 0.8698563059977236,
      "q3_ms": 0.8750400814519986,
      "tflops": 34.5607687915172
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 1024,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.01517343521118164,
      "rel_err": 0.004465306177735329,
      "peak_mb": 152.501953125,
      "compile_s": 0.0009348479943582788,
      "inner": 53,
      "median_ms": 0.8967276791633405,
      "q1_ms": 0.8956912924471874,
      "q3_ms": 0.8972594529590218,
      "tflops": 33.527203152746274
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 1024,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.014698266983032227,
      "rel_err": 0.004325471352785826,
      "peak_mb": 156.517578125,
      "compile_s": 0.0022592429886572063,
      "inner": 22,
      "median_ms": 2.211267181502824,
      "q1_ms": 2.2106587950309566,
      "q3_ms": 2.212554636075352,
      "tflops": 13.596172965207822
    },
    {
      "impl": "torch_math",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 1024,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.010560214519500732,
      "rel_err": 0.002686998574063182,
      "peak_mb": 648.2509765625,
      "compile_s": 0.005890610991627909,
      "inner": 8,
      "median_ms": 6.029309624864254,
      "q1_ms": 6.028710937243886,
      "q3_ms": 6.029638250765856,
      "tflops": 4.986436746923059
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 1500,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0008845925331115723,
      "rel_err": 0.003496113931760192,
      "xla_temp_mb": 48.0,
      "xla_arg_mb": 35.15625,
      "xla_out_mb": 11.71875,
      "compile_s": 0.27849588300159667,
      "inner": 52,
      "median_ms": 0.8432145768892951,
      "q1_ms": 0.8427864711568243,
      "q3_ms": 0.8437058077106485,
      "tflops": 43.7184092997954
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 1500,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0008845925331115723,
      "rel_err": 0.003496113931760192,
      "xla_temp_mb": 48.0,
      "xla_arg_mb": 35.15625,
      "xla_out_mb": 11.71875,
      "compile_s": 0.2769466019963147,
      "inner": 53,
      "median_ms": 0.774716377425236,
      "q1_ms": 0.7741341320962539,
      "q3_ms": 0.7767418867860094,
      "tflops": 47.583865623852205
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 1500,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0008845925331115723,
      "rel_err": 0.003496113931760192,
      "xla_temp_mb": 48.0,
      "xla_arg_mb": 35.15625,
      "xla_out_mb": 11.71875,
      "compile_s": 0.27711942899622954,
      "inner": 57,
      "median_ms": 0.7780387018819487,
      "q1_ms": 0.7778569647997389,
      "q3_ms": 0.7780589210022681,
      "tflops": 47.38067645071125
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 1500,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0008845925331115723,
      "rel_err": 0.003496113931760192,
      "xla_temp_mb": 411.98731994628906,
      "xla_arg_mb": 35.15625,
      "xla_out_mb": 11.71875,
      "compile_s": 0.9280730159953237,
      "inner": 33,
      "median_ms": 1.3809238788036797,
      "q1_ms": 1.380681757644087,
      "q3_ms": 1.3809378332084496,
      "tflops": 26.695171664303444
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 1500,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0008845925331115723,
      "rel_err": 0.003496113931760192,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 35.15625,
      "xla_out_mb": 11.71875,
      "compile_s": 0.04258516200934537,
      "inner": 64,
      "median_ms": 0.623380562501552,
      "q1_ms": 0.6223608280606641,
      "q3_ms": 0.6263501718422049,
      "tflops": 59.13562632121405
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 1500,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0008845925331115723,
      "rel_err": 0.003496113931760192,
      "peak_mb": 99.18408203125,
      "compile_s": 0.0006212149892235175,
      "inner": 83,
      "median_ms": 0.5950125783586493,
      "q1_ms": 0.5950094397910534,
      "q3_ms": 0.5950478314252357,
      "tflops": 61.95499278635398
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 1500,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0008845925331115723,
      "rel_err": 0.003496113931760192,
      "peak_mb": 99.0009765625,
      "compile_s": 0.0009157719905488193,
      "inner": 55,
      "median_ms": 0.8930454001529142,
      "q1_ms": 0.8930169818499549,
      "q3_ms": 0.8931206182089888,
      "tflops": 41.27897640331371
    },
    {
      "impl": "torch_math",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 1500,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0008845925331115723,
      "rel_err": 0.003496113931760192,
      "peak_mb": 800.74609375,
      "compile_s": 0.00460565100365784,
      "inner": 10,
      "median_ms": 4.667991599126253,
      "q1_ms": 4.664754000259563,
      "q3_ms": 4.6697860496351495,
      "tflops": 7.897186448857393
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 1500,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0016873478889465332,
      "rel_err": 0.005069323815405369,
      "xla_temp_mb": 96.00057983398438,
      "xla_arg_mb": 58.59375,
      "xla_out_mb": 35.156272888183594,
      "compile_s": 0.5249000300100306,
      "inner": 14,
      "median_ms": 3.513698356982786,
      "q1_ms": 3.4953236073696257,
      "q3_ms": 3.5155729643260463,
      "tflops": 36.72028355638159
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 1500,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0016873478889465332,
      "rel_err": 0.005069323815405369,
      "xla_temp_mb": 96.00048828125,
      "xla_arg_mb": 58.59375,
      "xla_out_mb": 35.156272888183594,
      "compile_s": 0.5249584809935186,
      "inner": 13,
      "median_ms": 3.481713461671741,
      "q1_ms": 3.4802171928906596,
      "q3_ms": 3.48311957686047,
      "tflops": 37.05761586079783
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 1500,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0016873478889465332,
      "rel_err": 0.005069323815405369,
      "xla_temp_mb": 96.00048828125,
      "xla_arg_mb": 58.59375,
      "xla_out_mb": 35.156272888183594,
      "compile_s": 0.5245782020065235,
      "inner": 14,
      "median_ms": 3.480091999725638,
      "q1_ms": 3.479254642505631,
      "q3_ms": 3.4817160714218125,
      "tflops": 37.07488193133168
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 1500,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0016873478889465332,
      "rel_err": 0.005069323815405369,
      "xla_temp_mb": 686.8286285400391,
      "xla_arg_mb": 58.59375,
      "xla_out_mb": 35.156272888183594,
      "compile_s": 1.4579407620040001,
      "inner": 12,
      "median_ms": 3.8905900000827387,
      "q1_ms": 3.88833820882913,
      "q3_ms": 3.901824124720103,
      "tflops": 33.16309351467416
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 1500,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0016873478889465332,
      "rel_err": 0.005069323815405369,
      "xla_temp_mb": 70.741943359375,
      "xla_arg_mb": 58.59375,
      "xla_out_mb": 35.156272888183594,
      "compile_s": 0.10463473701383919,
      "inner": 17,
      "median_ms": 2.8222590001806727,
      "q1_ms": 2.8210758237572224,
      "q3_ms": 2.831433912112569,
      "tflops": 45.71656959610732
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 1500,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0016873478889465332,
      "rel_err": 0.004792788997292519,
      "peak_mb": 219.37255859375,
      "compile_s": 0.002561635003075935,
      "inner": 20,
      "median_ms": 2.5490906496997923,
      "q1_ms": 2.544935524929315,
      "q3_ms": 2.549506399736856,
      "tflops": 50.61569701932539
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 1500,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0018518567085266113,
      "rel_err": 0.006253216881304979,
      "peak_mb": 218.54833984375,
      "compile_s": 0.00885674099845346,
      "inner": 5,
      "median_ms": 8.798372399178334,
      "q1_ms": 8.78732260025572,
      "q3_ms": 8.820214800653048,
      "tflops": 14.66453045475199
    },
    {
      "impl": "torch_math",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 1500,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00137290358543396,
      "rel_err": 0.003899632953107357,
      "peak_mb": 1293.6337890625,
      "compile_s": 0.010839972004760057,
      "inner": 4,
      "median_ms": 11.141497001517564,
      "q1_ms": 11.119784625407192,
      "q3_ms": 11.14405950102082,
      "tflops": 11.580490483677899
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 1500,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.008003711700439453,
      "rel_err": 0.0027688515838235617,
      "xla_temp_mb": 48.0,
      "xla_arg_mb": 35.15625,
      "xla_out_mb": 11.71875,
      "compile_s": 0.2895186500099953,
      "inner": 74,
      "median_ms": 0.5523261080835821,
      "q1_ms": 0.551181175683883,
      "q3_ms": 0.5527545540275782,
      "tflops": 33.37158923005453
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 1500,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.008003711700439453,
      "rel_err": 0.0027688515838235617,
      "xla_temp_mb": 48.0,
      "xla_arg_mb": 35.15625,
      "xla_out_mb": 11.71875,
      "compile_s": 0.2873201119946316,
      "inner": 71,
      "median_ms": 0.5572948592114852,
      "q1_ms": 0.5571663802494558,
      "q3_ms": 0.5602367676492833,
      "tflops": 33.07405352003315
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 1500,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.008003711700439453,
      "rel_err": 0.0027688515838235617,
      "xla_temp_mb": 48.0,
      "xla_arg_mb": 35.15625,
      "xla_out_mb": 11.71875,
      "compile_s": 0.28352478399756365,
      "inner": 71,
      "median_ms": 0.5579252112869987,
      "q1_ms": 0.5574769647480448,
      "q3_ms": 0.5609659648387247,
      "tflops": 33.03668596993821
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 1500,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.009605884552001953,
      "rel_err": 0.0033231168054044247,
      "xla_temp_mb": 411.98731994628906,
      "xla_arg_mb": 35.15625,
      "xla_out_mb": 11.71875,
      "compile_s": 0.5387537409987999,
      "inner": 34,
      "median_ms": 1.4238553237245308,
      "q1_ms": 1.4232243971919696,
      "q3_ms": 1.424512147102231,
      "tflops": 12.945135431165466
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 1500,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.008003711700439453,
      "rel_err": 0.0027688515838235617,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 35.15625,
      "xla_out_mb": 11.71875,
      "compile_s": 0.04372845399484504,
      "inner": 89,
      "median_ms": 0.39895083147017385,
      "q1_ms": 0.3989322584847072,
      "q3_ms": 0.39905884837485806,
      "tflops": 46.20118206566015
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 1500,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.008003711700439453,
      "rel_err": 0.0027688515838235617,
      "peak_mb": 99.74658203125,
      "compile_s": 0.0003859900025418028,
      "inner": 135,
      "median_ms": 0.360226703708543,
      "q1_ms": 0.35969106301544584,
      "q3_ms": 0.3604873148207035,
      "tflops": 51.16777798603517
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 1500,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.008003711700439453,
      "rel_err": 0.0027688515838235617,
      "peak_mb": 99.5634765625,
      "compile_s": 0.0005506810121005401,
      "inner": 93,
      "median_ms": 0.5135183226973099,
      "q1_ms": 0.5124213010628498,
      "q3_ms": 0.5135518816826735,
      "tflops": 35.89355858459724
    },
    {
      "impl": "torch_math",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 1500,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.007621288299560547,
      "rel_err": 0.002636553719639778,
      "peak_mb": 810.29833984375,
      "compile_s": 0.0056416999868815765,
      "inner": 8,
      "median_ms": 5.690519999916432,
      "q1_ms": 5.689473687198188,
      "q3_ms": 5.692945937880722,
      "tflops": 3.2390712975739797
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 1500,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.014189004898071289,
      "rel_err": 0.0037457095459103584,
      "xla_temp_mb": 96.00048828125,
      "xla_arg_mb": 58.59375,
      "xla_out_mb": 35.156272888183594,
      "compile_s": 0.5354982440039748,
      "inner": 23,
      "median_ms": 2.069145608718669,
      "q1_ms": 2.0680565867223777,
      "q3_ms": 2.0695408915660507,
      "tflops": 31.178086127998238
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 1500,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.014189004898071289,
      "rel_err": 0.0037457095459103584,
      "xla_temp_mb": 96.00048828125,
      "xla_arg_mb": 58.59375,
      "xla_out_mb": 35.156272888183594,
      "compile_s": 0.5302512110065436,
      "inner": 23,
      "median_ms": 2.045064391416457,
      "q1_ms": 2.0449693911359645,
      "q3_ms": 2.0454261523453807,
      "tflops": 31.545216996965827
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 1500,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.014189004898071289,
      "rel_err": 0.0037457095459103584,
      "xla_temp_mb": 96.00048828125,
      "xla_arg_mb": 58.59375,
      "xla_out_mb": 35.156272888183594,
      "compile_s": 0.5283361089968821,
      "inner": 23,
      "median_ms": 2.045501304446193,
      "q1_ms": 2.0448686086835903,
      "q3_ms": 2.046624304215505,
      "tflops": 31.538479031899826
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 1500,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.014189004898071289,
      "rel_err": 0.005992697086185217,
      "xla_temp_mb": 686.8286285400391,
      "xla_arg_mb": 58.59375,
      "xla_out_mb": 35.156272888183594,
      "compile_s": 0.9584988910064567,
      "inner": 12,
      "median_ms": 4.061857250538499,
      "q1_ms": 4.059793708317253,
      "q3_ms": 4.0646500419825315,
      "tflops": 15.882389759376045
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 1500,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.014189004898071289,
      "rel_err": 0.0037457095459103584,
      "xla_temp_mb": 70.741943359375,
      "xla_arg_mb": 58.59375,
      "xla_out_mb": 35.156272888183594,
      "compile_s": 0.10396185399440583,
      "inner": 27,
      "median_ms": 1.736177740558223,
      "q1_ms": 1.7356747592674329,
      "q3_ms": 1.7370770184681061,
      "tflops": 37.15748594913896
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 1500,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.014189004898071289,
      "rel_err": 0.003783293068408966,
      "peak_mb": 218.81005859375,
      "compile_s": 0.0016372560057789087,
      "inner": 31,
      "median_ms": 1.630281741757156,
      "q1_ms": 1.6293648546824473,
      "q3_ms": 1.6305695645578746,
      "tflops": 39.571074341093606
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 1500,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.014189004898071289,
      "rel_err": 0.004624849185347557,
      "peak_mb": 218.82958984375,
      "compile_s": 0.004683927996666171,
      "inner": 10,
      "median_ms": 4.630364499462303,
      "q1_ms": 4.6277600988105405,
      "q3_ms": 4.632025550381513,
      "tflops": 13.93238048699868
    },
    {
      "impl": "torch_math",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 1500,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.012250661849975586,
      "rel_err": 0.0032025272957980633,
      "peak_mb": 1292.7900390625,
      "compile_s": 0.011776662999182008,
      "inner": 4,
      "median_ms": 12.01637650228804,
      "q1_ms": 12.013766750897048,
      "q3_ms": 12.019195875836886,
      "tflops": 5.36867332574976
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 2048,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0005794912576675415,
      "rel_err": 0.002374216914176941,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.2460062839963939,
      "inner": 35,
      "median_ms": 1.2610928288527896,
      "q1_ms": 1.1824193285844689,
      "q3_ms": 1.2834860429783084,
      "tflops": 54.49200500054686
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 2048,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0005794912576675415,
      "rel_err": 0.002374216914176941,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.24245253299886826,
      "inner": 38,
      "median_ms": 1.1760178948813853,
      "q1_ms": 1.1754274340714084,
      "q3_ms": 1.1768091447409366,
      "tflops": 58.43404002192597
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 2048,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0005794912576675415,
      "rel_err": 0.002374216914176941,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.24126363500545267,
      "inner": 40,
      "median_ms": 1.1774671998864505,
      "q1_ms": 1.1767407000661478,
      "q3_ms": 1.1781632751080906,
      "tflops": 58.36211551593707
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 2048,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0007002353668212891,
      "rel_err": 0.0028689142782241106,
      "xla_temp_mb": 768.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.8609542129997863,
      "inner": 20,
      "median_ms": 2.3381090999464504,
      "q1_ms": 2.3363464999420103,
      "q3_ms": 2.345180475458619,
      "tflops": 29.391047978716593
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 2048,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0005794912576675415,
      "rel_err": 0.002374216914176941,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.04436873500526417,
      "inner": 38,
      "median_ms": 1.1223730524713909,
      "q1_ms": 1.1222537237029573,
      "q3_ms": 1.137371920968259,
      "tflops": 61.22694819221139
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 2048,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0005985796451568604,
      "rel_err": 0.0024524235632270575,
      "peak_mb": 128.5009765625,
      "compile_s": 0.0010966639965772629,
      "inner": 46,
      "median_ms": 1.0873480436123866,
      "q1_ms": 1.0872940218415232,
      "q3_ms": 1.0878385216093864,
      "tflops": 63.19915425395922
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 2048,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0005794912576675415,
      "rel_err": 0.002374216914176941,
      "peak_mb": 128.2509765625,
      "compile_s": 0.0015460739959962666,
      "inner": 32,
      "median_ms": 1.5556262187601533,
      "q1_ms": 1.5532417344275018,
      "q3_ms": 1.5557110784811812,
      "tflops": 44.17479977341214
    },
    {
      "impl": "torch_math",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 2048,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00048814713954925537,
      "rel_err": 0.001999973552301526,
      "peak_mb": 1392.31298828125,
      "compile_s": 0.00827452000521589,
      "inner": 6,
      "median_ms": 8.40931800000059,
      "q1_ms": 8.39739958367621,
      "q3_ms": 8.414263999535857,
      "tflops": 8.171825198665953
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 2048,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0009521394968032837,
      "rel_err": 0.00357146211899817,
      "xla_temp_mb": 32.50048828125,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.3800784670020221,
      "inner": 10,
      "median_ms": 4.678205000527669,
      "q1_ms": 4.65465184970526,
      "q3_ms": 4.705029700562591,
      "tflops": 51.41249016425557
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 2048,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0009521394968032837,
      "rel_err": 0.00357146211899817,
      "xla_temp_mb": 32.50048828125,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.38092220399994403,
      "inner": 11,
      "median_ms": 4.670018909647214,
      "q1_ms": 4.667625682486687,
      "q3_ms": 4.677106318095783,
      "tflops": 51.50261128046897
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 2048,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0009521394968032837,
      "rel_err": 0.00357146211899817,
      "xla_temp_mb": 32.50048828125,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.37232585399760865,
      "inner": 10,
      "median_ms": 4.685724100272637,
      "q1_ms": 4.683938300149748,
      "q3_ms": 4.693728299753275,
      "tflops": 51.32998943791965
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 2048,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0012786388397216797,
      "rel_err": 0.004288523457944393,
      "xla_temp_mb": 1280.25,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 1.5738516249984968,
      "inner": 7,
      "median_ms": 6.840437571684431,
      "q1_ms": 6.819219357566908,
      "q3_ms": 6.841722641963445,
      "tflops": 35.16122558761593
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 2048,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0009521394968032837,
      "rel_err": 0.00357146211899817,
      "xla_temp_mb": 96.500732421875,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.10485072500887327,
      "inner": 10,
      "median_ms": 4.346432500460651,
      "q1_ms": 4.343264000635827,
      "q3_ms": 4.351835200213827,
      "tflops": 55.33691563149985
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 2048,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0012866407632827759,
      "rel_err": 0.004315361380577087,
      "peak_mb": 288.751953125,
      "compile_s": 0.004158925992669538,
      "inner": 12,
      "median_ms": 4.282903665928946,
      "q1_ms": 4.274982832915461,
      "q3_ms": 4.284036583461178,
      "tflops": 56.15773487724069
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 2048,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.001228049397468567,
      "rel_err": 0.004167243372648954,
      "peak_mb": 284.783203125,
      "compile_s": 0.014872829007799737,
      "inner": 3,
      "median_ms": 15.228738668762768,
      "q1_ms": 15.061871667664189,
      "q3_ms": 15.288901166059077,
      "tflops": 15.793702538828875
    },
    {
      "impl": "torch_math",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 2048,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0008591413497924805,
      "rel_err": 0.0031309486366808414,
      "peak_mb": 2304.2509765625,
      "compile_s": 0.01952612999593839,
      "inner": 2,
      "median_ms": 20.03994700498879,
      "q1_ms": 20.006749000458512,
      "q3_ms": 20.04313024735893,
      "tflops": 12.001936358221151
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 2048,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00552058219909668,
      "rel_err": 0.0019098230404779315,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.27741625100316014,
      "inner": 53,
      "median_ms": 0.7694942266065277,
      "q1_ms": 0.7169828301280263,
      "q3_ms": 0.7700117170110011,
      "tflops": 44.652366684447486
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 2048,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00552058219909668,
      "rel_err": 0.0019098230404779315,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.2691933890018845,
      "inner": 59,
      "median_ms": 0.6986457627104847,
      "q1_ms": 0.6981641779282777,
      "q3_ms": 0.6989452204079036,
      "tflops": 49.18048630925212
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 2048,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00552058219909668,
      "rel_err": 0.0019098230404779315,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.2627149059990188,
      "inner": 58,
      "median_ms": 0.6983989828611435,
      "q1_ms": 0.6980565345217875,
      "q3_ms": 0.6987717585208633,
      "tflops": 49.197864274140045
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 2048,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0068132877349853516,
      "rel_err": 0.0023570293560624123,
      "xla_temp_mb": 768.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.5261812329990789,
      "inner": 17,
      "median_ms": 2.8025169406592956,
      "q1_ms": 2.79772694065126,
      "q3_ms": 2.808520941043218,
      "tflops": 12.260314244493676
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 2048,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00552058219909668,
      "rel_err": 0.0019098230404779315,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.04381737200310454,
      "inner": 64,
      "median_ms": 0.6520289687159675,
      "q1_ms": 0.6519349688005605,
      "q3_ms": 0.6525280391542765,
      "tflops": 52.69664388633561
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 2048,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00552058219909668,
      "rel_err": 0.0019098230404779315,
      "peak_mb": 128.5009765625,
      "compile_s": 0.0006046530033927411,
      "inner": 83,
      "median_ms": 0.6090469638387818,
      "q1_ms": 0.6086615301529787,
      "q3_ms": 0.6090822651195451,
      "tflops": 56.415581076757846
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 2048,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00552058219909668,
      "rel_err": 0.0019098230404779315,
      "peak_mb": 128.2509765625,
      "compile_s": 0.0008898829983081669,
      "inner": 57,
      "median_ms": 0.8623230351715169,
      "q1_ms": 0.8617548419830487,
      "q3_ms": 0.862351333285404,
      "tflops": 39.8455531936078
    },
    {
      "impl": "torch_math",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 2048,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00552058219909668,
      "rel_err": 0.0019098230404779315,
      "peak_mb": 1408.31298828125,
      "compile_s": 0.010216342998319305,
      "inner": 4,
      "median_ms": 10.277204750309465,
      "q1_ms": 10.272363249896443,
      "q3_ms": 10.285036001732806,
      "tflops": 3.3432960812584147
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 2048,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.01994955539703369,
      "rel_err": 0.00563684431836009,
      "xla_temp_mb": 32.50048828125,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.4332833240041509,
      "inner": 17,
      "median_ms": 2.902543117456576,
      "q1_ms": 2.898385000127532,
      "q3_ms": 2.9060509700417136,
      "tflops": 41.432316221155716
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 2048,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.01994955539703369,
      "rel_err": 0.00563684431836009,
      "xla_temp_mb": 32.50048828125,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.42992472200421616,
      "inner": 17,
      "median_ms": 2.903154882540762,
      "q1_ms": 2.8979191470973413,
      "q3_ms": 2.903706558914992,
      "tflops": 41.42358542812312
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 2048,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.01994955539703369,
      "rel_err": 0.00563684431836009,
      "xla_temp_mb": 32.50048828125,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.42839635100972373,
      "inner": 17,
      "median_ms": 2.9033587642890564,
      "q1_ms": 2.901642882480861,
      "q3_ms": 2.9054439116461093,
      "tflops": 41.42067655130032
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 2048,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.01994955539703369,
      "rel_err": 0.00563684431836009,
      "xla_temp_mb": 1280.25,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 1.590788549001445,
      "inner": 6,
      "median_ms": 7.7374088335394235,
      "q1_ms": 7.734076750542348,
      "q3_ms": 7.74143974922481,
      "tflops": 15.542552665268474
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 2048,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.01994955539703369,
      "rel_err": 0.00563684431836009,
      "xla_temp_mb": 96.500732421875,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.10556582699064165,
      "inner": 18,
      "median_ms": 2.5992665557876333,
      "q1_ms": 2.5973034442156657,
      "q3_ms": 2.6214282220785714,
      "tflops": 46.26654546846155
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 2048,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.01994955539703369,
      "rel_err": 0.00563684431836009,
      "peak_mb": 288.751953125,
      "compile_s": 0.0025302950089098886,
      "inner": 19,
      "median_ms": 2.588244105266456,
      "q1_ms": 2.5855247632104073,
      "q3_ms": 2.5888848157544424,
      "tflops": 46.46357893496274
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 2048,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.01994955539703369,
      "rel_err": 0.007024108897894621,
      "peak_mb": 284.783203125,
      "compile_s": 0.007980502996360883,
      "inner": 6,
      "median_ms": 7.9833256677375175,
      "q1_ms": 7.963117415783927,
      "q3_ms": 7.984497083574146,
      "tflops": 15.063782851048533
    },
    {
      "impl": "torch_math",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 2048,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.012137055397033691,
      "rel_err": 0.0034293842036277056,
      "peak_mb": 2304.2509765625,
      "compile_s": 0.021226595010375604,
      "inner": 2,
      "median_ms": 21.700883000448812,
      "q1_ms": 21.650650749506894,
      "q3_ms": 21.726792248955462,
      "tflops": 5.541667787689231
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 3000,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0005877763032913208,
      "rel_err": 0.002183879492804408,
      "xla_temp_mb": 96.0,
      "xla_arg_mb": 70.3125,
      "xla_out_mb": 23.4375,
      "compile_s": 0.28030422299343627,
      "inner": 16,
      "median_ms": 3.007092312145687,
      "q1_ms": 2.9077833123665187,
      "q3_ms": 3.0087028435445973,
      "tflops": 49.03607361982976
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 3000,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0005877763032913208,
      "rel_err": 0.002183879492804408,
      "xla_temp_mb": 96.0,
      "xla_arg_mb": 70.3125,
      "xla_out_mb": 23.4375,
      "compile_s": 0.28171987101086415,
      "inner": 17,
      "median_ms": 2.7738685295095338,
      "q1_ms": 2.769215117629809,
      "q3_ms": 2.7767242649775548,
      "tflops": 53.15897218318154
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 3000,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0005877763032913208,
      "rel_err": 0.002183879492804408,
      "xla_temp_mb": 96.0,
      "xla_arg_mb": 70.3125,
      "xla_out_mb": 23.4375,
      "compile_s": 0.28171599299821537,
      "inner": 17,
      "median_ms": 2.7749175882062818,
      "q1_ms": 2.7716060888932965,
      "q3_ms": 2.7808621474404767,
      "tflops": 53.13887541262664
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 3000,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0005970597267150879,
      "rel_err": 0.002218371955677867,
      "xla_temp_mb": 1647.949234008789,
      "xla_arg_mb": 70.3125,
      "xla_out_mb": 23.4375,
      "compile_s": 1.0199152580025839,
      "inner": 9,
      "median_ms": 5.086632555402402,
      "q1_ms": 5.080058499717981,
      "q3_ms": 5.088245610952274,
      "tflops": 28.98892310265073
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 3000,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0005877763032913208,
      "rel_err": 0.002183879492804408,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 70.3125,
      "xla_out_mb": 23.4375,
      "compile_s": 0.0463988440023968,
      "inner": 20,
      "median_ms": 2.3591753997607157,
      "q1_ms": 2.3584883998410078,
      "q3_ms": 2.3600708751473576,
      "tflops": 62.5031949786167
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 3000,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0005711764097213745,
      "rel_err": 0.002122202655300498,
      "peak_mb": 180.6796875,
      "compile_s": 0.002248894001240842,
      "inner": 22,
      "median_ms": 2.291927454776173,
      "q1_ms": 2.2913876137780873,
      "q3_ms": 2.293386363321175,
      "tflops": 64.33711489982583
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 3000,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0005877763032913208,
      "rel_err": 0.002183879492804408,
      "peak_mb": 180.3134765625,
      "compile_s": 0.0032665480102878064,
      "inner": 15,
      "median_ms": 3.3409815330135944,
      "q1_ms": 3.33950906691219,
      "q3_ms": 3.3419269331110018,
      "tflops": 44.13553278966897
    },
    {
      "impl": "torch_math",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 3000,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00048798322677612305,
      "rel_err": 0.0018130989046767354,
      "peak_mb": 2816.39111328125,
      "compile_s": 0.01796525699319318,
      "inner": 2,
      "median_ms": 18.469050999556202,
      "q1_ms": 18.369848250586074,
      "q3_ms": 18.527508749684785,
      "tflops": 7.983951097625063
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 3000,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0009886473417282104,
      "rel_err": 0.004239012952893972,
      "xla_temp_mb": 192.00048828125,
      "xla_arg_mb": 117.1875,
      "xla_out_mb": 70.3125228881836,
      "compile_s": 0.4493178600096144,
      "inner": 4,
      "median_ms": 10.747296750196256,
      "q1_ms": 10.716324251916376,
      "q3_ms": 10.76031350021367,
      "tflops": 48.02100583949872
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 3000,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0009886473417282104,
      "rel_err": 0.004239012952893972,
      "xla_temp_mb": 192.00048828125,
      "xla_arg_mb": 117.1875,
      "xla_out_mb": 70.3125228881836,
      "compile_s": 0.4478287539968733,
      "inner": 4,
      "median_ms": 10.623870501149213,
      "q1_ms": 10.596994748993893,
      "q3_ms": 10.666875750757754,
      "tflops": 48.578905394617955
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 3000,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0009886473417282104,
      "rel_err": 0.004239012952893972,
      "xla_temp_mb": 192.00048828125,
      "xla_arg_mb": 117.1875,
      "xla_out_mb": 70.3125228881836,
      "compile_s": 0.4471380579925608,
      "inner": 4,
      "median_ms": 10.627118750562659,
      "q1_ms": 10.594414998195134,
      "q3_ms": 10.688782125725993,
      "tflops": 48.56405692960522
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 3000,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0009462237358093262,
      "rel_err": 0.00432322034612298,
      "xla_temp_mb": 2746.948257446289,
      "xla_arg_mb": 117.1875,
      "xla_out_mb": 70.3125228881836,
      "compile_s": 1.6308540049940348,
      "inner": 3,
      "median_ms": 14.294735000779232,
      "q1_ms": 14.255065997228181,
      "q3_ms": 14.32265383239913,
      "tflops": 36.1039221763724
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 3000,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0009886473417282104,
      "rel_err": 0.004239012952893972,
      "xla_temp_mb": 141.358154296875,
      "xla_arg_mb": 117.1875,
      "xla_out_mb": 70.3125228881836,
      "compile_s": 0.11427230499975849,
      "inner": 5,
      "median_ms": 9.96789839991834,
      "q1_ms": 9.950737000326626,
      "q3_ms": 9.984785299457144,
      "tflops": 51.7758086302553
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 3000,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0008276849985122681,
      "rel_err": 0.003578837960958481,
      "peak_mb": 416.5556640625,
      "compile_s": 0.008345152993570082,
      "inner": 6,
      "median_ms": 8.659253997999864,
      "q1_ms": 8.658916750088489,
      "q3_ms": 8.660644166714821,
      "tflops": 59.600515254456
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 3000,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0009874999523162842,
      "rel_err": 0.004759765230119228,
      "peak_mb": 405.1103515625,
      "compile_s": 0.03179256001021713,
      "inner": 1,
      "median_ms": 31.69889199489262,
      "q1_ms": 31.626059499103576,
      "q3_ms": 31.822942008147947,
      "tflops": 16.281199989045493
    },
    {
      "impl": "torch_math",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 3000,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0007945001125335693,
      "rel_err": 0.003578837960958481,
      "peak_mb": 4762.3447265625,
      "compile_s": 0.04205075099889655,
      "inner": 1,
      "median_ms": 42.76198700245004,
      "q1_ms": 42.6341245038202,
      "q3_ms": 42.77590249694185,
      "tflops": 12.069036922218567
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 3000,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.007696628570556641,
      "rel_err": 0.0026626174803823233,
      "xla_temp_mb": 96.0,
      "xla_arg_mb": 70.3125,
      "xla_out_mb": 23.4375,
      "compile_s": 0.2904075819969876,
      "inner": 28,
      "median_ms": 1.7029976786164167,
      "q1_ms": 1.702381893015367,
      "q3_ms": 1.7040138929067845,
      "tflops": 43.29307134458314
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 3000,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.007696628570556641,
      "rel_err": 0.0026626174803823233,
      "xla_temp_mb": 96.0,
      "xla_arg_mb": 70.3125,
      "xla_out_mb": 23.4375,
      "compile_s": 0.28896125601022504,
      "inner": 27,
      "median_ms": 1.72597696332054,
      "q1_ms": 1.7256470738276324,
      "q3_ms": 1.7276183886375898,
      "tflops": 42.71667673834856
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 3000,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.007696628570556641,
      "rel_err": 0.0026626174803823233,
      "xla_temp_mb": 96.0,
      "xla_arg_mb": 70.3125,
      "xla_out_mb": 23.4375,
      "compile_s": 0.2866874559986172,
      "inner": 27,
      "median_ms": 1.7267409997724894,
      "q1_ms": 1.7253227407410223,
      "q3_ms": 1.7290247590237953,
      "tflops": 42.6977757577507
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 3000,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.01217508316040039,
      "rel_err": 0.004211920779198408,
      "xla_temp_mb": 1647.949234008789,
      "xla_arg_mb": 70.3125,
      "xla_out_mb": 23.4375,
      "compile_s": 0.5049347200110788,
      "inner": 9,
      "median_ms": 5.062043778404283,
      "q1_ms": 5.05961877772481,
      "q3_ms": 5.078079888739416,
      "tflops": 14.564868110097894
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 3000,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.007696628570556641,
      "rel_err": 0.0026626174803823233,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 70.3125,
      "xla_out_mb": 23.4375,
      "compile_s": 0.04532156699860934,
      "inner": 35,
      "median_ms": 1.3309309142641723,
      "q1_ms": 1.3305122855984206,
      "q3_ms": 1.3315405141578855,
      "tflops": 55.39581296806963
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 3000,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.007696628570556641,
      "rel_err": 0.0026626174803823233,
      "peak_mb": 180.6796875,
      "compile_s": 0.0011839299986604601,
      "inner": 43,
      "median_ms": 1.2298944882701994,
      "q1_ms": 1.228700604643077,
      "q3_ms": 1.2324532906140395,
      "tflops": 59.94660574802288
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 3000,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.007696628570556641,
      "rel_err": 0.0026626174803823233,
      "peak_mb": 180.3134765625,
      "compile_s": 0.0017684339982224628,
      "inner": 28,
      "median_ms": 1.7737518568797637,
      "q1_ms": 1.7705706788027393,
      "q3_ms": 1.773858143029169,
      "tflops": 41.56612984732608
    },
    {
      "impl": "torch_math",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 3000,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.007696628570556641,
      "rel_err": 0.0026626174803823233,
      "peak_mb": 2850.7236328125,
      "compile_s": 0.02219038001203444,
      "inner": 2,
      "median_ms": 22.61581349739572,
      "q1_ms": 22.59266974942875,
      "q3_ms": 22.63652274996275,
      "tflops": 3.260019809081376
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 3000,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.015622377395629883,
      "rel_err": 0.004377692937850952,
      "xla_temp_mb": 192.00048828125,
      "xla_arg_mb": 117.1875,
      "xla_out_mb": 70.3125228881836,
      "compile_s": 0.44901416600623634,
      "inner": 7,
      "median_ms": 6.274577142903581,
      "q1_ms": 6.258541142285269,
      "q3_ms": 6.300506142516887,
      "tflops": 41.12595863003247
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 3000,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.015622377395629883,
      "rel_err": 0.004377692937850952,
      "xla_temp_mb": 192.00048828125,
      "xla_arg_mb": 117.1875,
      "xla_out_mb": 70.3125228881836,
      "compile_s": 0.4529000559996348,
      "inner": 8,
      "median_ms": 6.179090250952868,
      "q1_ms": 6.157604188047117,
      "q3_ms": 6.199157687660772,
      "tflops": 41.76148745524583
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 3000,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.015622377395629883,
      "rel_err": 0.004377692937850952,
      "xla_temp_mb": 192.00048828125,
      "xla_arg_mb": 117.1875,
      "xla_out_mb": 70.3125228881836,
      "compile_s": 0.45555221200629603,
      "inner": 8,
      "median_ms": 6.181829248816939,
      "q1_ms": 6.163149562780745,
      "q3_ms": 6.206969937920803,
      "tflops": 41.74298409316053
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 3000,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.015622377395629883,
      "rel_err": 0.004377692937850952,
      "xla_temp_mb": 2746.948257446289,
      "xla_arg_mb": 117.1875,
      "xla_out_mb": 70.3125228881836,
      "compile_s": 0.9419685680040857,
      "inner": 3,
      "median_ms": 14.248909001859525,
      "q1_ms": 14.230477667297237,
      "q3_ms": 14.282465834791461,
      "tflops": 18.110018104987823
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 3000,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.015622377395629883,
      "rel_err": 0.004377692937850952,
      "xla_temp_mb": 141.358154296875,
      "xla_arg_mb": 117.1875,
      "xla_out_mb": 70.3125228881836,
      "compile_s": 0.10864674500771798,
      "inner": 8,
      "median_ms": 5.50928349912283,
      "q1_ms": 5.496879499332863,
      "q3_ms": 5.544491436921817,
      "tflops": 46.83875862280196
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 3000,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.015622377395629883,
      "rel_err": 0.004377692937850952,
      "peak_mb": 417.6806640625,
      "compile_s": 0.00479468799312599,
      "inner": 10,
      "median_ms": 4.9661215001833625,
      "q1_ms": 4.963445900648367,
      "q3_ms": 4.96932804962853,
      "tflops": 51.96167673112149
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 3000,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.015622377395629883,
      "rel_err": 0.004377692937850952,
      "peak_mb": 406.7978515625,
      "compile_s": 0.016776438991655596,
      "inner": 2,
      "median_ms": 16.75418249942595,
      "q1_ms": 16.73778374606627,
      "q3_ms": 16.768554498412414,
      "tflops": 15.402004843199094
    },
    {
      "impl": "torch_math",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 3000,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.015622377395629883,
      "rel_err": 0.004377692937850952,
      "peak_mb": 4762.9072265625,
      "compile_s": 0.04604640800971538,
      "inner": 1,
      "median_ms": 46.65223399933893,
      "q1_ms": 46.59605800407007,
      "q3_ms": 46.680177001690026,
      "tflops": 5.531310676433128
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 4096,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0005183517932891846,
      "rel_err": 0.0037030023522675037,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.24958061499637552,
      "inner": 10,
      "median_ms": 4.402162799669895,
      "q1_ms": 4.369096749724122,
      "q3_ms": 4.715828200278338,
      "tflops": 62.4415587185036
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 4096,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0005183517932891846,
      "rel_err": 0.0037030023522675037,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.24287112499587238,
      "inner": 11,
      "median_ms": 4.35715800003064,
      "q1_ms": 4.350057727027557,
      "q3_ms": 4.390676954327236,
      "tflops": 63.08651348931277
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 4096,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0005183517932891846,
      "rel_err": 0.0037030023522675037,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.24483185399731155,
      "inner": 11,
      "median_ms": 4.353808908490464,
      "q1_ms": 4.350260863165286,
      "q3_ms": 4.369256818211976,
      "tflops": 63.13504168911828
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 4096,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0005831122398376465,
      "rel_err": 0.004165638238191605,
      "xla_temp_mb": 3072.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 1.1346965260017896,
      "inner": 5,
      "median_ms": 8.859906997531652,
      "q1_ms": 8.850567201443482,
      "q3_ms": 8.88660430064192,
      "tflops": 31.024920128459616
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 4096,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0005183517932891846,
      "rel_err": 0.0037030023522675037,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.0470641239953693,
      "inner": 11,
      "median_ms": 4.116308909777382,
      "q1_ms": 4.104720181864898,
      "q3_ms": 4.131868272617629,
      "tflops": 66.7777644897078
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 4096,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0005171298980712891,
      "rel_err": 0.0036942732986062765,
      "peak_mb": 240.7509765625,
      "compile_s": 0.0038937040080782026,
      "inner": 12,
      "median_ms": 4.018782499770168,
      "q1_ms": 4.017745958359835,
      "q3_ms": 4.035119583325771,
      "tflops": 68.39830395392637
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 4096,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0005183517932891846,
      "rel_err": 0.0037030023522675037,
      "peak_mb": 240.2509765625,
      "compile_s": 0.005728855001507327,
      "inner": 8,
      "median_ms": 5.924328499531839,
      "q1_ms": 5.923371749304351,
      "q3_ms": 5.9253879999232595,
      "tflops": 46.398154147887276
    },
    {
      "impl": "torch_math",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 4096,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00045943260192871094,
      "rel_err": 0.003282095305621624,
      "peak_mb": 5072.37548828125,
      "compile_s": 0.03204158000880852,
      "inner": 1,
      "median_ms": 32.751844002632424,
      "q1_ms": 32.718590497097466,
      "q3_ms": 32.77146550681209,
      "tflops": 8.392745975521462
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 4096,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0007895827293395996,
      "rel_err": 0.0042350213043391705,
      "xla_temp_mb": 65.00048828125,
      "xla_arg_mb": 160.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 0.39522070600651205,
      "inner": 3,
      "median_ms": 17.357461668628577,
      "q1_ms": 17.348441501477886,
      "q3_ms": 17.375008498978183,
      "tflops": 55.427037240291014
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 4096,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0007895827293395996,
      "rel_err": 0.0042350213043391705,
      "xla_temp_mb": 65.00048828125,
      "xla_arg_mb": 160.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 0.38657902300474234,
      "inner": 3,
      "median_ms": 17.404026332466554,
      "q1_ms": 17.37473250250332,
      "q3_ms": 17.454430334813274,
      "tflops": 55.27874159264456
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 4096,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0007895827293395996,
      "rel_err": 0.0042350213043391705,
      "xla_temp_mb": 65.00048828125,
      "xla_arg_mb": 160.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 0.3884452720085392,
      "inner": 3,
      "median_ms": 17.4478726645854,
      "q1_ms": 17.36115200280134,
      "q3_ms": 17.520252998413827,
      "tflops": 55.13982665959931
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 4096,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0009999126195907593,
      "rel_err": 0.005363150499761105,
      "xla_temp_mb": 5120.500015258789,
      "xla_arg_mb": 160.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 2.0783086820010794,
      "inner": 1,
      "median_ms": 25.538471003528684,
      "q1_ms": 25.512297499517445,
      "q3_ms": 25.548966499627568,
      "tflops": 37.67150641755605
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 4096,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0007895827293395996,
      "rel_err": 0.0042350213043391705,
      "xla_temp_mb": 193.000732421875,
      "xla_arg_mb": 160.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 0.11774569400586188,
      "inner": 3,
      "median_ms": 15.756450666231103,
      "q1_ms": 15.73624016843193,
      "q3_ms": 15.864045667209817,
      "tflops": 61.058971635400994
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 4096,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0009463578462600708,
      "rel_err": 0.004244851879775524,
      "peak_mb": 561.251953125,
      "compile_s": 0.014400585991097614,
      "inner": 3,
      "median_ms": 15.065847668059481,
      "q1_ms": 15.024288833956234,
      "q3_ms": 15.084449500136543,
      "tflops": 63.857852243100325
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 4096,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0009139552712440491,
      "rel_err": 0.004242294002324343,
      "peak_mb": 541.314453125,
      "compile_s": 0.05855276100919582,
      "inner": 1,
      "median_ms": 60.80296701111365,
      "q1_ms": 60.2625449973857,
      "q3_ms": 61.12285749259172,
      "tflops": 15.822791577393764
    },
    {
      "impl": "torch_math",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 4096,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "error": "OutOfMemoryError: CUDA out of memory. Tried to allocate 2.00 GiB. GPU 0 has a total capacity of 23.61 GiB of which 1.09 GiB is free. Including non-PyTorch memory, this process has 22.48 GiB memory in use. Of the allocated memory 4.39 GiB is allocated by PyTorch, and 45.75 MiB is reserved by PyTorch but unallocated. I"
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 4096,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0050841569900512695,
      "rel_err": 0.0017588435439392924,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.27733832300873473,
      "inner": 16,
      "median_ms": 2.7099629369331524,
      "q1_ms": 2.703513187043427,
      "q3_ms": 2.718358781294228,
      "tflops": 50.71617460109576
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 4096,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0050841569900512695,
      "rel_err": 0.0017588435439392924,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.27542562599410303,
      "inner": 19,
      "median_ms": 2.469636631614872,
      "q1_ms": 2.46795873681549,
      "q3_ms": 2.4705038418974974,
      "tflops": 55.65148804183795
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 4096,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0050841569900512695,
      "rel_err": 0.0017588435439392924,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.27395292099390645,
      "inner": 20,
      "median_ms": 2.468856849736767,
      "q1_ms": 2.467573724788963,
      "q3_ms": 2.469730025040917,
      "tflops": 55.669065416512076
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 4096,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.007138490676879883,
      "rel_err": 0.002469531958922744,
      "xla_temp_mb": 3072.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.4777839890011819,
      "inner": 5,
      "median_ms": 8.90622420120053,
      "q1_ms": 8.899855401250534,
      "q3_ms": 8.913648501038551,
      "tflops": 15.431786845594305
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 4096,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0050841569900512695,
      "rel_err": 0.0017588435439392924,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.047625084000173956,
      "inner": 20,
      "median_ms": 2.2781667503295466,
      "q1_ms": 2.2773281496483833,
      "q3_ms": 2.2833001999970293,
      "tflops": 60.32875049735445
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 4096,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0050841569900512695,
      "rel_err": 0.0017588435439392924,
      "peak_mb": 240.7509765625,
      "compile_s": 0.002103007005644031,
      "inner": 24,
      "median_ms": 2.222711333160987,
      "q1_ms": 2.216282312777669,
      "q3_ms": 2.225462375160229,
      "tflops": 61.8339194215309
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 4096,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0050841569900512695,
      "rel_err": 0.0017588435439392924,
      "peak_mb": 240.2509765625,
      "compile_s": 0.0030862569983582944,
      "inner": 16,
      "median_ms": 3.1685427502452512,
      "q1_ms": 3.1675474065195885,
      "q3_ms": 3.1687649998275447,
      "tflops": 43.376076734758264
    },
    {
      "impl": "torch_math",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 4096,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.004895687103271484,
      "rel_err": 0.0016936431638896465,
      "peak_mb": 5136.37548828125,
      "compile_s": 0.0400349880073918,
      "inner": 1,
      "median_ms": 40.35986299277283,
      "q1_ms": 40.35008449864108,
      "q3_ms": 40.41430649522226,
      "tflops": 3.4053374635243667
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 4096,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.014839410781860352,
      "rel_err": 0.003661602735519409,
      "xla_temp_mb": 65.00048828125,
      "xla_arg_mb": 160.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 0.44089455899666063,
      "inner": 5,
      "median_ms": 10.37721460161265,
      "q1_ms": 10.370808499283157,
      "q3_ms": 10.42071060219314,
      "tflops": 46.35505341454975
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 4096,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.014839410781860352,
      "rel_err": 0.003661602735519409,
      "xla_temp_mb": 65.00048828125,
      "xla_arg_mb": 160.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 0.43307648299378343,
      "inner": 5,
      "median_ms": 10.392487200442702,
      "q1_ms": 10.3371427001548,
      "q3_ms": 10.427435599558521,
      "tflops": 46.286930921744
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 4096,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.014839410781860352,
      "rel_err": 0.003661602735519409,
      "xla_temp_mb": 65.00048828125,
      "xla_arg_mb": 160.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 0.4443560770014301,
      "inner": 4,
      "median_ms": 10.383708748122444,
      "q1_ms": 10.366575250372989,
      "q3_ms": 10.393136624770705,
      "tflops": 46.32606218264546
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 4096,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.014839410781860352,
      "rel_err": 0.003661602735519409,
      "xla_temp_mb": 5120.500015258789,
      "xla_arg_mb": 160.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 1.133409671994741,
      "inner": 1,
      "median_ms": 30.214675003662705,
      "q1_ms": 30.178867498761974,
      "q3_ms": 30.283501000667457,
      "tflops": 15.920619271717715
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 4096,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.014839410781860352,
      "rel_err": 0.003661602735519409,
      "xla_temp_mb": 193.000732421875,
      "xla_arg_mb": 160.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 0.11123300698818639,
      "inner": 5,
      "median_ms": 8.804419601801783,
      "q1_ms": 8.767541099223308,
      "q3_ms": 8.82530249946285,
      "tflops": 54.63578054066826
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 4096,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.014839410781860352,
      "rel_err": 0.0037752636708319187,
      "peak_mb": 561.251953125,
      "compile_s": 0.00812840400612913,
      "inner": 6,
      "median_ms": 8.49538733261094,
      "q1_ms": 8.487247083394323,
      "q3_ms": 8.521812499869458,
      "tflops": 56.62323780170245
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 4096,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.014839410781860352,
      "rel_err": 0.003661602735519409,
      "peak_mb": 541.314453125,
      "compile_s": 0.030414664011914283,
      "inner": 1,
      "median_ms": 30.498952997731976,
      "q1_ms": 30.38497199304402,
      "q3_ms": 30.588998495659325,
      "tflops": 15.77222461334236
    },
    {
      "impl": "torch_math",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 4096,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "error": "OutOfMemoryError: CUDA out of memory. Tried to allocate 2.00 GiB. GPU 0 has a total capacity of 23.61 GiB of which 1.09 GiB is free. Including non-PyTorch memory, this process has 22.48 GiB memory in use. Of the allocated memory 4.39 GiB is allocated by PyTorch, and 45.75 MiB is reserved by PyTorch but unallocated. I"
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 8192,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0003293529152870178,
      "rel_err": 0.0030611662659794092,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 192.0,
      "xla_out_mb": 64.0,
      "compile_s": 0.2571625050040893,
      "inner": 3,
      "median_ms": 16.833417331023764,
      "q1_ms": 16.813184001269597,
      "q3_ms": 16.838228000172723,
      "tflops": 65.31719651182263
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 8192,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0003293529152870178,
      "rel_err": 0.0030611662659794092,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 192.0,
      "xla_out_mb": 64.0,
      "compile_s": 0.2573863680008799,
      "inner": 3,
      "median_ms": 16.802268335595727,
      "q1_ms": 16.755483331508003,
      "q3_ms": 16.86817466543289,
      "tflops": 65.43828522525597
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 8192,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0003293529152870178,
      "rel_err": 0.0030611662659794092,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 192.0,
      "xla_out_mb": 64.0,
      "compile_s": 0.26015664800070226,
      "inner": 2,
      "median_ms": 16.857743001310155,
      "q1_ms": 16.848643503180938,
      "q3_ms": 17.030740000336664,
      "tflops": 65.22294400208544
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 8192,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00033868104219436646,
      "rel_err": 0.0031478663440793753,
      "xla_temp_mb": 12288.0,
      "xla_arg_mb": 192.0,
      "xla_out_mb": 64.0,
      "compile_s": 3.0553781950002303,
      "inner": 1,
      "median_ms": 36.295187994255684,
      "q1_ms": 36.25091249705292,
      "q3_ms": 36.344858999655116,
      "tflops": 30.293592306231226
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 8192,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0003293529152870178,
      "rel_err": 0.0030611662659794092,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 192.0,
      "xla_out_mb": 64.0,
      "compile_s": 0.05659653199836612,
      "inner": 3,
      "median_ms": 15.84210233219589,
      "q1_ms": 15.81686150166206,
      "q3_ms": 15.87908066479334,
      "tflops": 69.40440130483589
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 8192,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0003225579857826233,
      "rel_err": 0.0029980107210576534,
      "peak_mb": 465.2509765625,
      "compile_s": 0.014868710990413092,
      "inner": 3,
      "median_ms": 15.501177331316285,
      "q1_ms": 15.462110333222276,
      "q3_ms": 15.503188165894244,
      "tflops": 70.93084636576019
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 8192,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0003293529152870178,
      "rel_err": 0.0030611662659794092,
      "peak_mb": 464.2509765625,
      "compile_s": 0.022575918002985418,
      "inner": 2,
      "median_ms": 23.501478499383666,
      "q1_ms": 23.494558248785324,
      "q3_ms": 23.512634750659345,
      "tflops": 46.78478538296367
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 8192,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0005331337451934814,
      "rel_err": 0.0036701797507703304,
      "xla_temp_mb": 130.00048828125,
      "xla_arg_mb": 320.0,
      "xla_out_mb": 192.0000228881836,
      "compile_s": 0.4402070590003859,
      "inner": 1,
      "median_ms": 71.22463000996504,
      "q1_ms": 70.9854649976478,
      "q3_ms": 71.78975699935108,
      "tflops": 54.0303360884793
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 8192,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0005331337451934814,
      "rel_err": 0.0036701797507703304,
      "xla_temp_mb": 130.00048828125,
      "xla_arg_mb": 320.0,
      "xla_out_mb": 192.0000228881836,
      "compile_s": 0.4401561739941826,
      "inner": 1,
      "median_ms": 70.5900580069283,
      "q1_ms": 70.42513049964327,
      "q3_ms": 71.20250799925998,
      "tflops": 54.516043843430424
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 8192,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0005331337451934814,
      "rel_err": 0.0036701797507703304,
      "xla_temp_mb": 130.00048828125,
      "xla_arg_mb": 320.0,
      "xla_out_mb": 192.0000228881836,
      "compile_s": 0.43973120000737254,
      "inner": 1,
      "median_ms": 70.93024101050105,
      "q1_ms": 70.2681990005658,
      "q3_ms": 71.64574950002134,
      "tflops": 54.254583692254336
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 8192,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 20.00GiB with allocator GPU_0_bfc on device 0. [tf-allocator-allocation-error=''] [executable_name='jit_loss']"
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 8192,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0005331337451934814,
      "rel_err": 0.0036701797507703304,
      "xla_temp_mb": 386.000732421875,
      "xla_arg_mb": 320.0,
      "xla_out_mb": 192.0000228881836,
      "compile_s": 0.17207006899116095,
      "inner": 1,
      "median_ms": 60.43937000504229,
      "q1_ms": 60.31071650068043,
      "q3_ms": 60.786512498452794,
      "tflops": 63.67191942760071
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 8192,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0004901066422462463,
      "rel_err": 0.0036701797507703304,
      "peak_mb": 1106.251953125,
      "compile_s": 0.05634329300664831,
      "inner": 1,
      "median_ms": 57.19808999856468,
      "q1_ms": 57.14039099984802,
      "q3_ms": 57.21297799755121,
      "tflops": 67.28005598285831
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 8192,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0005476623773574829,
      "rel_err": 0.004770925734192133,
      "peak_mb": 1054.376953125,
      "compile_s": 0.24590957199689,
      "inner": 1,
      "median_ms": 245.5522660020506,
      "q1_ms": 244.36368299939204,
      "q3_ms": 246.28402048983844,
      "tflops": 15.671982017807414
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 8192,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.005632519721984863,
      "rel_err": 0.0019485473167151213,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 192.0,
      "xla_out_mb": 64.0,
      "compile_s": 0.2789986440038774,
      "inner": 4,
      "median_ms": 10.103185999469133,
      "q1_ms": 9.740379624417983,
      "q3_ms": 10.127708375875955,
      "tflops": 54.41410401796885
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 8192,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.005632519721984863,
      "rel_err": 0.0019485473167151213,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 192.0,
      "xla_out_mb": 64.0,
      "compile_s": 0.27592675399500877,
      "inner": 5,
      "median_ms": 9.339768398785964,
      "q1_ms": 9.300252598768566,
      "q3_ms": 9.37362559925532,
      "tflops": 58.86182509187919
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 8192,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.005632519721984863,
      "rel_err": 0.0019485473167151213,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 192.0,
      "xla_out_mb": 64.0,
      "compile_s": 0.27416109399928246,
      "inner": 5,
      "median_ms": 9.348996001062915,
      "q1_ms": 9.296368400100619,
      "q3_ms": 9.394741299911402,
      "tflops": 58.80372756876746
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 8192,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.011533498764038086,
      "rel_err": 0.003989967051893473,
      "xla_temp_mb": 12288.0,
      "xla_arg_mb": 192.0,
      "xla_out_mb": 64.0,
      "compile_s": 0.47986349200073164,
      "inner": 1,
      "median_ms": 36.37802200682927,
      "q1_ms": 36.27666999818757,
      "q3_ms": 36.461616997257806,
      "tflops": 15.112306375118306
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 8192,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.005632519721984863,
      "rel_err": 0.0019485473167151213,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 192.0,
      "xla_out_mb": 64.0,
      "compile_s": 0.049673959001665935,
      "inner": 5,
      "median_ms": 8.526350202737376,
      "q1_ms": 8.419977100857068,
      "q3_ms": 8.594468800583854,
      "tflops": 64.47727348936495
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 8192,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.005632519721984863,
      "rel_err": 0.0019485473167151213,
      "peak_mb": 465.2509765625,
      "compile_s": 0.00787973300612066,
      "inner": 6,
      "median_ms": 8.516669334009444,
      "q1_ms": 8.510515249023836,
      "q3_ms": 8.553238500705143,
      "tflops": 64.5505645842878
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 8192,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.005632519721984863,
      "rel_err": 0.0019485473167151213,
      "peak_mb": 464.2509765625,
      "compile_s": 0.011554473006981425,
      "inner": 4,
      "median_ms": 12.067826497514034,
      "q1_ms": 12.06531550269574,
      "q3_ms": 12.073964377123048,
      "tflops": 45.55549534965965
    },
    {
      "impl": "probjax_mha",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 8192,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.018784761428833008,
      "rel_err": 0.00475410046055913,
      "xla_temp_mb": 130.00048828125,
      "xla_arg_mb": 320.0,
      "xla_out_mb": 192.0000228881836,
      "compile_s": 0.46639296600187663,
      "inner": 1,
      "median_ms": 38.72885800956283,
      "q1_ms": 38.651000504614785,
      "q3_ms": 38.8140950090019,
      "tflops": 49.68247057873215
    },
    {
      "impl": "probjax_flex",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 8192,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.018784761428833008,
      "rel_err": 0.00475410046055913,
      "xla_temp_mb": 130.00048828125,
      "xla_arg_mb": 320.0,
      "xla_out_mb": 192.0000228881836,
      "compile_s": 0.4614889399963431,
      "inner": 1,
      "median_ms": 38.9317730005132,
      "q1_ms": 38.64654700009851,
      "q3_ms": 39.148497497080825,
      "tflops": 49.42352223677651
    },
    {
      "impl": "probjax_default",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 8192,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.018784761428833008,
      "rel_err": 0.00475410046055913,
      "xla_temp_mb": 130.00048828125,
      "xla_arg_mb": 320.0,
      "xla_out_mb": 192.0000228881836,
      "compile_s": 0.46406399000261445,
      "inner": 1,
      "median_ms": 39.019919000566006,
      "q1_ms": 38.952567003434524,
      "q3_ms": 39.09205049421871,
      "tflops": 49.31187449620511
    },
    {
      "impl": "jax_xla",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 8192,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 20.00GiB with allocator GPU_0_bfc on device 0. [executable_name='jit_loss'] [tf-allocator-allocation-error='']"
    },
    {
      "impl": "jax_cudnn",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 8192,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.018784761428833008,
      "rel_err": 0.00475410046055913,
      "xla_temp_mb": 386.000732421875,
      "xla_arg_mb": 320.0,
      "xla_out_mb": 192.0000228881836,
      "compile_s": 0.1453436900046654,
      "inner": 1,
      "median_ms": 32.33725800237153,
      "q1_ms": 32.01779650407843,
      "q3_ms": 33.43917399615748,
      "tflops": 59.50242746205904
    },
    {
      "impl": "torch_flash",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 8192,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.018784761428833008,
      "rel_err": 0.005017808172851801,
      "peak_mb": 1106.251953125,
      "compile_s": 0.028885671999887563,
      "inner": 1,
      "median_ms": 30.505344999255612,
      "q1_ms": 30.48680050414987,
      "q3_ms": 30.537776001438033,
      "tflops": 63.07567898854948
    },
    {
      "impl": "torch_efficient",
      "dtype": "bfloat16",
      "head_dim": 128,
      "seq": 8192,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.018784761428833008,
      "rel_err": 0.006403984036296606,
      "peak_mb": 1054.376953125,
      "compile_s": 0.12453883199486881,
      "inner": 1,
      "median_ms": 123.82710599922575,
      "q1_ms": 123.73014200420585,
      "q3_ms": 123.96762650314486,
      "tflops": 15.538967280879769
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 128,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0003886828199028969,
      "rel_err": 0.0004159725795034319,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.21265564901113976,
      "inner": 200,
      "median_ms": 0.04257061504176818,
      "q1_ms": 0.04235209999023937,
      "q3_ms": 0.042744470010802615,
      "tflops": 3.1528256725516464
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 128,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0003886828199028969,
      "rel_err": 0.0004159725795034319,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.2067072880017804,
      "inner": 200,
      "median_ms": 0.044114685006206855,
      "q1_ms": 0.04343820754002081,
      "q3_ms": 0.045725654999841936,
      "tflops": 3.042472772527239
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 128,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0003886828199028969,
      "rel_err": 0.0004159725795034319,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.21364615199854597,
      "inner": 200,
      "median_ms": 0.04617598497134168,
      "q1_ms": 0.046070110038272105,
      "q3_ms": 0.04644501503207721,
      "tflops": 2.906656524669694
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 128,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00037732720375061035,
      "rel_err": 0.00040381966391578317,
      "xla_temp_mb": 4.0,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.6838444629975129,
      "inner": 200,
      "median_ms": 0.05266956504783593,
      "q1_ms": 0.051428014994598925,
      "q3_ms": 0.05283680500724586,
      "tflops": 2.5482976340909556
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 128,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 1.3709068298339844e-06,
      "rel_err": 1.4671593362436397e-06,
      "peak_mb": 22.2509765625,
      "compile_s": 0.00019310599600430578,
      "inner": 200,
      "median_ms": 0.03299660500488244,
      "q1_ms": 0.03294636248028837,
      "q3_ms": 0.03304064248368377,
      "tflops": 4.0676223502430044
    },
    {
      "impl": "torch_math",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 128,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 3.5762786865234375e-07,
      "rel_err": 3.827372268006002e-07,
      "peak_mb": 26.75439453125,
      "compile_s": 0.0009713170002214611,
      "inner": 200,
      "median_ms": 0.1299583899526624,
      "q1_ms": 0.12913969498185907,
      "q3_ms": 0.13035075753577985,
      "tflops": 1.0327746292400906
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 128,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0005737245082855225,
      "rel_err": 0.0007043913938105106,
      "xla_temp_mb": 1.031494140625,
      "xla_arg_mb": 4.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.33917160899727605,
      "inner": 200,
      "median_ms": 0.04332749500463251,
      "q1_ms": 0.043264452506264206,
      "q3_ms": 0.047367662518809084,
      "tflops": 10.842123412622255
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 128,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0005737245082855225,
      "rel_err": 0.0007043913938105106,
      "xla_temp_mb": 1.031494140625,
      "xla_arg_mb": 4.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.33726287000172306,
      "inner": 200,
      "median_ms": 0.04308474002755247,
      "q1_ms": 0.0428991375156329,
      "q3_ms": 0.04333628497988684,
      "tflops": 10.90321184947593
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 128,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0005737245082855225,
      "rel_err": 0.0007043913938105106,
      "xla_temp_mb": 1.031494140625,
      "xla_arg_mb": 4.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.33836018499278,
      "inner": 200,
      "median_ms": 0.05904557496251073,
      "q1_ms": 0.056200542458100244,
      "q3_ms": 0.05921872998442268,
      "tflops": 7.95592300182125
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 128,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.000576406717300415,
      "rel_err": 0.0007043913938105106,
      "xla_temp_mb": 6.0,
      "xla_arg_mb": 4.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 1.1281270410108846,
      "inner": 200,
      "median_ms": 0.0735492049716413,
      "q1_ms": 0.07354284251050558,
      "q3_ms": 0.07357595499343006,
      "tflops": 6.387044539517841
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 128,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 1.7881393432617188e-06,
      "rel_err": 2.1446617211040575e-06,
      "peak_mb": 31.283203125,
      "compile_s": 0.0005615730042336509,
      "inner": 167,
      "median_ms": 0.23216574251007646,
      "q1_ms": 0.2291929401132462,
      "q3_ms": 0.23520106285382453,
      "tflops": 2.0233908884279574
    },
    {
      "impl": "torch_math",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 128,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 4.76837158203125e-07,
      "rel_err": 5.916307941333798e-07,
      "peak_mb": 35.2509765625,
      "compile_s": 0.0007446679956046864,
      "inner": 86,
      "median_ms": 0.5166198720554496,
      "q1_ms": 0.5005966686016912,
      "q3_ms": 0.5195795407459183,
      "tflops": 0.9092992225230928
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 128,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0009524822235107422,
      "rel_err": 0.0003142662171740085,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.22950887300248723,
      "inner": 200,
      "median_ms": 0.03722526002093218,
      "q1_ms": 0.037032212494523264,
      "q3_ms": 0.04062503496243153,
      "tflops": 1.8027775752879611
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 128,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0009524822235107422,
      "rel_err": 0.0003142662171740085,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.22462177999841515,
      "inner": 200,
      "median_ms": 0.03702802998304833,
      "q1_ms": 0.03684360999614,
      "q3_ms": 0.03710708249855088,
      "tflops": 1.8123800815415476
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 128,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0009524822235107422,
      "rel_err": 0.0003142662171740085,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.22472976399876643,
      "inner": 200,
      "median_ms": 0.03785584500292316,
      "q1_ms": 0.037511044974962715,
      "q3_ms": 0.042178875009994954,
      "tflops": 1.7727477485925351
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 128,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0013991594314575195,
      "rel_err": 0.0004616449004970491,
      "xla_temp_mb": 4.0,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.5765544340101769,
      "inner": 200,
      "median_ms": 0.037589740022667684,
      "q1_ms": 0.037429914955282584,
      "q3_ms": 0.03944177249650238,
      "tflops": 1.7852973699613628
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 128,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 1.3709068298339844e-06,
      "rel_err": 4.523230927588884e-07,
      "peak_mb": 22.2509765625,
      "compile_s": 8.566200267523527e-05,
      "inner": 200,
      "median_ms": 0.032793974969536066,
      "q1_ms": 0.03277664003690006,
      "q3_ms": 0.03281406254245667,
      "tflops": 2.0463778502709937
    },
    {
      "impl": "torch_math",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 128,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 5.960464477539062e-07,
      "rel_err": 1.966622136251317e-07,
      "peak_mb": 26.81689453125,
      "compile_s": 0.0002708730025915429,
      "inner": 200,
      "median_ms": 0.17106939500081353,
      "q1_ms": 0.16892669002118055,
      "q3_ms": 0.17207803248311393,
      "tflops": 0.39229029832998974
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 128,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0016772747039794922,
      "rel_err": 0.0007810532115399837,
      "xla_temp_mb": 1.031494140625,
      "xla_arg_mb": 4.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.37092014600057155,
      "inner": 200,
      "median_ms": 0.046315244981087744,
      "q1_ms": 0.04602470002282644,
      "q3_ms": 0.04838288248720346,
      "tflops": 5.071354455663805
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 128,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0016772747039794922,
      "rel_err": 0.0007810532115399837,
      "xla_temp_mb": 1.031494140625,
      "xla_arg_mb": 4.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.36311303000547923,
      "inner": 200,
      "median_ms": 0.058056554989889264,
      "q1_ms": 0.05635077497572638,
      "q3_ms": 0.05822291750519071,
      "tflops": 4.0457278948243705
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 128,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0016772747039794922,
      "rel_err": 0.0007810532115399837,
      "xla_temp_mb": 1.031494140625,
      "xla_arg_mb": 4.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.36356123800214846,
      "inner": 200,
      "median_ms": 0.04853368998738006,
      "q1_ms": 0.04826671502087265,
      "q3_ms": 0.04884901252808049,
      "tflops": 4.839545974375218
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 128,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0016138553619384766,
      "rel_err": 0.0007290384965017438,
      "xla_temp_mb": 6.0,
      "xla_arg_mb": 4.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 1.2318408460123464,
      "inner": 200,
      "median_ms": 0.08779739997407887,
      "q1_ms": 0.08777240498602623,
      "q3_ms": 0.08781721502600703,
      "tflops": 2.6752617283580817
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 128,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 3.2186508178710938e-06,
      "rel_err": 1.4988227121648379e-06,
      "peak_mb": 31.283203125,
      "compile_s": 0.0004948459973093122,
      "inner": 172,
      "median_ms": 0.24087376740956032,
      "q1_ms": 0.23265700288310506,
      "q3_ms": 0.24653650866956775,
      "tflops": 0.9751208133869936
    },
    {
      "impl": "torch_math",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 128,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 1.430511474609375e-06,
      "rel_err": 6.661434213128814e-07,
      "peak_mb": 35.2509765625,
      "compile_s": 0.001454469995223917,
      "inner": 79,
      "median_ms": 0.5520059114846672,
      "q1_ms": 0.5455462278059944,
      "q3_ms": 0.5686864493741901,
      "tflops": 0.425504544631175
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 512,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0003453493118286133,
      "rel_err": 0.0008412703173235059,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.21320861400454305,
      "inner": 200,
      "median_ms": 0.10720996499003377,
      "q1_ms": 0.10712157247326104,
      "q3_ms": 0.1072280525113456,
      "tflops": 20.030634728773858
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 512,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0003453493118286133,
      "rel_err": 0.0008412703173235059,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.2131869139993796,
      "inner": 200,
      "median_ms": 0.10243072501907591,
      "q1_ms": 0.099759200020344,
      "q3_ms": 0.10243871747661615,
      "tflops": 20.96522940358051
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 512,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0003453493118286133,
      "rel_err": 0.0008412703173235059,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.20637159299803898,
      "inner": 188,
      "median_ms": 0.10241438831918516,
      "q1_ms": 0.09850020479519872,
      "q3_ms": 0.10242768613797454,
      "tflops": 20.968573686220168
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 512,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00041928887367248535,
      "rel_err": 0.0010213869391009212,
      "xla_temp_mb": 64.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.8759152599959634,
      "inner": 150,
      "median_ms": 0.23145247997793678,
      "q1_ms": 0.2311890166796123,
      "q3_ms": 0.23156275662283105,
      "tflops": 9.278291804022619
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 512,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 6.109476089477539e-07,
      "rel_err": 1.4882673440297367e-06,
      "peak_mb": 40.2509765625,
      "compile_s": 0.0003580670017981902,
      "inner": 145,
      "median_ms": 0.3196569172673505,
      "q1_ms": 0.3196280000024828,
      "q3_ms": 0.3197921068680569,
      "tflops": 6.718089088633472
    },
    {
      "impl": "torch_math",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 512,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 1.7136335372924805e-07,
      "rel_err": 4.1744084455785924e-07,
      "peak_mb": 112.26611328125,
      "compile_s": 0.0005311539862304926,
      "inner": 101,
      "median_ms": 0.4674951683633066,
      "q1_ms": 0.46725045046181246,
      "q3_ms": 0.4679418613731926,
      "tflops": 4.593595385206455
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 512,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00036150217056274414,
      "rel_err": 0.0007085130200721323,
      "xla_temp_mb": 4.125244140625,
      "xla_arg_mb": 16.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.3362821049959166,
      "inner": 93,
      "median_ms": 0.4444662151390046,
      "q1_ms": 0.4432608170988118,
      "q3_ms": 0.4446522688858652,
      "tflops": 16.910605377844856
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 512,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00036150217056274414,
      "rel_err": 0.0007085130200721323,
      "xla_temp_mb": 4.125244140625,
      "xla_arg_mb": 16.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.33599544200114906,
      "inner": 90,
      "median_ms": 0.444501688859115,
      "q1_ms": 0.4443326444339214,
      "q3_ms": 0.44465982775566065,
      "tflops": 16.909255816983546
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 512,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00036150217056274414,
      "rel_err": 0.0007085130200721323,
      "xla_temp_mb": 4.125244140625,
      "xla_arg_mb": 16.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.3427549070038367,
      "inner": 89,
      "median_ms": 0.4445827527500145,
      "q1_ms": 0.4444704606888418,
      "q3_ms": 0.44469881465650696,
      "tflops": 16.906172633796025
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 512,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0003629922866821289,
      "rel_err": 0.0007088001002557576,
      "xla_temp_mb": 96.0625,
      "xla_arg_mb": 16.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 1.537614590008161,
      "inner": 67,
      "median_ms": 0.6071433732060215,
      "q1_ms": 0.6065722165061092,
      "q3_ms": 0.607299410460913,
      "tflops": 12.379601095389928
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 512,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 1.430511474609375e-06,
      "rel_err": 2.4725475213926984e-06,
      "peak_mb": 72.439453125,
      "compile_s": 0.001320937997661531,
      "inner": 38,
      "median_ms": 1.2750726053031702,
      "q1_ms": 1.27051597366134,
      "q3_ms": 1.27629863173321,
      "tflops": 5.894717474706389
    },
    {
      "impl": "torch_math",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 512,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 7.748603820800781e-07,
      "rel_err": 1.329966380581027e-06,
      "peak_mb": 188.2509765625,
      "compile_s": 0.0012441430008038878,
      "inner": 40,
      "median_ms": 1.2288188499951502,
      "q1_ms": 1.2278010373847792,
      "q3_ms": 1.2289402999158483,
      "tflops": 6.116599503685726
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 512,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0012367963790893555,
      "rel_err": 0.0004327103088144213,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.22639365200302564,
      "inner": 200,
      "median_ms": 0.08214736502850428,
      "q1_ms": 0.08211440250306623,
      "q3_ms": 0.0821560800250154,
      "tflops": 13.07092228250319
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 512,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0012367963790893555,
      "rel_err": 0.0004327103088144213,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.22368636100145523,
      "inner": 200,
      "median_ms": 0.07542965999164153,
      "q1_ms": 0.07413353501760866,
      "q3_ms": 0.07545884251157986,
      "tflops": 14.235008140285705
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 512,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0012367963790893555,
      "rel_err": 0.0004327103088144213,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.2278052710025804,
      "inner": 200,
      "median_ms": 0.07522853004047647,
      "q1_ms": 0.07449181252013659,
      "q3_ms": 0.07527191250119358,
      "tflops": 14.273066659979621
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 512,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0017076730728149414,
      "rel_err": 0.0005974530358798802,
      "xla_temp_mb": 64.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.8194676889979746,
      "inner": 132,
      "median_ms": 0.2662100757306413,
      "q1_ms": 0.26588692421545135,
      "q3_ms": 0.2663379697575389,
      "tflops": 4.033437957045779
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 512,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 1.0281801223754883e-06,
      "rel_err": 3.5972303180642484e-07,
      "peak_mb": 40.2509765625,
      "compile_s": 0.000272245000815019,
      "inner": 200,
      "median_ms": 0.21453224500874057,
      "q1_ms": 0.2134428124918486,
      "q3_ms": 0.2151086875164765,
      "tflops": 5.005037000178006
    },
    {
      "impl": "torch_math",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 512,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 2.980232238769531e-07,
      "rel_err": 1.0426754215586698e-07,
      "peak_mb": 113.26611328125,
      "compile_s": 0.0006697770004393533,
      "inner": 80,
      "median_ms": 0.5614190749838599,
      "q1_ms": 0.5608662188024027,
      "q3_ms": 0.5614941437670495,
      "tflops": 1.9125495941349495
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 512,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.002021193504333496,
      "rel_err": 0.000897872494533658,
      "xla_temp_mb": 4.125244140625,
      "xla_arg_mb": 16.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.3701616630132776,
      "inner": 127,
      "median_ms": 0.2700578267585237,
      "q1_ms": 0.2695413779120848,
      "q3_ms": 0.270126066933749,
      "tflops": 13.915895084797372
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 512,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.002021193504333496,
      "rel_err": 0.000897872494533658,
      "xla_temp_mb": 4.125244140625,
      "xla_arg_mb": 16.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.3659984599944437,
      "inner": 126,
      "median_ms": 0.269770666602112,
      "q1_ms": 0.2681885039598297,
      "q3_ms": 0.27023173803057055,
      "tflops": 13.93070800222643
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 512,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.002021193504333496,
      "rel_err": 0.000897872494533658,
      "xla_temp_mb": 4.125244140625,
      "xla_arg_mb": 16.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.3666031859902432,
      "inner": 119,
      "median_ms": 0.27033357153169246,
      "q1_ms": 0.26942525632990555,
      "q3_ms": 0.27036733613271235,
      "tflops": 13.901700638610551
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 512,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0012693405151367188,
      "rel_err": 0.0006988496752455831,
      "xla_temp_mb": 96.0625,
      "xla_arg_mb": 16.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.9675974850106286,
      "inner": 68,
      "median_ms": 0.6265744116664043,
      "q1_ms": 0.626501691145035,
      "q3_ms": 0.627105992701699,
      "tflops": 5.997845290242806
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 512,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 3.0994415283203125e-06,
      "rel_err": 1.217992689817038e-06,
      "peak_mb": 72.439453125,
      "compile_s": 0.0008891030011000112,
      "inner": 57,
      "median_ms": 0.8111752982235008,
      "q1_ms": 0.806020114013288,
      "q3_ms": 0.8213682981716517,
      "tflops": 4.632902890694956
    },
    {
      "impl": "torch_math",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 512,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 2.9802322387695312e-06,
      "rel_err": 1.323905166827899e-06,
      "peak_mb": 188.2509765625,
      "compile_s": 0.0013834699930157512,
      "inner": 37,
      "median_ms": 1.3111194326680757,
      "q1_ms": 1.3105172297769427,
      "q3_ms": 1.3115462027820508,
      "tflops": 2.8663265072293402
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 600,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0002467334270477295,
      "rel_err": 0.0005619395524263382,
      "xla_temp_mb": 20.0,
      "xla_arg_mb": 14.0625,
      "xla_out_mb": 4.6875,
      "compile_s": 0.2357264029997168,
      "inner": 167,
      "median_ms": 0.17447762871314465,
      "q1_ms": 0.17373373951208024,
      "q3_ms": 0.17519821559251872,
      "tflops": 16.90256809283322
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 600,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0002467334270477295,
      "rel_err": 0.0005619395524263382,
      "xla_temp_mb": 20.0,
      "xla_arg_mb": 14.0625,
      "xla_out_mb": 4.6875,
      "compile_s": 0.23858948399720248,
      "inner": 175,
      "median_ms": 0.1749057085752221,
      "q1_ms": 0.17391505713541328,
      "q3_ms": 0.1750922342762351,
      "tflops": 16.861199237140195
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 600,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0002467334270477295,
      "rel_err": 0.0005619395524263382,
      "xla_temp_mb": 20.0,
      "xla_arg_mb": 14.0625,
      "xla_out_mb": 4.6875,
      "compile_s": 0.23349407500063535,
      "inner": 168,
      "median_ms": 0.17502749402150844,
      "q1_ms": 0.17395385414905226,
      "q3_ms": 0.1750403839434425,
      "tflops": 16.849467087939875
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 600,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0003236532211303711,
      "rel_err": 0.0007371256360784173,
      "xla_temp_mb": 87.890625,
      "xla_arg_mb": 14.0625,
      "xla_out_mb": 4.6875,
      "compile_s": 0.8592543809936615,
      "inner": 119,
      "median_ms": 0.31720157142849875,
      "q1_ms": 0.31693062607830924,
      "q3_ms": 0.3174438319577906,
      "tflops": 9.297305769069206
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 600,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 8.642673492431641e-07,
      "rel_err": 1.9683834580064286e-06,
      "peak_mb": 44.8759765625,
      "compile_s": 0.0004487799888011068,
      "inner": 114,
      "median_ms": 0.41814757896007215,
      "q1_ms": 0.4180607017923122,
      "q3_ms": 0.41828283773563607,
      "tflops": 7.052820937847888
    },
    {
      "impl": "torch_math",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 600,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 3.5762786865234375e-07,
      "rel_err": 8.145034939843754e-07,
      "peak_mb": 143.88037109375,
      "compile_s": 0.0006559820030815899,
      "inner": 78,
      "median_ms": 0.6122325126526471,
      "q1_ms": 0.6119305833589095,
      "q3_ms": 0.6128870064947898,
      "tflops": 4.816993444569312
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 600,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0005118250846862793,
      "rel_err": 0.0005855912459082901,
      "xla_temp_mb": 40.00048828125,
      "xla_arg_mb": 18.75,
      "xla_out_mb": 14.062522888183594,
      "compile_s": 0.3871121299889637,
      "inner": 66,
      "median_ms": 0.6968209848061881,
      "q1_ms": 0.6967800832309435,
      "q3_ms": 0.6971991288909604,
      "tflops": 14.812871921288233
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 600,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0005118250846862793,
      "rel_err": 0.0005855912459082901,
      "xla_temp_mb": 40.00048828125,
      "xla_arg_mb": 18.75,
      "xla_out_mb": 14.062522888183594,
      "compile_s": 0.3850838000071235,
      "inner": 66,
      "median_ms": 0.6769286364115858,
      "q1_ms": 0.6764724545946317,
      "q3_ms": 0.6772246363680372,
      "tflops": 15.248165677724513
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 600,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0005118250846862793,
      "rel_err": 0.0005855912459082901,
      "xla_temp_mb": 40.00048828125,
      "xla_arg_mb": 18.75,
      "xla_out_mb": 14.062522888183594,
      "compile_s": 0.3798480149998795,
      "inner": 65,
      "median_ms": 0.6822998308612463,
      "q1_ms": 0.6820840999493016,
      "q3_ms": 0.682394307715675,
      "tflops": 15.128129208197155
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 600,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.000516057014465332,
      "rel_err": 0.0005977088585495949,
      "xla_temp_mb": 131.9091796875,
      "xla_arg_mb": 18.75,
      "xla_out_mb": 14.062522888183594,
      "compile_s": 1.332710237998981,
      "inner": 55,
      "median_ms": 0.8207869456699965,
      "q1_ms": 0.8203847181655213,
      "q3_ms": 0.8209529727041214,
      "tflops": 12.575638604454614
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 600,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 1.6093254089355469e-06,
      "rel_err": 2.3016341401671525e-06,
      "peak_mb": 82.841796875,
      "compile_s": 0.001649493002332747,
      "inner": 30,
      "median_ms": 1.6420703667487637,
      "q1_ms": 1.641831400047522,
      "q3_ms": 1.6423896498357256,
      "tflops": 6.285918197547773
    },
    {
      "impl": "torch_math",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 600,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 6.556510925292969e-07,
      "rel_err": 9.377027936352533e-07,
      "peak_mb": 244.3134765625,
      "compile_s": 0.0016028949903557077,
      "inner": 32,
      "median_ms": 1.5774532812429243,
      "q1_ms": 1.5720452031473542,
      "q3_ms": 1.5810279999186605,
      "tflops": 6.543407733677563
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 600,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.001040816307067871,
      "rel_err": 0.00044011077261529863,
      "xla_temp_mb": 20.0,
      "xla_arg_mb": 14.0625,
      "xla_out_mb": 4.6875,
      "compile_s": 0.24329611600842327,
      "inner": 186,
      "median_ms": 0.14761737095610167,
      "q1_ms": 0.1468443225612939,
      "q3_ms": 0.14784945158935053,
      "tflops": 9.989068294940054
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 600,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.001040816307067871,
      "rel_err": 0.00044011077261529863,
      "xla_temp_mb": 20.0,
      "xla_arg_mb": 14.0625,
      "xla_out_mb": 4.6875,
      "compile_s": 0.24165892601013184,
      "inner": 182,
      "median_ms": 0.1490082032722677,
      "q1_ms": 0.14853843682570955,
      "q3_ms": 0.14904579670605977,
      "tflops": 9.895831018817702
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 600,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.001040816307067871,
      "rel_err": 0.00044011077261529863,
      "xla_temp_mb": 20.0,
      "xla_arg_mb": 14.0625,
      "xla_out_mb": 4.6875,
      "compile_s": 0.23940845399920363,
      "inner": 165,
      "median_ms": 0.14889581820391343,
      "q1_ms": 0.1483027605872573,
      "q3_ms": 0.14918182114894135,
      "tflops": 9.903300292695823
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 600,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.000980377197265625,
      "rel_err": 0.0004145540005993098,
      "xla_temp_mb": 87.890625,
      "xla_arg_mb": 14.0625,
      "xla_out_mb": 4.6875,
      "compile_s": 0.645204514003126,
      "inner": 113,
      "median_ms": 0.35102974324161595,
      "q1_ms": 0.3507272610147236,
      "q3_ms": 0.35103568145325326,
      "tflops": 4.200669682241288
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 600,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 1.3113021850585938e-06,
      "rel_err": 5.544861210182717e-07,
      "peak_mb": 44.8759765625,
      "compile_s": 0.0003440220025368035,
      "inner": 144,
      "median_ms": 0.29811957635198877,
      "q1_ms": 0.2976671214835531,
      "q3_ms": 0.2982948402758565,
      "tflops": 4.9462031914972
    },
    {
      "impl": "torch_math",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 600,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 3.2782554626464844e-07,
      "rel_err": 1.3862153025456792e-07,
      "peak_mb": 145.88037109375,
      "compile_s": 0.0008181300072465092,
      "inner": 62,
      "median_ms": 0.7418983870828074,
      "q1_ms": 0.7416112259248151,
      "q3_ms": 0.7419441129557878,
      "tflops": 1.9875498123106392
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 600,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0018803179264068604,
      "rel_err": 0.0007670256309211254,
      "xla_temp_mb": 40.00048828125,
      "xla_arg_mb": 18.75,
      "xla_out_mb": 14.062522888183594,
      "compile_s": 0.38066726599936374,
      "inner": 85,
      "median_ms": 0.4486361059688908,
      "q1_ms": 0.4482137765187551,
      "q3_ms": 0.4490882587988916,
      "tflops": 11.503666181423814
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 600,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0018803179264068604,
      "rel_err": 0.0007670256309211254,
      "xla_temp_mb": 40.00048828125,
      "xla_arg_mb": 18.75,
      "xla_out_mb": 14.062522888183594,
      "compile_s": 0.3865380829956848,
      "inner": 92,
      "median_ms": 0.4377768913716706,
      "q1_ms": 0.4375422662554268,
      "q3_ms": 0.4381072499793094,
      "tflops": 11.789018794092465
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 600,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0018803179264068604,
      "rel_err": 0.0007670256309211254,
      "xla_temp_mb": 40.00048828125,
      "xla_arg_mb": 18.75,
      "xla_out_mb": 14.062522888183594,
      "compile_s": 0.3812528249982279,
      "inner": 96,
      "median_ms": 0.43732664577570785,
      "q1_ms": 0.43725259888560686,
      "q3_ms": 0.43833000518134213,
      "tflops": 11.801156069156844
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 600,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.001466512680053711,
      "rel_err": 0.0005982248112559319,
      "xla_temp_mb": 131.9091796875,
      "xla_arg_mb": 18.75,
      "xla_out_mb": 14.062522888183594,
      "compile_s": 1.0896496680070413,
      "inner": 51,
      "median_ms": 0.8659510979972159,
      "q1_ms": 0.8655350098064553,
      "q3_ms": 0.8673808726027845,
      "tflops": 5.959874653356687
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 600,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 4.0531158447265625e-06,
      "rel_err": 1.3615911029773997e-06,
      "peak_mb": 82.841796875,
      "compile_s": 0.0011706369987223297,
      "inner": 42,
      "median_ms": 1.1199208571010135,
      "q1_ms": 1.1172726545261131,
      "q3_ms": 1.1205665117207295,
      "tflops": 4.6083256395094505
    },
    {
      "impl": "torch_math",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 600,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 2.5033950805664062e-06,
      "rel_err": 1.0211933840764686e-06,
      "peak_mb": 244.3134765625,
      "compile_s": 0.0017189649952342734,
      "inner": 29,
      "median_ms": 1.689127448179502,
      "q1_ms": 1.6873938105257773,
      "q3_ms": 1.6893041382868903,
      "tflops": 3.0553999969406394
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 1024,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00022730231285095215,
      "rel_err": 0.000759284826926887,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.21467783499974757,
      "inner": 114,
      "median_ms": 0.32791285086300614,
      "q1_ms": 0.32666321050400204,
      "q3_ms": 0.3289533772004417,
      "tflops": 26.195785158748354
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 1024,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00022730231285095215,
      "rel_err": 0.000759284826926887,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.21216886398906354,
      "inner": 113,
      "median_ms": 0.3314448673271912,
      "q1_ms": 0.3300826858514837,
      "q3_ms": 0.3320713141424153,
      "tflops": 25.91663181050351
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 1024,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00022730231285095215,
      "rel_err": 0.000759284826926887,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.20813366401125677,
      "inner": 109,
      "median_ms": 0.3275998715586841,
      "q1_ms": 0.3274711377068568,
      "q3_ms": 0.3285278210928023,
      "tflops": 26.22081184320994
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 1024,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0003216862678527832,
      "rel_err": 0.001074566738680005,
      "xla_temp_mb": 256.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.8875710099964635,
      "inner": 56,
      "median_ms": 0.7988487320419934,
      "q1_ms": 0.7985318035415341,
      "q3_ms": 0.7993677589703501,
      "tflops": 10.752892565833665
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 1024,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 7.152557373046875e-07,
      "rel_err": 2.3892532681202283e-06,
      "peak_mb": 64.2509765625,
      "compile_s": 0.001094923005439341,
      "inner": 45,
      "median_ms": 1.1183783777394436,
      "q1_ms": 1.1176867334547245,
      "q3_ms": 1.1188160999962646,
      "tflops": 7.680705173648535
    },
    {
      "impl": "torch_math",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 1024,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 3.2782554626464844e-07,
      "rel_err": 1.0950744808724266e-06,
      "peak_mb": 352.28173828125,
      "compile_s": 0.0016999189974740148,
      "inner": 30,
      "median_ms": 1.6993244328962949,
      "q1_ms": 1.696934416637911,
      "q3_ms": 1.6996762666773673,
      "tflops": 5.054911484653631
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 1024,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0003542304039001465,
      "rel_err": 0.000781299255322665,
      "xla_temp_mb": 0.25048828125,
      "xla_arg_mb": 32.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.3803192479972495,
      "inner": 33,
      "median_ms": 1.390973272712224,
      "q1_ms": 1.3902809241148077,
      "q3_ms": 1.3938173182065259,
      "tflops": 21.61419752758977
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 1024,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0003542304039001465,
      "rel_err": 0.0007813428528606892,
      "xla_temp_mb": 0.25048828125,
      "xla_arg_mb": 32.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.3759245670080418,
      "inner": 33,
      "median_ms": 1.394738000023856,
      "q1_ms": 1.3944138786014677,
      "q3_ms": 1.3963710909151952,
      "tflops": 21.55585570299638
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 1024,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0003542304039001465,
      "rel_err": 0.000781299255322665,
      "xla_temp_mb": 0.25048828125,
      "xla_arg_mb": 32.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.3759383719880134,
      "inner": 33,
      "median_ms": 1.3906460305654698,
      "q1_ms": 1.3893128484734474,
      "q3_ms": 1.3985210605438403,
      "tflops": 21.619283707856948
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 1024,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00035199522972106934,
      "rel_err": 0.0007834789576008916,
      "xla_temp_mb": 384.0,
      "xla_arg_mb": 32.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 1.4288771059946157,
      "inner": 24,
      "median_ms": 1.9408087500778493,
      "q1_ms": 1.9404247498944944,
      "q3_ms": 1.941205583534611,
      "tflops": 15.490846829083004
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 1024,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 8.940696716308594e-07,
      "rel_err": 2.435739816064597e-06,
      "peak_mb": 128.626953125,
      "compile_s": 0.004293339006835595,
      "inner": 11,
      "median_ms": 4.396124271972274,
      "q1_ms": 4.390054181302814,
      "q3_ms": 4.3972692722243005,
      "tflops": 6.838926566220969
    },
    {
      "impl": "torch_math",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 1024,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 5.364418029785156e-07,
      "rel_err": 1.5693992736487417e-06,
      "peak_mb": 616.2509765625,
      "compile_s": 0.004340629995567724,
      "inner": 11,
      "median_ms": 4.488259727209384,
      "q1_ms": 4.487795636206019,
      "q3_ms": 4.488758818214675,
      "tflops": 6.698536381425735
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 1024,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0009534358978271484,
      "rel_err": 0.0003247723507229239,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.22847220599942375,
      "inner": 148,
      "median_ms": 0.21196386483466453,
      "q1_ms": 0.21102743242384325,
      "q3_ms": 0.21255075675120727,
      "tflops": 20.26273345860224
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 1024,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0009534358978271484,
      "rel_err": 0.0003247723507229239,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.22476965399982873,
      "inner": 165,
      "median_ms": 0.2122727393835896,
      "q1_ms": 0.21218633637269677,
      "q3_ms": 0.21234171815167152,
      "tflops": 20.233249490593966
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 1024,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0009534358978271484,
      "rel_err": 0.0003247723507229239,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.2263263719942188,
      "inner": 155,
      "median_ms": 0.21262229025934734,
      "q1_ms": 0.2119273935433387,
      "q3_ms": 0.21299609352060925,
      "tflops": 20.199986044554343
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 1024,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.001054137945175171,
      "rel_err": 0.00035907485289499164,
      "xla_temp_mb": 256.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.6098749960074201,
      "inner": 52,
      "median_ms": 0.8183131730430222,
      "q1_ms": 0.8167301826828266,
      "q3_ms": 0.8185042307349459,
      "tflops": 5.248561843418101
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 1024,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 1.0728836059570312e-06,
      "rel_err": 3.654602664937556e-07,
      "peak_mb": 64.2509765625,
      "compile_s": 0.0007004560029599816,
      "inner": 74,
      "median_ms": 0.6618184594805604,
      "q1_ms": 0.6597636283667934,
      "q3_ms": 0.6625161419120755,
      "tflops": 6.489645664116077
    },
    {
      "impl": "torch_math",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 1024,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 4.172325134277344e-07,
      "rel_err": 1.4212344012776157e-07,
      "peak_mb": 356.28173828125,
      "compile_s": 0.0022083420044509694,
      "inner": 23,
      "median_ms": 2.170321826428256,
      "q1_ms": 2.164558086761921,
      "q3_ms": 2.171197391441598,
      "tflops": 1.9789541088790124
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 1024,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0013266801834106445,
      "rel_err": 0.0005563136655837297,
      "xla_temp_mb": 0.25048828125,
      "xla_arg_mb": 32.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.4120323189999908,
      "inner": 55,
      "median_ms": 0.7808066907190633,
      "q1_ms": 0.7803447090819563,
      "q3_ms": 0.7814380727888255,
      "tflops": 19.252377976111244
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 1024,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0013266801834106445,
      "rel_err": 0.0005563136655837297,
      "xla_temp_mb": 0.25048828125,
      "xla_arg_mb": 32.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.4141022179974243,
      "inner": 50,
      "median_ms": 0.7803610598784871,
      "q1_ms": 0.7800053899700288,
      "q3_ms": 0.7810947600228246,
      "tflops": 19.26337218612721
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 1024,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0013266801834106445,
      "rel_err": 0.0005563136655837297,
      "xla_temp_mb": 0.25048828125,
      "xla_arg_mb": 32.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.41385938800522126,
      "inner": 56,
      "median_ms": 0.7805469465113545,
      "q1_ms": 0.7795729642176801,
      "q3_ms": 0.7810269643024574,
      "tflops": 19.258784629402594
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 1024,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0013266801834106445,
      "rel_err": 0.0004979501245543361,
      "xla_temp_mb": 384.0,
      "xla_arg_mb": 32.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 1.369009874004405,
      "inner": 22,
      "median_ms": 2.130450045197821,
      "q1_ms": 2.1301588180490816,
      "q3_ms": 2.131905545618131,
      "tflops": 7.055967151111577
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 1024,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 3.6954879760742188e-06,
      "rel_err": 1.8905637944044429e-06,
      "peak_mb": 128.626953125,
      "compile_s": 0.002500204005627893,
      "inner": 19,
      "median_ms": 2.521234579035081,
      "q1_ms": 2.5175102369087794,
      "q3_ms": 2.5262309473442945,
      "tflops": 5.962311345798353
    },
    {
      "impl": "torch_math",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 1024,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 5.245208740234375e-06,
      "rel_err": 1.8295778545507346e-06,
      "peak_mb": 616.2509765625,
      "compile_s": 0.004822420989512466,
      "inner": 10,
      "median_ms": 4.916563800361473,
      "q1_ms": 4.911149050167296,
      "q3_ms": 4.916686049546115,
      "tflops": 3.0574983151636914
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 1500,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00014296174049377441,
      "rel_err": 0.0005673671257682145,
      "xla_temp_mb": 48.0,
      "xla_arg_mb": 35.15625,
      "xla_out_mb": 11.71875,
      "compile_s": 0.23751501900551375,
      "inner": 56,
      "median_ms": 0.8025505714093535,
      "q1_ms": 0.8020012232269178,
      "q3_ms": 0.802627419683891,
      "tflops": 22.966776994042498
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 1500,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00014296174049377441,
      "rel_err": 0.0005673671257682145,
      "xla_temp_mb": 48.0,
      "xla_arg_mb": 35.15625,
      "xla_out_mb": 11.71875,
      "compile_s": 0.24043760300264694,
      "inner": 58,
      "median_ms": 0.8038512758754097,
      "q1_ms": 0.8036776034087183,
      "q3_ms": 0.8044805777146771,
      "tflops": 22.929614660283015
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 1500,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00014296174049377441,
      "rel_err": 0.0005673671257682145,
      "xla_temp_mb": 48.0,
      "xla_arg_mb": 35.15625,
      "xla_out_mb": 11.71875,
      "compile_s": 0.24320790900674183,
      "inner": 58,
      "median_ms": 0.8047828618898699,
      "q1_ms": 0.8037331293774608,
      "q3_ms": 0.8049841206480652,
      "tflops": 22.903072210951624
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 1500,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00019609928131103516,
      "rel_err": 0.0007782522006891668,
      "xla_temp_mb": 549.3164215087891,
      "xla_arg_mb": 35.15625,
      "xla_out_mb": 11.71875,
      "compile_s": 0.8805203800002346,
      "inner": 28,
      "median_ms": 1.6623685358873834,
      "q1_ms": 1.6618255358480383,
      "q3_ms": 1.6623828037804094,
      "tflops": 11.08779407338871
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 1500,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 5.21540641784668e-07,
      "rel_err": 2.0698196294688387e-06,
      "peak_mb": 87.2822265625,
      "compile_s": 0.002307880000444129,
      "inner": 22,
      "median_ms": 2.3308554545193063,
      "q1_ms": 2.327813773180506,
      "q3_ms": 2.3371962729883804,
      "tflops": 7.9078262722221195
    },
    {
      "impl": "torch_math",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 1500,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 4.023313522338867e-07,
      "rel_err": 1.5967179933795705e-06,
      "peak_mb": 705.30859375,
      "compile_s": 0.0033961809967877343,
      "inner": 14,
      "median_ms": 3.4592559999769685,
      "q1_ms": 3.4534854998387994,
      "q3_ms": 3.459383035598356,
      "tflops": 5.328313371465633
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 1500,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.000298917293548584,
      "rel_err": 0.0005957730463705957,
      "xla_temp_mb": 84.00048828125,
      "xla_arg_mb": 46.875,
      "xla_out_mb": 35.156272888183594,
      "compile_s": 0.42874558600306045,
      "inner": 14,
      "median_ms": 3.1451312149459096,
      "q1_ms": 3.144526856756004,
      "q3_ms": 3.1529893572691696,
      "tflops": 20.51170383398757
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 1500,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.000298917293548584,
      "rel_err": 0.0005957730463705957,
      "xla_temp_mb": 84.00048828125,
      "xla_arg_mb": 46.875,
      "xla_out_mb": 35.156272888183594,
      "compile_s": 0.42937405600969214,
      "inner": 15,
      "median_ms": 3.178355066726605,
      "q1_ms": 3.1681791331114555,
      "q3_ms": 3.1809336331207305,
      "tflops": 20.297291726578887
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 1500,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.000298917293548584,
      "rel_err": 0.0005957730463705957,
      "xla_temp_mb": 84.00048828125,
      "xla_arg_mb": 46.875,
      "xla_out_mb": 35.156272888183594,
      "compile_s": 0.42283803000464104,
      "inner": 15,
      "median_ms": 3.169201733544469,
      "q1_ms": 3.1690749000214664,
      "q3_ms": 3.1728235995008918,
      "tflops": 20.355914651052867
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 1500,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00031751394271850586,
      "rel_err": 0.0006307705771178007,
      "xla_temp_mb": 824.15771484375,
      "xla_arg_mb": 46.875,
      "xla_out_mb": 35.156272888183594,
      "compile_s": 1.534782554008416,
      "inner": 11,
      "median_ms": 4.425253090713258,
      "q1_ms": 4.422422272909898,
      "q3_ms": 4.428759772219399,
      "tflops": 14.5781492442508
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 1500,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 1.5497207641601562e-06,
      "rel_err": 3.078662530242582e-06,
      "peak_mb": 183.66162109375,
      "compile_s": 0.008850087993778288,
      "inner": 5,
      "median_ms": 9.208833400043659,
      "q1_ms": 9.196014098415617,
      "q3_ms": 9.225515101570636,
      "tflops": 7.005447617251296
    },
    {
      "impl": "torch_math",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 1500,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 5.066394805908203e-07,
      "rel_err": 1.2071511719113914e-06,
      "peak_mb": 1245.3525390625,
      "compile_s": 0.008670298004290089,
      "inner": 5,
      "median_ms": 8.977224997943267,
      "q1_ms": 8.97157119907206,
      "q3_ms": 8.988036299706437,
      "tflops": 7.186185042123826
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 1500,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.001315295696258545,
      "rel_err": 0.0004427969397511333,
      "xla_temp_mb": 48.0,
      "xla_arg_mb": 35.15625,
      "xla_out_mb": 11.71875,
      "compile_s": 0.24867293798888568,
      "inner": 77,
      "median_ms": 0.5301722206701465,
      "q1_ms": 0.529871396108095,
      "q3_ms": 0.5309687921317754,
      "tflops": 17.38303072226384
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 1500,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.001315295696258545,
      "rel_err": 0.0004427969397511333,
      "xla_temp_mb": 48.0,
      "xla_arg_mb": 35.15625,
      "xla_out_mb": 11.71875,
      "compile_s": 0.23899830799200572,
      "inner": 78,
      "median_ms": 0.5404603205254086,
      "q1_ms": 0.5388196474926259,
      "q3_ms": 0.5408024423022886,
      "tflops": 17.05213065603163
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 1500,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.001315295696258545,
      "rel_err": 0.0004427969397511333,
      "xla_temp_mb": 48.0,
      "xla_arg_mb": 35.15625,
      "xla_out_mb": 11.71875,
      "compile_s": 0.2408659840002656,
      "inner": 73,
      "median_ms": 0.5382706165195711,
      "q1_ms": 0.5378483082781775,
      "q3_ms": 0.5398046849048315,
      "tflops": 17.121499329816963
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 1500,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0013722777366638184,
      "rel_err": 0.00046198006020858884,
      "xla_temp_mb": 549.3164215087891,
      "xla_arg_mb": 35.15625,
      "xla_out_mb": 11.71875,
      "compile_s": 0.4152384990011342,
      "inner": 28,
      "median_ms": 1.6513479999308143,
      "q1_ms": 1.6506228751365728,
      "q3_ms": 1.657121946534192,
      "tflops": 5.580895123490699
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 1500,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 1.5497207641601562e-06,
      "rel_err": 5.21716629009461e-07,
      "peak_mb": 88.1259765625,
      "compile_s": 0.0013009829999646172,
      "inner": 37,
      "median_ms": 1.3550813242988158,
      "q1_ms": 1.3525045675074532,
      "q3_ms": 1.3580068783060573,
      "tflops": 6.801067828728878
    },
    {
      "impl": "torch_math",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 1500,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 4.76837158203125e-07,
      "rel_err": 1.6052818807565927e-07,
      "peak_mb": 714.73583984375,
      "compile_s": 0.004448223000508733,
      "inner": 11,
      "median_ms": 4.457793817643754,
      "q1_ms": 4.455519545775711,
      "q3_ms": 4.465683364039499,
      "tflops": 2.067390367747264
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 1500,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0014875829219818115,
      "rel_err": 0.0005151138757355511,
      "xla_temp_mb": 84.00048828125,
      "xla_arg_mb": 46.875,
      "xla_out_mb": 35.156272888183594,
      "compile_s": 0.431774241005769,
      "inner": 26,
      "median_ms": 1.814126730515944,
      "q1_ms": 1.8120354423165106,
      "q3_ms": 1.8163189040653658,
      "tflops": 17.780455718673124
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 1500,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0014875829219818115,
      "rel_err": 0.0005151138757355511,
      "xla_temp_mb": 84.00048828125,
      "xla_arg_mb": 46.875,
      "xla_out_mb": 35.156272888183594,
      "compile_s": 0.43042169899854343,
      "inner": 27,
      "median_ms": 1.7716395555503874,
      "q1_ms": 1.7698147777183397,
      "q3_ms": 1.7726056111213335,
      "tflops": 18.20686374886181
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 1500,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0014875829219818115,
      "rel_err": 0.0005151138757355511,
      "xla_temp_mb": 84.00048828125,
      "xla_arg_mb": 46.875,
      "xla_out_mb": 35.156272888183594,
      "compile_s": 0.4317451159877237,
      "inner": 27,
      "median_ms": 1.7723211850453582,
      "q1_ms": 1.770453185330283,
      "q3_ms": 1.7738663146802638,
      "tflops": 18.199861442819962
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 1500,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0014894604682922363,
      "rel_err": 0.000543788424693048,
      "xla_temp_mb": 824.15771484375,
      "xla_arg_mb": 46.875,
      "xla_out_mb": 35.156272888183594,
      "compile_s": 0.9860538130014902,
      "inner": 10,
      "median_ms": 4.557299500447698,
      "q1_ms": 4.550151950388681,
      "q3_ms": 4.563285900076153,
      "tflops": 7.077875833447253
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 1500,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 5.0067901611328125e-06,
      "rel_err": 1.59887360950961e-06,
      "peak_mb": 184.50537109375,
      "compile_s": 0.004955373995471746,
      "inner": 9,
      "median_ms": 5.125273888956548,
      "q1_ms": 5.118973055181818,
      "q3_ms": 5.127004943561689,
      "tflops": 6.293517321972227
    },
    {
      "impl": "torch_math",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 1500,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 5.245208740234375e-06,
      "rel_err": 1.6750104805396404e-06,
      "peak_mb": 1245.6337890625,
      "compile_s": 0.009644902995205484,
      "inner": 5,
      "median_ms": 9.869308199267834,
      "q1_ms": 9.863340899755713,
      "q3_ms": 9.88598170079058,
      "tflops": 3.26831418664106
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 2048,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00013739615678787231,
      "rel_err": 0.0006897264975123107,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.2135079500003485,
      "inner": 39,
      "median_ms": 1.1798451540436452,
      "q1_ms": 1.1794902307790323,
      "q3_ms": 1.1801164101388377,
      "tflops": 29.122243923484348
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 2048,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00013739615678787231,
      "rel_err": 0.0006897264975123107,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.2084940060012741,
      "inner": 40,
      "median_ms": 1.190286625205772,
      "q1_ms": 1.1901201874934486,
      "q3_ms": 1.1914985501789488,
      "tflops": 28.866776825338206
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 2048,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00013739615678787231,
      "rel_err": 0.0006897264975123107,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.21704631099419203,
      "inner": 41,
      "median_ms": 1.1916716097536066,
      "q1_ms": 1.1907961951049113,
      "q3_ms": 1.1920532072885757,
      "tflops": 28.833227280713952
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 2048,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00015942752361297607,
      "rel_err": 0.000800323614384979,
      "xla_temp_mb": 1024.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 1.023730736997095,
      "inner": 16,
      "median_ms": 2.86743843662407,
      "q1_ms": 2.8668986874436087,
      "q3_ms": 2.872816718536342,
      "tflops": 11.982729229385951
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 2048,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 5.513429641723633e-07,
      "rel_err": 2.7677328944264445e-06,
      "peak_mb": 112.2509765625,
      "compile_s": 0.003934540000045672,
      "inner": 13,
      "median_ms": 4.097458769907602,
      "q1_ms": 4.093091692014311,
      "q3_ms": 4.098207538145989,
      "tflops": 8.385621502855248
    },
    {
      "impl": "torch_math",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 2048,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 2.7567148208618164e-07,
      "rel_err": 1.3838664472132223e-06,
      "peak_mb": 1264.31298828125,
      "compile_s": 0.006203857003129087,
      "inner": 8,
      "median_ms": 6.312617124422104,
      "q1_ms": 6.301494250692485,
      "q3_ms": 6.316880124359159,
      "tflops": 5.4430258782953675
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 2048,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0002512037754058838,
      "rel_err": 0.0006222709198482335,
      "xla_temp_mb": 0.50048828125,
      "xla_arg_mb": 64.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.38128508601221256,
      "inner": 10,
      "median_ms": 5.101007099437993,
      "q1_ms": 5.0907290999020915,
      "q3_ms": 5.110711450106464,
      "tflops": 23.575557128953932
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 2048,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0002512037754058838,
      "rel_err": 0.0006222709198482335,
      "xla_temp_mb": 0.50048828125,
      "xla_arg_mb": 64.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.3806725870090304,
      "inner": 9,
      "median_ms": 5.113135444440155,
      "q1_ms": 5.0967782784330975,
      "q3_ms": 5.1152842780639824,
      "tflops": 23.519635964027817
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 2048,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0002512037754058838,
      "rel_err": 0.0006222709198482335,
      "xla_temp_mb": 0.50048828125,
      "xla_arg_mb": 64.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.3832551770028658,
      "inner": 10,
      "median_ms": 5.092652200255543,
      "q1_ms": 5.089343549479963,
      "q3_ms": 5.096266100008506,
      "tflops": 23.614234697190895
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 2048,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0002518892288208008,
      "rel_err": 0.0006239688955247402,
      "xla_temp_mb": 1536.25,
      "xla_arg_mb": 64.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 1.5326878080086317,
      "inner": 6,
      "median_ms": 7.664102500712033,
      "q1_ms": 7.660645998839755,
      "q3_ms": 7.678781000625652,
      "tflops": 15.6912155437414
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 2048,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 7.748603820800781e-07,
      "rel_err": 2.472344249326852e-06,
      "peak_mb": 241.001953125,
      "compile_s": 0.01540448000014294,
      "inner": 3,
      "median_ms": 16.167080000741407,
      "q1_ms": 16.12901466675491,
      "q3_ms": 16.21425100165652,
      "tflops": 7.438516063660539
    },
    {
      "impl": "torch_math",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 2048,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 7.003545761108398e-07,
      "rel_err": 2.4208370632550213e-06,
      "peak_mb": 2240.2509765625,
      "compile_s": 0.015948019005008973,
      "inner": 3,
      "median_ms": 16.505240331753157,
      "q1_ms": 16.463501165465765,
      "q3_ms": 16.523825499461964,
      "tflops": 7.286115310701828
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 2048,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0008454322814941406,
      "rel_err": 0.00032341157202608883,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.2298401459993329,
      "inner": 55,
      "median_ms": 0.7099863999015228,
      "q1_ms": 0.7084944908538918,
      "q3_ms": 0.71107400911966,
      "tflops": 24.197462354747778
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 2048,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0008454322814941406,
      "rel_err": 0.00032341157202608883,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.22739433500100859,
      "inner": 63,
      "median_ms": 0.7019512697644089,
      "q1_ms": 0.7010857223085911,
      "q3_ms": 0.702688761867605,
      "tflops": 24.47444705067057
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 2048,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0008454322814941406,
      "rel_err": 0.00032341157202608883,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.2275160959979985,
      "inner": 64,
      "median_ms": 0.7029870780570491,
      "q1_ms": 0.7026265546983268,
      "q3_ms": 0.7049670936112307,
      "tflops": 24.438385455793274
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 2048,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0013530254364013672,
      "rel_err": 0.0005175861879251897,
      "xla_temp_mb": 1024.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.4980273960100021,
      "inner": 16,
      "median_ms": 2.8870104370071203,
      "q1_ms": 2.8847346247857786,
      "q3_ms": 2.889902531478583,
      "tflops": 5.9507471686904845
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 2048,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 1.9073486328125e-06,
      "rel_err": 7.296369517462153e-07,
      "peak_mb": 112.2509765625,
      "compile_s": 0.002223231000243686,
      "inner": 22,
      "median_ms": 2.3296080908039585,
      "q1_ms": 2.3189335452943025,
      "q3_ms": 2.3307559544760315,
      "tflops": 7.374574827335506
    },
    {
      "impl": "torch_math",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 2048,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 3.129243850708008e-07,
      "rel_err": 1.1970605839906057e-07,
      "peak_mb": 1280.31298828125,
      "compile_s": 0.00813159799145069,
      "inner": 6,
      "median_ms": 8.18241900075615,
      "q1_ms": 8.17817758312837,
      "q3_ms": 8.187761499963624,
      "tflops": 2.099607607776182
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 2048,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.001447916030883789,
      "rel_err": 0.0005637407884933054,
      "xla_temp_mb": 0.50048828125,
      "xla_arg_mb": 64.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.4050686720001977,
      "inner": 18,
      "median_ms": 2.752519278146792,
      "q1_ms": 2.7437075829463233,
      "q3_ms": 2.7630959998027214,
      "tflops": 21.845275570416288
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 2048,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.001447916030883789,
      "rel_err": 0.0005637407884933054,
      "xla_temp_mb": 0.50048828125,
      "xla_arg_mb": 64.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.40785546899132896,
      "inner": 17,
      "median_ms": 2.7555867058896077,
      "q1_ms": 2.7476911474128856,
      "q3_ms": 2.7632832938894243,
      "tflops": 21.820958134063837
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 2048,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.001447916030883789,
      "rel_err": 0.0005637407884933054,
      "xla_temp_mb": 0.50048828125,
      "xla_arg_mb": 64.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.40574514400213957,
      "inner": 17,
      "median_ms": 2.748917587964302,
      "q1_ms": 2.747113882351721,
      "q3_ms": 2.758707911519142,
      "tflops": 21.87389771423764
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 2048,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0014613419771194458,
      "rel_err": 0.0005936888628639281,
      "xla_temp_mb": 1536.25,
      "xla_arg_mb": 64.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 1.01772149799217,
      "inner": 6,
      "median_ms": 7.670644998142961,
      "q1_ms": 7.6671907493922244,
      "q3_ms": 7.679479251237353,
      "tflops": 7.838916044029826
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 2048,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 6.9141387939453125e-06,
      "rel_err": 2.808957106026355e-06,
      "peak_mb": 241.001953125,
      "compile_s": 0.008469778011203744,
      "inner": 5,
      "median_ms": 8.783469998161308,
      "q1_ms": 8.774223500222433,
      "q3_ms": 8.793742500711232,
      "tflops": 6.845761658727956
    },
    {
      "impl": "torch_math",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 2048,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 6.67572021484375e-06,
      "rel_err": 2.7120966024085646e-06,
      "peak_mb": 2240.2509765625,
      "compile_s": 0.01775761498720385,
      "inner": 2,
      "median_ms": 18.191987503087148,
      "q1_ms": 18.184370746894274,
      "q3_ms": 18.19512875226792,
      "tflops": 3.3052761351004736
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 3000,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0001064687967300415,
      "rel_err": 0.000586395210120827,
      "xla_temp_mb": 96.0,
      "xla_arg_mb": 70.3125,
      "xla_out_mb": 23.4375,
      "compile_s": 0.24237169600382913,
      "inner": 18,
      "median_ms": 2.8111837220623985,
      "q1_ms": 2.8011163057170863,
      "q3_ms": 2.8186456390055055,
      "tflops": 26.22667434411229
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 3000,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0001064687967300415,
      "rel_err": 0.000586395210120827,
      "xla_temp_mb": 96.0,
      "xla_arg_mb": 70.3125,
      "xla_out_mb": 23.4375,
      "compile_s": 0.23970626799564343,
      "inner": 16,
      "median_ms": 2.8081386244593887,
      "q1_ms": 2.806394718845695,
      "q3_ms": 2.809105749747687,
      "tflops": 26.255114102208474
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 3000,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0001064687967300415,
      "rel_err": 0.000586395210120827,
      "xla_temp_mb": 96.0,
      "xla_arg_mb": 70.3125,
      "xla_out_mb": 23.4375,
      "compile_s": 0.23920443801034708,
      "inner": 17,
      "median_ms": 2.8018978816416005,
      "q1_ms": 2.800600500075448,
      "q3_ms": 2.8111981471952543,
      "tflops": 26.313592826874757
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 3000,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0001172870397567749,
      "rel_err": 0.0006459785508923233,
      "xla_temp_mb": 2197.265640258789,
      "xla_arg_mb": 70.3125,
      "xla_out_mb": 23.4375,
      "compile_s": 1.219359220995102,
      "inner": 8,
      "median_ms": 5.978601999231614,
      "q1_ms": 5.976736686534423,
      "q3_ms": 5.987264125906222,
      "tflops": 12.33197995275078
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 3000,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 4.172325134277344e-07,
      "rel_err": 2.297979790455429e-06,
      "peak_mb": 156.8759765625,
      "compile_s": 0.008187062994693406,
      "inner": 6,
      "median_ms": 8.760129499326771,
      "q1_ms": 8.746001082423996,
      "q3_ms": 8.791969001322286,
      "tflops": 8.416313937558355
    },
    {
      "impl": "torch_math",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 3000,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 3.203749656677246e-07,
      "rel_err": 1.7645202206040267e-06,
      "peak_mb": 2628.89111328125,
      "compile_s": 0.013458865010761656,
      "inner": 3,
      "median_ms": 14.143421334059289,
      "q1_ms": 14.13165916649935,
      "q3_ms": 14.16758216752593,
      "tflops": 5.212882955162547
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 3000,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00015441328287124634,
      "rel_err": 0.0007109376019798219,
      "xla_temp_mb": 168.00048828125,
      "xla_arg_mb": 93.75,
      "xla_out_mb": 70.3125228881836,
      "compile_s": 0.4346250089874957,
      "inner": 4,
      "median_ms": 11.49464074842399,
      "q1_ms": 11.449203622760251,
      "q3_ms": 11.553597623787937,
      "tflops": 22.44941844184043
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 3000,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00015441328287124634,
      "rel_err": 0.0007109376019798219,
      "xla_temp_mb": 168.00048828125,
      "xla_arg_mb": 93.75,
      "xla_out_mb": 70.3125228881836,
      "compile_s": 0.42947652100701816,
      "inner": 4,
      "median_ms": 11.464366001746384,
      "q1_ms": 11.424936250477913,
      "q3_ms": 11.485569126307382,
      "tflops": 22.508702178619483
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 3000,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00015441328287124634,
      "rel_err": 0.0007109376019798219,
      "xla_temp_mb": 168.00048828125,
      "xla_arg_mb": 93.75,
      "xla_out_mb": 70.3125228881836,
      "compile_s": 0.4315280050068395,
      "inner": 4,
      "median_ms": 11.43423925168463,
      "q1_ms": 11.427690373238875,
      "q3_ms": 11.461183750725468,
      "tflops": 22.56800774585692
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 3000,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00015636533498764038,
      "rel_err": 0.0007197430240921676,
      "xla_temp_mb": 3296.264663696289,
      "xla_arg_mb": 93.75,
      "xla_out_mb": 70.3125228881836,
      "compile_s": 1.7513711789943045,
      "inner": 3,
      "median_ms": 16.151834337506443,
      "q1_ms": 16.150008001811024,
      "q3_ms": 16.15410716719149,
      "tflops": 15.97638971573541
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 3000,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 1.1920928955078125e-06,
      "rel_err": 4.159890522714704e-06,
      "peak_mb": 349.4140625,
      "compile_s": 0.03333326899155509,
      "inner": 1,
      "median_ms": 34.19577199383639,
      "q1_ms": 34.17421600170201,
      "q3_ms": 34.297449499717914,
      "tflops": 7.546196063259277
    },
    {
      "impl": "torch_math",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 3000,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 1.1026859283447266e-06,
      "rel_err": 3.847898824460572e-06,
      "peak_mb": 4670.2822265625,
      "compile_s": 0.034869327006163076,
      "inner": 1,
      "median_ms": 35.6837790022837,
      "q1_ms": 35.58348450314952,
      "q3_ms": 35.76920650084503,
      "tflops": 7.23152107806422
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 3000,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0009112358093261719,
      "rel_err": 0.00033203110797330737,
      "xla_temp_mb": 96.0,
      "xla_arg_mb": 70.3125,
      "xla_out_mb": 23.4375,
      "compile_s": 0.24369675500201993,
      "inner": 28,
      "median_ms": 1.6829038218669925,
      "q1_ms": 1.679805464456357,
      "q3_ms": 1.6836400354804937,
      "tflops": 21.90499511677592
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 3000,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0009112358093261719,
      "rel_err": 0.00033203110797330737,
      "xla_temp_mb": 96.0,
      "xla_arg_mb": 70.3125,
      "xla_out_mb": 23.4375,
      "compile_s": 0.24367861099017318,
      "inner": 29,
      "median_ms": 1.6638747926255495,
      "q1_ms": 1.6613294481368062,
      "q3_ms": 1.665717413837636,
      "tflops": 22.155513241371725
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 3000,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0009112358093261719,
      "rel_err": 0.00033203110797330737,
      "xla_temp_mb": 96.0,
      "xla_arg_mb": 70.3125,
      "xla_out_mb": 23.4375,
      "compile_s": 0.24009070599277038,
      "inner": 29,
      "median_ms": 1.656385068882032,
      "q1_ms": 1.65465320681286,
      "q3_ms": 1.660233051646581,
      "tflops": 22.255694459309
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 3000,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.001251220703125,
      "rel_err": 0.00045591293019242585,
      "xla_temp_mb": 2197.265640258789,
      "xla_arg_mb": 70.3125,
      "xla_out_mb": 23.4375,
      "compile_s": 0.5046180140052456,
      "inner": 8,
      "median_ms": 5.9894612513744505,
      "q1_ms": 5.979836375445302,
      "q3_ms": 6.003023499943083,
      "tflops": 6.154810667076362
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 3000,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 1.3709068298339844e-06,
      "rel_err": 4.995234803573112e-07,
      "peak_mb": 158.0009765625,
      "compile_s": 0.004673430012189783,
      "inner": 10,
      "median_ms": 4.846396899665706,
      "q1_ms": 4.839668199565494,
      "q3_ms": 4.860084900428774,
      "tflops": 7.606475648443651
    },
    {
      "impl": "torch_math",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 3000,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 4.76837158203125e-07,
      "rel_err": 1.7374730987285147e-07,
      "peak_mb": 2664.3486328125,
      "compile_s": 0.017827317002229393,
      "inner": 2,
      "median_ms": 18.296646005182993,
      "q1_ms": 18.27899249838083,
      "q3_ms": 18.30176049770671,
      "tflops": 2.0147954980140805
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 3000,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0013532638549804688,
      "rel_err": 0.0007203135173767805,
      "xla_temp_mb": 168.00048828125,
      "xla_arg_mb": 93.75,
      "xla_out_mb": 70.3125228881836,
      "compile_s": 0.435029043990653,
      "inner": 8,
      "median_ms": 6.341997748677386,
      "q1_ms": 6.331373249850003,
      "q3_ms": 6.363619248986652,
      "tflops": 20.34437808290105
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 3000,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0013532638549804688,
      "rel_err": 0.0007203135173767805,
      "xla_temp_mb": 168.00048828125,
      "xla_arg_mb": 93.75,
      "xla_out_mb": 70.3125228881836,
      "compile_s": 0.44039123799302615,
      "inner": 8,
      "median_ms": 6.138498376458301,
      "q1_ms": 6.125754624918045,
      "q3_ms": 6.142662749880401,
      "tflops": 21.018821230745743
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 3000,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0013532638549804688,
      "rel_err": 0.0007203135173767805,
      "xla_temp_mb": 168.00048828125,
      "xla_arg_mb": 93.75,
      "xla_out_mb": 70.3125228881836,
      "compile_s": 0.43818007099616807,
      "inner": 8,
      "median_ms": 6.129656874691136,
      "q1_ms": 6.113002186793892,
      "q3_ms": 6.185410624311771,
      "tflops": 21.049139068897933
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 3000,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0013537406921386719,
      "rel_err": 0.0005977512337267399,
      "xla_temp_mb": 3296.264663696289,
      "xla_arg_mb": 93.75,
      "xla_out_mb": 70.3125228881836,
      "compile_s": 0.9155783429887379,
      "inner": 3,
      "median_ms": 16.163373336894438,
      "q1_ms": 16.15421266615158,
      "q3_ms": 16.22572466537046,
      "tflops": 7.9824921017873445
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 3000,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 5.7220458984375e-06,
      "rel_err": 2.176755742766545e-06,
      "peak_mb": 349.9765625,
      "compile_s": 0.017519985005492344,
      "inner": 2,
      "median_ms": 18.136837999918498,
      "q1_ms": 18.113253248884575,
      "q3_ms": 18.190029004472308,
      "tflops": 7.113919195869743
    },
    {
      "impl": "torch_math",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 3000,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 7.3909759521484375e-06,
      "rel_err": 3.85596740670735e-06,
      "peak_mb": 4669.7197265625,
      "compile_s": 0.039042349002556875,
      "inner": 1,
      "median_ms": 39.516727993031964,
      "q1_ms": 39.47087549749995,
      "q3_ms": 39.624266006285325,
      "tflops": 3.2650476533064925
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 4096,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00011807680130004883,
      "rel_err": 0.0005248566158115864,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.21686756200506352,
      "inner": 10,
      "median_ms": 4.527687400695868,
      "q1_ms": 4.516599949420197,
      "q3_ms": 4.5329008004046045,
      "tflops": 30.355221398649736
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 4096,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00011807680130004883,
      "rel_err": 0.0005248566158115864,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.2118449310073629,
      "inner": 11,
      "median_ms": 4.533551818564196,
      "q1_ms": 4.512070682143199,
      "q3_ms": 4.5459315912342975,
      "tflops": 30.31595512137055
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 4096,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00011807680130004883,
      "rel_err": 0.0005248566158115864,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.21306842799822334,
      "inner": 11,
      "median_ms": 4.565392000420781,
      "q1_ms": 4.554925545314538,
      "q3_ms": 4.576681272407189,
      "tflops": 30.1045240932942
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 4096,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00012464821338653564,
      "rel_err": 0.000554066791664809,
      "xla_temp_mb": 4096.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 1.3063977139972849,
      "inner": 4,
      "median_ms": 10.82381324886228,
      "q1_ms": 10.806420123117277,
      "q3_ms": 10.83343862410402,
      "tflops": 12.69783119054152
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 4096,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 3.8743019104003906e-07,
      "rel_err": 1.722144361337996e-06,
      "peak_mb": 208.2509765625,
      "compile_s": 0.015181757000391372,
      "inner": 3,
      "median_ms": 16.123039997182786,
      "q1_ms": 16.10495066658283,
      "q3_ms": 16.14825933453782,
      "tflops": 8.524382095188935
    },
    {
      "impl": "torch_math",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 4096,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 4.172325134277344e-07,
      "rel_err": 1.8546169258115697e-06,
      "peak_mb": 4816.37548828125,
      "compile_s": 0.02404360800574068,
      "inner": 2,
      "median_ms": 24.843457496899646,
      "q1_ms": 24.83659975041519,
      "q3_ms": 24.980980499094585,
      "tflops": 5.5321991107377775
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 4096,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0001219138503074646,
      "rel_err": 0.0006636553443968296,
      "xla_temp_mb": 1.00048828125,
      "xla_arg_mb": 128.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 0.3949519460002193,
      "inner": 2,
      "median_ms": 19.522738999512512,
      "q1_ms": 19.45849050025572,
      "q3_ms": 19.551709501683945,
      "tflops": 24.639797579838138
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 4096,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0001219138503074646,
      "rel_err": 0.0006636553443968296,
      "xla_temp_mb": 1.00048828125,
      "xla_arg_mb": 128.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 0.3931225710111903,
      "inner": 2,
      "median_ms": 19.523865499650128,
      "q1_ms": 19.477802252367837,
      "q3_ms": 19.5395192495198,
      "tflops": 24.63837589746868
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 4096,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0001219138503074646,
      "rel_err": 0.0006636553443968296,
      "xla_temp_mb": 1.00048828125,
      "xla_arg_mb": 128.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 0.3946442129963543,
      "inner": 2,
      "median_ms": 19.531072503014002,
      "q1_ms": 19.520704001479317,
      "q3_ms": 19.55100049963221,
      "tflops": 24.629284289317305
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 4096,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00012190639972686768,
      "rel_err": 0.0006636147736571729,
      "xla_temp_mb": 6144.5,
      "xla_arg_mb": 128.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 2.0903819360100897,
      "inner": 1,
      "median_ms": 29.224236990557984,
      "q1_ms": 29.197362506238278,
      "q3_ms": 29.246033998788334,
      "tflops": 16.460184651096874
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 4096,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 8.642673492431641e-07,
      "rel_err": 4.622498181561241e-06,
      "peak_mb": 465.751953125,
      "compile_s": 0.06303136399947107,
      "inner": 1,
      "median_ms": 63.3526829915354,
      "q1_ms": 63.036593499418814,
      "q3_ms": 63.61834549170453,
      "tflops": 7.592990769092946
    },
    {
      "impl": "torch_math",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 4096,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "error": "OutOfMemoryError: CUDA out of memory. Tried to allocate 2.00 GiB. GPU 0 has a total capacity of 23.61 GiB of which 1.19 GiB is free. Including non-PyTorch memory, this process has 22.38 GiB memory in use. Of the allocated memory 4.27 GiB is allocated by PyTorch, and 63.75 MiB is reserved by PyTorch but unallocated. I"
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 4096,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0009582042694091797,
      "rel_err": 0.000337931647663936,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.23809564800467342,
      "inner": 16,
      "median_ms": 2.5347928749397397,
      "q1_ms": 2.530764968469157,
      "q3_ms": 2.536159531246085,
      "tflops": 27.11048993998521
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 4096,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0009582042694091797,
      "rel_err": 0.000337931647663936,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.2314239249972161,
      "inner": 18,
      "median_ms": 2.533892500246616,
      "q1_ms": 2.5318589723610785,
      "q3_ms": 2.5382571390461335,
      "tflops": 27.120123181749715
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 4096,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0009582042694091797,
      "rel_err": 0.000337931647663936,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.23399502399843186,
      "inner": 20,
      "median_ms": 2.5438699500227813,
      "q1_ms": 2.541034900059458,
      "q3_ms": 2.5538481499097543,
      "tflops": 27.01375388131952
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 4096,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.001003861427307129,
      "rel_err": 0.0003540336329024285,
      "xla_temp_mb": 4096.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.48861991199373733,
      "inner": 4,
      "median_ms": 10.823214252013713,
      "q1_ms": 10.815185751198442,
      "q3_ms": 10.834658502062666,
      "tflops": 6.349266967824683
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 4096,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 1.7881393432617188e-06,
      "rel_err": 6.306263458100148e-07,
      "peak_mb": 208.2509765625,
      "compile_s": 0.008339672000147402,
      "inner": 5,
      "median_ms": 8.882224999251775,
      "q1_ms": 8.85319400113076,
      "q3_ms": 8.893828000873327,
      "tflops": 7.73674127167335
    },
    {
      "impl": "torch_math",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 4096,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 3.5762786865234375e-07,
      "rel_err": 1.261252720041739e-07,
      "peak_mb": 4880.37548828125,
      "compile_s": 0.032342712991521694,
      "inner": 1,
      "median_ms": 32.74378199421335,
      "q1_ms": 32.661832505255006,
      "q3_ms": 32.788778000394814,
      "tflops": 2.0987030987484725
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 4096,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0013544559478759766,
      "rel_err": 0.0006111936527304351,
      "xla_temp_mb": 1.00048828125,
      "xla_arg_mb": 128.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 0.4210794389946386,
      "inner": 4,
      "median_ms": 10.28301124824793,
      "q1_ms": 10.27696887649654,
      "q3_ms": 10.305534249710036,
      "tflops": 23.389857578632977
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 4096,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0013544559478759766,
      "rel_err": 0.0006111936527304351,
      "xla_temp_mb": 1.00048828125,
      "xla_arg_mb": 128.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 0.4289012299996102,
      "inner": 5,
      "median_ms": 10.381704400060698,
      "q1_ms": 10.368822900636587,
      "q3_ms": 10.411743199802004,
      "tflops": 23.167503071518176
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 4096,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0013544559478759766,
      "rel_err": 0.0006111936527304351,
      "xla_temp_mb": 1.00048828125,
      "xla_arg_mb": 128.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 0.4163202069903491,
      "inner": 5,
      "median_ms": 10.321013399516232,
      "q1_ms": 10.279386601177976,
      "q3_ms": 10.351292999985162,
      "tflops": 23.303735715261602
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 4096,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0013545751571655273,
      "rel_err": 0.0005387319251894951,
      "xla_temp_mb": 6144.0,
      "xla_arg_mb": 128.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 1.0296488640015014,
      "inner": 1,
      "median_ms": 28.443070987123065,
      "q1_ms": 28.40855249814922,
      "q3_ms": 28.563981504703406,
      "tflops": 8.456125173153383
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 4096,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 8.58306884765625e-06,
      "rel_err": 2.1971550268062856e-06,
      "peak_mb": 465.751953125,
      "compile_s": 0.03270553999755066,
      "inner": 1,
      "median_ms": 33.05672499118373,
      "q1_ms": 32.93477399711264,
      "q3_ms": 33.1319725009962,
      "tflops": 7.2759224829485225
    },
    {
      "impl": "torch_math",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 4096,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "error": "OutOfMemoryError: CUDA out of memory. Tried to allocate 2.00 GiB. GPU 0 has a total capacity of 23.61 GiB of which 1.19 GiB is free. Including non-PyTorch memory, this process has 22.38 GiB memory in use. Of the allocated memory 4.27 GiB is allocated by PyTorch, and 63.75 MiB is reserved by PyTorch but unallocated. I"
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 8192,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 6.543099880218506e-05,
      "rel_err": 0.0005818834179081023,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 192.0,
      "xla_out_mb": 64.0,
      "compile_s": 0.22332489999826066,
      "inner": 2,
      "median_ms": 18.10219149774639,
      "q1_ms": 18.08403100221767,
      "q3_ms": 18.207997745776083,
      "tflops": 30.369572322579902
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 8192,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 6.543099880218506e-05,
      "rel_err": 0.0005818834179081023,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 192.0,
      "xla_out_mb": 64.0,
      "compile_s": 0.22386470199853648,
      "inner": 3,
      "median_ms": 18.095922666058566,
      "q1_ms": 18.043818667744443,
      "q3_ms": 18.12091150107638,
      "tflops": 30.380093020575508
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 8192,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 6.543099880218506e-05,
      "rel_err": 0.0005818834179081023,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 192.0,
      "xla_out_mb": 64.0,
      "compile_s": 0.2258884240000043,
      "inner": 3,
      "median_ms": 18.16897400325009,
      "q1_ms": 18.136831667410053,
      "q3_ms": 18.213299333486553,
      "tflops": 30.257944878431715
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 8192,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 7.790327072143555e-05,
      "rel_err": 0.00069280038587749,
      "xla_temp_mb": 16384.00001525879,
      "xla_arg_mb": 192.0,
      "xla_out_mb": 64.0,
      "compile_s": 2.6506363520020386,
      "inner": 1,
      "median_ms": 44.982676990912296,
      "q1_ms": 44.92161250527715,
      "q3_ms": 45.04327649920015,
      "tflops": 12.221500601199555
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 8192,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 2.905726432800293e-07,
      "rel_err": 2.584087269497104e-06,
      "peak_mb": 400.2509765625,
      "compile_s": 0.06807051600480918,
      "inner": 1,
      "median_ms": 68.21003899676725,
      "q1_ms": 68.10767149727326,
      "q3_ms": 68.3587265011738,
      "tflops": 8.059749297519902
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 8192,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00010244548320770264,
      "rel_err": 0.0006203833036124706,
      "xla_temp_mb": 2.00048828125,
      "xla_arg_mb": 256.0,
      "xla_out_mb": 192.0000228881836,
      "compile_s": 0.4471490739961155,
      "inner": 1,
      "median_ms": 78.7197199970251,
      "q1_ms": 78.06572749541374,
      "q3_ms": 78.7865270031034,
      "tflops": 24.442990253023204
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 8192,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00010244548320770264,
      "rel_err": 0.0006204370292834938,
      "xla_temp_mb": 2.00048828125,
      "xla_arg_mb": 256.0,
      "xla_out_mb": 192.0000228881836,
      "compile_s": 0.4485811760096112,
      "inner": 1,
      "median_ms": 78.66859399655368,
      "q1_ms": 78.16329200431937,
      "q3_ms": 78.87448999827029,
      "tflops": 24.45887552906174
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 8192,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00010244548320770264,
      "rel_err": 0.0006204370292834938,
      "xla_temp_mb": 2.00048828125,
      "xla_arg_mb": 256.0,
      "xla_out_mb": 192.0000228881836,
      "compile_s": 0.4494757900101831,
      "inner": 1,
      "median_ms": 78.21889400656801,
      "q1_ms": 77.85391100333072,
      "q3_ms": 78.51294799911557,
      "tflops": 24.599495723455643
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 8192,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 24.00GiB with allocator GPU_0_bfc on device 0. [tf-allocator-allocation-error=''] [executable_name='jit_loss']"
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 8192,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 5.885958671569824e-07,
      "rel_err": 4.683713541453471e-06,
      "peak_mb": 915.251953125,
      "compile_s": 0.250517911001225,
      "inner": 1,
      "median_ms": 254.23336600943003,
      "q1_ms": 253.38070150610292,
      "q3_ms": 254.428441003256,
      "tflops": 7.56842179612502
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 8192,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0007331371307373047,
      "rel_err": 0.0002778850612230599,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 192.0,
      "xla_out_mb": 64.0,
      "compile_s": 0.24412721999397036,
      "inner": 5,
      "median_ms": 9.714150399668142,
      "q1_ms": 9.701219799171668,
      "q3_ms": 9.748754699830897,
      "tflops": 28.296649283234327
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 8192,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0007331371307373047,
      "rel_err": 0.0002778850612230599,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 192.0,
      "xla_out_mb": 64.0,
      "compile_s": 0.24094066600082442,
      "inner": 5,
      "median_ms": 9.777512197615579,
      "q1_ms": 9.745100099826232,
      "q3_ms": 9.801072100526653,
      "tflops": 28.11327681197206
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 8192,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0007331371307373047,
      "rel_err": 0.0002778850612230599,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 192.0,
      "xla_out_mb": 64.0,
      "compile_s": 0.24588100099936128,
      "inner": 5,
      "median_ms": 9.742872999049723,
      "q1_ms": 9.702449198812246,
      "q3_ms": 9.776469398639165,
      "tflops": 28.2132289901357
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 8192,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0008221268653869629,
      "rel_err": 0.0003116153529845178,
      "xla_temp_mb": 16384.00001525879,
      "xla_arg_mb": 192.0,
      "xla_out_mb": 64.0,
      "compile_s": 0.4938220430049114,
      "inner": 1,
      "median_ms": 44.63177500292659,
      "q1_ms": 44.58849150250899,
      "q3_ms": 44.68401649501175,
      "tflops": 6.158793974158002
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 8192,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 1.7881393432617188e-06,
      "rel_err": 6.777684689041052e-07,
      "peak_mb": 400.2509765625,
      "compile_s": 0.035408137991908006,
      "inner": 1,
      "median_ms": 35.64409400860313,
      "q1_ms": 35.57000400178367,
      "q3_ms": 35.850995998771396,
      "tflops": 7.711737795261534
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 8192,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0016207695007324219,
      "rel_err": 0.0008843561518006027,
      "xla_temp_mb": 2.00048828125,
      "xla_arg_mb": 256.0,
      "xla_out_mb": 192.0000228881836,
      "compile_s": 0.4574556199950166,
      "inner": 1,
      "median_ms": 40.52458600199316,
      "q1_ms": 40.461963501002174,
      "q3_ms": 40.58076701039681,
      "tflops": 23.740468915751084
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 8192,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0016207695007324219,
      "rel_err": 0.0008843561518006027,
      "xla_temp_mb": 2.00048828125,
      "xla_arg_mb": 256.0,
      "xla_out_mb": 192.0000228881836,
      "compile_s": 0.46378448500763625,
      "inner": 1,
      "median_ms": 40.299500004039146,
      "q1_ms": 40.24332799599506,
      "q3_ms": 40.46820350049529,
      "tflops": 23.87306726404976
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 8192,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0016207695007324219,
      "rel_err": 0.0008843561518006027,
      "xla_temp_mb": 2.00048828125,
      "xla_arg_mb": 256.0,
      "xla_out_mb": 192.0000228881836,
      "compile_s": 0.45706602299469523,
      "inner": 1,
      "median_ms": 40.16194900032133,
      "q1_ms": 40.12184349994641,
      "q3_ms": 40.264793489768635,
      "tflops": 23.95483033695159
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 8192,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 24.00GiB with allocator GPU_0_bfc on device 0. [tf-allocator-allocation-error=''] [executable_name='jit_loss']"
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 64,
      "seq": 8192,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 1.1205673217773438e-05,
      "rel_err": 4.5531719479185995e-06,
      "peak_mb": 915.251953125,
      "compile_s": 0.13101380299485754,
      "inner": 1,
      "median_ms": 130.8842880098382,
      "q1_ms": 130.80623949645087,
      "q3_ms": 131.22062449838268,
      "tflops": 7.350558947393928
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 128,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0005145668983459473,
      "rel_err": 0.0006053378456272185,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.30827678099740297,
      "inner": 200,
      "median_ms": 0.03180262501700781,
      "q1_ms": 0.031716482480987906,
      "q3_ms": 0.031941032539180014,
      "tflops": 8.440669782964225
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 128,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0005145668983459473,
      "rel_err": 0.0006053378456272185,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.3054678220069036,
      "inner": 200,
      "median_ms": 0.031998039994505234,
      "q1_ms": 0.03167024497088278,
      "q3_ms": 0.03209141501429258,
      "tflops": 8.389121835152912
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 128,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0005145668983459473,
      "rel_err": 0.0006053378456272185,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.31166499700339045,
      "inner": 200,
      "median_ms": 0.031678135055699386,
      "q1_ms": 0.031365614995593205,
      "q3_ms": 0.033311802508251276,
      "tflops": 8.47384025379058
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 128,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00043076276779174805,
      "rel_err": 0.0005067504826001823,
      "xla_temp_mb": 4.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.8612026100017829,
      "inner": 200,
      "median_ms": 0.04285986004106235,
      "q1_ms": 0.04282273748685839,
      "q3_ms": 0.042896777486021165,
      "tflops": 6.26309688698989
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 128,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 1.0132789611816406e-06,
      "rel_err": 1.192024114970991e-06,
      "peak_mb": 28.2509765625,
      "compile_s": 0.0001010209962259978,
      "inner": 200,
      "median_ms": 0.0435637900227448,
      "q1_ms": 0.04354560500360094,
      "q3_ms": 0.04356847250164719,
      "tflops": 6.161893991772731
    },
    {
      "impl": "torch_math",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 128,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 1.0728836059570312e-06,
      "rel_err": 1.2621430869330652e-06,
      "peak_mb": 32.75439453125,
      "compile_s": 0.0002275909937452525,
      "inner": 200,
      "median_ms": 0.13098497998726089,
      "q1_ms": 0.13066955500107724,
      "q3_ms": 0.13156710501789348,
      "tflops": 2.0493605910090382
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 128,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0005626082420349121,
      "rel_err": 0.0008495131623931229,
      "xla_temp_mb": 2.031494140625,
      "xla_arg_mb": 8.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.7761601669917582,
      "inner": 200,
      "median_ms": 0.09565913504047785,
      "q1_ms": 0.09563887248077663,
      "q3_ms": 0.095670529990457,
      "tflops": 9.821582597442926
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 128,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0005626082420349121,
      "rel_err": 0.0008495131623931229,
      "xla_temp_mb": 2.031494140625,
      "xla_arg_mb": 8.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.7750665259954985,
      "inner": 196,
      "median_ms": 0.09568659189259349,
      "q1_ms": 0.09565827809749361,
      "q3_ms": 0.09570714539630881,
      "tflops": 9.818764336957463
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 128,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0005626082420349121,
      "rel_err": 0.0008495131623931229,
      "xla_temp_mb": 2.031494140625,
      "xla_arg_mb": 8.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.7877066809887765,
      "inner": 166,
      "median_ms": 0.0957962711473523,
      "q1_ms": 0.09573845481023147,
      "q3_ms": 0.09583121682664503,
      "tflops": 9.807522618023816
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 128,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0005522668361663818,
      "rel_err": 0.0008338980842381716,
      "xla_temp_mb": 4.0,
      "xla_arg_mb": 8.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 1.146629644004861,
      "inner": 200,
      "median_ms": 0.10328552998544183,
      "q1_ms": 0.1032536399725359,
      "q3_ms": 0.1032961724922643,
      "tflops": 9.096376773517322
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 128,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 1.2218952178955078e-06,
      "rel_err": 1.8000066575041274e-06,
      "peak_mb": 44.298828125,
      "compile_s": 0.0004057379992445931,
      "inner": 143,
      "median_ms": 0.23760776225985453,
      "q1_ms": 0.2264503671069649,
      "q3_ms": 0.2448891992806913,
      "tflops": 3.9540968151221847
    },
    {
      "impl": "torch_math",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 128,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 8.344650268554688e-07,
      "rel_err": 1.1250041325183702e-06,
      "peak_mb": 46.2509765625,
      "compile_s": 0.0006442800076911226,
      "inner": 87,
      "median_ms": 0.48083659771952947,
      "q1_ms": 0.4699685919933535,
      "q3_ms": 0.5003415633042493,
      "tflops": 1.953936327758524
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 128,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0009992122650146484,
      "rel_err": 0.0003183129010722041,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.3189657909970265,
      "inner": 200,
      "median_ms": 0.032205379975494,
      "q1_ms": 0.03202699503162876,
      "q3_ms": 0.032297379984811414,
      "tflops": 4.1675561071513565
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 128,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0009992122650146484,
      "rel_err": 0.0003183129010722041,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.32322002699947916,
      "inner": 200,
      "median_ms": 0.03195456003595609,
      "q1_ms": 0.031929737524478696,
      "q3_ms": 0.033506672516523395,
      "tflops": 4.200268376374914
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 128,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0009992122650146484,
      "rel_err": 0.0003183129010722041,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.32459786599793006,
      "inner": 200,
      "median_ms": 0.031721970008220524,
      "q1_ms": 0.03166741746099433,
      "q3_ms": 0.03234494499338325,
      "tflops": 4.231065345727847
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 128,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0011777877807617188,
      "rel_err": 0.0003752006159629673,
      "xla_temp_mb": 4.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.575973670987878,
      "inner": 200,
      "median_ms": 0.04933196498313919,
      "q1_ms": 0.049227917479583994,
      "q3_ms": 0.049372812536603305,
      "tflops": 2.7207050853513195
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 128,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 1.7285346984863281e-06,
      "rel_err": 5.506486786543974e-07,
      "peak_mb": 28.2509765625,
      "compile_s": 0.00010304599709343165,
      "inner": 200,
      "median_ms": 0.04316944003221579,
      "q1_ms": 0.04316217997256899,
      "q3_ms": 0.04317658000218216,
      "tflops": 3.1090912437093965
    },
    {
      "impl": "torch_math",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 128,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 9.5367431640625e-07,
      "rel_err": 3.03806160673048e-07,
      "peak_mb": 32.81689453125,
      "compile_s": 0.00026390000130049884,
      "inner": 200,
      "median_ms": 0.1796664099674672,
      "q1_ms": 0.1786199050184223,
      "q3_ms": 0.1810589450178668,
      "tflops": 0.7470385144574507
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 128,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0018055438995361328,
      "rel_err": 0.0005866009159944952,
      "xla_temp_mb": 2.031494140625,
      "xla_arg_mb": 8.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.8161001849948661,
      "inner": 200,
      "median_ms": 0.0921562549774535,
      "q1_ms": 0.0921435525015113,
      "q3_ms": 0.092187312511669,
      "tflops": 5.09745158497304
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 128,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0018055438995361328,
      "rel_err": 0.0005866009159944952,
      "xla_temp_mb": 2.031494140625,
      "xla_arg_mb": 8.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.8186188449908514,
      "inner": 200,
      "median_ms": 0.08444965998933185,
      "q1_ms": 0.08441812748060329,
      "q3_ms": 0.08445121253316756,
      "tflops": 5.562628056280427
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 128,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0018055438995361328,
      "rel_err": 0.0005866009159944952,
      "xla_temp_mb": 2.031494140625,
      "xla_arg_mb": 8.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.8180102329934016,
      "inner": 200,
      "median_ms": 0.08478043499053456,
      "q1_ms": 0.08477260002109688,
      "q3_ms": 0.08479584001179319,
      "tflops": 5.540925191672434
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 128,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0018055438995361328,
      "rel_err": 0.0006100927712395787,
      "xla_temp_mb": 4.0,
      "xla_arg_mb": 8.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 1.2340413630008698,
      "inner": 190,
      "median_ms": 0.10965616835930729,
      "q1_ms": 0.10913731844323736,
      "q3_ms": 0.10969850788425997,
      "tflops": 4.283954610384924
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 128,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 3.5762786865234375e-06,
      "rel_err": 1.659371719142655e-06,
      "peak_mb": 44.298828125,
      "compile_s": 0.00037737500679213554,
      "inner": 150,
      "median_ms": 0.24620313333192217,
      "q1_ms": 0.24165669330007705,
      "q3_ms": 0.24815663666231558,
      "tflops": 1.9080262774994166
    },
    {
      "impl": "torch_math",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 128,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 1.6689300537109375e-06,
      "rel_err": 6.590976795450842e-07,
      "peak_mb": 46.2509765625,
      "compile_s": 0.0007208449969766662,
      "inner": 79,
      "median_ms": 0.5414975063633645,
      "q1_ms": 0.5407320126254536,
      "q3_ms": 0.5485240443347352,
      "tflops": 0.8675239358992959
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 512,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0002974569797515869,
      "rel_err": 0.00043655920308083296,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.323316709997016,
      "inner": 149,
      "median_ms": 0.22461559059373384,
      "q1_ms": 0.21266783557358768,
      "q3_ms": 0.22467012414021179,
      "tflops": 19.121412207616444
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 512,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0002974569797515869,
      "rel_err": 0.00043655920308083296,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.30777482100529596,
      "inner": 139,
      "median_ms": 0.20956444600513274,
      "q1_ms": 0.20893098202931004,
      "q3_ms": 0.20962549997878885,
      "tflops": 20.494732660399873
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 512,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0002974569797515869,
      "rel_err": 0.00043655920308083296,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.308063928008778,
      "inner": 133,
      "median_ms": 0.20971583455006235,
      "q1_ms": 0.20950490224115634,
      "q3_ms": 0.20982498868501612,
      "tflops": 20.479938032407972
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 512,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00038889050483703613,
      "rel_err": 0.0005707505624741316,
      "xla_temp_mb": 64.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.9308899269963149,
      "inner": 118,
      "median_ms": 0.32209057616902576,
      "q1_ms": 0.3216120423561641,
      "q3_ms": 0.3221736610072894,
      "tflops": 13.334656813262676
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 512,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 1.3113021850585938e-06,
      "rel_err": 1.9245171642978676e-06,
      "peak_mb": 64.2509765625,
      "compile_s": 0.0005452820041682571,
      "inner": 95,
      "median_ms": 0.48925716837402433,
      "q1_ms": 0.4891751263609254,
      "q3_ms": 0.4897074473678673,
      "tflops": 8.77854750758932
    },
    {
      "impl": "torch_math",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 512,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 9.5367431640625e-07,
      "rel_err": 1.3996487950862502e-06,
      "peak_mb": 136.26611328125,
      "compile_s": 0.0005803700041724369,
      "inner": 90,
      "median_ms": 0.5471161110714698,
      "q1_ms": 0.5470377944422782,
      "q3_ms": 0.547493488961158,
      "tflops": 7.850193421628097
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 512,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0004993081092834473,
      "rel_err": 0.0006850222707726061,
      "xla_temp_mb": 8.125244140625,
      "xla_arg_mb": 32.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.7805289800016908,
      "inner": 43,
      "median_ms": 0.988171162790359,
      "q1_ms": 0.9879285931928343,
      "q3_ms": 0.9884849069897237,
      "tflops": 15.21232970769167
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 512,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0004993081092834473,
      "rel_err": 0.0006850222707726061,
      "xla_temp_mb": 8.125244140625,
      "xla_arg_mb": 32.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.7962909259949811,
      "inner": 43,
      "median_ms": 0.9878761863314309,
      "q1_ms": 0.9876645116343943,
      "q3_ms": 0.9887699768093392,
      "tflops": 15.216872057442892
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 512,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0004993081092834473,
      "rel_err": 0.0006850222707726061,
      "xla_temp_mb": 8.125244140625,
      "xla_arg_mb": 32.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.7750467180012492,
      "inner": 44,
      "median_ms": 0.9873238635588097,
      "q1_ms": 0.9870292272932023,
      "q3_ms": 0.9877699545044876,
      "tflops": 15.22538458841231
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 512,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0005005598068237305,
      "rel_err": 0.0006719224038533866,
      "xla_temp_mb": 96.0625,
      "xla_arg_mb": 32.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 1.6978769730048953,
      "inner": 52,
      "median_ms": 0.8399961538648663,
      "q1_ms": 0.8397569616555791,
      "q3_ms": 0.8400051153260462,
      "tflops": 17.895779006648073
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 512,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 1.6689300537109375e-06,
      "rel_err": 2.3817985947971465e-06,
      "peak_mb": 128.439453125,
      "compile_s": 0.002350820999708958,
      "inner": 21,
      "median_ms": 2.3168948093717474,
      "q1_ms": 2.316484261730996,
      "q3_ms": 2.317657666502609,
      "tflops": 6.488160565250782
    },
    {
      "impl": "torch_math",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 512,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 1.0728836059570312e-06,
      "rel_err": 1.4942368125048233e-06,
      "peak_mb": 232.2509765625,
      "compile_s": 0.0014055419887881726,
      "inner": 36,
      "median_ms": 1.4274632501635802,
      "q1_ms": 1.427149763710784,
      "q3_ms": 1.4297932221880827,
      "tflops": 10.530838909006844
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 512,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0009677410125732422,
      "rel_err": 0.0002774132299236953,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.31913721600722056,
      "inner": 177,
      "median_ms": 0.1463241356235566,
      "q1_ms": 0.1456134378420726,
      "q3_ms": 0.14666087290247126,
      "tflops": 14.676209354311595
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 512,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0009677410125732422,
      "rel_err": 0.0002774132299236953,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.31578734300273936,
      "inner": 200,
      "median_ms": 0.14686249502119608,
      "q1_ms": 0.14604163749027066,
      "q3_ms": 0.14705147499626037,
      "tflops": 14.622410219096864
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 512,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0009677410125732422,
      "rel_err": 0.0002774132299236953,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.3239951040013693,
      "inner": 189,
      "median_ms": 0.1462269418105414,
      "q1_ms": 0.14525369844040908,
      "q3_ms": 0.14676873282826303,
      "tflops": 14.685964305964777
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 512,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0014874935150146484,
      "rel_err": 0.0004264058079570532,
      "xla_temp_mb": 64.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.82660724199377,
      "inner": 109,
      "median_ms": 0.3556722017868054,
      "q1_ms": 0.35550178437431795,
      "q3_ms": 0.3561907982023288,
      "tflops": 6.037816948335
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 512,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 1.4901161193847656e-06,
      "rel_err": 4.271576017345069e-07,
      "peak_mb": 64.2509765625,
      "compile_s": 0.0003608429979067296,
      "inner": 145,
      "median_ms": 0.30842076551474246,
      "q1_ms": 0.3083791414171393,
      "q3_ms": 0.3088473033758105,
      "tflops": 6.962837422493041
    },
    {
      "impl": "torch_math",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 512,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 8.344650268554688e-07,
      "rel_err": 2.3920824787637684e-07,
      "peak_mb": 137.26611328125,
      "compile_s": 0.0007212260097730905,
      "inner": 73,
      "median_ms": 0.6421330410425794,
      "q1_ms": 0.6415638767664436,
      "q3_ms": 0.6424722397079602,
      "tflops": 3.3442970704533517
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 512,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0019669532775878906,
      "rel_err": 0.0009117966401390731,
      "xla_temp_mb": 8.125244140625,
      "xla_arg_mb": 32.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.8239233889908064,
      "inner": 65,
      "median_ms": 0.6171823694155767,
      "q1_ms": 0.6168737767998559,
      "q3_ms": 0.6185599769546221,
      "tflops": 12.178236353571222
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 512,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0019669532775878906,
      "rel_err": 0.0009117966401390731,
      "xla_temp_mb": 8.125244140625,
      "xla_arg_mb": 32.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.81890203099465,
      "inner": 63,
      "median_ms": 0.6187150635043277,
      "q1_ms": 0.6183177220731204,
      "q3_ms": 0.6193207460917562,
      "tflops": 12.148068167969257
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 512,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0019669532775878906,
      "rel_err": 0.0009117966401390731,
      "xla_temp_mb": 8.125244140625,
      "xla_arg_mb": 32.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.8159916199947475,
      "inner": 65,
      "median_ms": 0.6176258153013455,
      "q1_ms": 0.6173130692332052,
      "q3_ms": 0.6182552384696185,
      "tflops": 12.169492566195242
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 512,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0018992424011230469,
      "rel_err": 0.0004914960591122508,
      "xla_temp_mb": 96.0625,
      "xla_arg_mb": 32.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.9559442290046718,
      "inner": 51,
      "median_ms": 0.8559073726692255,
      "q1_ms": 0.8553953235854814,
      "q3_ms": 0.8561833922559505,
      "tflops": 8.781549275081092
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 512,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 3.814697265625e-06,
      "rel_err": 1.6431126823590603e-06,
      "peak_mb": 128.439453125,
      "compile_s": 0.001510740999947302,
      "inner": 33,
      "median_ms": 1.4498775457137856,
      "q1_ms": 1.449556954674225,
      "q3_ms": 1.451526302844286,
      "tflops": 5.184019016102303
    },
    {
      "impl": "torch_math",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 512,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 3.337860107421875e-06,
      "rel_err": 1.078292712008988e-06,
      "peak_mb": 232.2509765625,
      "compile_s": 0.0015364390128524974,
      "inner": 33,
      "median_ms": 1.5137108487331055,
      "q1_ms": 1.5133064545366461,
      "q3_ms": 1.5145037272525246,
      "tflops": 4.965408535117951
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 600,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00028915703296661377,
      "rel_err": 0.0006467371713370085,
      "xla_temp_mb": 40.0,
      "xla_arg_mb": 28.125,
      "xla_out_mb": 9.375,
      "compile_s": 0.3442224360042019,
      "inner": 103,
      "median_ms": 0.36009594171366754,
      "q1_ms": 0.3600000339962578,
      "q3_ms": 0.36015682523161424,
      "tflops": 16.37963474936916
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 600,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00028915703296661377,
      "rel_err": 0.0006467371713370085,
      "xla_temp_mb": 40.0,
      "xla_arg_mb": 28.125,
      "xla_out_mb": 9.375,
      "compile_s": 0.3348639150062809,
      "inner": 104,
      "median_ms": 0.360515451925144,
      "q1_ms": 0.3600870864969553,
      "q3_ms": 0.360718523972117,
      "tflops": 16.360574750689707
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 600,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00028915703296661377,
      "rel_err": 0.0006467371713370085,
      "xla_temp_mb": 40.0,
      "xla_arg_mb": 28.125,
      "xla_out_mb": 9.375,
      "compile_s": 0.3369067029998405,
      "inner": 103,
      "median_ms": 0.36059494172132495,
      "q1_ms": 0.36055988354384627,
      "q3_ms": 0.36069864559199805,
      "tflops": 16.356968214374675
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 600,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0003898888826370239,
      "rel_err": 0.0008720370242372155,
      "xla_temp_mb": 87.890625,
      "xla_arg_mb": 28.125,
      "xla_out_mb": 9.375,
      "compile_s": 0.9275454730086494,
      "inner": 97,
      "median_ms": 0.45178186603865017,
      "q1_ms": 0.4515756546734525,
      "q3_ms": 0.45201158242242545,
      "tflops": 13.05550409032
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 600,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 1.3113021850585938e-06,
      "rel_err": 2.9328973596420838e-06,
      "peak_mb": 73.0009765625,
      "compile_s": 0.0006567430100403726,
      "inner": 78,
      "median_ms": 0.6080543333574986,
      "q1_ms": 0.6079798396781254,
      "q3_ms": 0.6123496537978868,
      "tflops": 9.7001857834507
    },
    {
      "impl": "torch_math",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 600,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 1.9669532775878906e-06,
      "rel_err": 4.399345925776288e-06,
      "peak_mb": 172.00537109375,
      "compile_s": 0.0007368060032604262,
      "inner": 70,
      "median_ms": 0.7221863715260821,
      "q1_ms": 0.7218098643144393,
      "q3_ms": 0.7228370214371742,
      "tflops": 8.167199261232502
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 600,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00037339329719543457,
      "rel_err": 0.0008606364717707038,
      "xla_temp_mb": 80.00052642822266,
      "xla_arg_mb": 37.5,
      "xla_out_mb": 28.125022888183594,
      "compile_s": 0.8164959959976841,
      "inner": 32,
      "median_ms": 1.415747343799012,
      "q1_ms": 1.4147553281418368,
      "q3_ms": 1.420483843958209,
      "tflops": 14.581584836037472
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 600,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00037339329719543457,
      "rel_err": 0.0008606364717707038,
      "xla_temp_mb": 80.00048828125,
      "xla_arg_mb": 37.5,
      "xla_out_mb": 28.125022888183594,
      "compile_s": 0.8092613959888695,
      "inner": 33,
      "median_ms": 1.3963425453491227,
      "q1_ms": 1.3959718485376937,
      "q3_ms": 1.3967393485169546,
      "tflops": 14.784223304489007
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 600,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00037339329719543457,
      "rel_err": 0.0008606364717707038,
      "xla_temp_mb": 80.00048828125,
      "xla_arg_mb": 37.5,
      "xla_out_mb": 28.125022888183594,
      "compile_s": 0.816825098998379,
      "inner": 31,
      "median_ms": 1.3951310970550104,
      "q1_ms": 1.3950335320964036,
      "q3_ms": 1.396143854729035,
      "tflops": 14.79706103862009
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 600,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00037339329719543457,
      "rel_err": 0.0008606364717707038,
      "xla_temp_mb": 131.9091796875,
      "xla_arg_mb": 37.5,
      "xla_out_mb": 28.125022888183594,
      "compile_s": 1.5711316340020858,
      "inner": 40,
      "median_ms": 1.1822308752016397,
      "q1_ms": 1.1817567250545835,
      "q3_ms": 1.1836975250844262,
      "tflops": 17.46176692981311
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 600,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 1.2218952178955078e-06,
      "rel_err": 2.8163535716885235e-06,
      "peak_mb": 149.3994140625,
      "compile_s": 0.002947181012132205,
      "inner": 17,
      "median_ms": 2.9731274116784334,
      "q1_ms": 2.972616147256785,
      "q3_ms": 2.9781074703992894,
      "tflops": 6.943476394220804
    },
    {
      "impl": "torch_math",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 600,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 1.6093254089355469e-06,
      "rel_err": 3.709343900482054e-06,
      "peak_mb": 295.8759765625,
      "compile_s": 0.0017875749908853322,
      "inner": 28,
      "median_ms": 1.8508135711142262,
      "q1_ms": 1.8491557678186967,
      "q3_ms": 1.8527303927839966,
      "tflops": 11.15392729024134
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 600,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0012070536613464355,
      "rel_err": 0.0003694213228300214,
      "xla_temp_mb": 40.0,
      "xla_arg_mb": 28.125,
      "xla_out_mb": 9.375,
      "compile_s": 0.3443258819897892,
      "inner": 120,
      "median_ms": 0.2952645166563646,
      "q1_ms": 0.2950971624765467,
      "q3_ms": 0.2959711500807316,
      "tflops": 9.988060988148641
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 600,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0012070536613464355,
      "rel_err": 0.0003694213228300214,
      "xla_temp_mb": 40.0,
      "xla_arg_mb": 28.125,
      "xla_out_mb": 9.375,
      "compile_s": 0.3409127089980757,
      "inner": 118,
      "median_ms": 0.2967860508745662,
      "q1_ms": 0.29559960597198676,
      "q3_ms": 0.29700974155892224,
      "tflops": 9.936855156465615
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 600,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0012070536613464355,
      "rel_err": 0.0003694213228300214,
      "xla_temp_mb": 40.0,
      "xla_arg_mb": 28.125,
      "xla_out_mb": 9.375,
      "compile_s": 0.3382367320009507,
      "inner": 126,
      "median_ms": 0.2964519126902497,
      "q1_ms": 0.29573507148713346,
      "q3_ms": 0.29679676980964675,
      "tflops": 9.948055228374974
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 600,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.001227259635925293,
      "rel_err": 0.00037560539203695953,
      "xla_temp_mb": 87.890625,
      "xla_arg_mb": 28.125,
      "xla_out_mb": 9.375,
      "compile_s": 0.6239410499983933,
      "inner": 83,
      "median_ms": 0.4867917711843725,
      "q1_ms": 0.48666007231909736,
      "q3_ms": 0.48722699398514313,
      "tflops": 6.058278250728729
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 600,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 1.6689300537109375e-06,
      "rel_err": 5.107795573167095e-07,
      "peak_mb": 73.0009765625,
      "compile_s": 0.0004949880094500259,
      "inner": 104,
      "median_ms": 0.4418494519576108,
      "q1_ms": 0.4416682019067454,
      "q3_ms": 0.4444930288557841,
      "tflops": 6.674490568979876
    },
    {
      "impl": "torch_math",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 600,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 1.1920928955078125e-06,
      "rel_err": 3.648425490609952e-07,
      "peak_mb": 174.00537109375,
      "compile_s": 0.0009004159946925938,
      "inner": 58,
      "median_ms": 0.8496010173232583,
      "q1_ms": 0.8495019569166069,
      "q3_ms": 0.8501871292035754,
      "tflops": 3.471182284234379
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 600,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0019297003746032715,
      "rel_err": 0.0007219392573460937,
      "xla_temp_mb": 80.00048828125,
      "xla_arg_mb": 37.5,
      "xla_out_mb": 28.125022888183594,
      "compile_s": 0.8510726389940828,
      "inner": 46,
      "median_ms": 0.9579944564156112,
      "q1_ms": 0.9571075541949223,
      "q3_ms": 0.9584496631008654,
      "tflops": 10.774509112109092
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 600,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0019297003746032715,
      "rel_err": 0.0007219392573460937,
      "xla_temp_mb": 80.00048828125,
      "xla_arg_mb": 37.5,
      "xla_out_mb": 28.125022888183594,
      "compile_s": 0.8431989189994056,
      "inner": 48,
      "median_ms": 0.934455083249001,
      "q1_ms": 0.9327091561317502,
      "q3_ms": 0.936096302060226,
      "tflops": 11.045924180873179
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 600,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0019297003746032715,
      "rel_err": 0.0007219392573460937,
      "xla_temp_mb": 80.00048828125,
      "xla_arg_mb": 37.5,
      "xla_out_mb": 28.125022888183594,
      "compile_s": 0.846813183001359,
      "inner": 46,
      "median_ms": 0.9345916740763326,
      "q1_ms": 0.9328633804334854,
      "q3_ms": 0.936679858726996,
      "tflops": 11.044309816050168
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 600,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0017391443252563477,
      "rel_err": 0.0005386559059843421,
      "xla_temp_mb": 131.9091796875,
      "xla_arg_mb": 37.5,
      "xla_out_mb": 28.125022888183594,
      "compile_s": 1.0843617250066018,
      "inner": 38,
      "median_ms": 1.2229286053058643,
      "q1_ms": 1.2228662237344237,
      "q3_ms": 1.224023026320777,
      "tflops": 8.440329186198408
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 600,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 3.933906555175781e-06,
      "rel_err": 1.4717526255481062e-06,
      "peak_mb": 149.3994140625,
      "compile_s": 0.0021807300072396174,
      "inner": 23,
      "median_ms": 2.0325165216649035,
      "q1_ms": 2.0267325434479457,
      "q3_ms": 2.108169130224269,
      "tflops": 5.078394143406502
    },
    {
      "impl": "torch_math",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 600,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 3.814697265625e-06,
      "rel_err": 1.2933584230268025e-06,
      "peak_mb": 295.8759765625,
      "compile_s": 0.0019499720074236393,
      "inner": 26,
      "median_ms": 1.9615797310297447,
      "q1_ms": 1.9607462692455067,
      "q3_ms": 1.964489692349721,
      "tflops": 5.262044584127833
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 1024,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00017434358596801758,
      "rel_err": 0.00046585098607465625,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.31509357898903545,
      "inner": 59,
      "median_ms": 0.7054566781103643,
      "q1_ms": 0.7020251780316249,
      "q3_ms": 0.7056527287882783,
      "tflops": 24.35283372753375
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 1024,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00017434358596801758,
      "rel_err": 0.00046585098607465625,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.3134680710063549,
      "inner": 64,
      "median_ms": 0.7055685312025162,
      "q1_ms": 0.7055610938095924,
      "q3_ms": 0.7055829453292972,
      "tflops": 24.3489730965183
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 1024,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00017434358596801758,
      "rel_err": 0.00046585098607465625,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.3084892029874027,
      "inner": 60,
      "median_ms": 0.7053978833331105,
      "q1_ms": 0.7052450084302109,
      "q3_ms": 0.7055326333405294,
      "tflops": 24.35486353151862
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 1024,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00016224384307861328,
      "rel_err": 0.0004335201228968799,
      "xla_temp_mb": 256.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 1.0836363319976954,
      "inner": 46,
      "median_ms": 1.0527946306746858,
      "q1_ms": 1.0524237064680126,
      "q3_ms": 1.0530639347448216,
      "tflops": 16.318348026708914
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 1024,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 1.043081283569336e-06,
      "rel_err": 2.787142648230656e-06,
      "peak_mb": 112.2509765625,
      "compile_s": 0.0018232529982924461,
      "inner": 28,
      "median_ms": 1.7070928214317453,
      "q1_ms": 1.7000339285004884,
      "q3_ms": 1.713475482184939,
      "tflops": 10.063816664398585
    },
    {
      "impl": "torch_math",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 1024,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 7.152557373046875e-07,
      "rel_err": 1.9111835172225256e-06,
      "peak_mb": 400.28173828125,
      "compile_s": 0.001927158999023959,
      "inner": 26,
      "median_ms": 1.9533021538966122,
      "q1_ms": 1.9523522885440168,
      "q3_ms": 1.9542098461198083,
      "tflops": 8.795295264344098
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 1024,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0004487037658691406,
      "rel_err": 0.0006063930923119187,
      "xla_temp_mb": 16.250244140625,
      "xla_arg_mb": 64.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.7790155440015951,
      "inner": 14,
      "median_ms": 3.3025807143920765,
      "q1_ms": 3.2984160721493287,
      "q3_ms": 3.3137076071787828,
      "tflops": 18.20683500087245
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 1024,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0004487037658691406,
      "rel_err": 0.0006063930923119187,
      "xla_temp_mb": 16.250244140625,
      "xla_arg_mb": 64.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.770507983994321,
      "inner": 14,
      "median_ms": 3.319182214909233,
      "q1_ms": 3.3055174286086446,
      "q3_ms": 3.3276697145733385,
      "tflops": 18.115770165888378
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 1024,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0004487037658691406,
      "rel_err": 0.0006063930923119187,
      "xla_temp_mb": 16.250244140625,
      "xla_arg_mb": 64.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.7733171529980609,
      "inner": 14,
      "median_ms": 3.3111133566957767,
      "q1_ms": 3.298817250262281,
      "q3_ms": 3.320771286131016,
      "tflops": 18.159916519440586
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 1024,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0004436969757080078,
      "rel_err": 0.0005996267427690327,
      "xla_temp_mb": 384.0,
      "xla_arg_mb": 64.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 1.6776887360028923,
      "inner": 18,
      "median_ms": 2.616823444744417,
      "q1_ms": 2.604102139432345,
      "q3_ms": 2.624122277969339,
      "tflops": 22.978066122406208
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 1024,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 1.4603137969970703e-06,
      "rel_err": 2.4182752440538025e-06,
      "peak_mb": 240.626953125,
      "compile_s": 0.007833373005269095,
      "inner": 6,
      "median_ms": 7.990974833470924,
      "q1_ms": 7.988453333382495,
      "q3_ms": 8.02099574987854,
      "tflops": 7.524681706184576
    },
    {
      "impl": "torch_math",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 1024,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 1.1026859283447266e-06,
      "rel_err": 1.4902061593602411e-06,
      "peak_mb": 704.2509765625,
      "compile_s": 0.004921048006508499,
      "inner": 10,
      "median_ms": 5.093305499758571,
      "q1_ms": 5.080572849692544,
      "q3_ms": 5.094138599815778,
      "tflops": 11.805602893219385
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 1024,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0011317729949951172,
      "rel_err": 0.0003316882357466966,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.33170485399023164,
      "inner": 94,
      "median_ms": 0.42681202126488565,
      "q1_ms": 0.42612465953657486,
      "q3_ms": 0.4275352553219257,
      "tflops": 20.125802845344328
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 1024,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0011317729949951172,
      "rel_err": 0.0003316882357466966,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.3183819759869948,
      "inner": 89,
      "median_ms": 0.4301619437376686,
      "q1_ms": 0.42927178648742587,
      "q3_ms": 0.4316441404268282,
      "tflops": 19.969071455652788
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 1024,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0011317729949951172,
      "rel_err": 0.0003316882357466966,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.3198348470032215,
      "inner": 87,
      "median_ms": 0.43161429886989855,
      "q1_ms": 0.42938484484798156,
      "q3_ms": 0.43169197707701923,
      "tflops": 19.901876778621887
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 1024,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0013011693954467773,
      "rel_err": 0.000381333171389997,
      "xla_temp_mb": 256.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.6044734880124452,
      "inner": 44,
      "median_ms": 1.0647288634738121,
      "q1_ms": 1.064146954542428,
      "q3_ms": 1.0648694546530764,
      "tflops": 8.067720230645627
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 1024,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 2.9206275939941406e-06,
      "rel_err": 8.559471211810887e-07,
      "peak_mb": 112.2509765625,
      "compile_s": 0.001029849998303689,
      "inner": 48,
      "median_ms": 0.9909447502044106,
      "q1_ms": 0.9896328959560681,
      "q3_ms": 0.9939296769516659,
      "tflops": 8.66842938542041
    },
    {
      "impl": "torch_math",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 1024,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 1.1920928955078125e-06,
      "rel_err": 3.493661608899856e-07,
      "peak_mb": 404.28173828125,
      "compile_s": 0.002423148005618714,
      "inner": 20,
      "median_ms": 2.418809200025862,
      "q1_ms": 2.4185764752473915,
      "q3_ms": 2.4190546249883482,
      "tflops": 3.5513072266750743
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 1024,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0017703771591186523,
      "rel_err": 0.0008779477793723345,
      "xla_temp_mb": 16.250244140625,
      "xla_arg_mb": 64.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.8203301750036189,
      "inner": 25,
      "median_ms": 1.8892894801683724,
      "q1_ms": 1.8871389998821542,
      "q3_ms": 1.8909953202819452,
      "tflops": 15.913268658713244
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 1024,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0017703771591186523,
      "rel_err": 0.0008779477793723345,
      "xla_temp_mb": 16.250244140625,
      "xla_arg_mb": 64.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.8188229210063582,
      "inner": 24,
      "median_ms": 1.8883327496344766,
      "q1_ms": 1.8868485000590833,
      "q3_ms": 1.8891908330260776,
      "tflops": 15.921331173131229
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 1024,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0017703771591186523,
      "rel_err": 0.0008779477793723345,
      "xla_temp_mb": 16.250244140625,
      "xla_arg_mb": 64.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.8115010769979563,
      "inner": 25,
      "median_ms": 1.8874408397823572,
      "q1_ms": 1.886763479851652,
      "q3_ms": 1.8888416601112112,
      "tflops": 15.92885479550543
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 1024,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0017703771591186523,
      "rel_err": 0.0006765168509446084,
      "xla_temp_mb": 384.0,
      "xla_arg_mb": 64.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 1.3690041630034102,
      "inner": 17,
      "median_ms": 2.7925331181476767,
      "q1_ms": 2.79201270570732,
      "q3_ms": 2.795941617531592,
      "tflops": 10.766128744049542
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 1024,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 4.76837158203125e-06,
      "rel_err": 2.1831258436577627e-06,
      "peak_mb": 240.626953125,
      "compile_s": 0.004558521992294118,
      "inner": 11,
      "median_ms": 4.608716909636065,
      "q1_ms": 4.6041123185370285,
      "q3_ms": 4.609635909384286,
      "tflops": 6.523457973549976
    },
    {
      "impl": "torch_math",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 1024,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 4.291534423828125e-06,
      "rel_err": 1.6426589581897133e-06,
      "peak_mb": 704.2509765625,
      "compile_s": 0.0053535570041276515,
      "inner": 9,
      "median_ms": 5.4993734439550375,
      "q1_ms": 5.497741390071395,
      "q3_ms": 5.513196221727412,
      "tflops": 5.466944803511657
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 1500,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00017557293176651,
      "rel_err": 0.0007930556894280016,
      "xla_temp_mb": 96.0,
      "xla_arg_mb": 70.3125,
      "xla_out_mb": 23.4375,
      "compile_s": 0.34133414798998274,
      "inner": 27,
      "median_ms": 1.7240635185131665,
      "q1_ms": 1.723731388833753,
      "q3_ms": 1.7243789075615092,
      "tflops": 21.382042833196504
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 1500,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00017557293176651,
      "rel_err": 0.0007930556894280016,
      "xla_temp_mb": 96.0,
      "xla_arg_mb": 70.3125,
      "xla_out_mb": 23.4375,
      "compile_s": 0.3397911360079888,
      "inner": 27,
      "median_ms": 1.7265986669506808,
      "q1_ms": 1.7256309071887734,
      "q3_ms": 1.727805592828965,
      "tflops": 21.35064778261699
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 1500,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00017557293176651,
      "rel_err": 0.0007930556894280016,
      "xla_temp_mb": 96.0,
      "xla_arg_mb": 70.3125,
      "xla_out_mb": 23.4375,
      "compile_s": 0.3407057370059192,
      "inner": 27,
      "median_ms": 1.7265504074101854,
      "q1_ms": 1.725644851696281,
      "q3_ms": 1.7280094999242228,
      "tflops": 21.351244563601107
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 1500,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00016435235738754272,
      "rel_err": 0.0007423728238791227,
      "xla_temp_mb": 549.3164215087891,
      "xla_arg_mb": 70.3125,
      "xla_out_mb": 23.4375,
      "compile_s": 1.026975479995599,
      "inner": 22,
      "median_ms": 2.271715681920548,
      "q1_ms": 2.2705983865307644,
      "q3_ms": 2.2729959094828125,
      "tflops": 16.227382807356655
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 1500,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 5.066394805908203e-07,
      "rel_err": 2.2884696591063403e-06,
      "peak_mb": 159.6884765625,
      "compile_s": 0.004031274002045393,
      "inner": 12,
      "median_ms": 3.648445250291843,
      "q1_ms": 3.6477969582241103,
      "q3_ms": 3.7742054167514048,
      "tflops": 10.104029928104639
    },
    {
      "impl": "torch_math",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 1500,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 4.470348358154297e-07,
      "rel_err": 2.019237854256062e-06,
      "peak_mb": 777.71484375,
      "compile_s": 0.003923149997717701,
      "inner": 12,
      "median_ms": 4.013828833800896,
      "q1_ms": 4.010953332908684,
      "q3_ms": 4.015851750712803,
      "tflops": 9.1842481397224
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 1500,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00042682886123657227,
      "rel_err": 0.0010993176838383079,
      "xla_temp_mb": 192.00057983398438,
      "xla_arg_mb": 93.75,
      "xla_out_mb": 70.3125228881836,
      "compile_s": 0.8254231289902236,
      "inner": 6,
      "median_ms": 7.4731543330320465,
      "q1_ms": 7.4672078335424885,
      "q3_ms": 7.4918804166372865,
      "tflops": 17.264998720781364
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 1500,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00042682886123657227,
      "rel_err": 0.0010993176838383079,
      "xla_temp_mb": 192.00048828125,
      "xla_arg_mb": 93.75,
      "xla_out_mb": 70.3125228881836,
      "compile_s": 0.818079814009252,
      "inner": 6,
      "median_ms": 7.4110065009638975,
      "q1_ms": 7.408070917639028,
      "q3_ms": 7.4292627480948195,
      "tflops": 17.40978097687794
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 1500,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00042682886123657227,
      "rel_err": 0.0010993176838383079,
      "xla_temp_mb": 192.00048828125,
      "xla_arg_mb": 93.75,
      "xla_out_mb": 70.3125228881836,
      "compile_s": 0.8297970220010029,
      "inner": 6,
      "median_ms": 7.424631000806888,
      "q1_ms": 7.402699166656627,
      "q3_ms": 7.4255039168444155,
      "tflops": 17.377833320737157
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 1500,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0004245638847351074,
      "rel_err": 0.0011021577520295978,
      "xla_temp_mb": 824.1577301025391,
      "xla_arg_mb": 93.75,
      "xla_out_mb": 70.3125228881836,
      "compile_s": 1.6221394909953233,
      "inner": 8,
      "median_ms": 6.008497375660227,
      "q1_ms": 6.001888750688522,
      "q3_ms": 6.02942593741318,
      "tflops": 21.47358847532534
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 1500,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 1.2516975402832031e-06,
      "rel_err": 2.839966782630654e-06,
      "peak_mb": 352.07958984375,
      "compile_s": 0.01645895600086078,
      "inner": 2,
      "median_ms": 17.051593502401374,
      "q1_ms": 16.984348749247147,
      "q3_ms": 17.061637499864446,
      "tflops": 7.566682842974738
    },
    {
      "impl": "torch_math",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 1500,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 7.450580596923828e-07,
      "rel_err": 1.661755391069164e-06,
      "peak_mb": 1376.0712890625,
      "compile_s": 0.009875589996227063,
      "inner": 5,
      "median_ms": 10.258998000063002,
      "q1_ms": 10.258351900847629,
      "q3_ms": 10.270245299034286,
      "tflops": 12.576666843994673
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 1500,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0008330345153808594,
      "rel_err": 0.0003061060851905495,
      "xla_temp_mb": 96.0,
      "xla_arg_mb": 70.3125,
      "xla_out_mb": 23.4375,
      "compile_s": 0.34318403899669647,
      "inner": 41,
      "median_ms": 1.1012579758776488,
      "q1_ms": 1.0996752439163306,
      "q3_ms": 1.1024457197025894,
      "tflops": 16.73722270688718
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 1500,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0008330345153808594,
      "rel_err": 0.0003061060851905495,
      "xla_temp_mb": 96.0,
      "xla_arg_mb": 70.3125,
      "xla_out_mb": 23.4375,
      "compile_s": 0.33613332798995543,
      "inner": 40,
      "median_ms": 1.1104779998277081,
      "q1_ms": 1.1086518126830924,
      "q3_ms": 1.1143470874230843,
      "tflops": 16.598257689805422
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 1500,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0008330345153808594,
      "rel_err": 0.0003061060851905495,
      "xla_temp_mb": 96.0,
      "xla_arg_mb": 70.3125,
      "xla_out_mb": 23.4375,
      "compile_s": 0.34351848300138954,
      "inner": 41,
      "median_ms": 1.1100546338328545,
      "q1_ms": 1.1088712928495825,
      "q3_ms": 1.1120985120835871,
      "tflops": 16.604588133069655
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 1500,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0012900829315185547,
      "rel_err": 0.00047405267832800746,
      "xla_temp_mb": 549.3164215087891,
      "xla_arg_mb": 70.3125,
      "xla_out_mb": 23.4375,
      "compile_s": 0.41078149000531994,
      "inner": 22,
      "median_ms": 2.241206136452754,
      "q1_ms": 2.2405544315750543,
      "q3_ms": 2.2420482048553163,
      "tflops": 8.224143107680876
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 1500,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 1.5497207641601562e-06,
      "rel_err": 5.694589617633028e-07,
      "peak_mb": 159.6884765625,
      "compile_s": 0.0020678360015153885,
      "inner": 24,
      "median_ms": 2.0338637083720337,
      "q1_ms": 2.0330085832635327,
      "q3_ms": 2.047861499704595,
      "tflops": 9.062554154503072
    },
    {
      "impl": "torch_math",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 1500,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 1.1920928955078125e-06,
      "rel_err": 4.3804536176139663e-07,
      "peak_mb": 786.86083984375,
      "compile_s": 0.0049672049935907125,
      "inner": 10,
      "median_ms": 5.010840299655683,
      "q1_ms": 5.01025670018862,
      "q3_ms": 5.011201999877812,
      "tflops": 3.6784249542470033
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 1500,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.002349555492401123,
      "rel_err": 0.0008802466909401119,
      "xla_temp_mb": 192.00048828125,
      "xla_arg_mb": 93.75,
      "xla_out_mb": 70.3125228881836,
      "compile_s": 0.8473456799983978,
      "inner": 11,
      "median_ms": 4.302347544580698,
      "q1_ms": 4.293061454849191,
      "q3_ms": 4.308016545440875,
      "tflops": 14.99460453427578
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 1500,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.002349555492401123,
      "rel_err": 0.0008802466909401119,
      "xla_temp_mb": 192.00048828125,
      "xla_arg_mb": 93.75,
      "xla_out_mb": 70.3125228881836,
      "compile_s": 0.8412481150007807,
      "inner": 11,
      "median_ms": 4.238770000731826,
      "q1_ms": 4.2354842728607105,
      "q3_ms": 4.248916045154064,
      "tflops": 15.219509430533376
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 1500,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.002349555492401123,
      "rel_err": 0.0008802466909401119,
      "xla_temp_mb": 192.00048828125,
      "xla_arg_mb": 93.75,
      "xla_out_mb": 70.3125228881836,
      "compile_s": 0.8508221539959777,
      "inner": 11,
      "median_ms": 4.240338455102491,
      "q1_ms": 4.233355182143647,
      "q3_ms": 4.246746499451216,
      "tflops": 15.213879902056242
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 1500,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0017014145851135254,
      "rel_err": 0.0006374246440827847,
      "xla_temp_mb": 824.1577301025391,
      "xla_arg_mb": 93.75,
      "xla_out_mb": 70.3125228881836,
      "compile_s": 0.9855508809996536,
      "inner": 8,
      "median_ms": 6.094306250815862,
      "q1_ms": 6.078439875636832,
      "q3_ms": 6.105605500124511,
      "tflops": 10.58561833701147
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 1500,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 5.4836273193359375e-06,
      "rel_err": 2.0544077869999455e-06,
      "peak_mb": 352.64208984375,
      "compile_s": 0.009141560003627092,
      "inner": 5,
      "median_ms": 9.393503001774661,
      "q1_ms": 9.384362000855617,
      "q3_ms": 9.402651499840431,
      "tflops": 6.867725489395395
    },
    {
      "impl": "torch_math",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 1500,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 7.62939453125e-06,
      "rel_err": 2.8583062885445543e-06,
      "peak_mb": 1376.6337890625,
      "compile_s": 0.01085030598915182,
      "inner": 4,
      "median_ms": 11.146356249810196,
      "q1_ms": 11.14578650049225,
      "q3_ms": 11.16164774975914,
      "tflops": 5.787720987394291
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 2048,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00010418146848678589,
      "rel_err": 0.0005699748871847987,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.3148226660123328,
      "inner": 17,
      "median_ms": 2.7210760591140786,
      "q1_ms": 2.7134914117221554,
      "q3_ms": 2.7216497649479767,
      "tflops": 25.254522565008173
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 2048,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00010418146848678589,
      "rel_err": 0.0005699748871847987,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.3100450290075969,
      "inner": 19,
      "median_ms": 2.586973315613393,
      "q1_ms": 2.5857654997564263,
      "q3_ms": 2.593112815750493,
      "tflops": 26.563658898702645
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 2048,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00010418146848678589,
      "rel_err": 0.0005699748871847987,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.3156284319993574,
      "inner": 17,
      "median_ms": 2.5876492941546636,
      "q1_ms": 2.5865970002247622,
      "q3_ms": 2.5882071174290853,
      "tflops": 26.556719603090325
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 2048,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00011554360389709473,
      "rel_err": 0.0006321368855424225,
      "xla_temp_mb": 1024.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 1.1708237760030897,
      "inner": 13,
      "median_ms": 3.7963856145954478,
      "q1_ms": 3.785256846570589,
      "q3_ms": 3.7998741155696245,
      "tflops": 18.101289940569675
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 2048,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 6.407499313354492e-07,
      "rel_err": 3.505530912661925e-06,
      "peak_mb": 208.2509765625,
      "compile_s": 0.006384699998307042,
      "inner": 7,
      "median_ms": 6.3326025722614885,
      "q1_ms": 6.318175999954942,
      "q3_ms": 6.3668760703876615,
      "tflops": 10.851695799924329
    },
    {
      "impl": "torch_math",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 2048,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 4.023313522338867e-07,
      "rel_err": 2.2011474811733933e-06,
      "peak_mb": 1360.31298828125,
      "compile_s": 0.00711170700378716,
      "inner": 7,
      "median_ms": 7.28303228554848,
      "q1_ms": 7.265741000342782,
      "q3_ms": 7.2868917143620955,
      "tflops": 9.435558437981683
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 2048,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00021325796842575073,
      "rel_err": 0.0008774645393714309,
      "xla_temp_mb": 32.500244140625,
      "xla_arg_mb": 128.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 0.7870828700106358,
      "inner": 4,
      "median_ms": 12.106532500183675,
      "q1_ms": 12.097434499082738,
      "q3_ms": 12.141525750848814,
      "tflops": 19.86680897873532
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 2048,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00021325796842575073,
      "rel_err": 0.0008774645393714309,
      "xla_temp_mb": 32.500244140625,
      "xla_arg_mb": 128.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 0.7977755469910335,
      "inner": 4,
      "median_ms": 12.07720925231115,
      "q1_ms": 12.047029249515617,
      "q3_ms": 12.12201237467525,
      "tflops": 19.915045235303293
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 2048,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00021325796842575073,
      "rel_err": 0.0008774645393714309,
      "xla_temp_mb": 32.500244140625,
      "xla_arg_mb": 128.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 0.7900392679875949,
      "inner": 4,
      "median_ms": 12.102750250051031,
      "q1_ms": 12.082847000783659,
      "q3_ms": 12.116543497540988,
      "tflops": 19.873017587468258
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 2048,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00021577626466751099,
      "rel_err": 0.00088782625971362,
      "xla_temp_mb": 1536.250015258789,
      "xla_arg_mb": 128.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 1.8453795800014632,
      "inner": 5,
      "median_ms": 10.023415798787028,
      "q1_ms": 10.02025289926678,
      "q3_ms": 10.024086201156024,
      "tflops": 23.99562917514666
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 2048,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 8.046627044677734e-07,
      "rel_err": 3.310839929326903e-06,
      "peak_mb": 465.001953125,
      "compile_s": 0.02995310799451545,
      "inner": 1,
      "median_ms": 30.114654000499286,
      "q1_ms": 30.025530002603773,
      "q3_ms": 30.135624001559336,
      "tflops": 7.986748530201022
    },
    {
      "impl": "torch_math",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 2048,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 7.450580596923828e-07,
      "rel_err": 2.697721356526017e-06,
      "peak_mb": 2416.2509765625,
      "compile_s": 0.017964787999517284,
      "inner": 2,
      "median_ms": 18.642030001501553,
      "q1_ms": 18.630034501256887,
      "q3_ms": 18.664928251382662,
      "tflops": 12.90193013081875
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 2048,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0008929967880249023,
      "rel_err": 0.00025551050202921033,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.32323842200275976,
      "inner": 29,
      "median_ms": 1.5162701381306196,
      "q1_ms": 1.463247482426031,
      "q3_ms": 1.5359810515355061,
      "tflops": 22.66069713036851
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 2048,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0008929967880249023,
      "rel_err": 0.00025551050202921033,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.3233651920018019,
      "inner": 33,
      "median_ms": 1.4494989397159468,
      "q1_ms": 1.446722030364485,
      "q3_ms": 1.4581096666833302,
      "tflops": 23.70456260887873
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 2048,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0008929967880249023,
      "rel_err": 0.00025551050202921033,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.3171927939984016,
      "inner": 32,
      "median_ms": 1.4568404999408813,
      "q1_ms": 1.4526655468216632,
      "q3_ms": 1.4608900310122408,
      "tflops": 23.585106516049166
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 2048,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0011209845542907715,
      "rel_err": 0.00032074397313408554,
      "xla_temp_mb": 1024.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.5194271869986551,
      "inner": 13,
      "median_ms": 3.7896151535768205,
      "q1_ms": 3.7887256918251158,
      "q3_ms": 3.7907830388016572,
      "tflops": 9.066814696360296
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 2048,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 1.6689300537109375e-06,
      "rel_err": 4.775259867528803e-07,
      "peak_mb": 208.2509765625,
      "compile_s": 0.0035091249883407727,
      "inner": 14,
      "median_ms": 3.4909832860908603,
      "q1_ms": 3.4783854644047096,
      "q3_ms": 3.509064642587743,
      "tflops": 9.842424197474578
    },
    {
      "impl": "torch_math",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 2048,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 1.1920928955078125e-06,
      "rel_err": 3.410899864775274e-07,
      "peak_mb": 1376.31298828125,
      "compile_s": 0.009029157998156734,
      "inner": 5,
      "median_ms": 9.124235401395708,
      "q1_ms": 9.120807799627073,
      "q3_ms": 9.130326799640898,
      "tflops": 3.765766319744895
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 2048,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.002055048942565918,
      "rel_err": 0.001002775621600449,
      "xla_temp_mb": 32.500244140625,
      "xla_arg_mb": 128.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 0.8242339780117618,
      "inner": 7,
      "median_ms": 6.611008143019197,
      "q1_ms": 6.577385429409333,
      "q3_ms": 6.658118857428365,
      "tflops": 18.190733045002506
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 2048,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.002055048942565918,
      "rel_err": 0.001002775621600449,
      "xla_temp_mb": 32.500244140625,
      "xla_arg_mb": 128.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 0.8271738839976024,
      "inner": 7,
      "median_ms": 6.587956142279187,
      "q1_ms": 6.5821107865693715,
      "q3_ms": 6.602171570453461,
      "tflops": 18.2543844692923
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 2048,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.002055048942565918,
      "rel_err": 0.001002775621600449,
      "xla_temp_mb": 32.500244140625,
      "xla_arg_mb": 128.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 0.8182258509914391,
      "inner": 7,
      "median_ms": 6.570634856221399,
      "q1_ms": 6.567572573008615,
      "q3_ms": 6.582509357290942,
      "tflops": 18.3025060621247
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 2048,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0013349056243896484,
      "rel_err": 0.0006167421233840287,
      "xla_temp_mb": 1536.250015258789,
      "xla_arg_mb": 128.0,
      "xla_out_mb": 96.0000228881836,
      "compile_s": 1.0597510730003705,
      "inner": 5,
      "median_ms": 9.968147199833766,
      "q1_ms": 9.952591799083166,
      "q3_ms": 10.000575699086767,
      "tflops": 12.064336719466333
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 2048,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 7.3909759521484375e-06,
      "rel_err": 2.501267545085284e-06,
      "peak_mb": 465.001953125,
      "compile_s": 0.01558855699840933,
      "inner": 3,
      "median_ms": 16.09651333031555,
      "q1_ms": 16.079238664436467,
      "q3_ms": 16.11286566670363,
      "tflops": 7.471126312895892
    },
    {
      "impl": "torch_math",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 2048,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 5.7220458984375e-06,
      "rel_err": 2.67577456725121e-06,
      "peak_mb": 2416.2509765625,
      "compile_s": 0.019959304001531564,
      "inner": 2,
      "median_ms": 20.308374994783662,
      "q1_ms": 20.291348249884322,
      "q3_ms": 20.32152024912648,
      "tflops": 5.921649778423401
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 3000,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 9.348243474960327e-05,
      "rel_err": 0.0005278591997921467,
      "xla_temp_mb": 192.0,
      "xla_arg_mb": 140.625,
      "xla_out_mb": 46.875,
      "compile_s": 0.3454894150054315,
      "inner": 8,
      "median_ms": 6.110905000241473,
      "q1_ms": 6.103741562583309,
      "q3_ms": 6.125947999862547,
      "tflops": 24.129977473741327
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 3000,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 9.348243474960327e-05,
      "rel_err": 0.0005278591997921467,
      "xla_temp_mb": 192.0,
      "xla_arg_mb": 140.625,
      "xla_out_mb": 46.875,
      "compile_s": 0.3418167110066861,
      "inner": 8,
      "median_ms": 6.068339500416187,
      "q1_ms": 6.067198688469944,
      "q3_ms": 6.074785500459257,
      "tflops": 24.2992337508287
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 3000,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 9.348243474960327e-05,
      "rel_err": 0.0005278591997921467,
      "xla_temp_mb": 192.0,
      "xla_arg_mb": 140.625,
      "xla_out_mb": 46.875,
      "compile_s": 0.33766197400109377,
      "inner": 8,
      "median_ms": 6.065429124646471,
      "q1_ms": 6.056312375221751,
      "q3_ms": 6.108111749199452,
      "tflops": 24.310893255816357
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 3000,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00012782961130142212,
      "rel_err": 0.0007218044484034181,
      "xla_temp_mb": 2197.265640258789,
      "xla_arg_mb": 140.625,
      "xla_out_mb": 46.875,
      "compile_s": 1.1389008890109835,
      "inner": 6,
      "median_ms": 8.09374433447374,
      "q1_ms": 8.064031334167035,
      "q3_ms": 8.116008833894739,
      "tflops": 18.21851468323995
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 3000,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 3.5762786865234375e-07,
      "rel_err": 2.0193865566398017e-06,
      "peak_mb": 297.5009765625,
      "compile_s": 0.013633384995046072,
      "inner": 3,
      "median_ms": 13.719762331068827,
      "q1_ms": 13.703354665873727,
      "q3_ms": 13.776319666552201,
      "tflops": 10.747708046376381
    },
    {
      "impl": "torch_math",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 3000,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 3.5762786865234375e-07,
      "rel_err": 2.0193865566398017e-06,
      "peak_mb": 2769.51611328125,
      "compile_s": 0.01539433999278117,
      "inner": 3,
      "median_ms": 16.26209666816673,
      "q1_ms": 16.230495166382752,
      "q3_ms": 16.270530832116492,
      "tflops": 9.067465469483222
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 3000,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00014643371105194092,
      "rel_err": 0.0006584993680007756,
      "xla_temp_mb": 384.00048828125,
      "xla_arg_mb": 187.5,
      "xla_out_mb": 140.6250228881836,
      "compile_s": 0.8450326820020564,
      "inner": 1,
      "median_ms": 27.237647998845205,
      "q1_ms": 27.226345002418384,
      "q3_ms": 27.34970850724494,
      "tflops": 18.94789153681261
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 3000,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00014643371105194092,
      "rel_err": 0.0006584993680007756,
      "xla_temp_mb": 384.00048828125,
      "xla_arg_mb": 187.5,
      "xla_out_mb": 140.6250228881836,
      "compile_s": 0.8450902699987637,
      "inner": 1,
      "median_ms": 26.98335499735549,
      "q1_ms": 26.93537899904186,
      "q3_ms": 27.05708349822089,
      "tflops": 19.126457775564976
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 3000,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00014643371105194092,
      "rel_err": 0.0006584993680007756,
      "xla_temp_mb": 384.00048828125,
      "xla_arg_mb": 187.5,
      "xla_out_mb": 140.6250228881836,
      "compile_s": 0.845153840011335,
      "inner": 1,
      "median_ms": 27.132085990160704,
      "q1_ms": 27.084936496976297,
      "q3_ms": 27.145125495735556,
      "tflops": 19.02161154093199
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 3000,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00014573335647583008,
      "rel_err": 0.0006615787860937417,
      "xla_temp_mb": 3296.264663696289,
      "xla_arg_mb": 187.5,
      "xla_out_mb": 140.6250228881836,
      "compile_s": 2.2064223459892673,
      "inner": 2,
      "median_ms": 21.896562502661254,
      "q1_ms": 21.80135474918643,
      "q3_ms": 21.943639003438875,
      "tflops": 23.569726980537467
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 3000,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 9.834766387939453e-07,
      "rel_err": 3.751931444639922e-06,
      "peak_mb": 674.3603515625,
      "compile_s": 0.06378126399067696,
      "inner": 1,
      "median_ms": 64.162224996835,
      "q1_ms": 64.07988399587339,
      "q3_ms": 64.32138100353768,
      "tflops": 8.043611330895368
    },
    {
      "impl": "torch_math",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 3000,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 8.791685104370117e-07,
      "rel_err": 3.6811402424064e-06,
      "peak_mb": 4926.4072265625,
      "compile_s": 0.039493903997936286,
      "inner": 1,
      "median_ms": 40.38439999567345,
      "q1_ms": 40.363104999414645,
      "q3_ms": 40.46249849488959,
      "tflops": 12.77958815917264
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 3000,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0009948015213012695,
      "rel_err": 0.0002707532839849591,
      "xla_temp_mb": 192.0,
      "xla_arg_mb": 140.625,
      "xla_out_mb": 46.875,
      "compile_s": 0.34605928500241134,
      "inner": 13,
      "median_ms": 3.5662329229615772,
      "q1_ms": 3.5661442307173274,
      "q3_ms": 3.574028769575167,
      "tflops": 20.673916032038818
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 3000,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0009948015213012695,
      "rel_err": 0.0002707532839849591,
      "xla_temp_mb": 192.0,
      "xla_arg_mb": 140.625,
      "xla_out_mb": 46.875,
      "compile_s": 0.35051087400643155,
      "inner": 14,
      "median_ms": 3.4678280000142485,
      "q1_ms": 3.46299167878377,
      "q3_ms": 3.468175392689383,
      "tflops": 21.260570016649343
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 3000,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0009948015213012695,
      "rel_err": 0.0002707532839849591,
      "xla_temp_mb": 192.0,
      "xla_arg_mb": 140.625,
      "xla_out_mb": 46.875,
      "compile_s": 0.3422196739993524,
      "inner": 13,
      "median_ms": 3.4704615391092375,
      "q1_ms": 3.463162307385714,
      "q3_ms": 3.4738286536034697,
      "tflops": 21.244436559560246
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 3000,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0010943412780761719,
      "rel_err": 0.00029784481739625335,
      "xla_temp_mb": 2197.265640258789,
      "xla_arg_mb": 140.625,
      "xla_out_mb": 46.875,
      "compile_s": 0.3985445879952749,
      "inner": 6,
      "median_ms": 8.03251683343357,
      "q1_ms": 8.02447833242089,
      "q3_ms": 8.04023466601696,
      "tflops": 9.17869224912465
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 3000,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 1.6689300537109375e-06,
      "rel_err": 4.5422959260577045e-07,
      "peak_mb": 297.5009765625,
      "compile_s": 0.007132424987503327,
      "inner": 7,
      "median_ms": 7.233423999943105,
      "q1_ms": 7.197415929113049,
      "q3_ms": 7.257292857892545,
      "tflops": 10.192683299165084
    },
    {
      "impl": "torch_math",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 3000,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 1.3709068298339844e-06,
      "rel_err": 3.7311716027943476e-07,
      "peak_mb": 2803.8486328125,
      "compile_s": 0.019996654998976737,
      "inner": 2,
      "median_ms": 20.33231550012715,
      "q1_ms": 20.328991497081006,
      "q3_ms": 20.3463322504831,
      "tflops": 3.6261487285862217
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 3000,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0016725659370422363,
      "rel_err": 0.0007800427847541869,
      "xla_temp_mb": 384.00048828125,
      "xla_arg_mb": 187.5,
      "xla_out_mb": 140.6250228881836,
      "compile_s": 0.8582959659979679,
      "inner": 3,
      "median_ms": 14.966069002790997,
      "q1_ms": 14.9247041651203,
      "q3_ms": 15.049430502889056,
      "tflops": 17.242203009479447
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 3000,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0016725659370422363,
      "rel_err": 0.0007800427847541869,
      "xla_temp_mb": 384.00048828125,
      "xla_arg_mb": 187.5,
      "xla_out_mb": 140.6250228881836,
      "compile_s": 0.8524092489969917,
      "inner": 3,
      "median_ms": 14.686269666223476,
      "q1_ms": 14.646318668383174,
      "q3_ms": 14.712832834144743,
      "tflops": 17.57069738365741
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 3000,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0016725659370422363,
      "rel_err": 0.0007800427847541869,
      "xla_temp_mb": 384.00048828125,
      "xla_arg_mb": 187.5,
      "xla_out_mb": 140.6250228881836,
      "compile_s": 0.853489826011355,
      "inner": 3,
      "median_ms": 14.678487665757226,
      "q1_ms": 14.641437834749619,
      "q3_ms": 14.722648168875214,
      "tflops": 17.580012728558433
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 3000,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0015189647674560547,
      "rel_err": 0.0006280419183894992,
      "xla_temp_mb": 3296.264663696289,
      "xla_arg_mb": 187.5,
      "xla_out_mb": 140.6250228881836,
      "compile_s": 0.909212529004435,
      "inner": 2,
      "median_ms": 21.69179949851241,
      "q1_ms": 21.662110750185093,
      "q3_ms": 21.73853075146326,
      "tflops": 11.896108481811135
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 3000,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 1.0728836059570312e-05,
      "rel_err": 3.391369546079659e-06,
      "peak_mb": 674.3603515625,
      "compile_s": 0.033465699991211295,
      "inner": 1,
      "median_ms": 33.40711799683049,
      "q1_ms": 33.18177699838998,
      "q3_ms": 33.44064649718348,
      "tflops": 7.724341861051359
    },
    {
      "impl": "torch_math",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 3000,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 6.9141387939453125e-06,
      "rel_err": 3.2245807233266532e-06,
      "peak_mb": 4926.4072265625,
      "compile_s": 0.04338119499152526,
      "inner": 1,
      "median_ms": 44.03833900869358,
      "q1_ms": 43.97539549245266,
      "q3_ms": 44.234094995772466,
      "tflops": 5.859621543606785
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 4096,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00011709332466125488,
      "rel_err": 0.0007362766191363335,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 192.0,
      "xla_out_mb": 64.0,
      "compile_s": 0.31863106800301466,
      "inner": 5,
      "median_ms": 9.961713198572397,
      "q1_ms": 9.903077999479137,
      "q3_ms": 9.971006600244436,
      "tflops": 27.593437139245534
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 4096,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00011709332466125488,
      "rel_err": 0.0007362766191363335,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 192.0,
      "xla_out_mb": 64.0,
      "compile_s": 0.321469442002126,
      "inner": 5,
      "median_ms": 9.984568000072613,
      "q1_ms": 9.980969699972775,
      "q3_ms": 9.992399199109059,
      "tflops": 27.530275415220864
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 4096,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00011709332466125488,
      "rel_err": 0.0007362766191363335,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 192.0,
      "xla_out_mb": 64.0,
      "compile_s": 0.313156109987176,
      "inner": 5,
      "median_ms": 9.97349540120922,
      "q1_ms": 9.925286400539335,
      "q3_ms": 10.00956350035267,
      "tflops": 27.56083959397754
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 4096,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 8.15466046333313e-05,
      "rel_err": 0.0005127607146278024,
      "xla_temp_mb": 4096.000015258789,
      "xla_arg_mb": 192.0,
      "xla_out_mb": 64.0,
      "compile_s": 1.7788451719970908,
      "inner": 3,
      "median_ms": 14.330495663064843,
      "q1_ms": 14.304232832121972,
      "q3_ms": 14.355013666014807,
      "tflops": 19.181325852703427
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 4096,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 5.587935447692871e-07,
      "rel_err": 3.513664296406205e-06,
      "peak_mb": 400.2509765625,
      "compile_s": 0.023961833998328075,
      "inner": 2,
      "median_ms": 24.4457799999509,
      "q1_ms": 24.435257251752773,
      "q3_ms": 24.458571002469398,
      "tflops": 11.244390931463512
    },
    {
      "impl": "torch_math",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 4096,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 7.82310962677002e-07,
      "rel_err": 4.9191298785444815e-06,
      "peak_mb": 5008.37548828125,
      "compile_s": 0.027784762001829222,
      "inner": 1,
      "median_ms": 28.933277004398406,
      "q1_ms": 28.78601849806728,
      "q3_ms": 28.97021699027391,
      "tflops": 9.500406984739866
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 4096,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00019660592079162598,
      "rel_err": 0.0008184633916243911,
      "xla_temp_mb": 65.000244140625,
      "xla_arg_mb": 256.0,
      "xla_out_mb": 192.0000228881836,
      "compile_s": 0.8249089539895067,
      "inner": 1,
      "median_ms": 46.21592300827615,
      "q1_ms": 46.195593997254036,
      "q3_ms": 46.261419003712945,
      "tflops": 20.81690923129927
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 4096,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00019660592079162598,
      "rel_err": 0.0008184633916243911,
      "xla_temp_mb": 65.000244140625,
      "xla_arg_mb": 256.0,
      "xla_out_mb": 192.0000228881836,
      "compile_s": 0.8310846800013678,
      "inner": 1,
      "median_ms": 46.37624599854462,
      "q1_ms": 46.26061200542608,
      "q3_ms": 46.47627999656834,
      "tflops": 20.74494503790134
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 4096,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00019660592079162598,
      "rel_err": 0.0008184633916243911,
      "xla_temp_mb": 65.000244140625,
      "xla_arg_mb": 256.0,
      "xla_out_mb": 192.0000228881836,
      "compile_s": 0.8176688259991352,
      "inner": 1,
      "median_ms": 46.41466798784677,
      "q1_ms": 46.33374050172279,
      "q3_ms": 46.94547949475236,
      "tflops": 20.72777240496279
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 4096,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00019532442092895508,
      "rel_err": 0.0008162367739714682,
      "xla_temp_mb": 6144.500015258789,
      "xla_arg_mb": 256.0,
      "xla_out_mb": 192.0000228881836,
      "compile_s": 2.7382163900037995,
      "inner": 1,
      "median_ms": 38.403698010370135,
      "q1_ms": 38.325988498399965,
      "q3_ms": 38.41335750621511,
      "tflops": 25.051563368824844
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 4096,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 8.940696716308594e-07,
      "rel_err": 3.6119956803304376e-06,
      "peak_mb": 913.751953125,
      "compile_s": 0.11503923399141058,
      "inner": 1,
      "median_ms": 116.4089780068025,
      "q1_ms": 116.2523364982917,
      "q3_ms": 116.71293799736304,
      "tflops": 8.264591707417791
    },
    {
      "impl": "torch_math",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 4096,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "error": "OutOfMemoryError: CUDA out of memory. Tried to allocate 2.00 GiB. GPU 0 has a total capacity of 23.61 GiB of which 919.25 MiB is free. Including non-PyTorch memory, this process has 22.67 GiB memory in use. Of the allocated memory 4.52 GiB is allocated by PyTorch, and 109.75 MiB is reserved by PyTorch but unallocated"
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 4096,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0009087324142456055,
      "rel_err": 0.00032749667298048735,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 192.0,
      "xla_out_mb": 64.0,
      "compile_s": 0.3344513650081353,
      "inner": 8,
      "median_ms": 5.578082498686854,
      "q1_ms": 5.442155687887862,
      "q3_ms": 5.585553437413182,
      "tflops": 24.639103760898973
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 4096,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0009087324142456055,
      "rel_err": 0.00032749667298048735,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 192.0,
      "xla_out_mb": 64.0,
      "compile_s": 0.32500933500705287,
      "inner": 9,
      "median_ms": 5.395704888441187,
      "q1_ms": 5.370123722564637,
      "q3_ms": 5.422124888961359,
      "tflops": 25.47191818559706
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 4096,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0009087324142456055,
      "rel_err": 0.00032749667298048735,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 192.0,
      "xla_out_mb": 64.0,
      "compile_s": 0.324740235999343,
      "inner": 9,
      "median_ms": 5.348299444045147,
      "q1_ms": 5.347568722047274,
      "q3_ms": 5.354731223229059,
      "tflops": 25.697692305733927
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 4096,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0013453960418701172,
      "rel_err": 0.00048486521700397134,
      "xla_temp_mb": 4096.000015258789,
      "xla_arg_mb": 192.0,
      "xla_out_mb": 64.0,
      "compile_s": 0.5100997459958307,
      "inner": 3,
      "median_ms": 14.247207999384651,
      "q1_ms": 14.22858433458411,
      "q3_ms": 14.251305665917851,
      "tflops": 9.64672892246229
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 4096,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 2.6226043701171875e-06,
      "rel_err": 9.451563300899579e-07,
      "peak_mb": 400.2509765625,
      "compile_s": 0.012866072007454932,
      "inner": 3,
      "median_ms": 13.080642330654277,
      "q1_ms": 13.03046399940892,
      "q3_ms": 13.08918533322867,
      "tflops": 10.50704927157239
    },
    {
      "impl": "torch_math",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 4096,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 1.1920928955078125e-06,
      "rel_err": 4.2961653434758773e-07,
      "peak_mb": 5072.37548828125,
      "compile_s": 0.03597288600576576,
      "inner": 1,
      "median_ms": 36.716766000608914,
      "q1_ms": 36.54127350455383,
      "q3_ms": 36.89058950112667,
      "tflops": 3.743220562228185
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 4096,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0017998218536376953,
      "rel_err": 0.0005610198713839054,
      "xla_temp_mb": 65.000244140625,
      "xla_arg_mb": 256.0,
      "xla_out_mb": 192.0000228881836,
      "compile_s": 0.8431171539996285,
      "inner": 2,
      "median_ms": 24.540242498915177,
      "q1_ms": 24.495739249687176,
      "q3_ms": 24.604047248431016,
      "tflops": 19.60193902620419
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 4096,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0017998218536376953,
      "rel_err": 0.0005610198713839054,
      "xla_temp_mb": 65.000244140625,
      "xla_arg_mb": 256.0,
      "xla_out_mb": 192.0000228881836,
      "compile_s": 0.8462651449954137,
      "inner": 2,
      "median_ms": 24.469908996252343,
      "q1_ms": 24.443097747280262,
      "q3_ms": 24.50606374986819,
      "tflops": 19.658280593755887
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 4096,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0017998218536376953,
      "rel_err": 0.0005610198713839054,
      "xla_temp_mb": 65.000244140625,
      "xla_arg_mb": 256.0,
      "xla_out_mb": 192.0000228881836,
      "compile_s": 0.8438357450067997,
      "inner": 2,
      "median_ms": 24.412205995758995,
      "q1_ms": 24.35223775319173,
      "q3_ms": 24.43944275000831,
      "tflops": 19.704746766251603
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 4096,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0017995834350585938,
      "rel_err": 0.0004941088845953345,
      "xla_temp_mb": 6144.000015258789,
      "xla_arg_mb": 256.0,
      "xla_out_mb": 192.0000228881836,
      "compile_s": 1.041078586000367,
      "inner": 1,
      "median_ms": 37.18365100212395,
      "q1_ms": 37.09354400780285,
      "q3_ms": 37.306437494407874,
      "tflops": 12.936769902571507
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 4096,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 1.33514404296875e-05,
      "rel_err": 3.2256314170808764e-06,
      "peak_mb": 913.751953125,
      "compile_s": 0.06038904099841602,
      "inner": 1,
      "median_ms": 60.36412400135305,
      "q1_ms": 60.29262299853144,
      "q3_ms": 60.54204500105698,
      "tflops": 7.968911089328782
    },
    {
      "impl": "torch_math",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 4096,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "error": "OutOfMemoryError: CUDA out of memory. Tried to allocate 2.00 GiB. GPU 0 has a total capacity of 23.61 GiB of which 491.25 MiB is free. Including non-PyTorch memory, this process has 23.09 GiB memory in use. Of the allocated memory 4.52 GiB is allocated by PyTorch, and 537.75 MiB is reserved by PyTorch but unallocated"
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 8192,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 5.8241188526153564e-05,
      "rel_err": 0.0005545090534724295,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 384.0,
      "xla_out_mb": 128.0,
      "compile_s": 0.3535290579893626,
      "inner": 1,
      "median_ms": 41.26310400897637,
      "q1_ms": 40.52740049519343,
      "q3_ms": 41.34984849952161,
      "tflops": 26.646362511574804
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 8192,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 5.8241188526153564e-05,
      "rel_err": 0.0005545090534724295,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 384.0,
      "xla_out_mb": 128.0,
      "compile_s": 0.34943350400135387,
      "inner": 1,
      "median_ms": 39.80062500340864,
      "q1_ms": 39.724270507576875,
      "q3_ms": 39.91054801008431,
      "tflops": 27.62548647620068
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 8192,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 5.8241188526153564e-05,
      "rel_err": 0.0005545090534724295,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 384.0,
      "xla_out_mb": 128.0,
      "compile_s": 0.34668628200597595,
      "inner": 1,
      "median_ms": 39.88309999112971,
      "q1_ms": 39.746087495586835,
      "q3_ms": 40.12509950553067,
      "tflops": 27.568359230364223
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 8192,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 16.00GiB with allocator GPU_0_bfc on device 0. [tf-allocator-allocation-error=''] [executable_name='jit__lambda']"
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 8192,
      "causal": false,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 4.544854164123535e-07,
      "rel_err": 4.327114311308833e-06,
      "peak_mb": 784.2509765625,
      "compile_s": 0.10157280399289448,
      "inner": 1,
      "median_ms": 97.85466399625875,
      "q1_ms": 97.57101150171366,
      "q3_ms": 98.2643539973651,
      "tflops": 11.236169875541522
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 8192,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00011782348155975342,
      "rel_err": 0.0006405917229130864,
      "xla_temp_mb": 130.000244140625,
      "xla_arg_mb": 512.0,
      "xla_out_mb": 384.0000228881836,
      "compile_s": 0.9554990650067339,
      "inner": 1,
      "median_ms": 185.32452099316288,
      "q1_ms": 185.12094800098566,
      "q3_ms": 185.55897799524246,
      "tflops": 20.765145791786363
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 8192,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00011782348155975342,
      "rel_err": 0.0006405917229130864,
      "xla_temp_mb": 130.000244140625,
      "xla_arg_mb": 512.0,
      "xla_out_mb": 384.0000228881836,
      "compile_s": 0.9518720980122453,
      "inner": 1,
      "median_ms": 185.42870999954175,
      "q1_ms": 185.33895600558026,
      "q3_ms": 185.61292449885514,
      "tflops": 20.753478235519786
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 8192,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.00011782348155975342,
      "rel_err": 0.0006405917229130864,
      "xla_temp_mb": 130.000244140625,
      "xla_arg_mb": 512.0,
      "xla_out_mb": 384.0000228881836,
      "compile_s": 0.9616864630079363,
      "inner": 1,
      "median_ms": 185.66736299544573,
      "q1_ms": 185.54371999925934,
      "q3_ms": 185.76317399856634,
      "tflops": 20.72680214298296
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 8192,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 24.00GiB with allocator GPU_0_bfc on device 0. [tf-allocator-allocation-error=''] [executable_name='jit_loss']"
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 8192,
      "causal": false,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 1.0579824447631836e-06,
      "rel_err": 5.732731096941279e-06,
      "peak_mb": 1811.251953125,
      "compile_s": 0.4506645410001511,
      "inner": 1,
      "median_ms": 455.56085200223606,
      "q1_ms": 455.3055034921272,
      "q3_ms": 455.9303725036443,
      "tflops": 8.447369172092758
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 8192,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0012829303741455078,
      "rel_err": 0.00042788597056642175,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 384.0,
      "xla_out_mb": 128.0,
      "compile_s": 0.3437993150000693,
      "inner": 2,
      "median_ms": 21.229664504062384,
      "q1_ms": 21.19718349786126,
      "q3_ms": 21.326681249774992,
      "tflops": 25.895643041499874
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 8192,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0012829303741455078,
      "rel_err": 0.00042788597056642175,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 384.0,
      "xla_out_mb": 128.0,
      "compile_s": 0.34636650601169094,
      "inner": 2,
      "median_ms": 20.885026999167167,
      "q1_ms": 20.782693252840545,
      "q3_ms": 20.979201995942276,
      "tflops": 26.322964002389014
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 8192,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0012829303741455078,
      "rel_err": 0.00042788597056642175,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 384.0,
      "xla_out_mb": 128.0,
      "compile_s": 0.3542574879975291,
      "inner": 2,
      "median_ms": 20.715751001262106,
      "q1_ms": 20.639006997953402,
      "q3_ms": 20.75294649694115,
      "tflops": 26.538058593893417
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 8192,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 16.00GiB with allocator GPU_0_bfc on device 0. [executable_name='jit__lambda'] [tf-allocator-allocation-error='']"
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 8192,
      "causal": true,
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 1.430511474609375e-06,
      "rel_err": 4.771075623466459e-07,
      "peak_mb": 784.2509765625,
      "compile_s": 0.05591388599714264,
      "inner": 1,
      "median_ms": 51.38469999656081,
      "q1_ms": 51.12849000579445,
      "q3_ms": 51.68525400222279,
      "tflops": 10.698823072330777
    },
    {
      "impl": "probjax_mha",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 8192,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0013127326965332031,
      "rel_err": 0.0005491699557751417,
      "xla_temp_mb": 130.000244140625,
      "xla_arg_mb": 512.0,
      "xla_out_mb": 384.0000228881836,
      "compile_s": 0.9245191240042914,
      "inner": 1,
      "median_ms": 96.93727699050214,
      "q1_ms": 96.17370100022526,
      "q3_ms": 97.02595450653462,
      "tflops": 19.849385173017875
    },
    {
      "impl": "probjax_flex",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 8192,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0013127326965332031,
      "rel_err": 0.0005491699557751417,
      "xla_temp_mb": 130.000244140625,
      "xla_arg_mb": 512.0,
      "xla_out_mb": 384.0000228881836,
      "compile_s": 0.9191896819975227,
      "inner": 1,
      "median_ms": 96.5523269987898,
      "q1_ms": 96.44800549722277,
      "q3_ms": 96.80293649580562,
      "tflops": 19.928523821410515
    },
    {
      "impl": "probjax_default",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 8192,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 0.0013127326965332031,
      "rel_err": 0.0005491699557751417,
      "xla_temp_mb": 130.000244140625,
      "xla_arg_mb": 512.0,
      "xla_out_mb": 384.0000228881836,
      "compile_s": 0.9289685370022198,
      "inner": 1,
      "median_ms": 96.79520899953786,
      "q1_ms": 96.6079629943124,
      "q3_ms": 97.13123199617257,
      "tflops": 19.8785184566024
    },
    {
      "impl": "jax_xla",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 8192,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 24.00GiB with allocator GPU_0_bfc on device 0. [executable_name='jit_loss'] [tf-allocator-allocation-error='']"
    },
    {
      "impl": "torch_efficient",
      "dtype": "float32",
      "head_dim": 128,
      "seq": 8192,
      "causal": true,
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "tag": null,
      "max_abs_err": 1.5735626220703125e-05,
      "rel_err": 4.988353339285823e-06,
      "peak_mb": 1811.251953125,
      "compile_s": 0.2300619060115423,
      "inner": 1,
      "median_ms": 232.4385060055647,
      "q1_ms": 231.92800999822794,
      "q3_ms": 232.7487140064477,
      "tflops": 8.278083445270186
    }
  ]
}