{
  "device": {
    "jax": "0.11.2",
    "backend": "cpu",
    "devices": [
      "cpu:0"
    ],
    "python": "3.14.8",
    "device_kind": "cpu",
    "gpu": "NVIDIA GeForce RTX 3090",
    "driver": "615.78.08",
    "gpu_memory_gb": 24.0,
    "torch": "2.14.1+cu130",
    "torch_cuda": "13.0"
  },
  "config": {
    "device": "cpu",
    "configs": {
      "g2h8_d128": [
        4,
        2,
        8,
        128,
        128
      ],
      "g1h24_dk128_dv64": [
        2,
        1,
        24,
        128,
        64
      ],
      "g1h1_d128": [
        4,
        1,
        1,
        128,
        128
      ],
      "g2h4_dk48_dv40": [
        4,
        2,
        4,
        48,
        40
      ]
    },
    "seq_lens": [
      512,
      2048,
      8192
    ],
    "impls": [
      "chunked",
      "scan"
    ],
    "repeats": 3
  },
  "results": [
    {
      "impl": "chunked",
      "mode": "fwd",
      "config": "g2h8_d128",
      "batch": 4,
      "groups": 2,
      "heads": 8,
      "seq": 512,
      "dk": 128,
      "dv": 128,
      "compile_s": 0.11972837900248123,
      "inner": 3,
      "median_ms": 13.798804000543896,
      "q1_ms": 12.916895999903016,
      "q3_ms": 14.17080750009821,
      "xla_temp_mb": 30.063087463378906,
      "xla_arg_mb": 14.0625,
      "xla_out_mb": 8.0,
      "rel_err": 0.0
    },
    {
      "impl": "scan",
      "mode": "fwd",
      "config": "g2h8_d128",
      "batch": 4,
      "groups": 2,
      "heads": 8,
      "seq": 512,
      "dk": 128,
      "dv": 128,
      "compile_s": 0.14789589200518094,
      "inner": 1,
      "median_ms": 34.33739100000821,
      "q1_ms": 33.54457650129916,
      "q3_ms": 40.01802000129828,
      "xla_temp_mb": 34.062564849853516,
      "xla_arg_mb": 14.0625,
      "xla_out_mb": 8.0,
      "rel_err": 1.8334178289969834e-06
    },
    {
      "impl": "chunked",
      "mode": "fwd_bwd",
      "config": "g2h8_d128",
      "batch": 4,
      "groups": 2,
      "heads": 8,
      "seq": 512,
      "dk": 128,
      "dv": 128,
      "compile_s": 0.2975732220002101,
      "inner": 1,
      "median_ms": 47.228752002411056,
      "q1_ms": 46.21368749940302,
      "q3_ms": 47.3930680018384,
      "xla_temp_mb": 80.126708984375,
      "xla_arg_mb": 14.0625,
      "xla_out_mb": 14.062538146972656,
      "residual_mb": 0.0,
      "rel_err": 9.871577684429946e-08,
      "rel_err_grads": {
        "q": 0.0,
        "k": 0.0,
        "v": 0.0,
        "log_alpha": 9.871577684429946e-08,
        "h0": 0.0
      }
    },
    {
      "impl": "scan",
      "mode": "fwd_bwd",
      "config": "g2h8_d128",
      "batch": 4,
      "groups": 2,
      "heads": 8,
      "seq": 512,
      "dk": 128,
      "dv": 128,
      "compile_s": 1.1436803899996448,
      "inner": 1,
      "median_ms": 1126.6301439973176,
      "q1_ms": 1123.9483669960464,
      "q3_ms": 1156.3353995006764,
      "xla_temp_mb": 2092.1723403930664,
      "xla_arg_mb": 14.0625,
      "xla_out_mb": 14.062538146972656,
      "residual_mb": 2072.125,
      "rel_err": 1.9496365926749143e-06,
      "rel_err_grads": {
        "q": 1.2706571736107177e-06,
        "k": 1.3142707498380255e-06,
        "v": 1.937131754429109e-06,
        "log_alpha": 1.9496365926749143e-06,
        "h0": 3.1633659447364713e-07
      }
    },
    {
      "impl": "chunked",
      "mode": "fwd",
      "config": "g2h8_d128",
      "batch": 4,
      "groups": 2,
      "heads": 8,
      "seq": 2048,
      "dk": 128,
      "dv": 128,
      "compile_s": 0.17812974200205645,
      "inner": 1,
      "median_ms": 72.84338700264925,
      "q1_ms": 71.9736165010545,
      "q3_ms": 73.18867499998305,
      "xla_temp_mb": 114.2520523071289,
      "xla_arg_mb": 50.25,
      "xla_out_mb": 32.0,
      "rel_err": 0.0
    },
    {
      "impl": "scan",
      "mode": "fwd",
      "config": "g2h8_d128",
      "batch": 4,
      "groups": 2,
      "heads": 8,
      "seq": 2048,
      "dk": 128,
      "dv": 128,
      "compile_s": 0.23395699600223452,
      "inner": 1,
      "median_ms": 140.89613300166093,
      "q1_ms": 140.8420660009142,
      "q3_ms": 164.9957115005236,
      "xla_temp_mb": 130.25006484985352,
      "xla_arg_mb": 50.25,
      "xla_out_mb": 32.0,
      "rel_err": 1.6388536192067195e-06
    },
    {
      "impl": "chunked",
      "mode": "fwd_bwd",
      "config": "g2h8_d128",
      "batch": 4,
      "groups": 2,
      "heads": 8,
      "seq": 2048,
      "dk": 128,
      "dv": 128,
      "compile_s": 0.6790586239949334,
      "inner": 1,
      "median_ms": 204.17777600232512,
      "q1_ms": 204.0238049994514,
      "q3_ms": 205.81438100271043,
      "xla_temp_mb": 322.506103515625,
      "xla_arg_mb": 50.25,
      "xla_out_mb": 50.250038146972656,
      "residual_mb": 0.0,
      "rel_err": 1.4593799824144712e-07,
      "rel_err_grads": {
        "q": 0.0,
        "k": 0.0,
        "v": 0.0,
        "log_alpha": 1.4593799824144712e-07,
        "h0": 0.0
      }
    },
    {
      "impl": "scan",
      "mode": "fwd_bwd",
      "config": "g2h8_d128",
      "batch": 4,
      "groups": 2,
      "heads": 8,
      "seq": 2048,
      "dk": 128,
      "dv": 128,
      "compile_s": 4.277472622001369,
      "inner": 1,
      "median_ms": 4262.538326998765,
      "q1_ms": 4259.694323500298,
      "q3_ms": 4281.363365498692,
      "xla_temp_mb": 8356.547340393066,
      "xla_arg_mb": 50.25,
      "xla_out_mb": 50.250038146972656,
      "residual_mb": 8288.5,
      "rel_err": 2.1890699736217067e-06,
      "rel_err_grads": {
        "q": 1.4277492845011292e-06,
        "k": 1.417766237658065e-06,
        "v": 2.018123622835167e-06,
        "log_alpha": 2.1890699736217067e-06,
        "h0": 1.9108040491466445e-07
      }
    },
    {
      "impl": "chunked",
      "mode": "fwd",
      "config": "g2h8_d128",
      "batch": 4,
      "groups": 2,
      "heads": 8,
      "seq": 8192,
      "dk": 128,
      "dv": 128,
      "compile_s": 0.47884591800539056,
      "inner": 1,
      "median_ms": 295.5291260004742,
      "q1_ms": 289.20902800018666,
      "q3_ms": 298.5623275017133,
      "xla_temp_mb": 451.0079116821289,
      "xla_arg_mb": 195.0,
      "xla_out_mb": 128.0,
      "rel_err": 0.0
    },
    {
      "impl": "chunked",
      "mode": "fwd_bwd",
      "config": "g2h8_d128",
      "batch": 4,
      "groups": 2,
      "heads": 8,
      "seq": 8192,
      "dk": 128,
      "dv": 128,
      "compile_s": 2.3930787819990655,
      "inner": 1,
      "median_ms": 797.5451149977744,
      "q1_ms": 790.758218998235,
      "q3_ms": 801.7084035018343,
      "xla_temp_mb": 1284.023681640625,
      "xla_arg_mb": 195.0,
      "xla_out_mb": 195.00003814697266,
      "residual_mb": 0.0,
      "rel_err": 1.3290227808459377e-07,
      "rel_err_grads": {
        "q": 0.0,
        "k": 0.0,
        "v": 0.0,
        "log_alpha": 1.3290227808459377e-07,
        "h0": 0.0
      }
    },
    {
      "impl": "chunked",
      "mode": "fwd",
      "config": "g1h24_dk128_dv64",
      "batch": 2,
      "groups": 1,
      "heads": 24,
      "seq": 512,
      "dk": 128,
      "dv": 64,
      "compile_s": 0.12074688600114314,
      "inner": 5,
      "median_ms": 9.807480599556584,
      "q1_ms": 9.761269199952949,
      "q3_ms": 9.909730999788735,
      "xla_temp_mb": 26.094581604003906,
      "xla_arg_mb": 8.59375,
      "xla_out_mb": 6.0,
      "rel_err": 0.0
    },
    {
      "impl": "scan",
      "mode": "fwd",
      "config": "g1h24_dk128_dv64",
      "batch": 2,
      "groups": 1,
      "heads": 24,
      "seq": 512,
      "dk": 128,
      "dv": 64,
      "compile_s": 0.21255857600044692,
      "inner": 1,
      "median_ms": 24.983130002510734,
      "q1_ms": 24.71144100127276,
      "q3_ms": 41.74423000222305,
      "xla_temp_mb": 43.593814849853516,
      "xla_arg_mb": 8.59375,
      "xla_out_mb": 6.0,
      "rel_err": 2.212250888259699e-06
    },
    {
      "impl": "chunked",
      "mode": "fwd_bwd",
      "config": "g1h24_dk128_dv64",
      "batch": 2,
      "groups": 1,
      "heads": 24,
      "seq": 512,
      "dk": 128,
      "dv": 64,
      "compile_s": 0.2881069779978134,
      "inner": 1,
      "median_ms": 41.077658999711275,
      "q1_ms": 40.944701999251265,
      "q3_ms": 41.66034599984414,
      "xla_temp_mb": 75.596923828125,
      "xla_arg_mb": 8.59375,
      "xla_out_mb": 8.593788146972656,
      "residual_mb": 0.0,
      "rel_err": 2.023150097303404e-07,
      "rel_err_grads": {
        "q": 0.0,
        "k": 0.0,
        "v": 0.0,
        "log_alpha": 2.023150097303404e-07,
        "h0": 0.0
      }
    },
    {
      "impl": "scan",
      "mode": "fwd_bwd",
      "config": "g1h24_dk128_dv64",
      "batch": 2,
      "groups": 1,
      "heads": 24,
      "seq": 512,
      "dk": 128,
      "dv": 64,
      "compile_s": 0.9659117619958124,
      "inner": 1,
      "median_ms": 813.9433279939112,
      "q1_ms": 766.1308064962213,
      "q3_ms": 834.7564059986325,
      "xla_temp_mb": 1596.0474014282227,
      "xla_arg_mb": 8.59375,
      "xla_out_mb": 8.593788146972656,
      "residual_mb": 1566.1875,
      "rel_err": 3.102163482531886e-06,
      "rel_err_grads": {
        "q": 1.011956277860692e-06,
        "k": 9.25688789600195e-07,
        "v": 1.7935191864498018e-06,
        "log_alpha": 3.102163482531886e-06,
        "h0": 3.0536556303989587e-07
      }
    },
    {
      "impl": "chunked",
      "mode": "fwd",
      "config": "g1h24_dk128_dv64",
      "batch": 2,
      "groups": 1,
      "heads": 24,
      "seq": 2048,
      "dk": 128,
      "dv": 64,
      "compile_s": 0.18622828800289426,
      "inner": 1,
      "median_ms": 62.04341099510202,
      "q1_ms": 59.73549599730177,
      "q3_ms": 62.71732099776273,
      "xla_temp_mb": 99.8780288696289,
      "xla_arg_mb": 29.875,
      "xla_out_mb": 24.0,
      "rel_err": 0.0
    },
    {
      "impl": "scan",
      "mode": "fwd",
      "config": "g1h24_dk128_dv64",
      "batch": 2,
      "groups": 1,
      "heads": 24,
      "seq": 2048,
      "dk": 128,
      "dv": 64,
      "compile_s": 0.2323148810028215,
      "inner": 1,
      "median_ms": 133.45270499848993,
      "q1_ms": 130.6587915023556,
      "q3_ms": 142.47123650056892,
      "xla_temp_mb": 169.87506484985352,
      "xla_arg_mb": 29.875,
      "xla_out_mb": 24.0,
      "rel_err": 1.8809799950032319e-06
    },
    {
      "impl": "chunked",
      "mode": "fwd_bwd",
      "config": "g1h24_dk128_dv64",
      "batch": 2,
      "groups": 1,
      "heads": 24,
      "seq": 2048,
      "dk": 128,
      "dv": 64,
      "compile_s": 0.5710083120065974,
      "inner": 1,
      "median_ms": 175.46003299503354,
      "q1_ms": 173.92605650093174,
      "q3_ms": 178.6646184955316,
      "xla_temp_mb": 297.886962890625,
      "xla_arg_mb": 29.875,
      "xla_out_mb": 29.875038146972656,
      "residual_mb": 0.0,
      "rel_err": 1.6989601231406297e-07,
      "rel_err_grads": {
        "q": 0.0,
        "k": 0.0,
        "v": 0.0,
        "log_alpha": 1.6989601231406297e-07,
        "h0": 0.0
      }
    },
    {
      "impl": "scan",
      "mode": "fwd_bwd",
      "config": "g1h24_dk128_dv64",
      "batch": 2,
      "groups": 1,
      "heads": 24,
      "seq": 2048,
      "dk": 128,
      "dv": 64,
      "compile_s": 3.0192224490019726,
      "inner": 1,
      "median_ms": 3047.420921000594,
      "q1_ms": 3044.0624990005745,
      "q3_ms": 3071.398333500838,
      "xla_temp_mb": 6363.797401428223,
      "xla_arg_mb": 29.875,
      "xla_out_mb": 29.875038146972656,
      "residual_mb": 6264.75,
      "rel_err": 2.906737991832437e-06,
      "rel_err_grads": {
        "q": 1.165533053355893e-06,
        "k": 8.721573391346114e-07,
        "v": 2.906737991832437e-06,
        "log_alpha": 2.8882322093390707e-06,
        "h0": 3.3043555123052133e-07
      }
    },
    {
      "impl": "chunked",
      "mode": "fwd",
      "config": "g1h24_dk128_dv64",
      "batch": 2,
      "groups": 1,
      "heads": 24,
      "seq": 8192,
      "dk": 128,
      "dv": 64,
      "compile_s": 0.3497273300017696,
      "inner": 1,
      "median_ms": 241.85726699943189,
      "q1_ms": 238.00589849997777,
      "q3_ms": 243.64260899892543,
      "xla_temp_mb": 395.0118179321289,
      "xla_arg_mb": 115.0,
      "xla_out_mb": 96.0,
      "rel_err": 0.0
    },
    {
      "impl": "chunked",
      "mode": "fwd_bwd",
      "config": "g1h24_dk128_dv64",
      "batch": 2,
      "groups": 1,
      "heads": 24,
      "seq": 8192,
      "dk": 128,
      "dv": 64,
      "compile_s": 2.095402364997426,
      "inner": 1,
      "median_ms": 705.6498050005757,
      "q1_ms": 703.8886245027243,
      "q3_ms": 717.9170759991393,
      "xla_temp_mb": 1187.047119140625,
      "xla_arg_mb": 115.0,
      "xla_out_mb": 115.00003814697266,
      "residual_mb": 0.0,
      "rel_err": 6.901965244877931e-08,
      "rel_err_grads": {
        "q": 0.0,
        "k": 0.0,
        "v": 0.0,
        "log_alpha": 6.901965244877931e-08,
        "h0": 0.0
      }
    },
    {
      "impl": "chunked",
      "mode": "fwd",
      "config": "g1h1_d128",
      "batch": 4,
      "groups": 1,
      "heads": 1,
      "seq": 512,
      "dk": 128,
      "dv": 128,
      "compile_s": 0.1111150290016667,
      "inner": 23,
      "median_ms": 1.7137106523201193,
      "q1_ms": 1.7062286957377142,
      "q3_ms": 1.7205213478756258,
      "xla_temp_mb": 5.257972717285156,
      "xla_arg_mb": 3.2578125,
      "xla_out_mb": 1.0,
      "rel_err": 0.0
    },
    {
      "impl": "scan",
      "mode": "fwd",
      "config": "g1h1_d128",
      "batch": 4,
      "groups": 1,
      "heads": 1,
      "seq": 512,
      "dk": 128,
      "dv": 128,
      "compile_s": 0.0873073309994652,
      "inner": 13,
      "median_ms": 3.8208470768250455,
      "q1_ms": 3.778892692375498,
      "q3_ms": 4.457526192149649,
      "xla_temp_mb": 4.257877349853516,
      "xla_arg_mb": 3.2578125,
      "xla_out_mb": 1.0,
      "rel_err": 1.340497441339281e-06
    },
    {
      "impl": "chunked",
      "mode": "fwd_bwd",
      "config": "g1h1_d128",
      "batch": 4,
      "groups": 1,
      "heads": 1,
      "seq": 512,
      "dk": 128,
      "dv": 128,
      "compile_s": 0.2201892669982044,
      "inner": 9,
      "median_ms": 5.008367777918465,
      "q1_ms": 4.953118833428663,
      "q3_ms": 5.100030000095204,
      "xla_temp_mb": 12.27392578125,
      "xla_arg_mb": 3.2578125,
      "xla_out_mb": 3.2578506469726562,
      "residual_mb": 0.0,
      "rel_err": 4.557411951484709e-08,
      "rel_err_grads": {
        "q": 0.0,
        "k": 0.0,
        "v": 0.0,
        "log_alpha": 4.557411951484709e-08,
        "h0": 0.0
      }
    },
    {
      "impl": "scan",
      "mode": "fwd_bwd",
      "config": "g1h1_d128",
      "batch": 4,
      "groups": 1,
      "heads": 1,
      "seq": 512,
      "dk": 128,
      "dv": 128,
      "compile_s": 0.2794865239993669,
      "inner": 1,
      "median_ms": 119.2829170031473,
      "q1_ms": 119.23888849923969,
      "q3_ms": 123.96091900154715,
      "xla_temp_mb": 260.02188873291016,
      "xla_arg_mb": 3.2578125,
      "xla_out_mb": 3.2578506469726562,
      "residual_mb": 259.015625,
      "rel_err": 2.1304387306776192e-06,
      "rel_err_grads": {
        "q": 1.8725885891458484e-06,
        "k": 1.8332807345714447e-06,
        "v": 2.1304387306776192e-06,
        "log_alpha": 1.2305012269008715e-06,
        "h0": 2.523192979602424e-07
      }
    },
    {
      "impl": "chunked",
      "mode": "fwd",
      "config": "g1h1_d128",
      "batch": 4,
      "groups": 1,
      "heads": 1,
      "seq": 2048,
      "dk": 128,
      "dv": 128,
      "compile_s": 0.10988829800044186,
      "inner": 7,
      "median_ms": 5.999163143444873,
      "q1_ms": 5.956602000099208,
      "q3_ms": 6.021866143003406,
      "xla_temp_mb": 20.281593322753906,
      "xla_arg_mb": 12.28125,
      "xla_out_mb": 4.0,
      "rel_err": 0.0
    },
    {
      "impl": "scan",
      "mode": "fwd",
      "config": "g1h1_d128",
      "batch": 4,
      "groups": 1,
      "heads": 1,
      "seq": 2048,
      "dk": 128,
      "dv": 128,
      "compile_s": 0.12471259499579901,
      "inner": 1,
      "median_ms": 15.227839998260606,
      "q1_ms": 15.15366949752206,
      "q3_ms": 17.59974049855373,
      "xla_temp_mb": 16.281314849853516,
      "xla_arg_mb": 12.28125,
      "xla_out_mb": 4.0,
      "rel_err": 1.3735862792466566e-06
    },
    {
      "impl": "chunked",
      "mode": "fwd_bwd",
      "config": "g1h1_d128",
      "batch": 4,
      "groups": 1,
      "heads": 1,
      "seq": 2048,
      "dk": 128,
      "dv": 128,
      "compile_s": 0.2535005480021937,
      "inner": 1,
      "median_ms": 28.48770700074965,
      "q1_ms": 28.110449500672985,
      "q3_ms": 28.942101500433637,
      "xla_temp_mb": 52.282470703125,
      "xla_arg_mb": 12.28125,
      "xla_out_mb": 12.281288146972656,
      "residual_mb": 0.0,
      "rel_err": 9.114843011412877e-08,
      "rel_err_grads": {
        "q": 0.0,
        "k": 0.0,
        "v": 0.0,
        "log_alpha": 9.114843011412877e-08,
        "h0": 0.0
      }
    },
    {
      "impl": "scan",
      "mode": "fwd_bwd",
      "config": "g1h1_d128",
      "batch": 4,
      "groups": 1,
      "heads": 1,
      "seq": 2048,
      "dk": 128,
      "dv": 128,
      "compile_s": 0.6062401949966443,
      "inner": 1,
      "median_ms": 495.68058700242545,
      "q1_ms": 486.6001930022321,
      "q3_ms": 496.30627499936963,
      "xla_temp_mb": 1037.0687637329102,
      "xla_arg_mb": 12.28125,
      "xla_out_mb": 12.281288146972656,
      "residual_mb": 1036.0625,
      "rel_err": 3.1161290406347555e-06,
      "rel_err_grads": {
        "q": 2.82485483715257e-06,
        "k": 3.1161290406347555e-06,
        "v": 2.1304387306776192e-06,
        "log_alpha": 1.7318201721684467e-06,
        "h0": 2.523192979602424e-07
      }
    },
    {
      "impl": "chunked",
      "mode": "fwd",
      "config": "g1h1_d128",
      "batch": 4,
      "groups": 1,
      "heads": 1,
      "seq": 8192,
      "dk": 128,
      "dv": 128,
      "compile_s": 0.1562679440030479,
      "inner": 1,
      "median_ms": 40.94808400259353,
      "q1_ms": 39.9572550013545,
      "q3_ms": 41.63754250112106,
      "xla_temp_mb": 80.3760757446289,
      "xla_arg_mb": 48.375,
      "xla_out_mb": 16.0,
      "rel_err": 0.0
    },
    {
      "impl": "chunked",
      "mode": "fwd_bwd",
      "config": "g1h1_d128",
      "batch": 4,
      "groups": 1,
      "heads": 1,
      "seq": 8192,
      "dk": 128,
      "dv": 128,
      "compile_s": 0.39326584299851675,
      "inner": 1,
      "median_ms": 119.25790900568245,
      "q1_ms": 118.75185200551641,
      "q3_ms": 122.61991750347079,
      "xla_temp_mb": 208.379150390625,
      "xla_arg_mb": 48.375,
      "xla_out_mb": 48.375038146972656,
      "residual_mb": 0.0,
      "rel_err": 1.5850009458493143e-07,
      "rel_err_grads": {
        "q": 0.0,
        "k": 0.0,
        "v": 0.0,
        "log_alpha": 1.5850009458493143e-07,
        "h0": 0.0
      }
    },
    {
      "impl": "chunked",
      "mode": "fwd",
      "config": "g2h4_dk48_dv40",
      "batch": 4,
      "groups": 2,
      "heads": 4,
      "seq": 512,
      "dk": 48,
      "dv": 40,
      "compile_s": 0.4227013330018963,
      "inner": 21,
      "median_ms": 2.1829418571239025,
      "q1_ms": 2.1612095475748427,
      "q3_ms": 2.188824642785004,
      "xla_temp_mb": 8.088233947753906,
      "xla_arg_mb": 2.8984375,
      "xla_out_mb": 1.25,
      "rel_err": 0.0
    },
    {
      "impl": "scan",
      "mode": "fwd",
      "config": "g2h4_dk48_dv40",
      "batch": 4,
      "groups": 2,
      "heads": 4,
      "seq": 512,
      "dk": 48,
      "dv": 40,
      "skipped": "unsupported shape"
    },
    {
      "impl": "chunked",
      "mode": "fwd_bwd",
      "config": "g2h4_dk48_dv40",
      "batch": 4,
      "groups": 2,
      "heads": 4,
      "seq": 512,
      "dk": 48,
      "dv": 40,
      "compile_s": 0.5972753890018794,
      "inner": 7,
      "median_ms": 6.7101305716538,
      "q1_ms": 6.696139857272751,
      "q3_ms": 6.786837785641962,
      "xla_temp_mb": 19.587158203125,
      "xla_arg_mb": 2.8984375,
      "xla_out_mb": 2.8984756469726562,
      "residual_mb": 0.0,
      "rel_err": 1.7866940959944283e-07,
      "rel_err_grads": {
        "q": 0.0,
        "k": 0.0,
        "v": 0.0,
        "log_alpha": 1.7866940959944283e-07,
        "h0": 0.0
      }
    },
    {
      "impl": "scan",
      "mode": "fwd_bwd",
      "config": "g2h4_dk48_dv40",
      "batch": 4,
      "groups": 2,
      "heads": 4,
      "seq": 512,
      "dk": 48,
      "dv": 40,
      "skipped": "unsupported shape"
    },
    {
      "impl": "chunked",
      "mode": "fwd",
      "config": "g2h4_dk48_dv40",
      "batch": 4,
      "groups": 2,
      "heads": 4,
      "seq": 2048,
      "dk": 48,
      "dv": 40,
      "compile_s": 0.42613135000283364,
      "inner": 6,
      "median_ms": 8.465171500574797,
      "q1_ms": 8.386366667158047,
      "q3_ms": 8.819960916904773,
      "xla_temp_mb": 32.001075744628906,
      "xla_arg_mb": 11.2421875,
      "xla_out_mb": 5.0,
      "rel_err": 0.0
    },
    {
      "impl": "scan",
      "mode": "fwd",
      "config": "g2h4_dk48_dv40",
      "batch": 4,
      "groups": 2,
      "heads": 4,
      "seq": 2048,
      "dk": 48,
      "dv": 40,
      "skipped": "unsupported shape"
    },
    {
      "impl": "chunked",
      "mode": "fwd_bwd",
      "config": "g2h4_dk48_dv40",
      "batch": 4,
      "groups": 2,
      "heads": 4,
      "seq": 2048,
      "dk": 48,
      "dv": 40,
      "compile_s": 0.6552920949980034,
      "inner": 1,
      "median_ms": 44.16932699678,
      "q1_ms": 43.33276450051926,
      "q3_ms": 44.56651799773681,
      "xla_temp_mb": 77.996337890625,
      "xla_arg_mb": 11.2421875,
      "xla_out_mb": 11.242225646972656,
      "residual_mb": 0.0,
      "rel_err": 7.245602897197792e-08,
      "rel_err_grads": {
        "q": 0.0,
        "k": 0.0,
        "v": 0.0,
        "log_alpha": 7.245602897197792e-08,
        "h0": 0.0
      }
    },
    {
      "impl": "scan",
      "mode": "fwd_bwd",
      "config": "g2h4_dk48_dv40",
      "batch": 4,
      "groups": 2,
      "heads": 4,
      "seq": 2048,
      "dk": 48,
      "dv": 40,
      "skipped": "unsupported shape"
    },
    {
      "impl": "chunked",
      "mode": "fwd",
      "config": "g2h4_dk48_dv40",
      "batch": 4,
      "groups": 2,
      "heads": 4,
      "seq": 8192,
      "dk": 48,
      "dv": 40,
      "compile_s": 0.5816368029991281,
      "inner": 1,
      "median_ms": 59.18682800256647,
      "q1_ms": 58.64953700074693,
      "q3_ms": 59.907375998591306,
      "xla_temp_mb": 127.6524429321289,
      "xla_arg_mb": 44.6171875,
      "xla_out_mb": 20.0,
      "rel_err": 0.0
    },
    {
      "impl": "chunked",
      "mode": "fwd_bwd",
      "config": "g2h4_dk48_dv40",
      "batch": 4,
      "groups": 2,
      "heads": 4,
      "seq": 8192,
      "dk": 48,
      "dv": 40,
      "compile_s": 0.9904471949994331,
      "inner": 1,
      "median_ms": 176.48591300530825,
      "q1_ms": 175.92646000048262,
      "q3_ms": 176.84620400177664,
      "xla_temp_mb": 311.633056640625,
      "xla_arg_mb": 44.6171875,
      "xla_out_mb": 44.617225646972656,
      "residual_mb": 0.0,
      "rel_err": 7.003944901926562e-08,
      "rel_err_grads": {
        "q": 0.0,
        "k": 0.0,
        "v": 0.0,
        "log_alpha": 7.003944901926562e-08,
        "h0": 0.0
      }
    }
  ]
}