{
  "device": {
    "jax": "0.11.2",
    "backend": "gpu",
    "devices": [
      "cuda:0"
    ],
    "python": "3.14.8",
    "device_kind": "NVIDIA GeForce RTX 3090",
    "gpu": "NVIDIA GeForce RTX 3090",
    "driver": "615.78.08",
    "gpu_memory_gb": 24.0,
    "torch": "2.14.1+cu130",
    "torch_cuda": "13.0"
  },
  "config": {
    "suites": [
      "kernel",
      "layer",
      "peak"
    ],
    "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": 2.5027313232421875,
      "xla_arg_mb": 1.5625,
      "xla_out_mb": 0.5625152587890625,
      "compile_s": 0.006062699998437893,
      "inner": 10,
      "median_ms": 4.493603399896529,
      "q1_ms": 4.4813822001742665,
      "q3_ms": 4.5703616000537295,
      "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": 4.75,
      "xla_arg_mb": 1.5625,
      "xla_out_mb": 0.5625152587890625,
      "compile_s": 0.0012474289978854358,
      "inner": 200,
      "median_ms": 0.057597570012148935,
      "q1_ms": 0.057576075014367234,
      "q3_ms": 0.05902165998122655,
      "rel_err": 4.4158698553784665e-07
    },
    {
      "suite": "kernel",
      "impl": "pallas",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 512,
      "dim": 32,
      "dtype": "float32",
      "xla_temp_mb": 2.5,
      "xla_arg_mb": 1.5625,
      "xla_out_mb": 0.5625152587890625,
      "compile_s": 0.0009469110009376891,
      "inner": 200,
      "median_ms": 0.07848397999623558,
      "q1_ms": 0.07846414500818355,
      "q3_ms": 0.07852656501199817,
      "rel_err": 0.00032348705879353597
    },
    {
      "suite": "kernel",
      "impl": "pallas_scanbwd",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 512,
      "dim": 32,
      "dtype": "float32",
      "xla_temp_mb": 2.5,
      "xla_arg_mb": 1.5625,
      "xla_out_mb": 0.5625152587890625,
      "compile_s": 0.0010018850007327273,
      "inner": 200,
      "median_ms": 0.07998562501597917,
      "q1_ms": 0.07841504997486481,
      "q3_ms": 0.08045360998949036,
      "rel_err": 0.00032348705879353597
    },
    {
      "suite": "kernel",
      "impl": "scan",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 512,
      "dim": 32,
      "dtype": "float32",
      "xla_temp_mb": 66.61506652832031,
      "xla_arg_mb": 1.5625,
      "xla_out_mb": 1.562530517578125,
      "compile_s": 0.010277924993715715,
      "inner": 6,
      "median_ms": 7.34265133360168,
      "q1_ms": 7.323894333239878,
      "q3_ms": 7.9417514995535985,
      "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": 8.376754760742188,
      "xla_arg_mb": 1.5625,
      "xla_out_mb": 1.562530517578125,
      "compile_s": 0.001992931000131648,
      "inner": 196,
      "median_ms": 0.1548568469639131,
      "q1_ms": 0.15483906121962532,
      "q3_ms": 0.1548685510099357,
      "rel_err": 2.607905695675851e-05
    },
    {
      "suite": "kernel",
      "impl": "pallas",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 512,
      "dim": 32,
      "dtype": "float32",
      "xla_temp_mb": 9.127243041992188,
      "xla_arg_mb": 1.5625,
      "xla_out_mb": 1.562530517578125,
      "compile_s": 0.001937996996275615,
      "inner": 164,
      "median_ms": 0.2040788963424303,
      "q1_ms": 0.20402342683732191,
      "q3_ms": 0.2041736524250421,
      "rel_err": 0.0002411201375093003
    },
    {
      "suite": "kernel",
      "impl": "pallas_scanbwd",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 512,
      "dim": 32,
      "dtype": "float32",
      "xla_temp_mb": 66.12971496582031,
      "xla_arg_mb": 1.5625,
      "xla_out_mb": 1.562530517578125,
      "compile_s": 0.009727874006784987,
      "inner": 5,
      "median_ms": 7.272059800743591,
      "q1_ms": 7.261862599989399,
      "q3_ms": 7.284741800685879,
      "rel_err": 0.00024150013691744848
    },
    {
      "suite": "kernel",
      "impl": "scan",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 2048,
      "dim": 32,
      "dtype": "float32",
      "xla_temp_mb": 10.002731323242188,
      "xla_arg_mb": 6.0625,
      "xla_out_mb": 2.0625152587890625,
      "compile_s": 0.017663897000602446,
      "inner": 3,
      "median_ms": 16.446805665812764,
      "q1_ms": 16.366844332272496,
      "q3_ms": 16.502500999195036,
      "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": 24.000015258789062,
      "xla_arg_mb": 6.0625,
      "xla_out_mb": 2.0625152587890625,
      "compile_s": 0.0013513170051737688,
      "inner": 183,
      "median_ms": 0.149191923487532,
      "q1_ms": 0.14915398358306312,
      "q3_ms": 0.149212786918277,
      "rel_err": 8.12194378957968e-07
    },
    {
      "suite": "kernel",
      "impl": "pallas",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 2048,
      "dim": 32,
      "dtype": "float32",
      "xla_temp_mb": 10.0,
      "xla_arg_mb": 6.0625,
      "xla_out_mb": 2.0625152587890625,
      "compile_s": 0.0011363190060365014,
      "inner": 128,
      "median_ms": 0.2585575156217601,
      "q1_ms": 0.25844457815082933,
      "q3_ms": 0.2585742656719958,
      "rel_err": 0.0002935579085433888
    },
    {
      "suite": "kernel",
      "impl": "pallas_scanbwd",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 2048,
      "dim": 32,
      "dtype": "float32",
      "xla_temp_mb": 10.0,
      "xla_arg_mb": 6.0625,
      "xla_out_mb": 2.0625152587890625,
      "compile_s": 0.001181644998723641,
      "inner": 125,
      "median_ms": 0.2584396559977904,
      "q1_ms": 0.2583657599752769,
      "q3_ms": 0.25848317600321025,
      "rel_err": 0.0002935579085433888
    },
    {
      "suite": "kernel",
      "impl": "scan",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 2048,
      "dim": 32,
      "dtype": "float32",
      "xla_temp_mb": 266.4429473876953,
      "xla_arg_mb": 6.0625,
      "xla_out_mb": 6.062530517578125,
      "compile_s": 0.029622312002175022,
      "inner": 1,
      "median_ms": 28.243152002687566,
      "q1_ms": 28.2141780044185,
      "q3_ms": 28.293067000049632,
      "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": 33.50199890136719,
      "xla_arg_mb": 6.0625,
      "xla_out_mb": 6.062530517578125,
      "compile_s": 0.002232963997812476,
      "inner": 96,
      "median_ms": 0.41760766665296006,
      "q1_ms": 0.4149892708179929,
      "q3_ms": 0.41808147913495003,
      "rel_err": 8.818221776631865e-05
    },
    {
      "suite": "kernel",
      "impl": "pallas",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 2048,
      "dim": 32,
      "dtype": "float32",
      "xla_temp_mb": 36.50224304199219,
      "xla_arg_mb": 6.0625,
      "xla_out_mb": 6.062530517578125,
      "compile_s": 0.002508976001990959,
      "inner": 65,
      "median_ms": 0.65738816931396,
      "q1_ms": 0.6546135231754026,
      "q3_ms": 0.6579227230619066,
      "rel_err": 0.00023172943492737769
    },
    {
      "suite": "kernel",
      "impl": "pallas_scanbwd",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 2048,
      "dim": 32,
      "dtype": "float32",
      "xla_temp_mb": 264.5047149658203,
      "xla_arg_mb": 6.0625,
      "xla_out_mb": 6.062530517578125,
      "compile_s": 0.03426029600086622,
      "inner": 1,
      "median_ms": 33.49702199921012,
      "q1_ms": 33.47478000068804,
      "q3_ms": 33.62178800307447,
      "rel_err": 0.0002447676018769224
    },
    {
      "suite": "kernel",
      "impl": "scan",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 8192,
      "dim": 32,
      "dtype": "float32",
      "xla_temp_mb": 40.00273132324219,
      "xla_arg_mb": 24.0625,
      "xla_out_mb": 8.062515258789062,
      "compile_s": 0.06693534400255885,
      "inner": 1,
      "median_ms": 65.33812200359534,
      "q1_ms": 65.21687200438464,
      "q3_ms": 65.38399800047046,
      "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": 76.00001525878906,
      "xla_arg_mb": 24.0625,
      "xla_out_mb": 8.062515258789062,
      "compile_s": 0.0018263239981024526,
      "inner": 68,
      "median_ms": 0.6207138823596966,
      "q1_ms": 0.616913602939081,
      "q3_ms": 0.6234739411250784,
      "rel_err": 1.6374810849623156e-06
    },
    {
      "suite": "kernel",
      "impl": "pallas",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 8192,
      "dim": 32,
      "dtype": "float32",
      "xla_temp_mb": 40.0,
      "xla_arg_mb": 24.0625,
      "xla_out_mb": 8.062515258789062,
      "compile_s": 0.0019020879990421236,
      "inner": 42,
      "median_ms": 1.0767816427341175,
      "q1_ms": 1.0766237381696036,
      "q3_ms": 1.0768172143711818,
      "rel_err": 0.00029141859158091354
    },
    {
      "suite": "kernel",
      "impl": "pallas_scanbwd",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 8192,
      "dim": 32,
      "dtype": "float32",
      "xla_temp_mb": 40.0,
      "xla_arg_mb": 24.0625,
      "xla_out_mb": 8.062515258789062,
      "compile_s": 0.0024187859962694347,
      "inner": 41,
      "median_ms": 1.081535414680441,
      "q1_ms": 1.081293731680844,
      "q3_ms": 1.0818489268209173,
      "rel_err": 0.00029141859158091354
    },
    {
      "suite": "kernel",
      "impl": "scan",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 8192,
      "dim": 32,
      "dtype": "float32",
      "xla_temp_mb": 1066.0054473876953,
      "xla_arg_mb": 24.0625,
      "xla_out_mb": 24.062530517578125,
      "compile_s": 0.11361734100501053,
      "inner": 1,
      "median_ms": 112.16459199931705,
      "q1_ms": 112.14537599880714,
      "q3_ms": 112.17189599847188,
      "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": 134.00177001953125,
      "xla_arg_mb": 24.0625,
      "xla_out_mb": 24.062530517578125,
      "compile_s": 0.003998724998382386,
      "inner": 27,
      "median_ms": 1.8394459259946383,
      "q1_ms": 1.677141814811707,
      "q3_ms": 1.927717777747333,
      "rel_err": 0.00033074114413433685
    },
    {
      "suite": "kernel",
      "impl": "pallas",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 8192,
      "dim": 32,
      "dtype": "float32",
      "xla_temp_mb": 138.00225830078125,
      "xla_arg_mb": 24.0625,
      "xla_out_mb": 24.062530517578125,
      "compile_s": 0.004838212997128721,
      "inner": 18,
      "median_ms": 2.61731549997219,
      "q1_ms": 2.616413277994272,
      "q3_ms": 2.619369444599013,
      "rel_err": 0.0003923770203544209
    },
    {
      "suite": "kernel",
      "impl": "pallas_scanbwd",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 8192,
      "dim": 32,
      "dtype": "float32",
      "xla_temp_mb": 1058.0049591064453,
      "xla_arg_mb": 24.0625,
      "xla_out_mb": 24.062530517578125,
      "compile_s": 0.13231369200366316,
      "inner": 1,
      "median_ms": 131.1060459993314,
      "q1_ms": 130.40539200301282,
      "q3_ms": 131.53911599511048,
      "rel_err": 0.00046536266475680725
    },
    {
      "suite": "kernel",
      "impl": "scan",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 512,
      "dim": 64,
      "dtype": "float32",
      "xla_temp_mb": 6.0017547607421875,
      "xla_arg_mb": 3.25,
      "xla_out_mb": 1.2500152587890625,
      "compile_s": 0.006350304000079632,
      "inner": 11,
      "median_ms": 4.112074999614957,
      "q1_ms": 4.1035752728930674,
      "q3_ms": 4.115251908925447,
      "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": 12.000259399414062,
      "xla_arg_mb": 3.25,
      "xla_out_mb": 1.2500152587890625,
      "compile_s": 0.0012573490021168254,
      "inner": 200,
      "median_ms": 0.08795279998594197,
      "q1_ms": 0.08795270001428435,
      "q3_ms": 0.08800780000456143,
      "rel_err": 4.408408178587597e-07
    },
    {
      "suite": "kernel",
      "impl": "pallas",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 512,
      "dim": 64,
      "dtype": "float32",
      "xla_temp_mb": 5.0,
      "xla_arg_mb": 3.25,
      "xla_out_mb": 1.2500152587890625,
      "compile_s": 0.001034887995047029,
      "inner": 150,
      "median_ms": 0.20222264662152156,
      "q1_ms": 0.20221643329326372,
      "q3_ms": 0.20226052001817152,
      "rel_err": 0.00030438823639147963
    },
    {
      "suite": "kernel",
      "impl": "pallas_scanbwd",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 512,
      "dim": 64,
      "dtype": "float32",
      "xla_temp_mb": 5.0,
      "xla_arg_mb": 3.25,
      "xla_out_mb": 1.2500152587890625,
      "compile_s": 0.0011078059978899546,
      "inner": 178,
      "median_ms": 0.17897257865197286,
      "q1_ms": 0.1789384157010328,
      "q3_ms": 0.1790133932588381,
      "rel_err": 0.00030438823639147963
    },
    {
      "suite": "kernel",
      "impl": "scan",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 512,
      "dim": 64,
      "dtype": "float32",
      "xla_temp_mb": 261.1150665283203,
      "xla_arg_mb": 3.25,
      "xla_out_mb": 3.250030517578125,
      "compile_s": 0.010028944001533091,
      "inner": 5,
      "median_ms": 8.552037399203982,
      "q1_ms": 8.55080319888657,
      "q3_ms": 8.552650600904599,
      "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": 21.128463745117188,
      "xla_arg_mb": 3.25,
      "xla_out_mb": 3.250030517578125,
      "compile_s": 0.003037836999283172,
      "inner": 118,
      "median_ms": 0.27921357628903193,
      "q1_ms": 0.27874038981173144,
      "q3_ms": 0.2793651355960272,
      "rel_err": 4.549870161393564e-05
    },
    {
      "suite": "kernel",
      "impl": "pallas",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 512,
      "dim": 64,
      "dtype": "float32",
      "xla_temp_mb": 19.128707885742188,
      "xla_arg_mb": 3.25,
      "xla_out_mb": 3.250030517578125,
      "compile_s": 0.002512523002224043,
      "inner": 86,
      "median_ms": 0.4452289535419008,
      "q1_ms": 0.44385297672327634,
      "q3_ms": 0.44552951167556887,
      "rel_err": 0.00022892755271634958
    },
    {
      "suite": "kernel",
      "impl": "pallas_scanbwd",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 512,
      "dim": 64,
      "dtype": "float32",
      "xla_temp_mb": 260.1297149658203,
      "xla_arg_mb": 3.25,
      "xla_out_mb": 3.250030517578125,
      "compile_s": 0.010886216005019378,
      "inner": 5,
      "median_ms": 8.811788599996362,
      "q1_ms": 8.772081599454395,
      "q3_ms": 8.817557799920905,
      "rel_err": 0.0002289407920796243
    },
    {
      "suite": "kernel",
      "impl": "scan",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 2048,
      "dim": 64,
      "dtype": "float32",
      "xla_temp_mb": 24.001754760742188,
      "xla_arg_mb": 12.25,
      "xla_out_mb": 4.2500152587890625,
      "compile_s": 0.019201948001864366,
      "inner": 1,
      "median_ms": 30.046503998164553,
      "q1_ms": 28.53508499538293,
      "q3_ms": 30.261831998359412,
      "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": 40.00025939941406,
      "xla_arg_mb": 12.25,
      "xla_out_mb": 4.2500152587890625,
      "compile_s": 0.001489927999500651,
      "inner": 108,
      "median_ms": 0.3137010833058782,
      "q1_ms": 0.313464981512947,
      "q3_ms": 0.31381165744851053,
      "rel_err": 8.389118616682799e-07
    },
    {
      "suite": "kernel",
      "impl": "pallas",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 2048,
      "dim": 64,
      "dtype": "float32",
      "xla_temp_mb": 20.0,
      "xla_arg_mb": 12.25,
      "xla_out_mb": 4.2500152587890625,
      "compile_s": 0.0015901480001048185,
      "inner": 57,
      "median_ms": 0.6994165613146482,
      "q1_ms": 0.6992931578823022,
      "q3_ms": 0.6995160351435883,
      "rel_err": 0.00030224716132580895
    },
    {
      "suite": "kernel",
      "impl": "pallas_scanbwd",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 2048,
      "dim": 64,
      "dtype": "float32",
      "xla_temp_mb": 20.0,
      "xla_arg_mb": 12.25,
      "xla_out_mb": 4.2500152587890625,
      "compile_s": 0.0015965999991749413,
      "inner": 59,
      "median_ms": 0.7864516102649682,
      "q1_ms": 0.7858234576060106,
      "q3_ms": 0.7884187117747773,
      "rel_err": 0.00030224716132580895
    },
    {
      "suite": "kernel",
      "impl": "scan",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 2048,
      "dim": 64,
      "dtype": "float32",
      "xla_temp_mb": 1044.4429473876953,
      "xla_arg_mb": 12.25,
      "xla_out_mb": 12.250030517578125,
      "compile_s": 0.03562617999705253,
      "inner": 1,
      "median_ms": 34.384380996925756,
      "q1_ms": 34.315750999667216,
      "q3_ms": 34.99969500262523,
      "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": 68.50346374511719,
      "xla_arg_mb": 12.25,
      "xla_out_mb": 12.250030517578125,
      "compile_s": 0.003045281002414413,
      "inner": 44,
      "median_ms": 0.9702622044807703,
      "q1_ms": 0.9697232271719258,
      "q3_ms": 0.9746010908367515,
      "rel_err": 0.0001380479798488296
    },
    {
      "suite": "kernel",
      "impl": "pallas",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 2048,
      "dim": 64,
      "dtype": "float32",
      "xla_temp_mb": 60.50346374511719,
      "xla_arg_mb": 12.25,
      "xla_out_mb": 12.250030517578125,
      "compile_s": 0.0034527919997344725,
      "inner": 29,
      "median_ms": 1.6239201722783037,
      "q1_ms": 1.6237560688249029,
      "q3_ms": 1.6239678621661433,
      "rel_err": 0.0002529603690558912
    },
    {
      "suite": "kernel",
      "impl": "pallas_scanbwd",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 2048,
      "dim": 64,
      "dtype": "float32",
      "xla_temp_mb": 1040.5047149658203,
      "xla_arg_mb": 12.25,
      "xla_out_mb": 12.250030517578125,
      "compile_s": 0.038054805001593195,
      "inner": 1,
      "median_ms": 36.400315002538264,
      "q1_ms": 36.36560800077859,
      "q3_ms": 36.431382002774626,
      "rel_err": 0.0002735496905407979
    },
    {
      "suite": "kernel",
      "impl": "scan",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 8192,
      "dim": 64,
      "dtype": "float32",
      "xla_temp_mb": 96.00175476074219,
      "xla_arg_mb": 48.25,
      "xla_out_mb": 16.250015258789062,
      "compile_s": 0.06596854500094196,
      "inner": 1,
      "median_ms": 64.54389899590751,
      "q1_ms": 64.54176500119502,
      "q3_ms": 64.86234200565377,
      "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": 160.00001525878906,
      "xla_arg_mb": 48.25,
      "xla_out_mb": 16.250015258789062,
      "compile_s": 0.0025712950009619817,
      "inner": 37,
      "median_ms": 1.22556159456852,
      "q1_ms": 1.2220514594495449,
      "q3_ms": 1.2261505404682327,
      "rel_err": 1.6235340539277496e-06
    },
    {
      "suite": "kernel",
      "impl": "pallas",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 8192,
      "dim": 64,
      "dtype": "float32",
      "xla_temp_mb": 80.0,
      "xla_arg_mb": 48.25,
      "xla_out_mb": 16.250015258789062,
      "compile_s": 0.0036501359936664812,
      "inner": 17,
      "median_ms": 2.7944850585567216,
      "q1_ms": 2.79082047071664,
      "q3_ms": 2.796058647029841,
      "rel_err": 0.00030792716813046863
    },
    {
      "suite": "kernel",
      "impl": "pallas_scanbwd",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 8192,
      "dim": 64,
      "dtype": "float32",
      "xla_temp_mb": 80.0,
      "xla_arg_mb": 48.25,
      "xla_out_mb": 16.250015258789062,
      "compile_s": 0.00367384100536583,
      "inner": 17,
      "median_ms": 2.793413058765616,
      "q1_ms": 2.78839299982578,
      "q3_ms": 2.796588470707374,
      "rel_err": 0.00030792716813046863
    },
    {
      "suite": "kernel",
      "impl": "scan",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 8192,
      "dim": 64,
      "dtype": "float32",
      "xla_temp_mb": 4178.005447387695,
      "xla_arg_mb": 48.25,
      "xla_out_mb": 48.250030517578125,
      "compile_s": 0.13773225100158015,
      "inner": 1,
      "median_ms": 136.32693499675952,
      "q1_ms": 136.31615299527766,
      "q3_ms": 136.32874799804995,
      "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": 274.00250244140625,
      "xla_arg_mb": 48.25,
      "xla_out_mb": 48.250030517578125,
      "compile_s": 0.005927983998844866,
      "inner": 12,
      "median_ms": 3.7948753336725836,
      "q1_ms": 3.7916641667834483,
      "q3_ms": 3.7976839169762875,
      "rel_err": 0.0004934131429459607
    },
    {
      "suite": "kernel",
      "impl": "pallas",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 8192,
      "dim": 64,
      "dtype": "float32",
      "xla_temp_mb": 242.00299072265625,
      "xla_arg_mb": 48.25,
      "xla_out_mb": 48.250030517578125,
      "compile_s": 0.008612473000539467,
      "inner": 7,
      "median_ms": 6.364486572108165,
      "q1_ms": 6.363996999945292,
      "q3_ms": 6.367433571930243,
      "rel_err": 0.000540315949249991
    },
    {
      "suite": "kernel",
      "impl": "pallas_scanbwd",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 8192,
      "dim": 64,
      "dtype": "float32",
      "xla_temp_mb": 4162.004959106445,
      "xla_arg_mb": 48.25,
      "xla_out_mb": 48.250030517578125,
      "compile_s": 0.14656016400113003,
      "inner": 1,
      "median_ms": 144.61072599806357,
      "q1_ms": 144.42539599986048,
      "q3_ms": 144.68645899614785,
      "rel_err": 0.0006429402658098163
    },
    {
      "suite": "kernel",
      "impl": "scan",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 512,
      "dim": 128,
      "dtype": "float32",
      "xla_temp_mb": 10.014450073242188,
      "xla_arg_mb": 7.0,
      "xla_out_mb": 3.0000152587890625,
      "compile_s": 0.006252148996281903,
      "inner": 9,
      "median_ms": 5.226777332912509,
      "q1_ms": 5.224500777937162,
      "q3_ms": 5.227150333262721,
      "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.000259399414062,
      "xla_arg_mb": 7.0,
      "xla_out_mb": 3.0000152587890625,
      "compile_s": 0.0012418089972925372,
      "inner": 178,
      "median_ms": 0.1826234269930374,
      "q1_ms": 0.18258937081064391,
      "q3_ms": 0.1826520168465782,
      "rel_err": 4.487826631008555e-07
    },
    {
      "suite": "kernel",
      "impl": "scan",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 512,
      "dim": 128,
      "dtype": "float32",
      "xla_temp_mb": 1034.1172637939453,
      "xla_arg_mb": 7.0,
      "xla_out_mb": 7.000030517578125,
      "compile_s": 0.01702857700001914,
      "inner": 3,
      "median_ms": 15.48819200009651,
      "q1_ms": 15.475187666500764,
      "q3_ms": 15.519423999649007,
      "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": 38.12895202636719,
      "xla_arg_mb": 7.0,
      "xla_out_mb": 7.000030517578125,
      "compile_s": 0.0028764919989043847,
      "inner": 74,
      "median_ms": 0.5517665810436262,
      "q1_ms": 0.5407989325200057,
      "q3_ms": 0.5521634189495377,
      "rel_err": 8.111006037745094e-05
    },
    {
      "suite": "kernel",
      "impl": "scan",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 2048,
      "dim": 128,
      "dtype": "float32",
      "xla_temp_mb": 40.01445007324219,
      "xla_arg_mb": 25.0,
      "xla_out_mb": 9.000015258789062,
      "compile_s": 0.022555171999556478,
      "inner": 2,
      "median_ms": 21.538563498324947,
      "q1_ms": 21.4316659985343,
      "q3_ms": 21.667576998879667,
      "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": 68.00025939941406,
      "xla_arg_mb": 25.0,
      "xla_out_mb": 9.000015258789062,
      "compile_s": 0.0018061470036627725,
      "inner": 63,
      "median_ms": 0.6666006507686457,
      "q1_ms": 0.6645644286421999,
      "q3_ms": 0.6669629365205765,
      "rel_err": 8.388496478637273e-07
    },
    {
      "suite": "kernel",
      "impl": "scan",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 2048,
      "dim": 128,
      "dtype": "float32",
      "xla_temp_mb": 4136.50764465332,
      "xla_arg_mb": 25.0,
      "xla_out_mb": 25.000030517578125,
      "compile_s": 0.06330167900159722,
      "inner": 1,
      "median_ms": 62.10618699697079,
      "q1_ms": 62.01965400396148,
      "q3_ms": 62.31022300198674,
      "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": 128.5037078857422,
      "xla_arg_mb": 25.0,
      "xla_out_mb": 25.000030517578125,
      "compile_s": 0.004114575000130571,
      "inner": 22,
      "median_ms": 2.1247534089277247,
      "q1_ms": 2.1208259544933257,
      "q3_ms": 2.1267376363208648,
      "rel_err": 0.00022429702404776398
    },
    {
      "suite": "kernel",
      "impl": "scan",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 8192,
      "dim": 128,
      "dtype": "float32",
      "xla_temp_mb": 192.0066375732422,
      "xla_arg_mb": 97.0,
      "xla_out_mb": 33.00001525878906,
      "compile_s": 0.08460381099575898,
      "inner": 1,
      "median_ms": 83.53411599819083,
      "q1_ms": 83.49842900497606,
      "q3_ms": 83.54857400263427,
      "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": 260.00025939941406,
      "xla_arg_mb": 97.0,
      "xla_out_mb": 33.00001525878906,
      "compile_s": 0.0038918130012461916,
      "inner": 17,
      "median_ms": 2.694800000052269,
      "q1_ms": 2.694571882517159,
      "q3_ms": 2.69763947075904,
      "rel_err": 1.6181954142038886e-06
    },
    {
      "suite": "kernel",
      "impl": "scan",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 8192,
      "dim": 128,
      "dtype": "float32",
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 8.02GiB with allocator GPU_0_bfc on device 0. [tf-allocator-allocation-error='']"
    },
    {
      "suite": "kernel",
      "impl": "chunked",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 8192,
      "dim": 128,
      "dtype": "float32",
      "xla_temp_mb": 483.00299072265625,
      "xla_arg_mb": 97.0,
      "xla_out_mb": 97.00003051757812,
      "compile_s": 0.010489664004126098,
      "inner": 5,
      "median_ms": 8.3879365993198,
      "q1_ms": 8.38142680004239,
      "q3_ms": 8.423834200948477
    },
    {
      "suite": "kernel",
      "impl": "scan",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 512,
      "dim": 32,
      "dtype": "bfloat16",
      "xla_temp_mb": 1.5012664794921875,
      "xla_arg_mb": 0.8125,
      "xla_out_mb": 0.3125152587890625,
      "compile_s": 0.005328119994373992,
      "inner": 11,
      "median_ms": 4.204052636685612,
      "q1_ms": 4.168259363557974,
      "q3_ms": 4.22044445440406,
      "rel_err": 0.0017985250276885763
    },
    {
      "suite": "kernel",
      "impl": "chunked",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 512,
      "dim": 32,
      "dtype": "bfloat16",
      "xla_temp_mb": 4.0,
      "xla_arg_mb": 0.8125,
      "xla_out_mb": 0.3125152587890625,
      "compile_s": 0.0009658670023782179,
      "inner": 200,
      "median_ms": 0.048663915003999136,
      "q1_ms": 0.04816847500478616,
      "q3_ms": 0.04874336500506615,
      "rel_err": 0.001959559217579589
    },
    {
      "suite": "kernel",
      "impl": "pallas",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 512,
      "dim": 32,
      "dtype": "bfloat16",
      "xla_temp_mb": 1.25,
      "xla_arg_mb": 0.8125,
      "xla_out_mb": 0.3125152587890625,
      "compile_s": 0.001007236001896672,
      "inner": 200,
      "median_ms": 0.06318959498457843,
      "q1_ms": 0.05529310499696294,
      "q3_ms": 0.06345995501760626,
      "rel_err": 0.001798524487114004
    },
    {
      "suite": "kernel",
      "impl": "scan",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 512,
      "dim": 32,
      "dtype": "bfloat16",
      "xla_temp_mb": 65.36433410644531,
      "xla_arg_mb": 0.8125,
      "xla_out_mb": 0.812530517578125,
      "compile_s": 0.009749714998179115,
      "inner": 6,
      "median_ms": 7.9749243341211695,
      "q1_ms": 7.964864834017741,
      "q3_ms": 7.996670000769276,
      "rel_err": 0.10056415598715934
    },
    {
      "suite": "kernel",
      "impl": "chunked",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 512,
      "dim": 32,
      "dtype": "bfloat16",
      "xla_temp_mb": 7.876487731933594,
      "xla_arg_mb": 0.8125,
      "xla_out_mb": 0.812530517578125,
      "compile_s": 0.0019288590046926402,
      "inner": 200,
      "median_ms": 0.11055479997594375,
      "q1_ms": 0.11054574002628215,
      "q3_ms": 0.11058020001655677,
      "rel_err": 0.13335402167322177
    },
    {
      "suite": "kernel",
      "impl": "pallas",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 512,
      "dim": 32,
      "dtype": "bfloat16",
      "xla_temp_mb": 8.361358642578125,
      "xla_arg_mb": 0.8125,
      "xla_out_mb": 0.812530517578125,
      "compile_s": 0.0017910489987116307,
      "inner": 181,
      "median_ms": 0.1321977237721864,
      "q1_ms": 0.13215299998828112,
      "q3_ms": 0.13221814915395047,
      "rel_err": 0.13324929462152427
    },
    {
      "suite": "kernel",
      "impl": "scan",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 2048,
      "dim": 32,
      "dtype": "bfloat16",
      "xla_temp_mb": 6.0012664794921875,
      "xla_arg_mb": 3.0625,
      "xla_out_mb": 1.0625152587890625,
      "compile_s": 0.017851543001597747,
      "inner": 3,
      "median_ms": 16.278389666695148,
      "q1_ms": 16.26857133427014,
      "q3_ms": 16.28227033264314,
      "rel_err": 0.0016874089472429937
    },
    {
      "suite": "kernel",
      "impl": "chunked",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 2048,
      "dim": 32,
      "dtype": "bfloat16",
      "xla_temp_mb": 16.0,
      "xla_arg_mb": 3.0625,
      "xla_out_mb": 1.0625152587890625,
      "compile_s": 0.0010586429998511449,
      "inner": 200,
      "median_ms": 0.08804226501524681,
      "q1_ms": 0.08756886498304084,
      "q3_ms": 0.08807442500256002,
      "rel_err": 0.0017367527579602996
    },
    {
      "suite": "kernel",
      "impl": "pallas",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 2048,
      "dim": 32,
      "dtype": "bfloat16",
      "xla_temp_mb": 5.0,
      "xla_arg_mb": 3.0625,
      "xla_out_mb": 1.0625152587890625,
      "compile_s": 0.0010454970033606514,
      "inner": 158,
      "median_ms": 0.21510555059280648,
      "q1_ms": 0.21509813924712023,
      "q3_ms": 0.21521557591998466,
      "rel_err": 0.0016874214349422059
    },
    {
      "suite": "kernel",
      "impl": "scan",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 2048,
      "dim": 32,
      "dtype": "bfloat16",
      "xla_temp_mb": 261.4424591064453,
      "xla_arg_mb": 3.0625,
      "xla_out_mb": 3.062530517578125,
      "compile_s": 0.029598205001093447,
      "inner": 1,
      "median_ms": 28.22391600057017,
      "q1_ms": 28.2231229939498,
      "q3_ms": 28.225989000929985,
      "rel_err": 0.21071913998821246
    },
    {
      "suite": "kernel",
      "impl": "chunked",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 2048,
      "dim": 32,
      "dtype": "bfloat16",
      "xla_temp_mb": 31.501731872558594,
      "xla_arg_mb": 3.0625,
      "xla_out_mb": 3.062530517578125,
      "compile_s": 0.001971238998521585,
      "inner": 129,
      "median_ms": 0.2875425969024889,
      "q1_ms": 0.287488620126324,
      "q3_ms": 0.28761637985052957,
      "rel_err": 0.26505476498262587
    },
    {
      "suite": "kernel",
      "impl": "pallas",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 2048,
      "dim": 32,
      "dtype": "bfloat16",
      "xla_temp_mb": 33.439483642578125,
      "xla_arg_mb": 3.0625,
      "xla_out_mb": 3.062530517578125,
      "compile_s": 0.0020079290043213405,
      "inner": 95,
      "median_ms": 0.4162971579046969,
      "q1_ms": 0.4161582525978845,
      "q3_ms": 0.4163082316058296,
      "rel_err": 0.265301867451842
    },
    {
      "suite": "kernel",
      "impl": "scan",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 8192,
      "dim": 32,
      "dtype": "bfloat16",
      "xla_temp_mb": 24.001266479492188,
      "xla_arg_mb": 12.0625,
      "xla_out_mb": 4.0625152587890625,
      "compile_s": 0.0660295790003147,
      "inner": 1,
      "median_ms": 64.85354300093604,
      "q1_ms": 64.7835909985588,
      "q3_ms": 64.96844100183807,
      "rel_err": 0.001696689177711925
    },
    {
      "suite": "kernel",
      "impl": "chunked",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 8192,
      "dim": 32,
      "dtype": "bfloat16",
      "xla_temp_mb": 64.00001525878906,
      "xla_arg_mb": 12.0625,
      "xla_out_mb": 4.0625152587890625,
      "compile_s": 0.0014002090028952807,
      "inner": 105,
      "median_ms": 0.38558451432479723,
      "q1_ms": 0.38555494289279785,
      "q3_ms": 0.3856007428285444,
      "rel_err": 0.0017110637096450123
    },
    {
      "suite": "kernel",
      "impl": "pallas",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 8192,
      "dim": 32,
      "dtype": "bfloat16",
      "xla_temp_mb": 20.0,
      "xla_arg_mb": 12.0625,
      "xla_out_mb": 4.0625152587890625,
      "compile_s": 0.001653197999985423,
      "inner": 55,
      "median_ms": 0.762809308997186,
      "q1_ms": 0.7626238544947807,
      "q3_ms": 0.7628726909222843,
      "rel_err": 0.001696942800263777
    },
    {
      "suite": "kernel",
      "impl": "scan",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 8192,
      "dim": 32,
      "dtype": "bfloat16",
      "xla_temp_mb": 1045.7549591064453,
      "xla_arg_mb": 12.0625,
      "xla_out_mb": 12.062530517578125,
      "compile_s": 0.11452713000471704,
      "inner": 1,
      "median_ms": 113.28361699997913,
      "q1_ms": 113.25247700006003,
      "q3_ms": 113.37953900510911,
      "rel_err": 0.4235486742275538
    },
    {
      "suite": "kernel",
      "impl": "chunked",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 8192,
      "dim": 32,
      "dtype": "bfloat16",
      "xla_temp_mb": 110.0017318725586,
      "xla_arg_mb": 12.0625,
      "xla_out_mb": 12.062530517578125,
      "compile_s": 0.002954290001071058,
      "inner": 40,
      "median_ms": 1.1381083999367547,
      "q1_ms": 1.1377925500710262,
      "q3_ms": 1.1385640000298736,
      "rel_err": 0.47431376928206087
    },
    {
      "suite": "kernel",
      "impl": "pallas",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 8192,
      "dim": 32,
      "dtype": "bfloat16",
      "xla_temp_mb": 134.00173950195312,
      "xla_arg_mb": 12.0625,
      "xla_out_mb": 12.062530517578125,
      "compile_s": 0.003386607000720687,
      "inner": 28,
      "median_ms": 1.6763506786706526,
      "q1_ms": 1.6751877857521842,
      "q3_ms": 1.6763689285913383,
      "rel_err": 0.4743156520204893
    },
    {
      "suite": "kernel",
      "impl": "scan",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 512,
      "dim": 64,
      "dtype": "bfloat16",
      "xla_temp_mb": 2.5022430419921875,
      "xla_arg_mb": 1.75,
      "xla_out_mb": 0.7500152587890625,
      "compile_s": 0.004965104002621956,
      "inner": 11,
      "median_ms": 4.1203969998274586,
      "q1_ms": 4.112559454700782,
      "q3_ms": 4.133141090939965,
      "rel_err": 0.0017020240532270275
    },
    {
      "suite": "kernel",
      "impl": "chunked",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 512,
      "dim": 64,
      "dtype": "bfloat16",
      "xla_temp_mb": 9.000015258789062,
      "xla_arg_mb": 1.75,
      "xla_out_mb": 0.7500152587890625,
      "compile_s": 0.0010535929977777414,
      "inner": 200,
      "median_ms": 0.05289597000228241,
      "q1_ms": 0.04762284501339309,
      "q3_ms": 0.060292160014796536,
      "rel_err": 0.001841924556376581
    },
    {
      "suite": "kernel",
      "impl": "pallas",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 512,
      "dim": 64,
      "dtype": "bfloat16",
      "xla_temp_mb": 2.5,
      "xla_arg_mb": 1.75,
      "xla_out_mb": 0.7500152587890625,
      "compile_s": 0.001094219995138701,
      "inner": 158,
      "median_ms": 0.17356361394941136,
      "q1_ms": 0.17354509493341977,
      "q3_ms": 0.1735922721504894,
      "rel_err": 0.001702024990601871
    },
    {
      "suite": "kernel",
      "impl": "scan",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 512,
      "dim": 64,
      "dtype": "bfloat16",
      "xla_temp_mb": 258.6143341064453,
      "xla_arg_mb": 1.75,
      "xla_out_mb": 1.750030517578125,
      "compile_s": 0.01004247899982147,
      "inner": 5,
      "median_ms": 8.572982798796147,
      "q1_ms": 8.572406000166666,
      "q3_ms": 8.57355399930384,
      "rel_err": 0.14944267018330054
    },
    {
      "suite": "kernel",
      "impl": "chunked",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 512,
      "dim": 64,
      "dtype": "bfloat16",
      "xla_temp_mb": 17.127708435058594,
      "xla_arg_mb": 1.75,
      "xla_out_mb": 1.750030517578125,
      "compile_s": 0.0020521319966064766,
      "inner": 179,
      "median_ms": 0.19280912286950555,
      "q1_ms": 0.18871568155399518,
      "q3_ms": 0.19318486592152245,
      "rel_err": 0.20059601161659485
    },
    {
      "suite": "kernel",
      "impl": "pallas",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 512,
      "dim": 64,
      "dtype": "bfloat16",
      "xla_temp_mb": 20.393341064453125,
      "xla_arg_mb": 1.75,
      "xla_out_mb": 1.750030517578125,
      "compile_s": 0.002089694004098419,
      "inner": 126,
      "median_ms": 0.30835280157921335,
      "q1_ms": 0.308263349166188,
      "q3_ms": 0.3087533968648622,
      "rel_err": 0.20039747623616772
    },
    {
      "suite": "kernel",
      "impl": "scan",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 2048,
      "dim": 64,
      "dtype": "bfloat16",
      "xla_temp_mb": 10.002243041992188,
      "xla_arg_mb": 6.25,
      "xla_out_mb": 2.2500152587890625,
      "compile_s": 0.017255365004530177,
      "inner": 3,
      "median_ms": 16.336446334510885,
      "q1_ms": 16.242237999297988,
      "q3_ms": 16.345716998330317,
      "rel_err": 0.0017766021498748983
    },
    {
      "suite": "kernel",
      "impl": "chunked",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 2048,
      "dim": 64,
      "dtype": "bfloat16",
      "xla_temp_mb": 24.000015258789062,
      "xla_arg_mb": 6.25,
      "xla_out_mb": 2.2500152587890625,
      "compile_s": 0.001199479003844317,
      "inner": 187,
      "median_ms": 0.18966582355980965,
      "q1_ms": 0.1890297593840866,
      "q3_ms": 0.18973413371278053,
      "rel_err": 0.0017766020321543373
    },
    {
      "suite": "kernel",
      "impl": "pallas",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 2048,
      "dim": 64,
      "dtype": "bfloat16",
      "xla_temp_mb": 10.0,
      "xla_arg_mb": 6.25,
      "xla_out_mb": 2.2500152587890625,
      "compile_s": 0.001403565998771228,
      "inner": 66,
      "median_ms": 0.5941211212220052,
      "q1_ms": 0.5940918181846482,
      "q3_ms": 0.5941464545630327,
      "rel_err": 0.0017766061385388794
    },
    {
      "suite": "kernel",
      "impl": "scan",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 2048,
      "dim": 64,
      "dtype": "bfloat16",
      "xla_temp_mb": 1034.4424591064453,
      "xla_arg_mb": 6.25,
      "xla_out_mb": 6.250030517578125,
      "compile_s": 0.035505831001501065,
      "inner": 1,
      "median_ms": 34.19203399971593,
      "q1_ms": 34.13403699960327,
      "q3_ms": 34.19346099690301,
      "rel_err": 0.3005157388735735
    },
    {
      "suite": "kernel",
      "impl": "chunked",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 2048,
      "dim": 64,
      "dtype": "bfloat16",
      "xla_temp_mb": 52.502708435058594,
      "xla_arg_mb": 6.25,
      "xla_out_mb": 6.250030517578125,
      "compile_s": 0.0026871040026890114,
      "inner": 64,
      "median_ms": 0.688942453166419,
      "q1_ms": 0.688432562469643,
      "q3_ms": 0.6895125937944613,
      "rel_err": 0.3734220486236524
    },
    {
      "suite": "kernel",
      "impl": "pallas",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 2048,
      "dim": 64,
      "dtype": "bfloat16",
      "xla_temp_mb": 70.50271606445312,
      "xla_arg_mb": 6.25,
      "xla_out_mb": 6.250030517578125,
      "compile_s": 0.002976401003252249,
      "inner": 40,
      "median_ms": 1.1523483250130084,
      "q1_ms": 1.1500825499751954,
      "q3_ms": 1.154424250125885,
      "rel_err": 0.37350300740683934
    },
    {
      "suite": "kernel",
      "impl": "scan",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 8192,
      "dim": 64,
      "dtype": "bfloat16",
      "xla_temp_mb": 40.00224304199219,
      "xla_arg_mb": 24.25,
      "xla_out_mb": 8.250015258789062,
      "compile_s": 0.06831130500358995,
      "inner": 1,
      "median_ms": 66.49014000140596,
      "q1_ms": 65.6975500023691,
      "q3_ms": 68.06879699433921,
      "rel_err": 0.0017459182702584792
    },
    {
      "suite": "kernel",
      "impl": "chunked",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 8192,
      "dim": 64,
      "dtype": "bfloat16",
      "xla_temp_mb": 84.00001525878906,
      "xla_arg_mb": 24.25,
      "xla_out_mb": 8.250015258789062,
      "compile_s": 0.0018105650015058927,
      "inner": 59,
      "median_ms": 0.7654620339192245,
      "q1_ms": 0.7653139660853135,
      "q3_ms": 0.7657064066980807,
      "rel_err": 0.0017459152939059961
    },
    {
      "suite": "kernel",
      "impl": "pallas",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 8192,
      "dim": 64,
      "dtype": "bfloat16",
      "xla_temp_mb": 40.0,
      "xla_arg_mb": 24.25,
      "xla_out_mb": 8.250015258789062,
      "compile_s": 0.0031558810005662963,
      "inner": 19,
      "median_ms": 2.374532842383671,
      "q1_ms": 2.374121525989674,
      "q3_ms": 2.3746266842384727,
      "rel_err": 0.0017460077277075747
    },
    {
      "suite": "kernel",
      "impl": "scan",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 8192,
      "dim": 64,
      "dtype": "bfloat16",
      "xla_temp_mb": 4138.004959106445,
      "xla_arg_mb": 24.25,
      "xla_out_mb": 24.250030517578125,
      "compile_s": 0.13764462299877778,
      "inner": 1,
      "median_ms": 136.3218219994451,
      "q1_ms": 136.3124740018975,
      "q3_ms": 136.3355169960414,
      "rel_err": 0.6109625968079669
    },
    {
      "suite": "kernel",
      "impl": "chunked",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 8192,
      "dim": 64,
      "dtype": "bfloat16",
      "xla_temp_mb": 194.00271606445312,
      "xla_arg_mb": 24.25,
      "xla_out_mb": 24.250030517578125,
      "compile_s": 0.005014948001189623,
      "inner": 17,
      "median_ms": 2.762382470417138,
      "q1_ms": 2.7622575880206353,
      "q3_ms": 2.7673094707150834,
      "rel_err": 0.670858991073284
    },
    {
      "suite": "kernel",
      "impl": "pallas",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 8192,
      "dim": 64,
      "dtype": "bfloat16",
      "xla_temp_mb": 266.0027313232422,
      "xla_arg_mb": 24.25,
      "xla_out_mb": 24.250030517578125,
      "compile_s": 0.006733699003234506,
      "inner": 10,
      "median_ms": 4.739959300059127,
      "q1_ms": 4.739211000560317,
      "q3_ms": 4.747610699996585,
      "rel_err": 0.6708721665439686
    },
    {
      "suite": "kernel",
      "impl": "scan",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 512,
      "dim": 128,
      "dtype": "bfloat16",
      "xla_temp_mb": 5.0105438232421875,
      "xla_arg_mb": 4.0,
      "xla_out_mb": 2.0000152587890625,
      "compile_s": 0.005296260002069175,
      "inner": 11,
      "median_ms": 4.383672454754229,
      "q1_ms": 4.383364454869562,
      "q3_ms": 4.383776272747623,
      "rel_err": 0.0016502312442814198
    },
    {
      "suite": "kernel",
      "impl": "chunked",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 512,
      "dim": 128,
      "dtype": "bfloat16",
      "xla_temp_mb": 18.000015258789062,
      "xla_arg_mb": 4.0,
      "xla_out_mb": 2.0000152587890625,
      "compile_s": 0.0011252190015511587,
      "inner": 200,
      "median_ms": 0.10842983501788694,
      "q1_ms": 0.10722029499447672,
      "q3_ms": 0.10857155000849161,
      "rel_err": 0.00168554945180945
    },
    {
      "suite": "kernel",
      "impl": "scan",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 512,
      "dim": 128,
      "dtype": "bfloat16",
      "xla_temp_mb": 1029.1302032470703,
      "xla_arg_mb": 4.0,
      "xla_out_mb": 4.000030517578125,
      "compile_s": 0.015300506005587522,
      "inner": 3,
      "median_ms": 13.800979332396915,
      "q1_ms": 13.800802332601355,
      "q3_ms": 13.807428333772501,
      "rel_err": 0.20815048001416347
    },
    {
      "suite": "kernel",
      "impl": "chunked",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 512,
      "dim": 128,
      "dtype": "bfloat16",
      "xla_temp_mb": 32.127708435058594,
      "xla_arg_mb": 4.0,
      "xla_out_mb": 4.000030517578125,
      "compile_s": 0.002360214995860588,
      "inner": 104,
      "median_ms": 0.3968342115023496,
      "q1_ms": 0.39662670196076555,
      "q3_ms": 0.39698613461671184,
      "rel_err": 0.2718040412082461
    },
    {
      "suite": "kernel",
      "impl": "scan",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 2048,
      "dim": 128,
      "dtype": "bfloat16",
      "xla_temp_mb": 20.010543823242188,
      "xla_arg_mb": 13.0,
      "xla_out_mb": 5.0000152587890625,
      "compile_s": 0.018575032001535874,
      "inner": 2,
      "median_ms": 17.793137001717696,
      "q1_ms": 17.7877720016113,
      "q3_ms": 17.89227449989994,
      "rel_err": 0.0016681342538245186
    },
    {
      "suite": "kernel",
      "impl": "chunked",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 2048,
      "dim": 128,
      "dtype": "bfloat16",
      "xla_temp_mb": 56.00001525878906,
      "xla_arg_mb": 13.0,
      "xla_out_mb": 5.0000152587890625,
      "compile_s": 0.00139837599999737,
      "inner": 100,
      "median_ms": 0.42090128001291305,
      "q1_ms": 0.4208306500368053,
      "q3_ms": 0.421030029974645,
      "rel_err": 0.0016721306416086004
    },
    {
      "suite": "kernel",
      "impl": "scan",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 2048,
      "dim": 128,
      "dtype": "bfloat16",
      "xla_temp_mb": 4116.44270324707,
      "xla_arg_mb": 13.0,
      "xla_out_mb": 13.000030517578125,
      "compile_s": 0.056705687995417975,
      "inner": 1,
      "median_ms": 55.41358699701959,
      "q1_ms": 55.36488199868472,
      "q3_ms": 55.472337000537664,
      "rel_err": 0.4248054230742187
    },
    {
      "suite": "kernel",
      "impl": "chunked",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 2048,
      "dim": 128,
      "dtype": "bfloat16",
      "xla_temp_mb": 104.5031967163086,
      "xla_arg_mb": 13.0,
      "xla_out_mb": 13.000030517578125,
      "compile_s": 0.0035831500063068233,
      "inner": 31,
      "median_ms": 1.4775909031384564,
      "q1_ms": 1.4774574193589,
      "q3_ms": 1.4776836129209596,
      "rel_err": 0.5077394987550908
    },
    {
      "suite": "kernel",
      "impl": "scan",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 8192,
      "dim": 128,
      "dtype": "bfloat16",
      "xla_temp_mb": 80.01054382324219,
      "xla_arg_mb": 49.0,
      "xla_out_mb": 17.000015258789062,
      "compile_s": 0.07172476399864536,
      "inner": 1,
      "median_ms": 70.99840699811466,
      "q1_ms": 70.86510899534915,
      "q3_ms": 71.00059399817837,
      "rel_err": 0.0017216861287381094
    },
    {
      "suite": "kernel",
      "impl": "chunked",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 8192,
      "dim": 128,
      "dtype": "bfloat16",
      "xla_temp_mb": 224.00001525878906,
      "xla_arg_mb": 49.0,
      "xla_out_mb": 17.000015258789062,
      "compile_s": 0.0027513249951880425,
      "inner": 27,
      "median_ms": 1.734489555356386,
      "q1_ms": 1.7333670370539442,
      "q3_ms": 1.7346008888260391,
      "rel_err": 0.0017216847058450556
    },
    {
      "suite": "kernel",
      "impl": "scan",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 8192,
      "dim": 128,
      "dtype": "bfloat16",
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 8.02GiB with allocator GPU_0_bfc on device 0. [tf-allocator-allocation-error='']"
    },
    {
      "suite": "kernel",
      "impl": "chunked",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 4,
      "seq": 8192,
      "dim": 128,
      "dtype": "bfloat16",
      "xla_temp_mb": 419.0029602050781,
      "xla_arg_mb": 49.0,
      "xla_out_mb": 49.000030517578125,
      "compile_s": 0.008263994997832924,
      "inner": 8,
      "median_ms": 5.898234750020492,
      "q1_ms": 5.89710525036935,
      "q3_ms": 5.901753999751236
    },
    {
      "suite": "layer",
      "impl": "scan",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 8,
      "seq": 512,
      "dim": 64,
      "window": 128,
      "dtype": "float32",
      "xla_temp_mb": 34.00102233886719,
      "xla_arg_mb": 6.0118408203125,
      "xla_out_mb": 2.0,
      "compile_s": 0.005919344999711029,
      "inner": 11,
      "median_ms": 4.160324272974817,
      "q1_ms": 4.158788636480246,
      "q3_ms": 4.196898000373039
    },
    {
      "suite": "layer",
      "impl": "scan",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 8,
      "seq": 512,
      "dim": 64,
      "window": 128,
      "dtype": "float32",
      "xla_temp_mb": 532.3528594970703,
      "xla_arg_mb": 6.0118408203125,
      "xla_out_mb": 6.011932373046875,
      "compile_s": 0.016913235995161813,
      "inner": 4,
      "median_ms": 12.192149750262615,
      "q1_ms": 12.177038750451175,
      "q3_ms": 12.192600999696879
    },
    {
      "suite": "layer",
      "impl": "chunked",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 8,
      "seq": 512,
      "dim": 64,
      "window": 128,
      "dtype": "float32",
      "xla_temp_mb": 36.00050354003906,
      "xla_arg_mb": 6.0118408203125,
      "xla_out_mb": 2.0,
      "compile_s": 0.0021695329996873625,
      "inner": 101,
      "median_ms": 0.37180979207093545,
      "q1_ms": 0.3716324158946495,
      "q3_ms": 0.3718231782480835
    },
    {
      "suite": "layer",
      "impl": "chunked",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 8,
      "seq": 512,
      "dim": 64,
      "window": 128,
      "dtype": "float32",
      "xla_temp_mb": 74.39604187011719,
      "xla_arg_mb": 6.0118408203125,
      "xla_out_mb": 6.011932373046875,
      "compile_s": 0.005806922999909148,
      "inner": 36,
      "median_ms": 1.2111852776999716,
      "q1_ms": 1.2109025277216763,
      "q3_ms": 1.2114596665924182
    },
    {
      "suite": "layer",
      "impl": "pallas",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 8,
      "seq": 512,
      "dim": 64,
      "window": 128,
      "dtype": "float32",
      "xla_temp_mb": 34.00050354003906,
      "xla_arg_mb": 6.0118408203125,
      "xla_out_mb": 2.0,
      "compile_s": 0.0019563999958336353,
      "inner": 84,
      "median_ms": 0.4366136904753235,
      "q1_ms": 0.43661320237983353,
      "q3_ms": 0.43685103576405837
    },
    {
      "suite": "layer",
      "impl": "pallas",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 8,
      "seq": 512,
      "dim": 64,
      "window": 128,
      "dtype": "float32",
      "xla_temp_mb": 80.33378601074219,
      "xla_arg_mb": 6.0118408203125,
      "xla_out_mb": 6.011932373046875,
      "compile_s": 0.005587617000855971,
      "inner": 33,
      "median_ms": 1.3271471212977883,
      "q1_ms": 1.3245440605907897,
      "q3_ms": 1.3286933939194898
    },
    {
      "suite": "layer",
      "impl": "scan",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 8,
      "seq": 2048,
      "dim": 64,
      "window": 128,
      "dtype": "float32",
      "xla_temp_mb": 520.1260223388672,
      "xla_arg_mb": 12.0118408203125,
      "xla_out_mb": 8.0,
      "compile_s": 0.020315008994657546,
      "inner": 2,
      "median_ms": 17.858037499536294,
      "q1_ms": 17.806905998440925,
      "q3_ms": 17.8617995006789
    },
    {
      "suite": "layer",
      "impl": "scan",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 8,
      "seq": 2048,
      "dim": 64,
      "window": 128,
      "dtype": "float32",
      "xla_temp_mb": 2465.7591094970703,
      "xla_arg_mb": 12.0118408203125,
      "xla_out_mb": 12.011932373046875,
      "compile_s": 0.049555789999431,
      "inner": 1,
      "median_ms": 46.378330000152346,
      "q1_ms": 46.36450399993919,
      "q3_ms": 46.39578099886421
    },
    {
      "suite": "layer",
      "impl": "chunked",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 8,
      "seq": 2048,
      "dim": 64,
      "window": 128,
      "dtype": "float32",
      "xla_temp_mb": 520.1255035400391,
      "xla_arg_mb": 12.0118408203125,
      "xla_out_mb": 8.0,
      "compile_s": 0.004333575998316519,
      "inner": 18,
      "median_ms": 2.4472337778812894,
      "q1_ms": 2.446126610771494,
      "q3_ms": 2.4478777223622377
    },
    {
      "suite": "layer",
      "impl": "chunked",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 8,
      "seq": 2048,
      "dim": 64,
      "window": 128,
      "dtype": "float32",
      "xla_temp_mb": 862.5054168701172,
      "xla_arg_mb": 12.0118408203125,
      "xla_out_mb": 12.011932373046875,
      "compile_s": 0.012465503998100758,
      "inner": 6,
      "median_ms": 7.59994583252895,
      "q1_ms": 7.585305167594925,
      "q3_ms": 7.603250833199127
    },
    {
      "suite": "layer",
      "impl": "pallas",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 8,
      "seq": 2048,
      "dim": 64,
      "window": 128,
      "dtype": "float32",
      "xla_temp_mb": 520.0005035400391,
      "xla_arg_mb": 12.0118408203125,
      "xla_out_mb": 8.0,
      "compile_s": 0.003977612002927344,
      "inner": 17,
      "median_ms": 2.5722793528556767,
      "q1_ms": 2.5461204118876006,
      "q3_ms": 2.633005941307227
    },
    {
      "suite": "layer",
      "impl": "pallas",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 8,
      "seq": 2048,
      "dim": 64,
      "window": 128,
      "dtype": "float32",
      "xla_temp_mb": 874.6304168701172,
      "xla_arg_mb": 12.0118408203125,
      "xla_out_mb": 12.011932373046875,
      "compile_s": 0.012289172998862341,
      "inner": 5,
      "median_ms": 8.022107199940365,
      "q1_ms": 8.013659001153428,
      "q3_ms": 8.039816599921323
    },
    {
      "suite": "layer",
      "impl": "scan",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 8,
      "seq": 8192,
      "dim": 64,
      "window": 128,
      "dtype": "float32",
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 8.03GiB with allocator GPU_0_bfc on device 0. [executable_name='jit_fwd'] [tf-allocator-allocation-error='']"
    },
    {
      "suite": "layer",
      "impl": "scan",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 8,
      "seq": 8192,
      "dim": 64,
      "window": 128,
      "dtype": "float32",
      "error": "JaxRuntimeError: NOT_FOUND: All configs failed during profiling or were excluded from selection."
    },
    {
      "suite": "layer",
      "impl": "chunked",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 8,
      "seq": 8192,
      "dim": 64,
      "window": 128,
      "dtype": "float32",
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 8.03GiB with allocator GPU_0_bfc on device 0. [tf-allocator-allocation-error=''] [executable_name='jit_fwd']"
    },
    {
      "suite": "layer",
      "impl": "chunked",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 8,
      "seq": 8192,
      "dim": 64,
      "window": 128,
      "dtype": "float32",
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 12.35GiB with allocator GPU_0_bfc on device 0. [tf-allocator-allocation-error=''] [executable_name='jit_grad']"
    },
    {
      "suite": "layer",
      "impl": "pallas",
      "mode": "fwd",
      "chunk": 128,
      "batch": 2,
      "heads": 8,
      "seq": 8192,
      "dim": 64,
      "window": 128,
      "dtype": "float32",
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 8.06GiB with allocator GPU_0_bfc on device 0. [executable_name='jit_fwd'] [tf-allocator-allocation-error='']"
    },
    {
      "suite": "layer",
      "impl": "pallas",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 2,
      "heads": 8,
      "seq": 8192,
      "dim": 64,
      "window": 128,
      "dtype": "float32",
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 12.38GiB with allocator GPU_0_bfc on device 0. [tf-allocator-allocation-error=''] [executable_name='jit_grad']"
    },
    {
      "suite": "peak",
      "impl": "scan",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 4,
      "heads": 8,
      "seq": 2048,
      "dim": 64,
      "dtype": "float32",
      "peak_mb": 6225.00048828125,
      "inputs_mb": 51.0
    },
    {
      "suite": "peak",
      "impl": "pallas_scanbwd",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 4,
      "heads": 8,
      "seq": 2048,
      "dim": 64,
      "dtype": "float32",
      "peak_mb": 6225.00048828125,
      "inputs_mb": 51.0
    },
    {
      "suite": "peak",
      "impl": "pallas",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 4,
      "heads": 8,
      "seq": 2048,
      "dim": 64,
      "dtype": "float32",
      "peak_mb": 465.00048828125,
      "inputs_mb": 51.0
    },
    {
      "suite": "peak",
      "impl": "chunked",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 4,
      "heads": 8,
      "seq": 2048,
      "dim": 64,
      "dtype": "float32",
      "peak_mb": 389.000732421875,
      "inputs_mb": 51.0
    },
    {
      "suite": "peak",
      "impl": "scan",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 4,
      "heads": 8,
      "seq": 8192,
      "dim": 64,
      "dtype": "float32",
      "error": "JaxRuntimeError: NOT_FOUND: All configs failed during profiling or were excluded from selection."
    },
    {
      "suite": "peak",
      "impl": "pallas_scanbwd",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 4,
      "heads": 8,
      "seq": 8192,
      "dim": 64,
      "dtype": "float32",
      "error": "JaxRuntimeError: NOT_FOUND: All configs failed during profiling or were excluded from selection."
    },
    {
      "suite": "peak",
      "impl": "pallas",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 4,
      "heads": 8,
      "seq": 8192,
      "dim": 64,
      "dtype": "float32",
      "peak_mb": 1164.00341796875,
      "inputs_mb": 195.0
    },
    {
      "suite": "peak",
      "impl": "chunked",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 4,
      "heads": 8,
      "seq": 8192,
      "dim": 64,
      "dtype": "float32",
      "peak_mb": 1411.0,
      "inputs_mb": 195.0
    },
    {
      "suite": "peak",
      "impl": "scan",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 4,
      "heads": 8,
      "seq": 32768,
      "dim": 64,
      "dtype": "float32",
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 32.02GiB with allocator GPU_0_bfc on device 0. [tf-allocator-allocation-error='']"
    },
    {
      "suite": "peak",
      "impl": "pallas_scanbwd",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 4,
      "heads": 8,
      "seq": 32768,
      "dim": 64,
      "dtype": "float32",
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 32.02GiB with allocator GPU_0_bfc on device 0. [tf-allocator-allocation-error='']"
    },
    {
      "suite": "peak",
      "impl": "pallas",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 4,
      "heads": 8,
      "seq": 32768,
      "dim": 64,
      "dtype": "float32",
      "peak_mb": 4644.003173828125,
      "inputs_mb": 770.999267578125
    },
    {
      "suite": "peak",
      "impl": "chunked",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 4,
      "heads": 8,
      "seq": 32768,
      "dim": 64,
      "dtype": "float32",
      "peak_mb": 5155.0029296875,
      "inputs_mb": 770.999267578125
    },
    {
      "suite": "peak",
      "impl": "scan",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 4,
      "heads": 8,
      "seq": 65536,
      "dim": 64,
      "dtype": "float32",
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 64.02GiB with allocator GPU_0_bfc on device 0. [tf-allocator-allocation-error='']"
    },
    {
      "suite": "peak",
      "impl": "pallas_scanbwd",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 4,
      "heads": 8,
      "seq": 65536,
      "dim": 64,
      "dtype": "float32",
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 64.02GiB with allocator GPU_0_bfc on device 0. [tf-allocator-allocation-error='']"
    },
    {
      "suite": "peak",
      "impl": "pallas",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 4,
      "heads": 8,
      "seq": 65536,
      "dim": 64,
      "dtype": "float32",
      "peak_mb": 11843.004150390625,
      "inputs_mb": 1538.999267578125
    },
    {
      "suite": "peak",
      "impl": "chunked",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 4,
      "heads": 8,
      "seq": 65536,
      "dim": 64,
      "dtype": "float32",
      "peak_mb": 12867.00390625,
      "inputs_mb": 1538.999267578125
    },
    {
      "suite": "peak",
      "impl": "scan",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 4,
      "heads": 8,
      "seq": 131072,
      "dim": 64,
      "dtype": "float32",
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 128.02GiB with allocator GPU_0_bfc on device 0. [tf-allocator-allocation-error='']"
    },
    {
      "suite": "peak",
      "impl": "pallas_scanbwd",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 4,
      "heads": 8,
      "seq": 131072,
      "dim": 64,
      "dtype": "float32",
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 128.02GiB with allocator GPU_0_bfc on device 0. [tf-allocator-allocation-error='']"
    },
    {
      "suite": "peak",
      "impl": "pallas",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 4,
      "heads": 8,
      "seq": 131072,
      "dim": 64,
      "dtype": "float32",
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 10.59GiB with allocator GPU_0_bfc on device 0. [executable_name='jit_loss'] [tf-allocator-allocation-error='']"
    },
    {
      "suite": "peak",
      "impl": "chunked",
      "mode": "fwd_bwd",
      "chunk": 128,
      "batch": 4,
      "heads": 8,
      "seq": 131072,
      "dim": 64,
      "dtype": "float32",
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 12.62GiB with allocator GPU_0_bfc on device 0. [tf-allocator-allocation-error=''] [executable_name='jit_loss']"
    },
    {
      "suite": "peak",
      "impl": "scan",
      "mode": "fwd",
      "chunk": 128,
      "batch": 4,
      "heads": 8,
      "seq": 8192,
      "dim": 64,
      "dtype": "float32",
      "peak_mb": 771.0,
      "inputs_mb": 195.0
    },
    {
      "suite": "peak",
      "impl": "pallas",
      "mode": "fwd",
      "chunk": 128,
      "batch": 4,
      "heads": 8,
      "seq": 8192,
      "dim": 64,
      "dtype": "float32",
      "peak_mb": 579.0,
      "inputs_mb": 195.0
    },
    {
      "suite": "peak",
      "impl": "chunked",
      "mode": "fwd",
      "chunk": 128,
      "batch": 4,
      "heads": 8,
      "seq": 8192,
      "dim": 64,
      "dtype": "float32",
      "peak_mb": 981.000732421875,
      "inputs_mb": 195.0
    },
    {
      "suite": "peak",
      "impl": "scan",
      "mode": "fwd",
      "chunk": 128,
      "batch": 4,
      "heads": 8,
      "seq": 65536,
      "dim": 64,
      "dtype": "float32",
      "peak_mb": 5123.01806640625,
      "inputs_mb": 1538.999267578125
    },
    {
      "suite": "peak",
      "impl": "pallas",
      "mode": "fwd",
      "chunk": 128,
      "batch": 4,
      "heads": 8,
      "seq": 65536,
      "dim": 64,
      "dtype": "float32",
      "peak_mb": 4610.999755859375,
      "inputs_mb": 1538.999267578125
    },
    {
      "suite": "peak",
      "impl": "chunked",
      "mode": "fwd",
      "chunk": 128,
      "batch": 4,
      "heads": 8,
      "seq": 65536,
      "dim": 64,
      "dtype": "float32",
      "peak_mb": 6147.00048828125,
      "inputs_mb": 1538.999267578125
    }
  ]
}