{
  "device": {
    "jax": "0.11.2",
    "backend": "cpu",
    "devices": [
      "cpu:0"
    ],
    "python": "3.14.8",
    "device_kind": "cpu",
    "gpu": "NVIDIA GeForce RTX 3090",
    "driver": "615.78.08",
    "gpu_memory_gb": 24.0,
    "torch": "2.14.1+cu130",
    "torch_cuda": "13.0"
  },
  "config": {
    "device": "cpu",
    "append": true,
    "seqs": "128,256,512,1024,2048",
    "dims": "64",
    "dtypes": "float32",
    "variants": "plain",
    "modes": "fwd,fwdbwd",
    "batch": 1,
    "heads": 8,
    "impls": "probjax_mha,xla_dense,xla_tiled,mha_interpret,jax_xla,torch_sdpa",
    "interpret_max_seq": 1024,
    "no_ref": false,
    "peaks": false,
    "peak_seqs": "4096,8192,16384",
    "peak_impls": "xla_tiled,xla_dense,jax_xla,torch_sdpa",
    "peak_timeout": 300,
    "max_fit": false,
    "fit_seqs": "4096,8192,16384,32768,65536,131072",
    "fit_impls": "xla_tiled,xla_dense,jax_xla"
  },
  "results": [
    {
      "impl": "probjax_mha",
      "mode": "fwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 128,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.061776973998348694,
      "xla_temp_mb": 1.25,
      "xla_arg_mb": 0.75,
      "xla_out_mb": 0.25,
      "compile_s": 0.003265559993451461,
      "inner": 101,
      "median_ms": 0.21959727724822756,
      "q1_ms": 0.21784980193536618,
      "q3_ms": 0.29809236636565134,
      "rel_err": 4.013325498109296e-08
    },
    {
      "impl": "xla_dense",
      "mode": "fwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 128,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.06622154600336216,
      "xla_temp_mb": 1.25,
      "xla_arg_mb": 0.75,
      "xla_out_mb": 0.25,
      "compile_s": 0.002711199995246716,
      "inner": 148,
      "median_ms": 0.2794714729699646,
      "q1_ms": 0.23234234460097788,
      "q3_ms": 0.34131249324402874,
      "rel_err": 4.013325498109296e-08
    },
    {
      "impl": "xla_tiled",
      "mode": "fwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 128,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.07407216299907304,
      "xla_temp_mb": 1.25390625,
      "xla_arg_mb": 0.75,
      "xla_out_mb": 0.25,
      "compile_s": 0.0026885580009547994,
      "inner": 137,
      "median_ms": 0.3091864598329505,
      "q1_ms": 0.23872346716489387,
      "q3_ms": 0.3349994671238709,
      "rel_err": 2.934683051465238e-07
    },
    {
      "impl": "mha_interpret",
      "mode": "fwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 128,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.14870582000003196,
      "xla_temp_mb": 1.5314788818359375,
      "xla_arg_mb": 0.75,
      "xla_out_mb": 0.25,
      "compile_s": 0.0027278029956505634,
      "inner": 35,
      "median_ms": 0.6679166000269886,
      "q1_ms": 0.6651983430076922,
      "q3_ms": 0.7030746001484138,
      "rel_err": 3.026228455251298e-07
    },
    {
      "impl": "jax_xla",
      "mode": "fwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 128,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.0637509400039562,
      "xla_temp_mb": 1.00390625,
      "xla_arg_mb": 0.75,
      "xla_out_mb": 0.25,
      "compile_s": 0.0035656079999171197,
      "inner": 111,
      "median_ms": 0.2929736486597297,
      "q1_ms": 0.2918470990980413,
      "q3_ms": 0.29485314414020514,
      "rel_err": 0.0
    },
    {
      "impl": "torch_sdpa",
      "mode": "fwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 128,
      "heads": 8,
      "head_dim": 64,
      "compile_s": 0.0006497300055343658,
      "inner": 200,
      "median_ms": 0.06325461497908691,
      "q1_ms": 0.06289527998887934,
      "q3_ms": 0.07213456501631299,
      "rel_err": 3.039969419313544e-07
    },
    {
      "impl": "probjax_mha",
      "mode": "fwdbwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 128,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.11120137100078864,
      "xla_temp_mb": 2.7578125,
      "xla_arg_mb": 0.75,
      "xla_out_mb": 0.7500228881835938,
      "compile_s": 0.005482004999066703,
      "inner": 41,
      "median_ms": 0.4412858781018635,
      "q1_ms": 0.4344597805245984,
      "q3_ms": 0.48231817061427934,
      "rel_err": 4.4167406165351345e-08
    },
    {
      "impl": "xla_dense",
      "mode": "fwdbwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 128,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.1071613449967117,
      "xla_temp_mb": 2.7578125,
      "xla_arg_mb": 0.75,
      "xla_out_mb": 0.7500228881835938,
      "compile_s": 0.005438742999103852,
      "inner": 70,
      "median_ms": 0.7470954571285152,
      "q1_ms": 0.7240196571469174,
      "q3_ms": 0.7580855571307308,
      "rel_err": 4.4167406165351345e-08
    },
    {
      "impl": "xla_tiled",
      "mode": "fwdbwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 128,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.1350269619942992,
      "xla_temp_mb": 2.25,
      "xla_arg_mb": 0.75,
      "xla_out_mb": 0.7500228881835938,
      "compile_s": 0.007805201996234246,
      "inner": 52,
      "median_ms": 0.8350983269058185,
      "q1_ms": 0.8301158076602195,
      "q3_ms": 0.8388348269363632,
      "rel_err": 2.664876369864099e-07
    },
    {
      "impl": "mha_interpret",
      "mode": "fwdbwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 128,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.3056952850020025,
      "xla_temp_mb": 3.040008544921875,
      "xla_arg_mb": 0.75,
      "xla_out_mb": 0.7500228881835938,
      "compile_s": 0.017991209002502728,
      "inner": 11,
      "median_ms": 4.019236454041675,
      "q1_ms": 3.965709727543237,
      "q3_ms": 4.049356272761625,
      "rel_err": 3.143545391967955e-07
    },
    {
      "impl": "jax_xla",
      "mode": "fwdbwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 128,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.10633513200446032,
      "xla_temp_mb": 2.75390625,
      "xla_arg_mb": 0.75,
      "xla_out_mb": 0.7500228881835938,
      "compile_s": 0.005385882002883591,
      "inner": 62,
      "median_ms": 0.6887812741033192,
      "q1_ms": 0.6873285323484695,
      "q3_ms": 0.6893544677070402,
      "rel_err": 0.0
    },
    {
      "impl": "torch_sdpa",
      "mode": "fwdbwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 128,
      "heads": 8,
      "head_dim": 64,
      "compile_s": 0.316034421994118,
      "inner": 92,
      "median_ms": 0.40748192393790383,
      "q1_ms": 0.4072263369459466,
      "q3_ms": 0.4087536847731352,
      "rel_err": 3.0365244388380266e-07
    },
    {
      "impl": "probjax_mha",
      "mode": "fwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 256,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.07808349499828182,
      "xla_temp_mb": 4.5,
      "xla_arg_mb": 1.5,
      "xla_out_mb": 0.5,
      "compile_s": 0.004495187000429723,
      "inner": 26,
      "median_ms": 1.0934624613303012,
      "q1_ms": 1.0489343462373095,
      "q3_ms": 1.1312994616813599,
      "rel_err": 4.2272482880613725e-08
    },
    {
      "impl": "xla_dense",
      "mode": "fwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 256,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.07294024099974195,
      "xla_temp_mb": 4.5,
      "xla_arg_mb": 1.5,
      "xla_out_mb": 0.5,
      "compile_s": 0.0037493560012080707,
      "inner": 38,
      "median_ms": 1.0995696579093899,
      "q1_ms": 1.0555803158862738,
      "q3_ms": 1.1213814474391064,
      "rel_err": 4.2272482880613725e-08
    },
    {
      "impl": "xla_tiled",
      "mode": "fwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 256,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.08265009699971415,
      "xla_temp_mb": 4.5078125,
      "xla_arg_mb": 1.5,
      "xla_out_mb": 0.5,
      "compile_s": 0.0037845319966436364,
      "inner": 36,
      "median_ms": 0.9119734168052673,
      "q1_ms": 0.8969024165708106,
      "q3_ms": 0.9130935832217801,
      "rel_err": 4.091313574264293e-07
    },
    {
      "impl": "mha_interpret",
      "mode": "fwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 256,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.135244914003124,
      "xla_temp_mb": 3.2203369140625,
      "xla_arg_mb": 1.5,
      "xla_out_mb": 0.5,
      "compile_s": 0.00546078499610303,
      "inner": 15,
      "median_ms": 2.9930891333303102,
      "q1_ms": 2.8190357334096916,
      "q3_ms": 3.0692973334225826,
      "rel_err": 3.685833922991945e-07
    },
    {
      "impl": "jax_xla",
      "mode": "fwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 256,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.06253134299913654,
      "xla_temp_mb": 4.0078125,
      "xla_arg_mb": 1.5,
      "xla_out_mb": 0.5,
      "compile_s": 0.004520023998338729,
      "inner": 51,
      "median_ms": 0.7747792353221288,
      "q1_ms": 0.7029089020387106,
      "q3_ms": 0.8660032549147111,
      "rel_err": 0.0
    },
    {
      "impl": "torch_sdpa",
      "mode": "fwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 256,
      "heads": 8,
      "head_dim": 64,
      "compile_s": 0.0005014189955545589,
      "inner": 200,
      "median_ms": 0.1795985949866008,
      "q1_ms": 0.16211950998695102,
      "q3_ms": 0.20052609997947002,
      "rel_err": 3.6643008399589345e-07
    },
    {
      "impl": "probjax_mha",
      "mode": "fwdbwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 256,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.12598405700555304,
      "xla_temp_mb": 10.515625,
      "xla_arg_mb": 1.5,
      "xla_out_mb": 1.5000228881835938,
      "compile_s": 0.008316307998029515,
      "inner": 12,
      "median_ms": 2.092667916561671,
      "q1_ms": 2.030096166587706,
      "q3_ms": 2.160318500197415,
      "rel_err": 4.764114103042833e-08
    },
    {
      "impl": "xla_dense",
      "mode": "fwdbwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 256,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.12889550699765095,
      "xla_temp_mb": 10.515625,
      "xla_arg_mb": 1.5,
      "xla_out_mb": 1.5000228881835938,
      "compile_s": 0.0078652950032847,
      "inner": 22,
      "median_ms": 2.066980454401346,
      "q1_ms": 2.036774636177473,
      "q3_ms": 2.074796136267568,
      "rel_err": 4.764114103042833e-08
    },
    {
      "impl": "xla_tiled",
      "mode": "fwdbwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 256,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.14943397699971683,
      "xla_temp_mb": 8.5,
      "xla_arg_mb": 1.5,
      "xla_out_mb": 1.5000228881835938,
      "compile_s": 0.008858043998770881,
      "inner": 19,
      "median_ms": 2.2121845265102915,
      "q1_ms": 2.203504894599686,
      "q3_ms": 2.232664841836882,
      "rel_err": 3.024212408359534e-07
    },
    {
      "impl": "mha_interpret",
      "mode": "fwdbwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 256,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.3289115320003475,
      "xla_temp_mb": 6.236846923828125,
      "xla_arg_mb": 1.5,
      "xla_out_mb": 1.5000228881835938,
      "compile_s": 0.020303295001212973,
      "inner": 3,
      "median_ms": 13.592096666495005,
      "q1_ms": 13.199029332705928,
      "q3_ms": 14.097489666407151,
      "rel_err": 4.3491910313762544e-07
    },
    {
      "impl": "jax_xla",
      "mode": "fwdbwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 256,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.12061926499882247,
      "xla_temp_mb": 10.5078125,
      "xla_arg_mb": 1.5,
      "xla_out_mb": 1.5000228881835938,
      "compile_s": 0.006487237995315809,
      "inner": 28,
      "median_ms": 1.881024142806252,
      "q1_ms": 1.8565327856257292,
      "q3_ms": 1.9001188928606487,
      "rel_err": 0.0
    },
    {
      "impl": "torch_sdpa",
      "mode": "fwdbwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 256,
      "heads": 8,
      "head_dim": 64,
      "compile_s": 0.002280166001582984,
      "inner": 34,
      "median_ms": 0.961232088122736,
      "q1_ms": 0.948429352908587,
      "q3_ms": 0.9789506175901111,
      "rel_err": 3.878633629555442e-07
    },
    {
      "impl": "xla_dense",
      "mode": "fwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 512,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.061018308995699044,
      "xla_temp_mb": 17.0,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.00565718599682441,
      "inner": 16,
      "median_ms": 3.1350681874755537,
      "q1_ms": 3.1235576875587867,
      "q3_ms": 3.167246624798281,
      "rel_err": 5.0695175675676665e-08
    },
    {
      "impl": "xla_tiled",
      "mode": "fwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 512,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.11152318000677042,
      "xla_temp_mb": 8.53167724609375,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.0070528180003748275,
      "inner": 12,
      "median_ms": 4.163896083506795,
      "q1_ms": 4.102741999910601,
      "q3_ms": 4.2342584165453445,
      "rel_err": 5.029622353720825e-07
    },
    {
      "impl": "mha_interpret",
      "mode": "fwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 512,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.13440507499763044,
      "xla_temp_mb": 6.3453369140625,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.013758166998741217,
      "inner": 4,
      "median_ms": 11.591206499360851,
      "q1_ms": 11.533625000083703,
      "q3_ms": 11.771034500270616,
      "rel_err": 4.6668447000148565e-07
    },
    {
      "impl": "jax_xla",
      "mode": "fwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 512,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.06100354099908145,
      "xla_temp_mb": 16.0,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.006622413005970884,
      "inner": 19,
      "median_ms": 2.166741684241539,
      "q1_ms": 2.1426255261758342,
      "q3_ms": 2.2412426842493005,
      "rel_err": 0.0
    },
    {
      "impl": "torch_sdpa",
      "mode": "fwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 512,
      "heads": 8,
      "head_dim": 64,
      "compile_s": 0.0016285219971905462,
      "inner": 40,
      "median_ms": 0.6390636748619727,
      "q1_ms": 0.6331309748929925,
      "q3_ms": 0.653922824858455,
      "rel_err": 4.836162691068112e-07
    },
    {
      "impl": "xla_dense",
      "mode": "fwdbwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 512,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.1407336909978767,
      "xla_temp_mb": 41.015625,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.017836506995081436,
      "inner": 2,
      "median_ms": 16.83893300287309,
      "q1_ms": 16.373286500311224,
      "q3_ms": 17.47683999928995,
      "rel_err": 9.832460469616295e-08
    },
    {
      "impl": "xla_tiled",
      "mode": "fwdbwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 512,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.23397910000494448,
      "xla_temp_mb": 16.031936645507812,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.01915688499866519,
      "inner": 3,
      "median_ms": 13.528610000018185,
      "q1_ms": 13.244114668244341,
      "q3_ms": 13.586843667629486,
      "rel_err": 4.832077495913222e-07
    },
    {
      "impl": "mha_interpret",
      "mode": "fwdbwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 512,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.3502027470021858,
      "xla_temp_mb": 12.377471923828125,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.06800923899572808,
      "inner": 1,
      "median_ms": 51.933315000496805,
      "q1_ms": 51.66139100037981,
      "q3_ms": 52.21980700298445,
      "rel_err": 5.421284527647368e-07
    },
    {
      "impl": "jax_xla",
      "mode": "fwdbwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 512,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.13537763499334687,
      "xla_temp_mb": 25.03125,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.014340890003950335,
      "inner": 9,
      "median_ms": 5.706051111499416,
      "q1_ms": 5.653711999507828,
      "q3_ms": 5.758003888735806,
      "rel_err": 0.0
    },
    {
      "impl": "torch_sdpa",
      "mode": "fwdbwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 512,
      "heads": 8,
      "head_dim": 64,
      "compile_s": 0.005309018000843935,
      "inner": 10,
      "median_ms": 3.3094242004153784,
      "q1_ms": 3.2662924997566734,
      "q3_ms": 3.5123190995363984,
      "rel_err": 4.973698664009482e-07
    },
    {
      "impl": "probjax_mha",
      "mode": "fwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 1024,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.11661051800183486,
      "xla_temp_mb": 11.53167724609375,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.018450468000082765,
      "inner": 3,
      "median_ms": 15.359438000208078,
      "q1_ms": 15.200300333769215,
      "q3_ms": 15.461898665914001,
      "rel_err": 6.524791666936522e-07
    },
    {
      "impl": "xla_dense",
      "mode": "fwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 1024,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.05872318500041729,
      "xla_temp_mb": 66.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.02259042300283909,
      "inner": 1,
      "median_ms": 26.58764800435165,
      "q1_ms": 26.221544001600705,
      "q3_ms": 26.60870800173143,
      "rel_err": 6.206117587227691e-08
    },
    {
      "impl": "xla_tiled",
      "mode": "fwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 1024,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.11516527200001292,
      "xla_temp_mb": 11.53167724609375,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.018658380999113433,
      "inner": 3,
      "median_ms": 15.587423001610054,
      "q1_ms": 15.575252999648606,
      "q3_ms": 15.685414999704033,
      "rel_err": 6.524791666936522e-07
    },
    {
      "impl": "mha_interpret",
      "mode": "fwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 1024,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.14554404599766713,
      "xla_temp_mb": 12.5953369140625,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.07662212899595033,
      "inner": 1,
      "median_ms": 41.2950309982989,
      "q1_ms": 41.28213799413061,
      "q3_ms": 41.328204999445006,
      "rel_err": 6.231357629984253e-07
    },
    {
      "impl": "jax_xla",
      "mode": "fwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 1024,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.05902275200060103,
      "xla_temp_mb": 64.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.01971795599820325,
      "inner": 1,
      "median_ms": 46.838802998536266,
      "q1_ms": 41.026733997568954,
      "q3_ms": 49.27672700432595,
      "rel_err": 0.0
    },
    {
      "impl": "torch_sdpa",
      "mode": "fwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 1024,
      "heads": 8,
      "head_dim": 64,
      "compile_s": 0.0030998359943623655,
      "inner": 19,
      "median_ms": 2.1550411053651355,
      "q1_ms": 2.152526841918619,
      "q3_ms": 2.1731769999237978,
      "rel_err": 6.347056830136409e-07
    },
    {
      "impl": "probjax_mha",
      "mode": "fwdbwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 1024,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.3145039250011905,
      "xla_temp_mb": 19.063186645507812,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.057192047002899926,
      "inner": 1,
      "median_ms": 48.511508997762576,
      "q1_ms": 44.38670200033812,
      "q3_ms": 48.71034499956295,
      "rel_err": 6.660117743954283e-07
    },
    {
      "impl": "xla_dense",
      "mode": "fwdbwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 1024,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.18593605799833313,
      "xla_temp_mb": 162.03125,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.07123658600175986,
      "inner": 1,
      "median_ms": 84.62097700248705,
      "q1_ms": 80.96268400549889,
      "q3_ms": 106.47827199863968,
      "rel_err": 1.1551242810726563e-07
    },
    {
      "impl": "xla_tiled",
      "mode": "fwdbwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 1024,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.29433829199842876,
      "xla_temp_mb": 19.063186645507812,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.05751971699646674,
      "inner": 1,
      "median_ms": 48.935382001218386,
      "q1_ms": 47.75599899585359,
      "q3_ms": 50.47750900121173,
      "rel_err": 6.660117743954283e-07
    },
    {
      "impl": "jax_xla",
      "mode": "fwdbwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 1024,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.18388433499785606,
      "xla_temp_mb": 98.0625,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.039138790998549666,
      "inner": 1,
      "median_ms": 43.48560700600501,
      "q1_ms": 43.206157999520656,
      "q3_ms": 44.40846400393639,
      "rel_err": 0.0
    },
    {
      "impl": "torch_sdpa",
      "mode": "fwdbwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 1024,
      "heads": 8,
      "head_dim": 64,
      "compile_s": 0.013800168002489954,
      "inner": 3,
      "median_ms": 11.086210998958753,
      "q1_ms": 10.67025366743716,
      "q3_ms": 12.92706166714197,
      "rel_err": 6.627456577369124e-07
    },
    {
      "impl": "probjax_mha",
      "mode": "fwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 2048,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.11596489600196946,
      "xla_temp_mb": 17.53167724609375,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.06536264099850086,
      "inner": 1,
      "median_ms": 62.11271999927703,
      "q1_ms": 60.38052199437516,
      "q3_ms": 63.08866699691862,
      "rel_err": 6.459818006076694e-07
    },
    {
      "impl": "xla_dense",
      "mode": "fwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 2048,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.05982904700067593,
      "xla_temp_mb": 260.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.09105764799460303,
      "inner": 1,
      "median_ms": 111.4522060015588,
      "q1_ms": 109.3245099982596,
      "q3_ms": 111.52276000211714,
      "rel_err": 6.570094639754267e-08
    },
    {
      "impl": "xla_tiled",
      "mode": "fwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 2048,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.11331534200144233,
      "xla_temp_mb": 17.53167724609375,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.0634181399946101,
      "inner": 1,
      "median_ms": 59.78016500012018,
      "q1_ms": 59.540080997976474,
      "q3_ms": 60.51180999929784,
      "rel_err": 6.459818006076694e-07
    },
    {
      "impl": "jax_xla",
      "mode": "fwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 2048,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.05733065000094939,
      "xla_temp_mb": 256.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.06792779399984283,
      "inner": 1,
      "median_ms": 124.42170100257499,
      "q1_ms": 122.65169199963566,
      "q3_ms": 126.02284099557437,
      "rel_err": 0.0
    },
    {
      "impl": "torch_sdpa",
      "mode": "fwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 2048,
      "heads": 8,
      "head_dim": 64,
      "compile_s": 0.01618735500233015,
      "inner": 3,
      "median_ms": 9.313830002308046,
      "q1_ms": 8.87933766595476,
      "q3_ms": 9.566130999398107,
      "rel_err": 6.29828968042282e-07
    },
    {
      "impl": "probjax_mha",
      "mode": "fwdbwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 2048,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.46272572899761144,
      "xla_temp_mb": 25.125686645507812,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.21593063500040444,
      "inner": 1,
      "median_ms": 207.05104799708351,
      "q1_ms": 205.25207399623469,
      "q3_ms": 216.47508400201332,
      "rel_err": 7.161710527060266e-07
    },
    {
      "impl": "xla_dense",
      "mode": "fwdbwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 2048,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.30442275699897436,
      "xla_temp_mb": 644.0625,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.281409152004926,
      "inner": 1,
      "median_ms": 336.73574000567896,
      "q1_ms": 330.84439099911833,
      "q3_ms": 339.74965400557267,
      "rel_err": 1.2083650737615875e-07
    },
    {
      "impl": "xla_tiled",
      "mode": "fwdbwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 2048,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.4805111479945481,
      "xla_temp_mb": 25.125686645507812,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.23213395900529576,
      "inner": 1,
      "median_ms": 229.5582429942442,
      "q1_ms": 227.38000199751696,
      "q3_ms": 237.035484999069,
      "rel_err": 7.161710527060266e-07
    },
    {
      "impl": "jax_xla",
      "mode": "fwdbwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 2048,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.32421493400033796,
      "xla_temp_mb": 388.125,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.16116131199669326,
      "inner": 1,
      "median_ms": 192.52507799683372,
      "q1_ms": 188.82760000269627,
      "q3_ms": 194.82067200442543,
      "rel_err": 0.0
    },
    {
      "impl": "torch_sdpa",
      "mode": "fwdbwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 2048,
      "heads": 8,
      "head_dim": 64,
      "compile_s": 0.08375293399876682,
      "inner": 1,
      "median_ms": 73.91651899524732,
      "q1_ms": 71.29526799690211,
      "q3_ms": 74.42122400243534,
      "rel_err": 7.111362158237429e-07
    },
    {
      "impl": "probjax_mha",
      "mode": "fwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 128,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.06904585599841084,
      "xla_temp_mb": 1.25,
      "xla_arg_mb": 0.75,
      "xla_out_mb": 0.25,
      "compile_s": 0.003017668997927103,
      "inner": 94,
      "median_ms": 0.29994792551386784,
      "q1_ms": 0.2748926489538404,
      "q3_ms": 0.3275924680836281,
      "rel_err": 2.1822928654703168e-08
    },
    {
      "impl": "xla_dense",
      "mode": "fwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 128,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.06896971099922666,
      "xla_temp_mb": 1.25,
      "xla_arg_mb": 0.75,
      "xla_out_mb": 0.25,
      "compile_s": 0.0027957300044363365,
      "inner": 145,
      "median_ms": 0.23120606206718766,
      "q1_ms": 0.22659326897500534,
      "q3_ms": 0.24041970343738472,
      "rel_err": 2.1822928654703168e-08
    },
    {
      "impl": "xla_tiled",
      "mode": "fwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 128,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.08042135799769312,
      "xla_temp_mb": 1.25390625,
      "xla_arg_mb": 0.75,
      "xla_out_mb": 0.25,
      "compile_s": 0.0028860610036645085,
      "inner": 137,
      "median_ms": 0.2348853503443117,
      "q1_ms": 0.22630839415079826,
      "q3_ms": 0.24192648906648215,
      "rel_err": 1.4159266742518304e-07
    },
    {
      "impl": "mha_interpret",
      "mode": "fwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 128,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.16874173100222833,
      "xla_temp_mb": 1.5317535400390625,
      "xla_arg_mb": 0.75,
      "xla_out_mb": 0.25,
      "compile_s": 0.0027773550027632155,
      "inner": 51,
      "median_ms": 0.7415569609079017,
      "q1_ms": 0.7391834509440278,
      "q3_ms": 0.7448431568656244,
      "rel_err": 1.5594517710562076e-07
    },
    {
      "impl": "jax_xla",
      "mode": "fwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 128,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.0696404310001526,
      "xla_temp_mb": 1.00390625,
      "xla_arg_mb": 0.75,
      "xla_out_mb": 0.25,
      "compile_s": 0.00308198100537993,
      "inner": 124,
      "median_ms": 0.29981936290090333,
      "q1_ms": 0.28020708874288586,
      "q3_ms": 0.31097450001623417,
      "rel_err": 0.0
    },
    {
      "impl": "torch_sdpa",
      "mode": "fwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 128,
      "heads": 8,
      "head_dim": 64,
      "compile_s": 0.0006351320043904707,
      "inner": 200,
      "median_ms": 0.07043265999527648,
      "q1_ms": 0.06911611500981962,
      "q3_ms": 0.07063493998430204,
      "rel_err": 1.5435091948127467e-07
    },
    {
      "impl": "probjax_mha",
      "mode": "fwdbwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 128,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.11780099000316113,
      "xla_temp_mb": 2.7578125,
      "xla_arg_mb": 0.75,
      "xla_out_mb": 0.7500228881835938,
      "compile_s": 0.005665889999363571,
      "inner": 39,
      "median_ms": 0.7466991795435882,
      "q1_ms": 0.7430080769177622,
      "q3_ms": 0.7546883845442118,
      "rel_err": 3.311274867213489e-08
    },
    {
      "impl": "xla_dense",
      "mode": "fwdbwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 128,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.11783610700513236,
      "xla_temp_mb": 2.7578125,
      "xla_arg_mb": 0.75,
      "xla_out_mb": 0.7500228881835938,
      "compile_s": 0.006060446998162661,
      "inner": 59,
      "median_ms": 0.5030463389476161,
      "q1_ms": 0.4395614915835213,
      "q3_ms": 0.7583459831085608,
      "rel_err": 3.311274867213489e-08
    },
    {
      "impl": "xla_tiled",
      "mode": "fwdbwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 128,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.14440029700199375,
      "xla_temp_mb": 2.25,
      "xla_arg_mb": 0.75,
      "xla_out_mb": 0.7500228881835938,
      "compile_s": 0.007716230000369251,
      "inner": 47,
      "median_ms": 0.7103077234128153,
      "q1_ms": 0.692446936111392,
      "q3_ms": 0.8418515319638746,
      "rel_err": 2.2572715585192996e-07
    },
    {
      "impl": "mha_interpret",
      "mode": "fwdbwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 128,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.34364618200197583,
      "xla_temp_mb": 3.04058837890625,
      "xla_arg_mb": 0.75,
      "xla_out_mb": 0.7500228881835938,
      "compile_s": 0.017190364000271074,
      "inner": 13,
      "median_ms": 3.8223850765811784,
      "q1_ms": 3.709079153710403,
      "q3_ms": 3.969857153536465,
      "rel_err": 2.829457422080454e-07
    },
    {
      "impl": "jax_xla",
      "mode": "fwdbwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 128,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.12047274599899538,
      "xla_temp_mb": 2.75390625,
      "xla_arg_mb": 0.75,
      "xla_out_mb": 0.7500228881835938,
      "compile_s": 0.005283425998641178,
      "inner": 63,
      "median_ms": 0.7073421110694725,
      "q1_ms": 0.7052239682433004,
      "q3_ms": 0.708617698465888,
      "rel_err": 0.0
    },
    {
      "impl": "torch_sdpa",
      "mode": "fwdbwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 128,
      "heads": 8,
      "head_dim": 64,
      "compile_s": 0.3168784270019387,
      "inner": 89,
      "median_ms": 0.42517455054376424,
      "q1_ms": 0.4220046179847834,
      "q3_ms": 0.4306578538786055,
      "rel_err": 2.913166414163691e-07
    },
    {
      "impl": "probjax_mha",
      "mode": "fwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 256,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.06441836999874795,
      "xla_temp_mb": 4.5,
      "xla_arg_mb": 1.5,
      "xla_out_mb": 0.5,
      "compile_s": 0.004311948003305588,
      "inner": 41,
      "median_ms": 1.071904878034086,
      "q1_ms": 1.0634383659457742,
      "q3_ms": 1.0794298537004534,
      "rel_err": 2.49299923772152e-08
    },
    {
      "impl": "xla_dense",
      "mode": "fwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 256,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.0651356670059613,
      "xla_temp_mb": 4.5,
      "xla_arg_mb": 1.5,
      "xla_out_mb": 0.5,
      "compile_s": 0.0034026879948214628,
      "inner": 40,
      "median_ms": 1.0731431499152677,
      "q1_ms": 1.0612346250127302,
      "q3_ms": 1.0786608248963603,
      "rel_err": 2.49299923772152e-08
    },
    {
      "impl": "xla_tiled",
      "mode": "fwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 256,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.07364328399853548,
      "xla_temp_mb": 4.5078125,
      "xla_arg_mb": 1.5,
      "xla_out_mb": 0.5,
      "compile_s": 0.0036482720024650916,
      "inner": 54,
      "median_ms": 0.9073840000955561,
      "q1_ms": 0.9049286111096282,
      "q3_ms": 0.9153602036854459,
      "rel_err": 1.8707507464028955e-07
    },
    {
      "impl": "mha_interpret",
      "mode": "fwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 256,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.15255802100000437,
      "xla_temp_mb": 3.2225570678710938,
      "xla_arg_mb": 1.5,
      "xla_out_mb": 0.5,
      "compile_s": 0.005179439998755697,
      "inner": 15,
      "median_ms": 2.7592104001087137,
      "q1_ms": 2.6866585331542105,
      "q3_ms": 2.8586178663924025,
      "rel_err": 1.9232041376997993e-07
    },
    {
      "impl": "jax_xla",
      "mode": "fwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 256,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.06724431699694833,
      "xla_temp_mb": 4.0078125,
      "xla_arg_mb": 1.5,
      "xla_out_mb": 0.5,
      "compile_s": 0.004405915999086574,
      "inner": 44,
      "median_ms": 1.1581566817436726,
      "q1_ms": 1.133269159088906,
      "q3_ms": 1.1608481590056645,
      "rel_err": 0.0
    },
    {
      "impl": "torch_sdpa",
      "mode": "fwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 256,
      "heads": 8,
      "head_dim": 64,
      "compile_s": 0.0005118879998917691,
      "inner": 200,
      "median_ms": 0.1793515849931282,
      "q1_ms": 0.17541318000439787,
      "q3_ms": 0.18026196001301287,
      "rel_err": 1.8985504206001347e-07
    },
    {
      "impl": "probjax_mha",
      "mode": "fwdbwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 256,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.12792209000326693,
      "xla_temp_mb": 10.515625,
      "xla_arg_mb": 1.5,
      "xla_out_mb": 1.5000228881835938,
      "compile_s": 0.00812194700120017,
      "inner": 11,
      "median_ms": 2.18140363590051,
      "q1_ms": 2.0028558184011755,
      "q3_ms": 2.183703454317186,
      "rel_err": 3.7654143454632616e-08
    },
    {
      "impl": "xla_dense",
      "mode": "fwdbwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 256,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.12829967399738962,
      "xla_temp_mb": 10.515625,
      "xla_arg_mb": 1.5,
      "xla_out_mb": 1.5000228881835938,
      "compile_s": 0.007820125996659044,
      "inner": 26,
      "median_ms": 2.102291038541947,
      "q1_ms": 2.0709058461272014,
      "q3_ms": 2.1084761924262585,
      "rel_err": 3.7654143454632616e-08
    },
    {
      "impl": "xla_tiled",
      "mode": "fwdbwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 256,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.14738034499896457,
      "xla_temp_mb": 8.5,
      "xla_arg_mb": 1.5,
      "xla_out_mb": 1.5000228881835938,
      "compile_s": 0.009632233995944262,
      "inner": 20,
      "median_ms": 2.2772269996494288,
      "q1_ms": 2.2547304499312304,
      "q3_ms": 2.3642586998903425,
      "rel_err": 2.508540689013853e-07
    },
    {
      "impl": "mha_interpret",
      "mode": "fwdbwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 256,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.36503166400507325,
      "xla_temp_mb": 6.23931884765625,
      "xla_arg_mb": 1.5,
      "xla_out_mb": 1.5000228881835938,
      "compile_s": 0.018353374995058402,
      "inner": 3,
      "median_ms": 19.14809166676908,
      "q1_ms": 13.2793883337096,
      "q3_ms": 20.027492333611008,
      "rel_err": 3.729404753096071e-07
    },
    {
      "impl": "jax_xla",
      "mode": "fwdbwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 256,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.12923589599813567,
      "xla_temp_mb": 10.5078125,
      "xla_arg_mb": 1.5,
      "xla_out_mb": 1.5000228881835938,
      "compile_s": 0.006865249000838958,
      "inner": 22,
      "median_ms": 1.941669136365156,
      "q1_ms": 1.9272914545250717,
      "q3_ms": 1.95177368170169,
      "rel_err": 0.0
    },
    {
      "impl": "torch_sdpa",
      "mode": "fwdbwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 256,
      "heads": 8,
      "head_dim": 64,
      "compile_s": 0.0018274670001119375,
      "inner": 41,
      "median_ms": 0.9988351705875929,
      "q1_ms": 0.9845215120926363,
      "q3_ms": 1.0243285365371464,
      "rel_err": 3.608001096293087e-07
    },
    {
      "impl": "xla_dense",
      "mode": "fwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 512,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.07020022000506287,
      "xla_temp_mb": 17.0,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.005904962003114633,
      "inner": 14,
      "median_ms": 3.2596580001284434,
      "q1_ms": 3.2592665717986944,
      "q3_ms": 3.301792785350699,
      "rel_err": 2.872570186144982e-08
    },
    {
      "impl": "xla_tiled",
      "mode": "fwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 512,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.12482806600019103,
      "xla_temp_mb": 8.532119750976562,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.00654333100101212,
      "inner": 13,
      "median_ms": 3.259759845744926,
      "q1_ms": 3.241375769273593,
      "q3_ms": 3.331722076556896,
      "rel_err": 2.3767681197166188e-07
    },
    {
      "impl": "mha_interpret",
      "mode": "fwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 512,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.1466199670030619,
      "xla_temp_mb": 6.348533630371094,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.013664733996847644,
      "inner": 5,
      "median_ms": 9.757982000883203,
      "q1_ms": 9.619223600020632,
      "q3_ms": 9.863435800070874,
      "rel_err": 2.3371727271052317e-07
    },
    {
      "impl": "jax_xla",
      "mode": "fwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 512,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.0634314729977632,
      "xla_temp_mb": 16.0,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.006662175001110882,
      "inner": 14,
      "median_ms": 2.714519499672211,
      "q1_ms": 2.7125457142704232,
      "q3_ms": 2.7571300711965057,
      "rel_err": 0.0
    },
    {
      "impl": "torch_sdpa",
      "mode": "fwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 512,
      "heads": 8,
      "head_dim": 64,
      "compile_s": 0.001709545002086088,
      "inner": 40,
      "median_ms": 0.6939241498912452,
      "q1_ms": 0.6818670999564347,
      "q3_ms": 0.7645816000149352,
      "rel_err": 2.3985823104779036e-07
    },
    {
      "impl": "xla_dense",
      "mode": "fwdbwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 512,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.14942764799343422,
      "xla_temp_mb": 41.015625,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.019032709999009967,
      "inner": 3,
      "median_ms": 16.236503002194997,
      "q1_ms": 16.188488666860696,
      "q3_ms": 16.426435666896094,
      "rel_err": 7.952522422288965e-08
    },
    {
      "impl": "xla_tiled",
      "mode": "fwdbwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 512,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.2614205409990973,
      "xla_temp_mb": 16.032623291015625,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.018218249002529774,
      "inner": 4,
      "median_ms": 10.4573504995642,
      "q1_ms": 10.433502999148914,
      "q3_ms": 10.51182575065468,
      "rel_err": 4.269848701565202e-07
    },
    {
      "impl": "mha_interpret",
      "mode": "fwdbwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 512,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.3834165180014679,
      "xla_temp_mb": 12.38128662109375,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.0518144539964851,
      "inner": 1,
      "median_ms": 46.55843900400214,
      "q1_ms": 42.59198500221828,
      "q3_ms": 46.75983000197448,
      "rel_err": 5.117985612168145e-07
    },
    {
      "impl": "jax_xla",
      "mode": "fwdbwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 512,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.15194874900043942,
      "xla_temp_mb": 25.03125,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.01581510299729416,
      "inner": 7,
      "median_ms": 7.258467428203273,
      "q1_ms": 7.245122142194305,
      "q3_ms": 7.304186857904175,
      "rel_err": 0.0
    },
    {
      "impl": "torch_sdpa",
      "mode": "fwdbwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 512,
      "heads": 8,
      "head_dim": 64,
      "compile_s": 0.006096495002566371,
      "inner": 13,
      "median_ms": 3.326117615730394,
      "q1_ms": 3.3259534617312827,
      "q3_ms": 3.333909307772186,
      "rel_err": 4.937880398332202e-07
    },
    {
      "impl": "probjax_mha",
      "mode": "fwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 1024,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.12376674799452303,
      "xla_temp_mb": 11.532119750976562,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.014131186006125063,
      "inner": 4,
      "median_ms": 10.10656174912583,
      "q1_ms": 9.972317749998183,
      "q3_ms": 10.216742750344565,
      "rel_err": 2.962723390657541e-07
    },
    {
      "impl": "xla_dense",
      "mode": "fwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 1024,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.06933711600140668,
      "xla_temp_mb": 66.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.02331902000150876,
      "inner": 1,
      "median_ms": 26.766862996737473,
      "q1_ms": 25.902086999849416,
      "q3_ms": 27.693846001056954,
      "rel_err": 3.255044381850416e-08
    },
    {
      "impl": "xla_tiled",
      "mode": "fwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 1024,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.12629309800104238,
      "xla_temp_mb": 11.532119750976562,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.013096017006319016,
      "inner": 4,
      "median_ms": 10.004050000134157,
      "q1_ms": 9.524783499728073,
      "q3_ms": 10.099428249304765,
      "rel_err": 2.962723390657541e-07
    },
    {
      "impl": "mha_interpret",
      "mode": "fwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 1024,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.15228693900280632,
      "xla_temp_mb": 12.600486755371094,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.06210542400367558,
      "inner": 1,
      "median_ms": 35.14922299655154,
      "q1_ms": 34.137438000470866,
      "q3_ms": 54.457254002045374,
      "rel_err": 2.8789885948488557e-07
    },
    {
      "impl": "jax_xla",
      "mode": "fwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 1024,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.06799595899792621,
      "xla_temp_mb": 64.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.02253126900177449,
      "inner": 1,
      "median_ms": 42.57083500124281,
      "q1_ms": 27.76633400208084,
      "q3_ms": 44.58864199841628,
      "rel_err": 0.0
    },
    {
      "impl": "torch_sdpa",
      "mode": "fwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 1024,
      "heads": 8,
      "head_dim": 64,
      "compile_s": 0.0027091059964732267,
      "inner": 22,
      "median_ms": 1.7622501363803167,
      "q1_ms": 1.760438454353293,
      "q3_ms": 1.7675797726869031,
      "rel_err": 2.945259098176026e-07
    },
    {
      "impl": "probjax_mha",
      "mode": "fwdbwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 1024,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.31368466499407077,
      "xla_temp_mb": 19.063873291015625,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.038458012997580227,
      "inner": 1,
      "median_ms": 32.46350200060988,
      "q1_ms": 32.022336999943946,
      "q3_ms": 33.30865099997027,
      "rel_err": 6.356668284179638e-07
    },
    {
      "impl": "xla_dense",
      "mode": "fwdbwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 1024,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.1975247180016595,
      "xla_temp_mb": 162.03125,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.07118148599693086,
      "inner": 1,
      "median_ms": 84.67633899999782,
      "q1_ms": 83.71835599973565,
      "q3_ms": 96.20139400067274,
      "rel_err": 8.819854702062318e-08
    },
    {
      "impl": "xla_tiled",
      "mode": "fwdbwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 1024,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.30762073100049747,
      "xla_temp_mb": 19.063873291015625,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.03981253599340562,
      "inner": 1,
      "median_ms": 31.433563002792653,
      "q1_ms": 31.357407999166753,
      "q3_ms": 33.4873690007953,
      "rel_err": 6.356668284179638e-07
    },
    {
      "impl": "jax_xla",
      "mode": "fwdbwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 1024,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.19834411900228588,
      "xla_temp_mb": 98.0625,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.056311531996470876,
      "inner": 1,
      "median_ms": 55.128182997577824,
      "q1_ms": 54.353086998162325,
      "q3_ms": 55.919049998919945,
      "rel_err": 0.0
    },
    {
      "impl": "torch_sdpa",
      "mode": "fwdbwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 1024,
      "heads": 8,
      "head_dim": 64,
      "compile_s": 0.010840278999239672,
      "inner": 4,
      "median_ms": 9.402797501024907,
      "q1_ms": 9.030190249177394,
      "q3_ms": 10.07870924877352,
      "rel_err": 6.624466273365493e-07
    },
    {
      "impl": "probjax_mha",
      "mode": "fwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 2048,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.12530497799889417,
      "xla_temp_mb": 17.532119750976562,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.03834508900035871,
      "inner": 1,
      "median_ms": 35.41739000502275,
      "q1_ms": 34.76226099883206,
      "q3_ms": 35.46832699794322,
      "rel_err": 3.448090669826248e-07
    },
    {
      "impl": "xla_dense",
      "mode": "fwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 2048,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.06800808199477615,
      "xla_temp_mb": 260.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.09094498900230974,
      "inner": 1,
      "median_ms": 113.4385370023665,
      "q1_ms": 112.47816000104649,
      "q3_ms": 115.52968400064856,
      "rel_err": 3.6978503559462416e-08
    },
    {
      "impl": "xla_tiled",
      "mode": "fwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 2048,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.12896158700459637,
      "xla_temp_mb": 17.532119750976562,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.038056694000260904,
      "inner": 1,
      "median_ms": 34.9475210023229,
      "q1_ms": 33.80392699909862,
      "q3_ms": 34.965575003298,
      "rel_err": 3.448090669826248e-07
    },
    {
      "impl": "jax_xla",
      "mode": "fwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 2048,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.06438389600225491,
      "xla_temp_mb": 256.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.08134470500226598,
      "inner": 1,
      "median_ms": 143.9195480052149,
      "q1_ms": 135.22090899641626,
      "q3_ms": 144.81423999677645,
      "rel_err": 0.0
    },
    {
      "impl": "torch_sdpa",
      "mode": "fwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 2048,
      "heads": 8,
      "head_dim": 64,
      "compile_s": 0.007745505005004816,
      "inner": 7,
      "median_ms": 6.494959857913533,
      "q1_ms": 6.36791671422543,
      "q3_ms": 6.597964857064653,
      "rel_err": 3.4046575015725216e-07
    },
    {
      "impl": "probjax_mha",
      "mode": "fwdbwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 2048,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.42783288400096353,
      "xla_temp_mb": 25.126373291015625,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.11667978999321349,
      "inner": 1,
      "median_ms": 106.9171790004475,
      "q1_ms": 104.19746299885446,
      "q3_ms": 108.30405200249515,
      "rel_err": 6.432054839581624e-07
    },
    {
      "impl": "xla_dense",
      "mode": "fwdbwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 2048,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.2845395019976422,
      "xla_temp_mb": 644.0625,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.2650486149941571,
      "inner": 1,
      "median_ms": 344.48356599750696,
      "q1_ms": 344.26558400446083,
      "q3_ms": 357.1689370000968,
      "rel_err": 9.40667218996439e-08
    },
    {
      "impl": "xla_tiled",
      "mode": "fwdbwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 2048,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.4324952660026611,
      "xla_temp_mb": 25.126373291015625,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.11649945899989689,
      "inner": 1,
      "median_ms": 109.57655999663984,
      "q1_ms": 109.39147000317462,
      "q3_ms": 110.33103799854871,
      "rel_err": 6.432054839581624e-07
    },
    {
      "impl": "jax_xla",
      "mode": "fwdbwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 2048,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.29923532599786995,
      "xla_temp_mb": 388.125,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.24290294499951415,
      "inner": 1,
      "median_ms": 237.56361300183926,
      "q1_ms": 234.86369499732973,
      "q3_ms": 238.3930119976867,
      "rel_err": 0.0
    },
    {
      "impl": "torch_sdpa",
      "mode": "fwdbwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 2048,
      "heads": 8,
      "head_dim": 64,
      "compile_s": 0.03216738099581562,
      "inner": 1,
      "median_ms": 32.48716700181831,
      "q1_ms": 29.65671099809697,
      "q3_ms": 32.790129000204615,
      "rel_err": 6.674964410825458e-07
    },
    {
      "impl": "probjax_mha",
      "mode": "fwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 512,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.06743726100103231,
      "xla_temp_mb": 17.015625,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.006698723002045881,
      "inner": 8,
      "median_ms": 2.5743539999893983,
      "q1_ms": 2.5350458745379,
      "q3_ms": 2.755014624199248,
      "rel_err": 5.745409798625413e-07
    },
    {
      "impl": "probjax_mha",
      "mode": "fwdbwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 512,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.16971529299917165,
      "xla_temp_mb": 33.0,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.017521647998364642,
      "inner": 3,
      "median_ms": 13.4903750004014,
      "q1_ms": 13.340708333998919,
      "q3_ms": 13.671106998420631,
      "rel_err": 3.5330077567049384e-07
    },
    {
      "impl": "probjax_mha",
      "mode": "fwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 512,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.07655628200154752,
      "xla_temp_mb": 17.015625,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.0061354569988907315,
      "inner": 15,
      "median_ms": 2.6387267333727022,
      "q1_ms": 2.6114017333990582,
      "q3_ms": 2.6691782666603103,
      "rel_err": 2.486659629315213e-07
    },
    {
      "impl": "probjax_mha",
      "mode": "fwdbwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 512,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.17498645600426244,
      "xla_temp_mb": 33.0,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.019136132003040984,
      "inner": 3,
      "median_ms": 15.307782998813005,
      "q1_ms": 15.040240332988711,
      "q3_ms": 15.541972999926656,
      "rel_err": 2.821393163656735e-07
    },
    {
      "impl": "probjax_mha",
      "mode": "fwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 4096,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.20679386999836424,
      "xla_temp_mb": 32.000244140625,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.23346182199748,
      "inner": 1,
      "median_ms": 243.93156800215365,
      "q1_ms": 243.60472099942854,
      "q3_ms": 245.32917400210863,
      "rel_err": 6.497361972594765e-07
    },
    {
      "impl": "xla_tiled",
      "mode": "fwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 4096,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.12009534099343,
      "xla_temp_mb": 32.000244140625,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.21588886700192234,
      "inner": 1,
      "median_ms": 241.00402999465587,
      "q1_ms": 240.6255840032827,
      "q3_ms": 243.13859900576063,
      "rel_err": 6.497361972594765e-07
    },
    {
      "impl": "jax_xla",
      "mode": "fwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 4096,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.05622364599548746,
      "xla_temp_mb": 1024.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.30379660399921704,
      "inner": 1,
      "median_ms": 406.4529030001722,
      "q1_ms": 402.26642200286733,
      "q3_ms": 406.6682610064163,
      "rel_err": 0.0
    },
    {
      "impl": "torch_sdpa",
      "mode": "fwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 4096,
      "heads": 8,
      "head_dim": 64,
      "compile_s": 0.0366282949980814,
      "inner": 1,
      "median_ms": 35.27513500011992,
      "q1_ms": 34.761052003887016,
      "q3_ms": 38.43819900066592,
      "rel_err": 6.342996378153116e-07
    },
    {
      "impl": "probjax_mha",
      "mode": "fwdbwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 4096,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.8875499200003105,
      "xla_temp_mb": 40.12550354003906,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.5922225489994162,
      "inner": 1,
      "median_ms": 721.1920480040135,
      "q1_ms": 706.800061998365,
      "q3_ms": 721.8787480014726,
      "rel_err": 7.313026030666727e-07
    },
    {
      "impl": "xla_tiled",
      "mode": "fwdbwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 4096,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.6751667159987846,
      "xla_temp_mb": 40.12550354003906,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.7038237500019022,
      "inner": 1,
      "median_ms": 707.0225019997451,
      "q1_ms": 702.093646999856,
      "q3_ms": 724.4709820006392,
      "rel_err": 7.313026030666727e-07
    },
    {
      "impl": "jax_xla",
      "mode": "fwdbwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 4096,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.5186329599964665,
      "xla_temp_mb": 1544.25,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.662046253994049,
      "inner": 1,
      "median_ms": 778.6298319988418,
      "q1_ms": 775.9024010010762,
      "q3_ms": 789.6306729962816,
      "rel_err": 0.0
    },
    {
      "impl": "torch_sdpa",
      "mode": "fwdbwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 4096,
      "heads": 8,
      "head_dim": 64,
      "compile_s": 0.48040752400265774,
      "inner": 1,
      "median_ms": 169.23324199888157,
      "q1_ms": 165.6551429987303,
      "q3_ms": 171.13748499832582,
      "rel_err": 7.264272585620225e-07
    },
    {
      "impl": "probjax_mha",
      "mode": "fwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 4096,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.13115878099779366,
      "xla_temp_mb": 32.00044250488281,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.13129645199660445,
      "inner": 1,
      "median_ms": 130.54880699928617,
      "q1_ms": 129.75501600158168,
      "q3_ms": 134.03452200145693,
      "rel_err": 3.827393591071548e-07
    },
    {
      "impl": "xla_tiled",
      "mode": "fwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 4096,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.12799291199917207,
      "xla_temp_mb": 32.00044250488281,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.12953728400316322,
      "inner": 1,
      "median_ms": 133.33980699826498,
      "q1_ms": 131.2511860014638,
      "q3_ms": 134.86092599487165,
      "rel_err": 3.827393591071548e-07
    },
    {
      "impl": "jax_xla",
      "mode": "fwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 4096,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.06265523399633821,
      "xla_temp_mb": 1024.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 8.0,
      "compile_s": 0.31357703599496745,
      "inner": 1,
      "median_ms": 415.1447670010384,
      "q1_ms": 412.422034998599,
      "q3_ms": 416.3155310016009,
      "rel_err": 0.0
    },
    {
      "impl": "torch_sdpa",
      "mode": "fwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 4096,
      "heads": 8,
      "head_dim": 64,
      "compile_s": 0.032711754996853415,
      "inner": 1,
      "median_ms": 28.637206996791065,
      "q1_ms": 26.715512001828756,
      "q3_ms": 29.099311999743804,
      "rel_err": 3.766782693380433e-07
    },
    {
      "impl": "probjax_mha",
      "mode": "fwdbwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 4096,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.9617830970018986,
      "xla_temp_mb": 40.125946044921875,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.3637768469998264,
      "inner": 1,
      "median_ms": 391.8708959972719,
      "q1_ms": 378.8007990006008,
      "q3_ms": 394.60322500235634,
      "rel_err": 6.569159430239878e-07
    },
    {
      "impl": "xla_tiled",
      "mode": "fwdbwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 4096,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.7024737859974266,
      "xla_temp_mb": 40.125946044921875,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.3617958179966081,
      "inner": 1,
      "median_ms": 386.13464400259545,
      "q1_ms": 385.65823300450575,
      "q3_ms": 398.04315299988957,
      "rel_err": 6.569159430239878e-07
    },
    {
      "impl": "jax_xla",
      "mode": "fwdbwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 4096,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.5424347489970387,
      "xla_temp_mb": 1544.25,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 24.000022888183594,
      "compile_s": 0.9558474180012126,
      "inner": 1,
      "median_ms": 1010.8238840039121,
      "q1_ms": 1002.3803009971743,
      "q3_ms": 1020.0746359987534,
      "rel_err": 0.0
    },
    {
      "impl": "torch_sdpa",
      "mode": "fwdbwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 4096,
      "heads": 8,
      "head_dim": 64,
      "compile_s": 0.09815043400158174,
      "inner": 1,
      "median_ms": 101.30212800140725,
      "q1_ms": 100.30300699872896,
      "q3_ms": 105.60635099682258,
      "rel_err": 6.795112885071446e-07
    },
    {
      "impl": "probjax_mha",
      "mode": "fwd",
      "variant": "masked_bias",
      "dtype": "float32",
      "batch": 1,
      "seq": 128,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.09979373299574945,
      "xla_temp_mb": 1.75,
      "xla_arg_mb": 0.7509765625,
      "xla_out_mb": 0.25,
      "compile_s": 0.003483960004814435,
      "inner": 75,
      "median_ms": 0.49414669328446814,
      "q1_ms": 0.4780936666065827,
      "q3_ms": 0.49783473339630296,
      "rel_err": 2.100135621176691e-08
    },
    {
      "impl": "xla_dense",
      "mode": "fwd",
      "variant": "masked_bias",
      "dtype": "float32",
      "batch": 1,
      "seq": 128,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.10238982399459928,
      "xla_temp_mb": 1.75,
      "xla_arg_mb": 0.7509765625,
      "xla_out_mb": 0.25,
      "compile_s": 0.003169584997522179,
      "inner": 75,
      "median_ms": 0.4825037333648652,
      "q1_ms": 0.48184154661915574,
      "q3_ms": 0.5064448933505143,
      "rel_err": 2.100135621176691e-08
    },
    {
      "impl": "xla_tiled",
      "mode": "fwd",
      "variant": "masked_bias",
      "dtype": "float32",
      "batch": 1,
      "seq": 128,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.12113660399336368,
      "xla_temp_mb": 1.75390625,
      "xla_arg_mb": 0.7509765625,
      "xla_out_mb": 0.25,
      "compile_s": 0.0032374739967053756,
      "inner": 99,
      "median_ms": 0.46976150505625724,
      "q1_ms": 0.4187493534779853,
      "q3_ms": 0.503821181770734,
      "rel_err": 1.4636642489439472e-07
    },
    {
      "impl": "jax_xla",
      "mode": "fwd",
      "variant": "masked_bias",
      "dtype": "float32",
      "batch": 1,
      "seq": 128,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.10134986699995352,
      "xla_temp_mb": 1.00390625,
      "xla_arg_mb": 0.7509765625,
      "xla_out_mb": 0.25,
      "compile_s": 0.0038608230024692602,
      "inner": 55,
      "median_ms": 0.6380742181632244,
      "q1_ms": 0.6334048182576556,
      "q3_ms": 0.7008921272989193,
      "rel_err": 0.0
    },
    {
      "impl": "torch_sdpa",
      "mode": "fwd",
      "variant": "masked_bias",
      "dtype": "float32",
      "batch": 1,
      "seq": 128,
      "heads": 8,
      "head_dim": 64,
      "compile_s": 0.0007272660004673526,
      "inner": 200,
      "median_ms": 0.08517040001606802,
      "q1_ms": 0.08515387002262287,
      "q3_ms": 0.08672274001582991,
      "rel_err": 1.565683681121507e-07
    },
    {
      "impl": "probjax_mha",
      "mode": "fwdbwd",
      "variant": "masked_bias",
      "dtype": "float32",
      "batch": 1,
      "seq": 128,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.1985307169961743,
      "xla_temp_mb": 3.321319580078125,
      "xla_arg_mb": 0.7509765625,
      "xla_out_mb": 0.751007080078125,
      "compile_s": 0.006048319999536034,
      "inner": 29,
      "median_ms": 0.9206165516937698,
      "q1_ms": 0.9205972069296343,
      "q3_ms": 0.9232183794414721,
      "rel_err": 4.512603661174836e-08
    },
    {
      "impl": "xla_dense",
      "mode": "fwdbwd",
      "variant": "masked_bias",
      "dtype": "float32",
      "batch": 1,
      "seq": 128,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.19628924800053937,
      "xla_temp_mb": 3.321319580078125,
      "xla_arg_mb": 0.7509765625,
      "xla_out_mb": 0.751007080078125,
      "compile_s": 0.005671998995239846,
      "inner": 49,
      "median_ms": 0.9166416122604218,
      "q1_ms": 0.8986011632760911,
      "q3_ms": 0.9414133672514569,
      "rel_err": 4.512603661174836e-08
    },
    {
      "impl": "xla_tiled",
      "mode": "fwdbwd",
      "variant": "masked_bias",
      "dtype": "float32",
      "batch": 1,
      "seq": 128,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.23956335199909518,
      "xla_temp_mb": 2.813507080078125,
      "xla_arg_mb": 0.7509765625,
      "xla_out_mb": 0.751007080078125,
      "compile_s": 0.0078009060016484,
      "inner": 38,
      "median_ms": 1.1798535790723308,
      "q1_ms": 1.1786415525421035,
      "q3_ms": 1.2142731579124781,
      "rel_err": 2.6495419841468037e-07
    },
    {
      "impl": "jax_xla",
      "mode": "fwdbwd",
      "variant": "masked_bias",
      "dtype": "float32",
      "batch": 1,
      "seq": 128,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.18073856299452018,
      "xla_temp_mb": 3.317413330078125,
      "xla_arg_mb": 0.7509765625,
      "xla_out_mb": 0.751007080078125,
      "compile_s": 0.005669744998158421,
      "inner": 52,
      "median_ms": 0.9394577885429428,
      "q1_ms": 0.9295134038368885,
      "q3_ms": 0.9622865384769214,
      "rel_err": 0.0
    },
    {
      "impl": "torch_sdpa",
      "mode": "fwdbwd",
      "variant": "masked_bias",
      "dtype": "float32",
      "batch": 1,
      "seq": 128,
      "heads": 8,
      "head_dim": 64,
      "compile_s": 0.2126900899966131,
      "inner": 59,
      "median_ms": 0.5222901355887574,
      "q1_ms": 0.5182193898520889,
      "q3_ms": 0.5338856100965039,
      "rel_err": 2.8186239388644355e-07
    },
    {
      "impl": "probjax_mha",
      "mode": "fwd",
      "variant": "masked_bias",
      "dtype": "float32",
      "batch": 1,
      "seq": 512,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.1355870280021918,
      "xla_temp_mb": 25.015625,
      "xla_arg_mb": 3.0009765625,
      "xla_out_mb": 1.0,
      "compile_s": 0.011134009997476824,
      "inner": 4,
      "median_ms": 4.924332250084262,
      "q1_ms": 4.515714250374003,
      "q3_ms": 5.186320000575506,
      "rel_err": 2.4502309816617566e-07
    },
    {
      "impl": "xla_dense",
      "mode": "fwd",
      "variant": "masked_bias",
      "dtype": "float32",
      "batch": 1,
      "seq": 512,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.11473050699714804,
      "xla_temp_mb": 25.0,
      "xla_arg_mb": 3.0009765625,
      "xla_out_mb": 1.0,
      "compile_s": 0.008345535999978893,
      "inner": 9,
      "median_ms": 5.267471444514943,
      "q1_ms": 5.191269444670373,
      "q3_ms": 5.303497333241264,
      "rel_err": 2.8144779173344003e-08
    },
    {
      "impl": "xla_tiled",
      "mode": "fwd",
      "variant": "masked_bias",
      "dtype": "float32",
      "batch": 1,
      "seq": 512,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.1800634059982258,
      "xla_temp_mb": 10.532188415527344,
      "xla_arg_mb": 3.0009765625,
      "xla_out_mb": 1.0,
      "compile_s": 0.007748828000330832,
      "inner": 10,
      "median_ms": 4.585856300400337,
      "q1_ms": 4.5629299005668145,
      "q3_ms": 4.633376100537134,
      "rel_err": 2.3505544388751877e-07
    },
    {
      "impl": "jax_xla",
      "mode": "fwd",
      "variant": "masked_bias",
      "dtype": "float32",
      "batch": 1,
      "seq": 512,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.10317835499881767,
      "xla_temp_mb": 16.0,
      "xla_arg_mb": 3.0009765625,
      "xla_out_mb": 1.0,
      "compile_s": 0.009767716001078952,
      "inner": 7,
      "median_ms": 7.1780818574812395,
      "q1_ms": 7.16122285770585,
      "q3_ms": 7.743097713890685,
      "rel_err": 0.0
    },
    {
      "impl": "torch_sdpa",
      "mode": "fwd",
      "variant": "masked_bias",
      "dtype": "float32",
      "batch": 1,
      "seq": 512,
      "heads": 8,
      "head_dim": 64,
      "compile_s": 0.0018079090004903264,
      "inner": 42,
      "median_ms": 0.8949432858138433,
      "q1_ms": 0.8948915237851907,
      "q3_ms": 0.89692038091016,
      "rel_err": 2.3887160440689945e-07
    },
    {
      "impl": "probjax_mha",
      "mode": "fwdbwd",
      "variant": "masked_bias",
      "dtype": "float32",
      "batch": 1,
      "seq": 512,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.24850712999614188,
      "xla_temp_mb": 42.001007080078125,
      "xla_arg_mb": 3.0009765625,
      "xla_out_mb": 3.001007080078125,
      "compile_s": 0.02291170800162945,
      "inner": 2,
      "median_ms": 20.039285001985263,
      "q1_ms": 19.975188999524107,
      "q3_ms": 20.22764649882447,
      "rel_err": 6.145865650710189e-07
    },
    {
      "impl": "xla_dense",
      "mode": "fwdbwd",
      "variant": "masked_bias",
      "dtype": "float32",
      "batch": 1,
      "seq": 512,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.20670786499977112,
      "xla_temp_mb": 50.016632080078125,
      "xla_arg_mb": 3.0009765625,
      "xla_out_mb": 3.001007080078125,
      "compile_s": 0.02255616600450594,
      "inner": 2,
      "median_ms": 22.091001002991106,
      "q1_ms": 21.9417535008688,
      "q3_ms": 22.40372699816362,
      "rel_err": 1.8770742382315318e-07
    },
    {
      "impl": "xla_tiled",
      "mode": "fwdbwd",
      "variant": "masked_bias",
      "dtype": "float32",
      "batch": 1,
      "seq": 512,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.4889310189973912,
      "xla_temp_mb": 18.283737182617188,
      "xla_arg_mb": 3.0009765625,
      "xla_out_mb": 3.001007080078125,
      "compile_s": 0.025931650001439266,
      "inner": 3,
      "median_ms": 15.753344666639652,
      "q1_ms": 15.676666332486395,
      "q3_ms": 15.90855833395229,
      "rel_err": 9.6330509634173e-07
    },
    {
      "impl": "jax_xla",
      "mode": "fwdbwd",
      "variant": "masked_bias",
      "dtype": "float32",
      "batch": 1,
      "seq": 512,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.2853079289998277,
      "xla_temp_mb": 34.001007080078125,
      "xla_arg_mb": 3.0009765625,
      "xla_out_mb": 3.001007080078125,
      "compile_s": 0.02761745899624657,
      "inner": 2,
      "median_ms": 18.078285498631885,
      "q1_ms": 17.32480549981119,
      "q3_ms": 19.051851999392966,
      "rel_err": 0.0
    },
    {
      "impl": "torch_sdpa",
      "mode": "fwdbwd",
      "variant": "masked_bias",
      "dtype": "float32",
      "batch": 1,
      "seq": 512,
      "heads": 8,
      "head_dim": 64,
      "compile_s": 0.0056583640034659766,
      "inner": 10,
      "median_ms": 3.987639800470788,
      "q1_ms": 3.980306800076505,
      "q3_ms": 4.022031900240108,
      "rel_err": 4.88813683403208e-07
    },
    {
      "impl": "probjax_mha",
      "mode": "fwd",
      "variant": "masked_bias",
      "dtype": "float32",
      "batch": 1,
      "seq": 1024,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.17599335800332483,
      "xla_temp_mb": 13.532188415527344,
      "xla_arg_mb": 6.0009765625,
      "xla_out_mb": 2.0,
      "compile_s": 0.018371860998740885,
      "inner": 3,
      "median_ms": 13.975809000839945,
      "q1_ms": 13.832207332598045,
      "q3_ms": 14.251607666665222,
      "rel_err": 2.8992422019814456e-07
    },
    {
      "impl": "xla_dense",
      "mode": "fwd",
      "variant": "masked_bias",
      "dtype": "float32",
      "batch": 1,
      "seq": 1024,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.10709054600010859,
      "xla_temp_mb": 98.0,
      "xla_arg_mb": 6.0009765625,
      "xla_out_mb": 2.0,
      "compile_s": 0.03716096199786989,
      "inner": 1,
      "median_ms": 45.16499099554494,
      "q1_ms": 44.73779399995692,
      "q3_ms": 45.472903999325354,
      "rel_err": 3.244912021352727e-08
    },
    {
      "impl": "xla_tiled",
      "mode": "fwd",
      "variant": "masked_bias",
      "dtype": "float32",
      "batch": 1,
      "seq": 1024,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.1734792209972511,
      "xla_temp_mb": 13.532188415527344,
      "xla_arg_mb": 6.0009765625,
      "xla_out_mb": 2.0,
      "compile_s": 0.0173944920024951,
      "inner": 3,
      "median_ms": 14.431384001606299,
      "q1_ms": 14.139412000076845,
      "q3_ms": 14.61868166613082,
      "rel_err": 2.8992422019814456e-07
    },
    {
      "impl": "jax_xla",
      "mode": "fwd",
      "variant": "masked_bias",
      "dtype": "float32",
      "batch": 1,
      "seq": 1024,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.2096631750027882,
      "xla_temp_mb": 64.0,
      "xla_arg_mb": 6.0009765625,
      "xla_out_mb": 2.0,
      "compile_s": 0.11924804199952632,
      "inner": 1,
      "median_ms": 100.06038800202077,
      "q1_ms": 96.98443899833364,
      "q3_ms": 102.42470000230242,
      "rel_err": 0.0
    },
    {
      "impl": "torch_sdpa",
      "mode": "fwd",
      "variant": "masked_bias",
      "dtype": "float32",
      "batch": 1,
      "seq": 1024,
      "heads": 8,
      "head_dim": 64,
      "compile_s": 0.006575006002094597,
      "inner": 11,
      "median_ms": 3.2680561820267364,
      "q1_ms": 3.251217091350224,
      "q3_ms": 4.492698999952568,
      "rel_err": 2.880252890095381e-07
    },
    {
      "impl": "probjax_mha",
      "mode": "fwdbwd",
      "variant": "masked_bias",
      "dtype": "float32",
      "batch": 1,
      "seq": 1024,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.4271163340017665,
      "xla_temp_mb": 21.314987182617188,
      "xla_arg_mb": 6.0009765625,
      "xla_out_mb": 6.001007080078125,
      "compile_s": 0.055950811998627614,
      "inner": 1,
      "median_ms": 45.42871999728959,
      "q1_ms": 45.03044699958991,
      "q3_ms": 46.363849003682844,
      "rel_err": 6.298606968997291e-06
    },
    {
      "impl": "xla_dense",
      "mode": "fwdbwd",
      "variant": "masked_bias",
      "dtype": "float32",
      "batch": 1,
      "seq": 1024,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.27600598499702755,
      "xla_temp_mb": 198.03225708007812,
      "xla_arg_mb": 6.0009765625,
      "xla_out_mb": 6.001007080078125,
      "compile_s": 0.09981715699541382,
      "inner": 1,
      "median_ms": 104.10864499863237,
      "q1_ms": 102.94012499798555,
      "q3_ms": 114.32770499959588,
      "rel_err": 2.1348821380048186e-07
    },
    {
      "impl": "xla_tiled",
      "mode": "fwdbwd",
      "variant": "masked_bias",
      "dtype": "float32",
      "batch": 1,
      "seq": 1024,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.42019987200183095,
      "xla_temp_mb": 21.314987182617188,
      "xla_arg_mb": 6.0009765625,
      "xla_out_mb": 6.001007080078125,
      "compile_s": 0.05156821299897274,
      "inner": 1,
      "median_ms": 43.784129004052375,
      "q1_ms": 42.969779002305586,
      "q3_ms": 45.487271003366914,
      "rel_err": 6.298606968997291e-06
    },
    {
      "impl": "jax_xla",
      "mode": "fwdbwd",
      "variant": "masked_bias",
      "dtype": "float32",
      "batch": 1,
      "seq": 1024,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.2667964850043063,
      "xla_temp_mb": 134.00100708007812,
      "xla_arg_mb": 6.0009765625,
      "xla_out_mb": 6.001007080078125,
      "compile_s": 0.07268017699971097,
      "inner": 1,
      "median_ms": 72.45248700201046,
      "q1_ms": 71.06390000262763,
      "q3_ms": 74.11316799698398,
      "rel_err": 0.0
    },
    {
      "impl": "torch_sdpa",
      "mode": "fwdbwd",
      "variant": "masked_bias",
      "dtype": "float32",
      "batch": 1,
      "seq": 1024,
      "heads": 8,
      "head_dim": 64,
      "compile_s": 0.018158748003770597,
      "inner": 3,
      "median_ms": 17.36662299905826,
      "q1_ms": 15.625008665665518,
      "q3_ms": 29.408594001627836,
      "rel_err": 6.589567616929611e-07
    },
    {
      "impl": "probjax_mha",
      "mode": "fwd",
      "variant": "masked_bias",
      "dtype": "float32",
      "batch": 1,
      "seq": 2048,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.18660595099936472,
      "xla_temp_mb": 19.532188415527344,
      "xla_arg_mb": 12.0009765625,
      "xla_out_mb": 4.0,
      "compile_s": 0.07017480800277553,
      "inner": 1,
      "median_ms": 53.47560100199189,
      "q1_ms": 52.20595900027547,
      "q3_ms": 54.3944079981884,
      "rel_err": 3.3641972615702244e-07
    },
    {
      "impl": "xla_dense",
      "mode": "fwd",
      "variant": "masked_bias",
      "dtype": "float32",
      "batch": 1,
      "seq": 2048,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.11168454500148073,
      "xla_temp_mb": 388.0,
      "xla_arg_mb": 12.0009765625,
      "xla_out_mb": 4.0,
      "compile_s": 0.14723440800298704,
      "inner": 1,
      "median_ms": 188.45392700313823,
      "q1_ms": 184.93583300005412,
      "q3_ms": 210.70255099766655,
      "rel_err": 3.6363376507038376e-08
    },
    {
      "impl": "xla_tiled",
      "mode": "fwd",
      "variant": "masked_bias",
      "dtype": "float32",
      "batch": 1,
      "seq": 2048,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.19697078699391568,
      "xla_temp_mb": 19.532188415527344,
      "xla_arg_mb": 12.0009765625,
      "xla_out_mb": 4.0,
      "compile_s": 0.0574373449999257,
      "inner": 1,
      "median_ms": 49.616701995546464,
      "q1_ms": 49.50385799747892,
      "q3_ms": 49.845244000607636,
      "rel_err": 3.3641972615702244e-07
    },
    {
      "impl": "jax_xla",
      "mode": "fwd",
      "variant": "masked_bias",
      "dtype": "float32",
      "batch": 1,
      "seq": 2048,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.12987550400430337,
      "xla_temp_mb": 256.0,
      "xla_arg_mb": 12.0009765625,
      "xla_out_mb": 4.0,
      "compile_s": 0.13198465300229145,
      "inner": 1,
      "median_ms": 181.7531519991462,
      "q1_ms": 176.11084999953164,
      "q3_ms": 181.76683800265891,
      "rel_err": 0.0
    },
    {
      "impl": "torch_sdpa",
      "mode": "fwd",
      "variant": "masked_bias",
      "dtype": "float32",
      "batch": 1,
      "seq": 2048,
      "heads": 8,
      "head_dim": 64,
      "compile_s": 0.020477574005781207,
      "inner": 2,
      "median_ms": 13.572402498539304,
      "q1_ms": 13.410020499577513,
      "q3_ms": 15.466330001800088,
      "rel_err": 3.322602689174116e-07
    },
    {
      "impl": "probjax_mha",
      "mode": "fwdbwd",
      "variant": "masked_bias",
      "dtype": "float32",
      "batch": 1,
      "seq": 2048,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.5218833800026914,
      "xla_temp_mb": 27.377487182617188,
      "xla_arg_mb": 12.0009765625,
      "xla_out_mb": 12.001007080078125,
      "compile_s": 0.18157708900253056,
      "inner": 1,
      "median_ms": 446.2576110017835,
      "q1_ms": 186.89769299817272,
      "q3_ms": 502.06571399758104,
      "rel_err": 8.95620398640334e-06
    },
    {
      "impl": "xla_dense",
      "mode": "fwdbwd",
      "variant": "masked_bias",
      "dtype": "float32",
      "batch": 1,
      "seq": 2048,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.3898433720023604,
      "xla_temp_mb": 788.0635070800781,
      "xla_arg_mb": 12.0009765625,
      "xla_out_mb": 12.001007080078125,
      "compile_s": 0.32686232100240886,
      "inner": 1,
      "median_ms": 423.7612189972424,
      "q1_ms": 399.6804490016075,
      "q3_ms": 446.3771170048858,
      "rel_err": 3.6953213145965567e-07
    },
    {
      "impl": "xla_tiled",
      "mode": "fwdbwd",
      "variant": "masked_bias",
      "dtype": "float32",
      "batch": 1,
      "seq": 2048,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 1.9294204640027601,
      "xla_temp_mb": 27.377487182617188,
      "xla_arg_mb": 12.0009765625,
      "xla_out_mb": 12.001007080078125,
      "compile_s": 0.47323023100034334,
      "inner": 1,
      "median_ms": 188.20358400262194,
      "q1_ms": 177.163380998536,
      "q3_ms": 265.4204619975644,
      "rel_err": 8.95620398640334e-06
    },
    {
      "impl": "jax_xla",
      "mode": "fwdbwd",
      "variant": "masked_bias",
      "dtype": "float32",
      "batch": 1,
      "seq": 2048,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.36106582800130127,
      "xla_temp_mb": 532.0010070800781,
      "xla_arg_mb": 12.0009765625,
      "xla_out_mb": 12.001007080078125,
      "compile_s": 0.27972650199808413,
      "inner": 1,
      "median_ms": 301.6502610043972,
      "q1_ms": 297.0919489962398,
      "q3_ms": 303.1656680032029,
      "rel_err": 0.0
    },
    {
      "impl": "torch_sdpa",
      "mode": "fwdbwd",
      "variant": "masked_bias",
      "dtype": "float32",
      "batch": 1,
      "seq": 2048,
      "heads": 8,
      "head_dim": 64,
      "compile_s": 0.05632870699628256,
      "inner": 1,
      "median_ms": 78.23369200195884,
      "q1_ms": 56.94651700468967,
      "q3_ms": 109.30138800176792,
      "rel_err": 6.651810839232719e-07
    },
    {
      "impl": "probjax_mha",
      "mode": "fwd",
      "variant": "masked_bias",
      "dtype": "float32",
      "batch": 1,
      "seq": 4096,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.28963141800340964,
      "xla_temp_mb": 32.000511169433594,
      "xla_arg_mb": 24.0009765625,
      "xla_out_mb": 8.0,
      "compile_s": 0.2714693339949008,
      "inner": 1,
      "median_ms": 186.14452900510514,
      "q1_ms": 181.2954470005934,
      "q3_ms": 187.8513670017128,
      "rel_err": 3.7451015655948445e-07
    },
    {
      "impl": "xla_dense",
      "mode": "fwd",
      "variant": "masked_bias",
      "dtype": "float32",
      "batch": 1,
      "seq": 4096,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.11325498600490391,
      "xla_temp_mb": 1544.0,
      "xla_arg_mb": 24.0009765625,
      "xla_out_mb": 8.0,
      "compile_s": 0.5525579020031728,
      "inner": 1,
      "median_ms": 705.4353650019038,
      "q1_ms": 693.1983269969351,
      "q3_ms": 711.1072839979897,
      "rel_err": 4.009121355682119e-08
    },
    {
      "impl": "xla_tiled",
      "mode": "fwd",
      "variant": "masked_bias",
      "dtype": "float32",
      "batch": 1,
      "seq": 4096,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.16716486799850827,
      "xla_temp_mb": 32.000511169433594,
      "xla_arg_mb": 24.0009765625,
      "xla_out_mb": 8.0,
      "compile_s": 0.17420223000226542,
      "inner": 1,
      "median_ms": 171.14163000223925,
      "q1_ms": 171.12388600071426,
      "q3_ms": 173.8101180017111,
      "rel_err": 3.7451015655948445e-07
    },
    {
      "impl": "jax_xla",
      "mode": "fwd",
      "variant": "masked_bias",
      "dtype": "float32",
      "batch": 1,
      "seq": 4096,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.09804861300654011,
      "xla_temp_mb": 1024.0,
      "xla_arg_mb": 24.0009765625,
      "xla_out_mb": 8.0,
      "compile_s": 0.5081185739982175,
      "inner": 1,
      "median_ms": 609.5388430039748,
      "q1_ms": 604.5850629961933,
      "q3_ms": 613.5764709979412,
      "rel_err": 0.0
    },
    {
      "impl": "torch_sdpa",
      "mode": "fwd",
      "variant": "masked_bias",
      "dtype": "float32",
      "batch": 1,
      "seq": 4096,
      "heads": 8,
      "head_dim": 64,
      "compile_s": 0.05911432600260014,
      "inner": 1,
      "median_ms": 51.30970400205115,
      "q1_ms": 44.71168299642159,
      "q3_ms": 57.985942003142554,
      "rel_err": 3.688001411876452e-07
    },
    {
      "impl": "probjax_mha",
      "mode": "fwdbwd",
      "variant": "masked_bias",
      "dtype": "float32",
      "batch": 1,
      "seq": 4096,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 1.3459001149967662,
      "xla_temp_mb": 40.12602233886719,
      "xla_arg_mb": 24.0009765625,
      "xla_out_mb": 24.001007080078125,
      "compile_s": 0.5526680299954023,
      "inner": 1,
      "median_ms": 572.0149159969878,
      "q1_ms": 571.9319790005102,
      "q3_ms": 572.8793899979792,
      "rel_err": 1.3332394005935819e-05
    },
    {
      "impl": "xla_dense",
      "mode": "fwdbwd",
      "variant": "masked_bias",
      "dtype": "float32",
      "batch": 1,
      "seq": 4096,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.5763689919986064,
      "xla_temp_mb": 3144.126007080078,
      "xla_arg_mb": 24.0009765625,
      "xla_out_mb": 24.001007080078125,
      "compile_s": 1.2240898169984575,
      "inner": 1,
      "median_ms": 1505.1253760029795,
      "q1_ms": 1496.0085409984458,
      "q3_ms": 1509.5734389979043,
      "rel_err": 3.6176373872421734e-07
    },
    {
      "impl": "xla_tiled",
      "mode": "fwdbwd",
      "variant": "masked_bias",
      "dtype": "float32",
      "batch": 1,
      "seq": 4096,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.773779941999237,
      "xla_temp_mb": 40.12602233886719,
      "xla_arg_mb": 24.0009765625,
      "xla_out_mb": 24.001007080078125,
      "compile_s": 0.542305729999498,
      "inner": 1,
      "median_ms": 574.373066003318,
      "q1_ms": 571.0967990016798,
      "q3_ms": 577.2330949985189,
      "rel_err": 1.3332394005935819e-05
    },
    {
      "impl": "jax_xla",
      "mode": "fwdbwd",
      "variant": "masked_bias",
      "dtype": "float32",
      "batch": 1,
      "seq": 4096,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.6064442200004123,
      "xla_temp_mb": 2120.001007080078,
      "xla_arg_mb": 24.0009765625,
      "xla_out_mb": 24.001007080078125,
      "compile_s": 1.118564261996653,
      "inner": 1,
      "median_ms": 1201.9058059959207,
      "q1_ms": 1193.9988789963536,
      "q3_ms": 1205.1126920050592,
      "rel_err": 0.0
    },
    {
      "impl": "torch_sdpa",
      "mode": "fwdbwd",
      "variant": "masked_bias",
      "dtype": "float32",
      "batch": 1,
      "seq": 4096,
      "heads": 8,
      "head_dim": 64,
      "compile_s": 0.1805009939998854,
      "inner": 1,
      "median_ms": 184.89443400176242,
      "q1_ms": 179.00253900006646,
      "q3_ms": 186.61317499936558,
      "rel_err": 6.779501490871814e-07
    },
    {
      "impl": "probjax_mha",
      "mode": "fwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 256,
      "heads": 8,
      "head_dim": 128,
      "lower_compile_s": 0.07337839800311485,
      "xla_temp_mb": 5.0,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.004850413999520242,
      "inner": 20,
      "median_ms": 1.5061976002471056,
      "q1_ms": 1.4992479998909403,
      "q3_ms": 1.5174018000834621,
      "rel_err": 1.3738302035077561e-07
    },
    {
      "impl": "xla_dense",
      "mode": "fwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 256,
      "heads": 8,
      "head_dim": 128,
      "lower_compile_s": 0.07501418099855073,
      "xla_temp_mb": 5.0,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.004235029002302326,
      "inner": 43,
      "median_ms": 0.8391624185944323,
      "q1_ms": 0.8278303022623646,
      "q3_ms": 0.8810330001261021,
      "rel_err": 1.3738302035077561e-07
    },
    {
      "impl": "xla_tiled",
      "mode": "fwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 256,
      "heads": 8,
      "head_dim": 128,
      "lower_compile_s": 0.08185962800052948,
      "xla_temp_mb": 5.0078125,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.004051312003866769,
      "inner": 28,
      "median_ms": 1.2800620001403982,
      "q1_ms": 1.2656640355999116,
      "q3_ms": 1.3024536071627932,
      "rel_err": 4.119164093973956e-07
    },
    {
      "impl": "jax_xla",
      "mode": "fwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 256,
      "heads": 8,
      "head_dim": 128,
      "lower_compile_s": 0.05933490699680988,
      "xla_temp_mb": 4.0078125,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.00491776000126265,
      "inner": 33,
      "median_ms": 1.3291759393971667,
      "q1_ms": 1.3065070000293693,
      "q3_ms": 1.339367939429146,
      "rel_err": 0.0
    },
    {
      "impl": "torch_sdpa",
      "mode": "fwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 256,
      "heads": 8,
      "head_dim": 128,
      "compile_s": 0.0035178929974790663,
      "inner": 44,
      "median_ms": 0.27278386369289365,
      "q1_ms": 0.27098750003708777,
      "q3_ms": 0.2737076591239946,
      "rel_err": 3.643299245099848e-07
    },
    {
      "impl": "probjax_mha",
      "mode": "fwdbwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 256,
      "heads": 8,
      "head_dim": 128,
      "lower_compile_s": 0.155887252993125,
      "xla_temp_mb": 11.015625,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.009461052999540698,
      "inner": 11,
      "median_ms": 2.6888516361147845,
      "q1_ms": 2.6680250005235644,
      "q3_ms": 2.7030913638141514,
      "rel_err": 1.383557786454406e-07
    },
    {
      "impl": "xla_dense",
      "mode": "fwdbwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 256,
      "heads": 8,
      "head_dim": 128,
      "lower_compile_s": 0.1451374919997761,
      "xla_temp_mb": 11.015625,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.008146788000885863,
      "inner": 17,
      "median_ms": 2.6998036469997126,
      "q1_ms": 2.6717158235602683,
      "q3_ms": 2.7608721177196407,
      "rel_err": 1.383557786454406e-07
    },
    {
      "impl": "xla_tiled",
      "mode": "fwdbwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 256,
      "heads": 8,
      "head_dim": 128,
      "lower_compile_s": 0.17409937599586556,
      "xla_temp_mb": 9.0,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.010497342998860404,
      "inner": 13,
      "median_ms": 3.3422349230162443,
      "q1_ms": 3.3047879226684858,
      "q3_ms": 3.3783955385353273,
      "rel_err": 3.0467115611896387e-07
    },
    {
      "impl": "jax_xla",
      "mode": "fwdbwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 256,
      "heads": 8,
      "head_dim": 128,
      "lower_compile_s": 0.13638819399784552,
      "xla_temp_mb": 11.0078125,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.007767279996187426,
      "inner": 18,
      "median_ms": 2.4410026667788363,
      "q1_ms": 2.4264407220471185,
      "q3_ms": 2.474680277777629,
      "rel_err": 0.0
    },
    {
      "impl": "torch_sdpa",
      "mode": "fwdbwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 256,
      "heads": 8,
      "head_dim": 128,
      "compile_s": 0.32022095000138506,
      "inner": 21,
      "median_ms": 1.673171571469755,
      "q1_ms": 1.6712679044714953,
      "q3_ms": 1.678381857345812,
      "rel_err": 6.887719126530129e-07
    },
    {
      "impl": "probjax_mha",
      "mode": "fwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 256,
      "heads": 8,
      "head_dim": 128,
      "lower_compile_s": 0.06483375700190663,
      "xla_temp_mb": 5.0,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.0042157319985562935,
      "inner": 31,
      "median_ms": 1.4240212258300744,
      "q1_ms": 1.3868079676608285,
      "q3_ms": 1.4618178709849503,
      "rel_err": 2.5965927729561828e-08
    },
    {
      "impl": "xla_dense",
      "mode": "fwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 256,
      "heads": 8,
      "head_dim": 128,
      "lower_compile_s": 0.06670907299849205,
      "xla_temp_mb": 5.0,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.004132064997975249,
      "inner": 37,
      "median_ms": 0.9059522973886389,
      "q1_ms": 0.9002087297975211,
      "q3_ms": 1.0468092431610048,
      "rel_err": 2.5965927729561828e-08
    },
    {
      "impl": "xla_tiled",
      "mode": "fwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 256,
      "heads": 8,
      "head_dim": 128,
      "lower_compile_s": 0.07399706700380193,
      "xla_temp_mb": 5.0078125,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.004021234999527223,
      "inner": 40,
      "median_ms": 1.3233946248874418,
      "q1_ms": 1.3076909250230528,
      "q3_ms": 1.3404942250417662,
      "rel_err": 1.8989882994472704e-07
    },
    {
      "impl": "jax_xla",
      "mode": "fwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 256,
      "heads": 8,
      "head_dim": 128,
      "lower_compile_s": 0.06361474200093653,
      "xla_temp_mb": 4.0078125,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.004674880998209119,
      "inner": 28,
      "median_ms": 1.5672132857227032,
      "q1_ms": 1.5558807858074684,
      "q3_ms": 1.568494999966268,
      "rel_err": 0.0
    },
    {
      "impl": "torch_sdpa",
      "mode": "fwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 256,
      "heads": 8,
      "head_dim": 128,
      "compile_s": 0.0007471030039596371,
      "inner": 144,
      "median_ms": 0.3053534722110069,
      "q1_ms": 0.2921125972332245,
      "q3_ms": 0.31028989582814065,
      "rel_err": 1.9121080488179032e-07
    },
    {
      "impl": "probjax_mha",
      "mode": "fwdbwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 256,
      "heads": 8,
      "head_dim": 128,
      "lower_compile_s": 0.1486748120005359,
      "xla_temp_mb": 11.015625,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.008420274003583472,
      "inner": 16,
      "median_ms": 2.698323250115209,
      "q1_ms": 2.681713750007475,
      "q3_ms": 2.7295748127471597,
      "rel_err": 3.881335867006964e-08
    },
    {
      "impl": "xla_dense",
      "mode": "fwdbwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 256,
      "heads": 8,
      "head_dim": 128,
      "lower_compile_s": 0.14400260600086767,
      "xla_temp_mb": 11.015625,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.008163067999703344,
      "inner": 16,
      "median_ms": 2.7013889998670493,
      "q1_ms": 2.684828999917954,
      "q3_ms": 2.7215439999963564,
      "rel_err": 3.881335867006964e-08
    },
    {
      "impl": "xla_tiled",
      "mode": "fwdbwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 256,
      "heads": 8,
      "head_dim": 128,
      "lower_compile_s": 0.1803561099950457,
      "xla_temp_mb": 9.0,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.01058673200168414,
      "inner": 14,
      "median_ms": 3.418611928542045,
      "q1_ms": 3.386953428292015,
      "q3_ms": 3.423274999866927,
      "rel_err": 2.5649926571036723e-07
    },
    {
      "impl": "jax_xla",
      "mode": "fwdbwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 256,
      "heads": 8,
      "head_dim": 128,
      "lower_compile_s": 0.14433868099877145,
      "xla_temp_mb": 11.0078125,
      "xla_arg_mb": 3.0,
      "xla_out_mb": 3.0000228881835938,
      "compile_s": 0.0074452209955779836,
      "inner": 18,
      "median_ms": 2.5416045557196614,
      "q1_ms": 2.5066684999425584,
      "q3_ms": 2.5565389441908337,
      "rel_err": 0.0
    },
    {
      "impl": "torch_sdpa",
      "mode": "fwdbwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 256,
      "heads": 8,
      "head_dim": 128,
      "compile_s": 0.002903852000599727,
      "inner": 18,
      "median_ms": 1.6967444999964856,
      "q1_ms": 1.677070111327339,
      "q3_ms": 1.699176944384817,
      "rel_err": 6.138205173041401e-07
    },
    {
      "impl": "probjax_mha",
      "mode": "fwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 1024,
      "heads": 8,
      "head_dim": 128,
      "lower_compile_s": 0.1147216980025405,
      "xla_temp_mb": 19.03167724609375,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.028557990000990685,
      "inner": 1,
      "median_ms": 23.838184999476653,
      "q1_ms": 23.64831500017317,
      "q3_ms": 24.05218899366446,
      "rel_err": 6.457201554586343e-07
    },
    {
      "impl": "xla_dense",
      "mode": "fwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 1024,
      "heads": 8,
      "head_dim": 128,
      "lower_compile_s": 0.058946723002009094,
      "xla_temp_mb": 68.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.02631393599585863,
      "inner": 1,
      "median_ms": 29.715398995904252,
      "q1_ms": 29.20858200377552,
      "q3_ms": 30.42288700089557,
      "rel_err": 1.8566269426575043e-07
    },
    {
      "impl": "xla_tiled",
      "mode": "fwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 1024,
      "heads": 8,
      "head_dim": 128,
      "lower_compile_s": 0.11703404199943179,
      "xla_temp_mb": 19.03167724609375,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.02889340500405524,
      "inner": 2,
      "median_ms": 25.120063499343814,
      "q1_ms": 23.90581700092298,
      "q3_ms": 25.488760500593344,
      "rel_err": 6.457201554586343e-07
    },
    {
      "impl": "jax_xla",
      "mode": "fwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 1024,
      "heads": 8,
      "head_dim": 128,
      "lower_compile_s": 0.057972500995674636,
      "xla_temp_mb": 64.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.0255439699976705,
      "inner": 1,
      "median_ms": 28.807552000216674,
      "q1_ms": 27.834081003675237,
      "q3_ms": 29.17560800415231,
      "rel_err": 0.0
    },
    {
      "impl": "torch_sdpa",
      "mode": "fwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 1024,
      "heads": 8,
      "head_dim": 128,
      "compile_s": 0.00777023499540519,
      "inner": 11,
      "median_ms": 3.989276545656718,
      "q1_ms": 3.9038194544528695,
      "q3_ms": 4.403634727747307,
      "rel_err": 6.264540350179933e-07
    },
    {
      "impl": "probjax_mha",
      "mode": "fwdbwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 1024,
      "heads": 8,
      "head_dim": 128,
      "lower_compile_s": 0.41720991100009996,
      "xla_temp_mb": 28.063186645507812,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.08268838500225684,
      "inner": 1,
      "median_ms": 73.2518280055956,
      "q1_ms": 72.30166099907365,
      "q3_ms": 73.88465500116581,
      "rel_err": 6.665849985691146e-07
    },
    {
      "impl": "xla_dense",
      "mode": "fwdbwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 1024,
      "heads": 8,
      "head_dim": 128,
      "lower_compile_s": 0.2905398160000914,
      "xla_temp_mb": 164.03125,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.07986761999927694,
      "inner": 1,
      "median_ms": 93.43640199949732,
      "q1_ms": 92.2943419936928,
      "q3_ms": 94.18359499977669,
      "rel_err": 1.9501046070570158e-07
    },
    {
      "impl": "xla_tiled",
      "mode": "fwdbwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 1024,
      "heads": 8,
      "head_dim": 128,
      "lower_compile_s": 0.4128444340021815,
      "xla_temp_mb": 28.063186645507812,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.08055256600346183,
      "inner": 1,
      "median_ms": 71.79411099787103,
      "q1_ms": 71.02880300226388,
      "q3_ms": 71.87477400293574,
      "rel_err": 6.665849985691146e-07
    },
    {
      "impl": "jax_xla",
      "mode": "fwdbwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 1024,
      "heads": 8,
      "head_dim": 128,
      "lower_compile_s": 0.27479437600413803,
      "xla_temp_mb": 100.0625,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.05213225600164151,
      "inner": 1,
      "median_ms": 57.3404360038694,
      "q1_ms": 57.075875003647525,
      "q3_ms": 57.622699001512956,
      "rel_err": 0.0
    },
    {
      "impl": "torch_sdpa",
      "mode": "fwdbwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 1024,
      "heads": 8,
      "head_dim": 128,
      "compile_s": 0.024101733004499692,
      "inner": 2,
      "median_ms": 20.92841150079039,
      "q1_ms": 20.90151549782604,
      "q3_ms": 21.813239000039175,
      "rel_err": 8.736927987031707e-07
    },
    {
      "impl": "probjax_mha",
      "mode": "fwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 1024,
      "heads": 8,
      "head_dim": 128,
      "lower_compile_s": 0.11522558100114111,
      "xla_temp_mb": 19.032119750976562,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.01945290100411512,
      "inner": 3,
      "median_ms": 16.888550667014595,
      "q1_ms": 16.62026299891295,
      "q3_ms": 16.951977333519608,
      "rel_err": 2.987612670671099e-07
    },
    {
      "impl": "xla_dense",
      "mode": "fwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 1024,
      "heads": 8,
      "head_dim": 128,
      "lower_compile_s": 0.0656318069959525,
      "xla_temp_mb": 68.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.025304798000433948,
      "inner": 1,
      "median_ms": 29.431282004225068,
      "q1_ms": 28.638632997171953,
      "q3_ms": 29.55521599506028,
      "rel_err": 3.365201655727169e-08
    },
    {
      "impl": "xla_tiled",
      "mode": "fwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 1024,
      "heads": 8,
      "head_dim": 128,
      "lower_compile_s": 0.11236019199714065,
      "xla_temp_mb": 19.032119750976562,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.018716156999289524,
      "inner": 3,
      "median_ms": 17.164175333164167,
      "q1_ms": 16.950080333723843,
      "q3_ms": 17.2451850012294,
      "rel_err": 2.987612670671099e-07
    },
    {
      "impl": "jax_xla",
      "mode": "fwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 1024,
      "heads": 8,
      "head_dim": 128,
      "lower_compile_s": 0.06764554399705958,
      "xla_temp_mb": 64.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 4.0,
      "compile_s": 0.0248654870010796,
      "inner": 1,
      "median_ms": 27.55845000501722,
      "q1_ms": 27.492715998960193,
      "q3_ms": 28.29871999711031,
      "rel_err": 0.0
    },
    {
      "impl": "torch_sdpa",
      "mode": "fwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 1024,
      "heads": 8,
      "head_dim": 128,
      "compile_s": 0.006048969997209497,
      "inner": 9,
      "median_ms": 3.263936333597586,
      "q1_ms": 3.258146444245035,
      "q3_ms": 3.2771100005548863,
      "rel_err": 2.9656892958054394e-07
    },
    {
      "impl": "probjax_mha",
      "mode": "fwdbwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 1024,
      "heads": 8,
      "head_dim": 128,
      "lower_compile_s": 0.4052000669980771,
      "xla_temp_mb": 28.063873291015625,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.057026180998946074,
      "inner": 1,
      "median_ms": 48.06442200060701,
      "q1_ms": 47.95584800012875,
      "q3_ms": 49.31126100564143,
      "rel_err": 6.348066677375121e-07
    },
    {
      "impl": "xla_dense",
      "mode": "fwdbwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 1024,
      "heads": 8,
      "head_dim": 128,
      "lower_compile_s": 0.2902246189987636,
      "xla_temp_mb": 164.03125,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.08038037899677875,
      "inner": 1,
      "median_ms": 90.95909699681215,
      "q1_ms": 89.67749300063588,
      "q3_ms": 91.69222300261026,
      "rel_err": 8.827457108706153e-08
    },
    {
      "impl": "xla_tiled",
      "mode": "fwdbwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 1024,
      "heads": 8,
      "head_dim": 128,
      "lower_compile_s": 0.41567605899763294,
      "xla_temp_mb": 28.063873291015625,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.056446444999892265,
      "inner": 1,
      "median_ms": 46.02423399774125,
      "q1_ms": 45.85437299829209,
      "q3_ms": 47.27992900006939,
      "rel_err": 6.348066677375121e-07
    },
    {
      "impl": "jax_xla",
      "mode": "fwdbwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 1024,
      "heads": 8,
      "head_dim": 128,
      "lower_compile_s": 0.28858446699450724,
      "xla_temp_mb": 100.0625,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 12.000022888183594,
      "compile_s": 0.06689374300185591,
      "inner": 1,
      "median_ms": 64.25672600016696,
      "q1_ms": 63.21184999978868,
      "q3_ms": 65.61096599762095,
      "rel_err": 0.0
    },
    {
      "impl": "torch_sdpa",
      "mode": "fwdbwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 1024,
      "heads": 8,
      "head_dim": 128,
      "compile_s": 0.017410888998711016,
      "inner": 3,
      "median_ms": 17.06117033366657,
      "q1_ms": 17.052002998146538,
      "q3_ms": 17.457824001515593,
      "rel_err": 8.3982934767225e-07
    },
    {
      "impl": "probjax_mha",
      "mode": "fwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 4096,
      "heads": 8,
      "head_dim": 128,
      "lower_compile_s": 0.11065551699721254,
      "xla_temp_mb": 64.000244140625,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.3696889659986482,
      "inner": 1,
      "median_ms": 367.15607599762734,
      "q1_ms": 363.87247600214323,
      "q3_ms": 368.0191180028487,
      "rel_err": 6.524236687462414e-07
    },
    {
      "impl": "xla_dense",
      "mode": "fwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 4096,
      "heads": 8,
      "head_dim": 128,
      "lower_compile_s": 0.06300434800505172,
      "xla_temp_mb": 1040.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.40199327300069854,
      "inner": 1,
      "median_ms": 486.35511199972825,
      "q1_ms": 485.93242199422093,
      "q3_ms": 487.411059999431,
      "rel_err": 2.004682929632955e-07
    },
    {
      "impl": "xla_tiled",
      "mode": "fwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 4096,
      "heads": 8,
      "head_dim": 128,
      "lower_compile_s": 0.1206341400029487,
      "xla_temp_mb": 64.000244140625,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.3620424449982238,
      "inner": 1,
      "median_ms": 368.91465200460516,
      "q1_ms": 365.4959269988467,
      "q3_ms": 372.1248129950254,
      "rel_err": 6.524236687462414e-07
    },
    {
      "impl": "jax_xla",
      "mode": "fwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 4096,
      "heads": 8,
      "head_dim": 128,
      "lower_compile_s": 0.0594228629997815,
      "xla_temp_mb": 1024.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.3501184030028526,
      "inner": 1,
      "median_ms": 434.2696050007362,
      "q1_ms": 433.4918429958634,
      "q3_ms": 434.4948409998324,
      "rel_err": 0.0
    },
    {
      "impl": "torch_sdpa",
      "mode": "fwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 4096,
      "heads": 8,
      "head_dim": 128,
      "compile_s": 0.08695112600253196,
      "inner": 1,
      "median_ms": 76.17688999744132,
      "q1_ms": 71.39472500421107,
      "q3_ms": 78.88552299846197,
      "rel_err": 6.323630209628945e-07
    },
    {
      "impl": "probjax_mha",
      "mode": "fwdbwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 4096,
      "heads": 8,
      "head_dim": 128,
      "lower_compile_s": 1.347596268999041,
      "xla_temp_mb": 80.12550354003906,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.872266017999209,
      "inner": 1,
      "median_ms": 1057.8132060036296,
      "q1_ms": 1014.8651389972656,
      "q3_ms": 1064.9430700068478,
      "rel_err": 7.316858712818506e-07
    },
    {
      "impl": "xla_dense",
      "mode": "fwdbwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 4096,
      "heads": 8,
      "head_dim": 128,
      "lower_compile_s": 0.9671523819997674,
      "xla_temp_mb": 2576.125,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 1.1703864809969673,
      "inner": 1,
      "median_ms": 1416.2753490018076,
      "q1_ms": 1408.5721009978442,
      "q3_ms": 1422.1941029973095,
      "rel_err": 2.1013758706355202e-07
    },
    {
      "impl": "xla_tiled",
      "mode": "fwdbwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 4096,
      "heads": 8,
      "head_dim": 128,
      "lower_compile_s": 1.2285473470037687,
      "xla_temp_mb": 80.12550354003906,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 1.0151706469987403,
      "inner": 1,
      "median_ms": 1042.8868150047492,
      "q1_ms": 1033.2789550011512,
      "q3_ms": 1051.9143490018905,
      "rel_err": 7.316858712818506e-07
    },
    {
      "impl": "jax_xla",
      "mode": "fwdbwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 4096,
      "heads": 8,
      "head_dim": 128,
      "lower_compile_s": 0.9965776129974984,
      "xla_temp_mb": 1552.25,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.8089913190051448,
      "inner": 1,
      "median_ms": 926.2172829985502,
      "q1_ms": 925.6118400007836,
      "q3_ms": 929.3671900013578,
      "rel_err": 0.0
    },
    {
      "impl": "torch_sdpa",
      "mode": "fwdbwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 4096,
      "heads": 8,
      "head_dim": 128,
      "compile_s": 0.2998003889952088,
      "inner": 1,
      "median_ms": 302.48204999952577,
      "q1_ms": 301.03761999635026,
      "q3_ms": 306.67051300406456,
      "rel_err": 9.28965817388167e-07
    },
    {
      "impl": "probjax_mha",
      "mode": "fwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 4096,
      "heads": 8,
      "head_dim": 128,
      "lower_compile_s": 0.11635988600028213,
      "xla_temp_mb": 64.00044250488281,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.20234038699709345,
      "inner": 1,
      "median_ms": 203.1754470008309,
      "q1_ms": 202.45735900243744,
      "q3_ms": 203.7740189989563,
      "rel_err": 3.848323348332148e-07
    },
    {
      "impl": "xla_dense",
      "mode": "fwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 4096,
      "heads": 8,
      "head_dim": 128,
      "lower_compile_s": 0.06560079799965024,
      "xla_temp_mb": 1040.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.3924515659964527,
      "inner": 1,
      "median_ms": 484.62692400062224,
      "q1_ms": 483.2017779990565,
      "q3_ms": 486.71799799922155,
      "rel_err": 4.224104675042508e-08
    },
    {
      "impl": "xla_tiled",
      "mode": "fwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 4096,
      "heads": 8,
      "head_dim": 128,
      "lower_compile_s": 0.11383099299564492,
      "xla_temp_mb": 64.00044250488281,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.20456455400562845,
      "inner": 1,
      "median_ms": 205.09443599439692,
      "q1_ms": 203.9020910015097,
      "q3_ms": 206.80224500392796,
      "rel_err": 3.848323348332148e-07
    },
    {
      "impl": "jax_xla",
      "mode": "fwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 4096,
      "heads": 8,
      "head_dim": 128,
      "lower_compile_s": 0.06507031499495497,
      "xla_temp_mb": 1024.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.34818133000226226,
      "inner": 1,
      "median_ms": 442.4418609996792,
      "q1_ms": 439.5188019989291,
      "q3_ms": 444.18523799686227,
      "rel_err": 0.0
    },
    {
      "impl": "torch_sdpa",
      "mode": "fwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 4096,
      "heads": 8,
      "head_dim": 128,
      "compile_s": 0.04286535899882438,
      "inner": 1,
      "median_ms": 45.544116001110524,
      "q1_ms": 45.44857500150101,
      "q3_ms": 48.49162000027718,
      "rel_err": 3.784842371147991e-07
    },
    {
      "impl": "probjax_mha",
      "mode": "fwdbwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 4096,
      "heads": 8,
      "head_dim": 128,
      "lower_compile_s": 1.449155514004815,
      "xla_temp_mb": 80.12594604492188,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.5497821790049784,
      "inner": 1,
      "median_ms": 584.9706899971352,
      "q1_ms": 578.4035689939628,
      "q3_ms": 589.4648879984743,
      "rel_err": 6.558927652602838e-07
    },
    {
      "impl": "xla_dense",
      "mode": "fwdbwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 4096,
      "heads": 8,
      "head_dim": 128,
      "lower_compile_s": 0.9310824950007373,
      "xla_temp_mb": 2576.125,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 1.1636647269988316,
      "inner": 1,
      "median_ms": 1394.4837559974985,
      "q1_ms": 1385.2257059988915,
      "q3_ms": 1405.4123339956277,
      "rel_err": 1.0068948560017e-07
    },
    {
      "impl": "xla_tiled",
      "mode": "fwdbwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 4096,
      "heads": 8,
      "head_dim": 128,
      "lower_compile_s": 1.2152402900028392,
      "xla_temp_mb": 80.12594604492188,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 0.5509294579969719,
      "inner": 1,
      "median_ms": 584.6815620025154,
      "q1_ms": 573.5609100011061,
      "q3_ms": 587.0040750014596,
      "rel_err": 6.558927652602838e-07
    },
    {
      "impl": "jax_xla",
      "mode": "fwdbwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 4096,
      "heads": 8,
      "head_dim": 128,
      "lower_compile_s": 0.9982007430007798,
      "xla_temp_mb": 1552.25,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 48.000022888183594,
      "compile_s": 1.101557006004441,
      "inner": 1,
      "median_ms": 1155.1522580048186,
      "q1_ms": 1139.7926900026505,
      "q3_ms": 1156.5316269989125,
      "rel_err": 0.0
    },
    {
      "impl": "torch_sdpa",
      "mode": "fwdbwd",
      "variant": "causal",
      "dtype": "float32",
      "batch": 1,
      "seq": 4096,
      "heads": 8,
      "head_dim": 128,
      "compile_s": 0.17343236500164494,
      "inner": 1,
      "median_ms": 186.4500850060722,
      "q1_ms": 185.5596899986267,
      "q3_ms": 187.35186000412796,
      "rel_err": 8.62723446474624e-07
    },
    {
      "impl": "probjax_mha",
      "mode": "fwd",
      "variant": "causal",
      "dtype": "bfloat16",
      "batch": 1,
      "seq": 512,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.08632980199763551,
      "xla_temp_mb": 16.515625,
      "xla_arg_mb": 1.5,
      "xla_out_mb": 0.5,
      "compile_s": 0.006755997994332574,
      "inner": 8,
      "median_ms": 2.9018943750998005,
      "q1_ms": 2.8881347498099785,
      "q3_ms": 2.9716138751609833,
      "rel_err": 0.001973010642361036
    },
    {
      "impl": "xla_tiled",
      "mode": "fwd",
      "variant": "causal",
      "dtype": "bfloat16",
      "batch": 1,
      "seq": 512,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.13601630199991632,
      "xla_temp_mb": 7.0321197509765625,
      "xla_arg_mb": 1.5,
      "xla_out_mb": 0.5,
      "compile_s": 0.006056403995899018,
      "inner": 13,
      "median_ms": 3.6002867692821803,
      "q1_ms": 3.536520769277838,
      "q3_ms": 3.604604923076784,
      "rel_err": 0.0019711647188369987
    },
    {
      "impl": "jax_xla",
      "mode": "fwd",
      "variant": "causal",
      "dtype": "bfloat16",
      "batch": 1,
      "seq": 512,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.07622254599846201,
      "xla_temp_mb": 16.0,
      "xla_arg_mb": 1.5,
      "xla_out_mb": 0.5,
      "compile_s": 0.00670930899650557,
      "inner": 15,
      "median_ms": 2.766700533296292,
      "q1_ms": 2.751128466722245,
      "q3_ms": 2.7689655330808214,
      "rel_err": 0.002242159810983226
    },
    {
      "impl": "torch_sdpa",
      "mode": "fwd",
      "variant": "causal",
      "dtype": "bfloat16",
      "batch": 1,
      "seq": 512,
      "heads": 8,
      "head_dim": 64,
      "compile_s": 0.0016785250045359135,
      "inner": 52,
      "median_ms": 0.7079042114548009,
      "q1_ms": 0.7028947307284957,
      "q3_ms": 0.7457299808055933,
      "rel_err": 0.0019743659208738868
    },
    {
      "impl": "probjax_mha",
      "mode": "fwdbwd",
      "variant": "causal",
      "dtype": "bfloat16",
      "batch": 1,
      "seq": 512,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.15200653299689293,
      "xla_temp_mb": 30.0,
      "xla_arg_mb": 1.5,
      "xla_out_mb": 1.5000228881835938,
      "compile_s": 0.020076168999366928,
      "inner": 3,
      "median_ms": 8.730054333379181,
      "q1_ms": 8.44092733314028,
      "q3_ms": 8.731186334140753,
      "rel_err": 0.001977548985271394
    },
    {
      "impl": "xla_tiled",
      "mode": "fwdbwd",
      "variant": "causal",
      "dtype": "bfloat16",
      "batch": 1,
      "seq": 512,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.22373140400304692,
      "xla_temp_mb": 14.782623291015625,
      "xla_arg_mb": 1.5,
      "xla_out_mb": 1.5000228881835938,
      "compile_s": 0.01699466200079769,
      "inner": 4,
      "median_ms": 10.456766500283265,
      "q1_ms": 10.208509251242504,
      "q3_ms": 10.541945001023123,
      "rel_err": 0.0023029435255706144
    },
    {
      "impl": "jax_xla",
      "mode": "fwdbwd",
      "variant": "causal",
      "dtype": "bfloat16",
      "batch": 1,
      "seq": 512,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.14524642800097354,
      "xla_temp_mb": 26.015625,
      "xla_arg_mb": 1.5,
      "xla_out_mb": 1.5000228881835938,
      "compile_s": 0.014814428002864588,
      "inner": 7,
      "median_ms": 6.622022571848772,
      "q1_ms": 6.5735464283664315,
      "q3_ms": 6.622174142518945,
      "rel_err": 0.001954885513519406
    },
    {
      "impl": "torch_sdpa",
      "mode": "fwdbwd",
      "variant": "causal",
      "dtype": "bfloat16",
      "batch": 1,
      "seq": 512,
      "heads": 8,
      "head_dim": 64,
      "compile_s": 0.325732101999165,
      "inner": 2,
      "median_ms": 16.071876001660712,
      "q1_ms": 15.994199500710238,
      "q3_ms": 18.601189498440363,
      "rel_err": 0.004092898043326846
    },
    {
      "impl": "probjax_mha",
      "mode": "fwd",
      "variant": "causal",
      "dtype": "bfloat16",
      "batch": 1,
      "seq": 2048,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.12928910899790935,
      "xla_temp_mb": 13.032119750976562,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.04232075999607332,
      "inner": 1,
      "median_ms": 39.4657599972561,
      "q1_ms": 39.10469700349495,
      "q3_ms": 39.739166997605935,
      "rel_err": 0.0020319399214280944
    },
    {
      "impl": "xla_tiled",
      "mode": "fwd",
      "variant": "causal",
      "dtype": "bfloat16",
      "batch": 1,
      "seq": 2048,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.12797532400145428,
      "xla_temp_mb": 13.032119750976562,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.043346338999981526,
      "inner": 1,
      "median_ms": 40.82949800067581,
      "q1_ms": 40.50781999831088,
      "q3_ms": 40.97838099551154,
      "rel_err": 0.0020319399214280944
    },
    {
      "impl": "jax_xla",
      "mode": "fwd",
      "variant": "causal",
      "dtype": "bfloat16",
      "batch": 1,
      "seq": 2048,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.07502807700075209,
      "xla_temp_mb": 256.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 2.0,
      "compile_s": 0.0663152590059326,
      "inner": 1,
      "median_ms": 115.95554100495065,
      "q1_ms": 106.10653500043554,
      "q3_ms": 119.67716999788536,
      "rel_err": 0.0022670725800938538
    },
    {
      "impl": "torch_sdpa",
      "mode": "fwd",
      "variant": "causal",
      "dtype": "bfloat16",
      "batch": 1,
      "seq": 2048,
      "heads": 8,
      "head_dim": 64,
      "compile_s": 0.009430514997802675,
      "inner": 6,
      "median_ms": 7.0909993337409105,
      "q1_ms": 7.086902999920615,
      "q3_ms": 7.092528667271836,
      "rel_err": 0.0020370475914894158
    },
    {
      "impl": "probjax_mha",
      "mode": "fwdbwd",
      "variant": "causal",
      "dtype": "bfloat16",
      "batch": 1,
      "seq": 2048,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.2802650529993116,
      "xla_temp_mb": 23.876373291015625,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.11976391400094144,
      "inner": 1,
      "median_ms": 110.98905900144018,
      "q1_ms": 108.07275299885077,
      "q3_ms": 112.8180970044923,
      "rel_err": 0.0022949345432637207
    },
    {
      "impl": "xla_tiled",
      "mode": "fwdbwd",
      "variant": "causal",
      "dtype": "bfloat16",
      "batch": 1,
      "seq": 2048,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.28646983799990267,
      "xla_temp_mb": 23.876373291015625,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.1131498959948658,
      "inner": 1,
      "median_ms": 109.21742799837375,
      "q1_ms": 108.55165900284192,
      "q3_ms": 109.89384799904656,
      "rel_err": 0.0022949345432637207
    },
    {
      "impl": "jax_xla",
      "mode": "fwdbwd",
      "variant": "causal",
      "dtype": "bfloat16",
      "batch": 1,
      "seq": 2048,
      "heads": 8,
      "head_dim": 64,
      "lower_compile_s": 0.20756529900245368,
      "xla_temp_mb": 392.0625,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 6.000022888183594,
      "compile_s": 0.197779891997925,
      "inner": 1,
      "median_ms": 219.2717090001679,
      "q1_ms": 218.52409499842906,
      "q3_ms": 219.34977700584568,
      "rel_err": 0.001917751766425642
    },
    {
      "impl": "torch_sdpa",
      "mode": "fwdbwd",
      "variant": "causal",
      "dtype": "bfloat16",
      "batch": 1,
      "seq": 2048,
      "heads": 8,
      "head_dim": 64,
      "compile_s": 0.1447762980023981,
      "inner": 1,
      "median_ms": 153.98654699674807,
      "q1_ms": 153.22588900016854,
      "q3_ms": 166.03231999761192,
      "rel_err": 0.004207182137042932
    }
  ],
  "configs": [
    {
      "device": "cpu",
      "append": true,
      "seqs": "128,256,512,1024,2048",
      "dims": "64",
      "dtypes": "float32",
      "variants": "causal",
      "modes": "fwd,fwdbwd",
      "batch": 1,
      "heads": 8,
      "impls": "probjax_mha,xla_dense,xla_tiled,mha_interpret,jax_xla,torch_sdpa",
      "interpret_max_seq": 1024,
      "no_ref": false,
      "peaks": false,
      "peak_seqs": "4096,8192,16384",
      "peak_impls": "xla_tiled,xla_dense,jax_xla,torch_sdpa",
      "peak_timeout": 300,
      "max_fit": false,
      "fit_seqs": "4096,8192,16384,32768,65536,131072",
      "fit_impls": "xla_tiled,xla_dense,jax_xla"
    },
    {
      "device": "cpu",
      "append": true,
      "seqs": "512",
      "dims": "64",
      "dtypes": "float32",
      "variants": "plain,causal",
      "modes": "fwd,fwdbwd",
      "batch": 1,
      "heads": 8,
      "impls": "probjax_mha",
      "interpret_max_seq": 1024,
      "no_ref": false,
      "peaks": false,
      "peak_seqs": "4096,8192,16384",
      "peak_impls": "xla_tiled,xla_dense,jax_xla,torch_sdpa",
      "peak_timeout": 300,
      "max_fit": false,
      "fit_seqs": "4096,8192,16384,32768,65536,131072",
      "fit_impls": "xla_tiled,xla_dense,jax_xla"
    },
    {
      "device": "cpu",
      "append": true,
      "seqs": "4096",
      "dims": "64",
      "dtypes": "float32",
      "variants": "plain,causal",
      "modes": "fwd,fwdbwd",
      "batch": 1,
      "heads": 8,
      "impls": "probjax_mha,xla_tiled,jax_xla,torch_sdpa",
      "interpret_max_seq": 1024,
      "no_ref": false,
      "peaks": false,
      "peak_seqs": "4096,8192,16384",
      "peak_impls": "xla_tiled,xla_dense,jax_xla,torch_sdpa",
      "peak_timeout": 300,
      "max_fit": false,
      "fit_seqs": "4096,8192,16384,32768,65536,131072",
      "fit_impls": "xla_tiled,xla_dense,jax_xla"
    },
    {
      "device": "cpu",
      "append": true,
      "seqs": "128,512,1024,2048,4096",
      "dims": "64",
      "dtypes": "float32",
      "variants": "masked_bias",
      "modes": "fwd,fwdbwd",
      "batch": 1,
      "heads": 8,
      "impls": "probjax_mha,xla_dense,xla_tiled,jax_xla,torch_sdpa",
      "interpret_max_seq": 1024,
      "no_ref": false,
      "peaks": false,
      "peak_seqs": "4096,8192,16384",
      "peak_impls": "xla_tiled,xla_dense,jax_xla,torch_sdpa",
      "peak_timeout": 300,
      "max_fit": false,
      "fit_seqs": "4096,8192,16384,32768,65536,131072",
      "fit_impls": "xla_tiled,xla_dense,jax_xla"
    },
    {
      "device": "cpu",
      "append": true,
      "seqs": "256,1024,4096",
      "dims": "128",
      "dtypes": "float32",
      "variants": "plain,causal",
      "modes": "fwd,fwdbwd",
      "batch": 1,
      "heads": 8,
      "impls": "probjax_mha,xla_dense,xla_tiled,jax_xla,torch_sdpa",
      "interpret_max_seq": 1024,
      "no_ref": false,
      "peaks": false,
      "peak_seqs": "4096,8192,16384",
      "peak_impls": "xla_tiled,xla_dense,jax_xla,torch_sdpa",
      "peak_timeout": 300,
      "max_fit": false,
      "fit_seqs": "4096,8192,16384,32768,65536,131072",
      "fit_impls": "xla_tiled,xla_dense,jax_xla"
    },
    {
      "device": "cpu",
      "append": true,
      "seqs": "512,2048",
      "dims": "64",
      "dtypes": "bfloat16",
      "variants": "causal",
      "modes": "fwd,fwdbwd",
      "batch": 1,
      "heads": 8,
      "impls": "probjax_mha,xla_tiled,jax_xla,torch_sdpa",
      "interpret_max_seq": 1024,
      "no_ref": false,
      "peaks": false,
      "peak_seqs": "4096,8192,16384",
      "peak_impls": "xla_tiled,xla_dense,jax_xla,torch_sdpa",
      "peak_timeout": 300,
      "max_fit": false,
      "fit_seqs": "4096,8192,16384,32768,65536,131072",
      "fit_impls": "xla_tiled,xla_dense,jax_xla"
    }
  ],
  "note": "probjax_mha rows at seq=512 were re-run after the auto-strategy change (single tile instead of dense); all other rows are from the same code. Timings on a shared 24-core host under a bench lock (other agents' unlocked test runs were active), so expect ~10-30% run-to-run noise."
}