{
  "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": {
    "suites": [
      "kernel",
      "layer"
    ],
    "window": 64,
    "chunk": 128,
    "batch": 2,
    "heads": 4,
    "layer": {
      "batch": 2,
      "heads": 8,
      "head_dim": 64,
      "window": 128
    }
  },
  "results": [
    {
      "suite": "kernel",
      "impl": "scan",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 512,
      "dim": 32,
      "dtype": "float32",
      "xla_temp_mb": 4.0001068115234375,
      "xla_arg_mb": 1.5625,
      "xla_out_mb": 0.5625152587890625,
      "compile_s": 0.009579293997376226,
      "inner": 9,
      "median_ms": 4.42225411178596,
      "q1_ms": 4.406850332290762,
      "q3_ms": 4.473452554925138,
      "rel_err": 0.0
    },
    {
      "suite": "kernel",
      "impl": "chunked",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 512,
      "dim": 32,
      "dtype": "float32",
      "xla_temp_mb": 5.0,
      "xla_arg_mb": 1.5625,
      "xla_out_mb": 0.5625152587890625,
      "compile_s": 0.019588004986871965,
      "inner": 23,
      "median_ms": 1.2955098266146428,
      "q1_ms": 1.2894491734661406,
      "q3_ms": 1.304185869546769,
      "rel_err": 4.295214368434808e-07
    },
    {
      "suite": "kernel",
      "impl": "scan",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 512,
      "dim": 32,
      "dtype": "float32",
      "xla_temp_mb": 44.07838439941406,
      "xla_arg_mb": 1.5625,
      "xla_out_mb": 1.562530517578125,
      "compile_s": 0.029158222998376004,
      "inner": 2,
      "median_ms": 24.437911495624576,
      "q1_ms": 24.207619004300795,
      "q3_ms": 24.74989449547138,
      "rel_err": 0.0
    },
    {
      "suite": "kernel",
      "impl": "chunked",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 512,
      "dim": 32,
      "dtype": "float32",
      "xla_temp_mb": 14.953125,
      "xla_arg_mb": 1.5625,
      "xla_out_mb": 1.562530517578125,
      "compile_s": 0.022382415001629852,
      "inner": 5,
      "median_ms": 5.74986900028307,
      "q1_ms": 5.689680599607527,
      "q3_ms": 5.798715798300691,
      "rel_err": 2.6317341870939302e-05
    },
    {
      "suite": "kernel",
      "impl": "scan",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 2048,
      "dim": 32,
      "dtype": "float32",
      "xla_temp_mb": 16.000106811523438,
      "xla_arg_mb": 6.0625,
      "xla_out_mb": 2.0625152587890625,
      "compile_s": 0.023095751006621867,
      "inner": 2,
      "median_ms": 17.238665503100492,
      "q1_ms": 16.80445000238251,
      "q3_ms": 17.287658498389646,
      "rel_err": 0.0
    },
    {
      "suite": "kernel",
      "impl": "chunked",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 2048,
      "dim": 32,
      "dtype": "float32",
      "xla_temp_mb": 20.0,
      "xla_arg_mb": 6.0625,
      "xla_out_mb": 2.0625152587890625,
      "compile_s": 0.04322059599508066,
      "inner": 2,
      "median_ms": 16.99314350116765,
      "q1_ms": 16.568105995247606,
      "q3_ms": 17.502581496955827,
      "rel_err": 8.168978251241978e-07
    },
    {
      "suite": "kernel",
      "impl": "scan",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 2048,
      "dim": 32,
      "dtype": "float32",
      "xla_temp_mb": 176.31275939941406,
      "xla_arg_mb": 6.0625,
      "xla_out_mb": 6.062530517578125,
      "compile_s": 0.12543134899169672,
      "inner": 1,
      "median_ms": 103.36283000651747,
      "q1_ms": 101.46823999821208,
      "q3_ms": 103.65222000109497,
      "rel_err": 0.0
    },
    {
      "suite": "kernel",
      "impl": "chunked",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 2048,
      "dim": 32,
      "dtype": "float32",
      "xla_temp_mb": 59.4375,
      "xla_arg_mb": 6.0625,
      "xla_out_mb": 6.062530517578125,
      "compile_s": 0.04404215699469205,
      "inner": 1,
      "median_ms": 32.50821899564471,
      "q1_ms": 31.94928599987179,
      "q3_ms": 32.68611599924043,
      "rel_err": 8.834376408183848e-05
    },
    {
      "suite": "kernel",
      "impl": "scan",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 8192,
      "dim": 32,
      "dtype": "float32",
      "xla_temp_mb": 64.00010681152344,
      "xla_arg_mb": 24.0625,
      "xla_out_mb": 8.062515258789062,
      "compile_s": 0.07240623100369703,
      "inner": 1,
      "median_ms": 74.51066000794526,
      "q1_ms": 73.72500700876117,
      "q3_ms": 74.74709900270682,
      "rel_err": 0.0
    },
    {
      "suite": "kernel",
      "impl": "chunked",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 8192,
      "dim": 32,
      "dtype": "float32",
      "xla_temp_mb": 80.0,
      "xla_arg_mb": 24.0625,
      "xla_out_mb": 8.062515258789062,
      "compile_s": 0.042437158001121134,
      "inner": 1,
      "median_ms": 37.15971800556872,
      "q1_ms": 36.492751009063795,
      "q3_ms": 37.57186100119725,
      "rel_err": 1.6341763575699853e-06
    },
    {
      "suite": "kernel",
      "impl": "scan",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 8192,
      "dim": 32,
      "dtype": "float32",
      "xla_temp_mb": 705.2502593994141,
      "xla_arg_mb": 24.0625,
      "xla_out_mb": 24.062530517578125,
      "compile_s": 0.34105528000509366,
      "inner": 1,
      "median_ms": 393.14951999404,
      "q1_ms": 392.82678699237294,
      "q3_ms": 395.3846579970559,
      "rel_err": 0.0
    },
    {
      "suite": "kernel",
      "impl": "chunked",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 8192,
      "dim": 32,
      "dtype": "float32",
      "xla_temp_mb": 237.375,
      "xla_arg_mb": 24.0625,
      "xla_out_mb": 24.062530517578125,
      "compile_s": 0.1498893599928124,
      "inner": 1,
      "median_ms": 143.08248399174772,
      "q1_ms": 139.46915698761586,
      "q3_ms": 143.7044960039202,
      "rel_err": 0.00033065989864435936
    },
    {
      "suite": "kernel",
      "impl": "scan",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 512,
      "dim": 64,
      "dtype": "float32",
      "xla_temp_mb": 9.000106811523438,
      "xla_arg_mb": 3.25,
      "xla_out_mb": 1.2500152587890625,
      "compile_s": 0.01505130299483426,
      "inner": 6,
      "median_ms": 6.135127499874216,
      "q1_ms": 5.41352133344238,
      "q3_ms": 6.309908167168032,
      "rel_err": 0.0
    },
    {
      "suite": "kernel",
      "impl": "chunked",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 512,
      "dim": 64,
      "dtype": "float32",
      "xla_temp_mb": 8.0,
      "xla_arg_mb": 3.25,
      "xla_out_mb": 1.2500152587890625,
      "compile_s": 0.01115412799117621,
      "inner": 16,
      "median_ms": 2.6885142497121706,
      "q1_ms": 2.650982937666413,
      "q3_ms": 2.701547749893507,
      "rel_err": 4.357646828394869e-07
    },
    {
      "suite": "kernel",
      "impl": "scan",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 512,
      "dim": 64,
      "dtype": "float32",
      "xla_temp_mb": 154.07838439941406,
      "xla_arg_mb": 3.25,
      "xla_out_mb": 3.250030517578125,
      "compile_s": 0.05802696400496643,
      "inner": 1,
      "median_ms": 66.28946599084884,
      "q1_ms": 65.6679859966971,
      "q3_ms": 66.49851400288753,
      "rel_err": 0.0
    },
    {
      "suite": "kernel",
      "impl": "chunked",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 512,
      "dim": 64,
      "dtype": "float32",
      "xla_temp_mb": 28.328125,
      "xla_arg_mb": 3.25,
      "xla_out_mb": 3.250030517578125,
      "compile_s": 0.025316561994259246,
      "inner": 4,
      "median_ms": 10.43965200005914,
      "q1_ms": 10.2829492498131,
      "q3_ms": 10.721106751589105,
      "rel_err": 4.686784295496889e-05
    },
    {
      "suite": "kernel",
      "impl": "scan",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 2048,
      "dim": 64,
      "dtype": "float32",
      "xla_temp_mb": 36.00010681152344,
      "xla_arg_mb": 12.25,
      "xla_out_mb": 4.2500152587890625,
      "compile_s": 0.02555030600342434,
      "inner": 2,
      "median_ms": 19.91921399894636,
      "q1_ms": 19.755282999540213,
      "q3_ms": 22.65268900373485,
      "rel_err": 0.0
    },
    {
      "suite": "kernel",
      "impl": "chunked",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 2048,
      "dim": 64,
      "dtype": "float32",
      "xla_temp_mb": 32.0,
      "xla_arg_mb": 12.25,
      "xla_out_mb": 4.2500152587890625,
      "compile_s": 0.019393896000110544,
      "inner": 4,
      "median_ms": 12.974048251635395,
      "q1_ms": 12.95302849757718,
      "q3_ms": 13.128359249094501,
      "rel_err": 8.327394768453762e-07
    },
    {
      "suite": "kernel",
      "impl": "scan",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 2048,
      "dim": 64,
      "dtype": "float32",
      "xla_temp_mb": 616.3127593994141,
      "xla_arg_mb": 12.25,
      "xla_out_mb": 12.250030517578125,
      "compile_s": 0.23916235299839173,
      "inner": 1,
      "median_ms": 265.54490100534167,
      "q1_ms": 263.6689659993863,
      "q3_ms": 268.7442910100799,
      "rel_err": 0.0
    },
    {
      "suite": "kernel",
      "impl": "chunked",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 2048,
      "dim": 64,
      "dtype": "float32",
      "xla_temp_mb": 112.5625,
      "xla_arg_mb": 12.25,
      "xla_out_mb": 12.250030517578125,
      "compile_s": 0.07212518699816428,
      "inner": 1,
      "median_ms": 66.86760499724187,
      "q1_ms": 66.68735199491493,
      "q3_ms": 66.94962100300472,
      "rel_err": 0.00013942034059385747
    },
    {
      "suite": "kernel",
      "impl": "scan",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 8192,
      "dim": 64,
      "dtype": "float32",
      "xla_temp_mb": 144.00010681152344,
      "xla_arg_mb": 48.25,
      "xla_out_mb": 16.250015258789062,
      "compile_s": 0.09852700099872891,
      "inner": 1,
      "median_ms": 104.77157699642703,
      "q1_ms": 104.41554100543726,
      "q3_ms": 105.0161929888418,
      "rel_err": 0.0
    },
    {
      "suite": "kernel",
      "impl": "chunked",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 8192,
      "dim": 64,
      "dtype": "float32",
      "xla_temp_mb": 144.0,
      "xla_arg_mb": 48.25,
      "xla_out_mb": 16.250015258789062,
      "compile_s": 0.07823547800944652,
      "inner": 1,
      "median_ms": 81.52057099505328,
      "q1_ms": 81.28649699210655,
      "q3_ms": 84.26025799417403,
      "rel_err": 1.6221775620670304e-06
    },
    {
      "suite": "kernel",
      "impl": "scan",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 8192,
      "dim": 64,
      "dtype": "float32",
      "xla_temp_mb": 2465.250259399414,
      "xla_arg_mb": 48.25,
      "xla_out_mb": 48.250030517578125,
      "compile_s": 0.8326925760047743,
      "inner": 1,
      "median_ms": 1094.040480995318,
      "q1_ms": 1090.8530239976244,
      "q3_ms": 1121.385097008897,
      "rel_err": 0.0
    },
    {
      "suite": "kernel",
      "impl": "chunked",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 8192,
      "dim": 64,
      "dtype": "float32",
      "xla_temp_mb": 449.5,
      "xla_arg_mb": 48.25,
      "xla_out_mb": 48.250030517578125,
      "compile_s": 0.2710021520033479,
      "inner": 1,
      "median_ms": 295.9559029986849,
      "q1_ms": 292.6157059991965,
      "q3_ms": 297.4077130056685,
      "rel_err": 0.0004938367489051214
    },
    {
      "suite": "kernel",
      "impl": "scan",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 512,
      "dim": 128,
      "dtype": "float32",
      "xla_temp_mb": 18.000106811523438,
      "xla_arg_mb": 7.0,
      "xla_out_mb": 3.0000152587890625,
      "compile_s": 0.05354307898960542,
      "inner": 1,
      "median_ms": 37.027262995252386,
      "q1_ms": 17.371144000208005,
      "q3_ms": 58.01265500485897,
      "rel_err": 0.0
    },
    {
      "suite": "kernel",
      "impl": "chunked",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 512,
      "dim": 128,
      "dtype": "float32",
      "xla_temp_mb": 20.0,
      "xla_arg_mb": 7.0,
      "xla_out_mb": 3.0000152587890625,
      "compile_s": 0.02587885998946149,
      "inner": 6,
      "median_ms": 5.658054501206304,
      "q1_ms": 5.630605834691475,
      "q3_ms": 5.750718332516651,
      "rel_err": 4.460460020963777e-07
    },
    {
      "suite": "kernel",
      "impl": "scan",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 512,
      "dim": 128,
      "dtype": "float32",
      "xla_temp_mb": 564.0783843994141,
      "xla_arg_mb": 7.0,
      "xla_out_mb": 7.000030517578125,
      "compile_s": 0.2098907840118045,
      "inner": 1,
      "median_ms": 271.3200529979076,
      "q1_ms": 266.32182698813267,
      "q3_ms": 278.7923699943349,
      "rel_err": 0.0
    },
    {
      "suite": "kernel",
      "impl": "chunked",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 512,
      "dim": 128,
      "dtype": "float32",
      "xla_temp_mb": 61.078125,
      "xla_arg_mb": 7.0,
      "xla_out_mb": 7.000030517578125,
      "compile_s": 0.03930968999338802,
      "inner": 1,
      "median_ms": 28.814004006562755,
      "q1_ms": 28.78113099723123,
      "q3_ms": 29.017909007961862,
      "rel_err": 8.256220714945254e-05
    },
    {
      "suite": "kernel",
      "impl": "scan",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 2048,
      "dim": 128,
      "dtype": "float32",
      "xla_temp_mb": 72.00010681152344,
      "xla_arg_mb": 25.0,
      "xla_out_mb": 9.000015258789062,
      "compile_s": 0.12292628599971067,
      "inner": 1,
      "median_ms": 113.34906599950045,
      "q1_ms": 112.94081399682909,
      "q3_ms": 130.04900098894723,
      "rel_err": 0.0
    },
    {
      "suite": "kernel",
      "impl": "chunked",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 2048,
      "dim": 128,
      "dtype": "float32",
      "xla_temp_mb": 80.0,
      "xla_arg_mb": 25.0,
      "xla_out_mb": 9.000015258789062,
      "compile_s": 0.040044751003733836,
      "inner": 1,
      "median_ms": 39.964538998901844,
      "q1_ms": 39.23810500418767,
      "q3_ms": 40.19600700121373,
      "rel_err": 8.367493819730074e-07
    },
    {
      "suite": "kernel",
      "impl": "scan",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 2048,
      "dim": 128,
      "dtype": "float32",
      "xla_temp_mb": 2256.312759399414,
      "xla_arg_mb": 25.0,
      "xla_out_mb": 25.000030517578125,
      "compile_s": 0.7541006309911609,
      "inner": 1,
      "median_ms": 946.6367429995444,
      "q1_ms": 940.7537619990762,
      "q3_ms": 948.1849309959216,
      "rel_err": 0.0
    },
    {
      "suite": "kernel",
      "impl": "chunked",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 2048,
      "dim": 128,
      "dtype": "float32",
      "xla_temp_mb": 241.3125,
      "xla_arg_mb": 25.0,
      "xla_out_mb": 25.000030517578125,
      "compile_s": 0.13104046700755134,
      "inner": 1,
      "median_ms": 144.89819499431178,
      "q1_ms": 140.1209259929601,
      "q3_ms": 145.84440500766505,
      "rel_err": 0.00022540633689397324
    },
    {
      "suite": "kernel",
      "impl": "scan",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 8192,
      "dim": 128,
      "dtype": "float32",
      "xla_temp_mb": 288.00010681152344,
      "xla_arg_mb": 97.0,
      "xla_out_mb": 33.00001525878906,
      "compile_s": 0.32121674400696065,
      "inner": 1,
      "median_ms": 396.8124029925093,
      "q1_ms": 379.9582770006964,
      "q3_ms": 410.5415469966829,
      "rel_err": 0.0
    },
    {
      "suite": "kernel",
      "impl": "chunked",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 8192,
      "dim": 128,
      "dtype": "float32",
      "xla_temp_mb": 352.0,
      "xla_arg_mb": 97.0,
      "xla_out_mb": 33.00001525878906,
      "compile_s": 0.16959997799131088,
      "inner": 1,
      "median_ms": 202.7488239982631,
      "q1_ms": 197.1163760026684,
      "q3_ms": 204.9132579995785,
      "rel_err": 1.6169468770726746e-06
    },
    {
      "suite": "kernel",
      "impl": "scan",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 8192,
      "dim": 128,
      "dtype": "float32",
      "xla_temp_mb": 9025.250259399414,
      "xla_arg_mb": 97.0,
      "xla_out_mb": 97.00003051757812,
      "compile_s": 2.9111353689950192,
      "inner": 1,
      "median_ms": 3684.503863012651,
      "q1_ms": 3653.7678530003177,
      "q3_ms": 3696.0261829954106,
      "rel_err": 0.0
    },
    {
      "suite": "kernel",
      "impl": "chunked",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 8192,
      "dim": 128,
      "dtype": "float32",
      "xla_temp_mb": 930.25,
      "xla_arg_mb": 97.0,
      "xla_out_mb": 97.00003051757812,
      "compile_s": 0.5085814010089962,
      "inner": 1,
      "median_ms": 592.529873989406,
      "q1_ms": 591.04964400467,
      "q3_ms": 593.9364649966592,
      "rel_err": 0.0007240942716000266
    },
    {
      "suite": "layer",
      "impl": "scan",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 8,
      "seq": 512,
      "dim": 64,
      "window": 128,
      "dtype": "float32",
      "xla_temp_mb": 52.50010681152344,
      "xla_arg_mb": 6.0118408203125,
      "xla_out_mb": 2.0,
      "compile_s": 0.04721152001002338,
      "inner": 1,
      "median_ms": 40.02815000421833,
      "q1_ms": 38.177417998667806,
      "q3_ms": 44.4815860100789
    },
    {
      "suite": "layer",
      "impl": "scan",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 8,
      "seq": 512,
      "dim": 64,
      "window": 128,
      "dtype": "float32",
      "xla_temp_mb": 387.03138732910156,
      "xla_arg_mb": 6.0118408203125,
      "xla_out_mb": 6.011932373046875,
      "compile_s": 0.18237209200742655,
      "inner": 1,
      "median_ms": 195.30205400951672,
      "q1_ms": 194.06701799016446,
      "q3_ms": 213.04249700915534
    },
    {
      "suite": "layer",
      "impl": "chunked",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 8,
      "seq": 512,
      "dim": 64,
      "window": 128,
      "dtype": "float32",
      "xla_temp_mb": 50.0,
      "xla_arg_mb": 6.0118408203125,
      "xla_out_mb": 2.0,
      "compile_s": 0.042413092000060715,
      "inner": 2,
      "median_ms": 20.30542099964805,
      "q1_ms": 19.776214998273645,
      "q3_ms": 20.735409503686242
    },
    {
      "suite": "layer",
      "impl": "chunked",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 8,
      "seq": 512,
      "dim": 64,
      "window": 128,
      "dtype": "float32",
      "xla_temp_mb": 127.03125,
      "xla_arg_mb": 6.0118408203125,
      "xla_out_mb": 6.011932373046875,
      "compile_s": 0.11834686200018041,
      "inner": 1,
      "median_ms": 81.06306400441099,
      "q1_ms": 80.09972200670745,
      "q3_ms": 81.17093800683506
    },
    {
      "suite": "layer",
      "impl": "scan",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 8,
      "seq": 2048,
      "dim": 64,
      "window": 128,
      "dtype": "float32",
      "xla_temp_mb": 592.5001068115234,
      "xla_arg_mb": 12.0118408203125,
      "xla_out_mb": 8.0,
      "compile_s": 0.2879061929997988,
      "inner": 1,
      "median_ms": 333.25963700190187,
      "q1_ms": 313.0805710097775,
      "q3_ms": 337.2544750018278
    },
    {
      "suite": "layer",
      "impl": "scan",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 8,
      "seq": 2048,
      "dim": 64,
      "window": 128,
      "dtype": "float32",
      "xla_temp_mb": 2308.3751373291016,
      "xla_arg_mb": 12.0118408203125,
      "xla_out_mb": 12.011932373046875,
      "compile_s": 1.2073441380052827,
      "inner": 1,
      "median_ms": 1285.8953690010821,
      "q1_ms": 1270.5148700042628,
      "q3_ms": 1303.3982530032517
    },
    {
      "suite": "layer",
      "impl": "chunked",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 8,
      "seq": 2048,
      "dim": 64,
      "window": 128,
      "dtype": "float32",
      "xla_temp_mb": 584.0,
      "xla_arg_mb": 12.0118408203125,
      "xla_out_mb": 8.0,
      "compile_s": 0.2280723330040928,
      "inner": 1,
      "median_ms": 294.5144349941984,
      "q1_ms": 271.85807199566625,
      "q3_ms": 337.27991399064194
    },
    {
      "suite": "layer",
      "impl": "chunked",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 8,
      "seq": 2048,
      "dim": 64,
      "window": 128,
      "dtype": "float32",
      "xla_temp_mb": 1228.375,
      "xla_arg_mb": 12.0118408203125,
      "xla_out_mb": 12.011932373046875,
      "compile_s": 0.7824682109931018,
      "inner": 1,
      "median_ms": 864.5361249946291,
      "q1_ms": 846.7187659989577,
      "q3_ms": 887.6267809973797
    }
  ]
}