{
  "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": {
    "device": "gpu",
    "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",
      "pallas"
    ],
    "repeats": 5,
    "merged_from": [
      "gpu_a.json",
      "gpu_b.json",
      "gpu_peak.json",
      "gpu_maxlen.json"
    ]
  },
  "results": [
    {
      "impl": "chunked",
      "mode": "fwd",
      "config": "g2h8_d128",
      "batch": 4,
      "groups": 2,
      "heads": 8,
      "seq": 512,
      "dk": 128,
      "dv": 128,
      "compile_s": 0.54906745399785,
      "inner": 104,
      "median_ms": 0.3323198846443959,
      "q1_ms": 0.33225823076813515,
      "q3_ms": 0.33233163461633153,
      "xla_temp_mb": 28.063270568847656,
      "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.1993586959943059,
      "inner": 9,
      "median_ms": 4.902800000208017,
      "q1_ms": 4.900997888600816,
      "q3_ms": 4.905002221575059,
      "xla_temp_mb": 26.12749481201172,
      "xla_arg_mb": 14.0625,
      "xla_out_mb": 8.0,
      "rel_err": 4.549593008348435e-06
    },
    {
      "impl": "pallas",
      "mode": "fwd",
      "config": "g2h8_d128",
      "batch": 4,
      "groups": 2,
      "heads": 8,
      "seq": 512,
      "dk": 128,
      "dv": 128,
      "compile_s": 1.0372557579976274,
      "inner": 97,
      "median_ms": 0.362678226797611,
      "q1_ms": 0.36250150511365326,
      "q3_ms": 0.3627115979397036,
      "xla_temp_mb": 16.140884399414062,
      "xla_arg_mb": 14.0625,
      "xla_out_mb": 8.0,
      "rel_err": 0.00048587276676097225
    },
    {
      "impl": "chunked",
      "mode": "fwd_bwd",
      "config": "g2h8_d128",
      "batch": 4,
      "groups": 2,
      "heads": 8,
      "seq": 512,
      "dk": 128,
      "dv": 128,
      "compile_s": 0.7122757950055529,
      "inner": 48,
      "median_ms": 0.901018208272338,
      "q1_ms": 0.9001862291976673,
      "q3_ms": 0.9017719374545171,
      "xla_temp_mb": 48.19146728515625,
      "xla_arg_mb": 14.0625,
      "xla_out_mb": 14.062538146972656,
      "residual_mb": 0.0,
      "rel_err": 9.87157963339125e-08,
      "rel_err_grads": {
        "q": 0.0,
        "k": 0.0,
        "v": 0.0,
        "log_alpha": 9.87157963339125e-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": 0.5401538700025412,
      "inner": 1,
      "median_ms": 27.65891900344286,
      "q1_ms": 27.640675005386584,
      "q3_ms": 27.661083993734792,
      "xla_temp_mb": 3080.069435119629,
      "xla_arg_mb": 14.0625,
      "xla_out_mb": 14.062538146972656,
      "residual_mb": 2072.125,
      "rel_err": 4.787716122194756e-06,
      "rel_err_grads": {
        "q": 2.5984219412347745e-06,
        "k": 2.3105357840761994e-06,
        "v": 3.937778526379115e-06,
        "log_alpha": 4.787716122194756e-06,
        "h0": 3.1633659447364713e-07
      }
    },
    {
      "impl": "pallas",
      "mode": "fwd_bwd",
      "config": "g2h8_d128",
      "batch": 4,
      "groups": 2,
      "heads": 8,
      "seq": 512,
      "dk": 128,
      "dv": 128,
      "compile_s": 2.8616631220065756,
      "inner": 29,
      "median_ms": 1.4587744138633094,
      "q1_ms": 1.457392103349838,
      "q3_ms": 1.4589381723359611,
      "xla_temp_mb": 48.000518798828125,
      "xla_arg_mb": 14.0625,
      "xla_out_mb": 14.062538146972656,
      "residual_mb": 8.0,
      "rel_err": 0.002601583860401649,
      "rel_err_grads": {
        "q": 0.0004575078632245514,
        "k": 0.0005007748943250819,
        "v": 0.00046132345687668874,
        "log_alpha": 0.002601583860401649,
        "h0": 0.0004707879367254054
      }
    },
    {
      "impl": "chunked",
      "mode": "fwd",
      "config": "g2h8_d128",
      "batch": 4,
      "groups": 2,
      "heads": 8,
      "seq": 2048,
      "dk": 128,
      "dv": 128,
      "compile_s": 0.30792829299753066,
      "inner": 38,
      "median_ms": 1.1700650263121228,
      "q1_ms": 1.1681824999916564,
      "q3_ms": 1.1700990262876363,
      "xla_temp_mb": 112.50101470947266,
      "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.2152617929968983,
      "inner": 2,
      "median_ms": 21.751092001068173,
      "q1_ms": 21.750710999185685,
      "q3_ms": 21.834674498677487,
      "xla_temp_mb": 98.31523895263672,
      "xla_arg_mb": 50.25,
      "xla_out_mb": 32.0,
      "rel_err": 4.621030373710249e-06
    },
    {
      "impl": "pallas",
      "mode": "fwd",
      "config": "g2h8_d128",
      "batch": 4,
      "groups": 2,
      "heads": 8,
      "seq": 2048,
      "dk": 128,
      "dv": 128,
      "compile_s": 1.0529757700060145,
      "inner": 31,
      "median_ms": 1.4564900967900853,
      "q1_ms": 1.4553705806329458,
      "q3_ms": 1.4565495805600808,
      "xla_temp_mb": 64.51564025878906,
      "xla_arg_mb": 50.25,
      "xla_out_mb": 32.0,
      "rel_err": 0.0004677449930598572
    },
    {
      "impl": "chunked",
      "mode": "fwd_bwd",
      "config": "g2h8_d128",
      "batch": 4,
      "groups": 2,
      "heads": 8,
      "seq": 2048,
      "dk": 128,
      "dv": 128,
      "compile_s": 1.3325461099957465,
      "inner": 14,
      "median_ms": 3.210066071395496,
      "q1_ms": 3.202280000030961,
      "q3_ms": 3.2131061426688183,
      "xla_temp_mb": 224.75469970703125,
      "xla_arg_mb": 50.25,
      "xla_out_mb": 50.250038146972656,
      "residual_mb": 0.0,
      "rel_err": 1.4593746579590645e-07,
      "rel_err_grads": {
        "q": 0.0,
        "k": 0.0,
        "v": 0.0,
        "log_alpha": 1.4593746579590645e-07,
        "h0": 0.0
      }
    },
    {
      "impl": "scan",
      "mode": "fwd_bwd",
      "config": "g2h8_d128",
      "batch": 4,
      "groups": 2,
      "heads": 8,
      "seq": 2048,
      "dk": 128,
      "dv": 128,
      "error": "JaxRuntimeError: NOT_FOUND: All configs failed during profiling or were excluded from selection."
    },
    {
      "impl": "pallas",
      "mode": "fwd_bwd",
      "config": "g2h8_d128",
      "batch": 4,
      "groups": 2,
      "heads": 8,
      "seq": 2048,
      "dk": 128,
      "dv": 128,
      "compile_s": 4.513113451997924,
      "inner": 8,
      "median_ms": 5.566713125517708,
      "q1_ms": 5.564602874983393,
      "q3_ms": 5.5749849998392165,
      "xla_temp_mb": 192.00027465820312,
      "xla_arg_mb": 50.25,
      "xla_out_mb": 50.250038146972656,
      "residual_mb": 32.0,
      "rel_err": 0.003862453938487358,
      "rel_err_grads": {
        "q": 0.0005433829826763994,
        "k": 0.0005830750200557958,
        "v": 0.0004927495347678596,
        "log_alpha": 0.003862453938487358,
        "h0": 0.00037916726511995553
      }
    },
    {
      "impl": "chunked",
      "mode": "fwd",
      "config": "g2h8_d128",
      "batch": 4,
      "groups": 2,
      "heads": 8,
      "seq": 8192,
      "dk": 128,
      "dv": 128,
      "compile_s": 0.6730233520065667,
      "inner": 11,
      "median_ms": 4.182355454593727,
      "q1_ms": 4.175189272659322,
      "q3_ms": 4.186062454457648,
      "xla_temp_mb": 292.00150299072266,
      "xla_arg_mb": 195.0,
      "xla_out_mb": 128.0,
      "rel_err": 0.0
    },
    {
      "impl": "scan",
      "mode": "fwd",
      "config": "g2h8_d128",
      "batch": 4,
      "groups": 2,
      "heads": 8,
      "seq": 8192,
      "dk": 128,
      "dv": 128,
      "compile_s": 0.4174369270040188,
      "inner": 1,
      "median_ms": 90.59587200317765,
      "q1_ms": 86.9035160067142,
      "q3_ms": 92.24598400032846,
      "xla_temp_mb": 387.0652389526367,
      "xla_arg_mb": 195.0,
      "xla_out_mb": 128.0,
      "rel_err": 4.185940538449718e-06
    },
    {
      "impl": "pallas",
      "mode": "fwd",
      "config": "g2h8_d128",
      "batch": 4,
      "groups": 2,
      "heads": 8,
      "seq": 8192,
      "dk": 128,
      "dv": 128,
      "compile_s": 1.9883567659999244,
      "inner": 7,
      "median_ms": 5.998996856630713,
      "q1_ms": 5.964113856732313,
      "q3_ms": 6.0614808576896655,
      "xla_temp_mb": 258.01564025878906,
      "xla_arg_mb": 195.0,
      "xla_out_mb": 128.0,
      "rel_err": 0.0005217589428115744
    },
    {
      "impl": "chunked",
      "mode": "fwd_bwd",
      "config": "g2h8_d128",
      "batch": 4,
      "groups": 2,
      "heads": 8,
      "seq": 8192,
      "dk": 128,
      "dv": 128,
      "compile_s": 2.7715248960012104,
      "inner": 4,
      "median_ms": 11.883384999237023,
      "q1_ms": 11.865401000250131,
      "q3_ms": 11.888096749316901,
      "xla_temp_mb": 804.0076293945312,
      "xla_arg_mb": 195.0,
      "xla_out_mb": 195.00003814697266,
      "residual_mb": 0.0,
      "rel_err": 1.3290226042158058e-07,
      "rel_err_grads": {
        "q": 0.0,
        "k": 0.0,
        "v": 0.0,
        "log_alpha": 1.3290226042158058e-07,
        "h0": 0.0
      }
    },
    {
      "impl": "scan",
      "mode": "fwd_bwd",
      "config": "g2h8_d128",
      "batch": 4,
      "groups": 2,
      "heads": 8,
      "seq": 8192,
      "dk": 128,
      "dv": 128,
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 16.02GiB with allocator GPU_0_bfc on device 0. [tf-allocator-allocation-error='']"
    },
    {
      "impl": "pallas",
      "mode": "fwd_bwd",
      "config": "g2h8_d128",
      "batch": 4,
      "groups": 2,
      "heads": 8,
      "seq": 8192,
      "dk": 128,
      "dv": 128,
      "compile_s": 9.36640527999407,
      "inner": 2,
      "median_ms": 22.22943200104055,
      "q1_ms": 22.15461049854639,
      "q3_ms": 22.344133500155294,
      "xla_temp_mb": 768.0002746582031,
      "xla_arg_mb": 195.0,
      "xla_out_mb": 195.00003814697266,
      "residual_mb": 128.0,
      "rel_err": 0.010721524378294855,
      "rel_err_grads": {
        "q": 0.0005465899913527634,
        "k": 0.00048241417102928006,
        "v": 0.0005377775152933869,
        "log_alpha": 0.010721524378294855,
        "h0": 0.00040029473742534565
      }
    },
    {
      "impl": "chunked",
      "mode": "fwd",
      "config": "g1h24_dk128_dv64",
      "batch": 2,
      "groups": 1,
      "heads": 24,
      "seq": 512,
      "dk": 128,
      "dv": 64,
      "compile_s": 0.333817241997167,
      "inner": 130,
      "median_ms": 0.25800490005237336,
      "q1_ms": 0.25666744617270665,
      "q3_ms": 0.2580523769071218,
      "xla_temp_mb": 31.094520568847656,
      "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.22389018700050656,
      "inner": 9,
      "median_ms": 4.712633777267507,
      "q1_ms": 4.712251999687093,
      "q3_ms": 4.713319554866757,
      "xla_temp_mb": 31.66655731201172,
      "xla_arg_mb": 8.59375,
      "xla_out_mb": 6.0,
      "rel_err": 4.014825686100935e-06
    },
    {
      "impl": "pallas",
      "mode": "fwd",
      "config": "g1h24_dk128_dv64",
      "batch": 2,
      "groups": 1,
      "heads": 24,
      "seq": 512,
      "dk": 128,
      "dv": 64,
      "compile_s": 1.2218494590051705,
      "inner": 105,
      "median_ms": 0.32054702862210216,
      "q1_ms": 0.32051229521000224,
      "q3_ms": 0.32063558097607237,
      "xla_temp_mb": 12.203384399414062,
      "xla_arg_mb": 8.59375,
      "xla_out_mb": 6.0,
      "rel_err": 0.0007534762654976368
    },
    {
      "impl": "chunked",
      "mode": "fwd_bwd",
      "config": "g1h24_dk128_dv64",
      "batch": 2,
      "groups": 1,
      "heads": 24,
      "seq": 512,
      "dk": 128,
      "dv": 64,
      "compile_s": 0.7715464919965598,
      "inner": 57,
      "median_ms": 0.7252080174287568,
      "q1_ms": 0.7237422455147138,
      "q3_ms": 0.7287260525851723,
      "xla_temp_mb": 55.37896728515625,
      "xla_arg_mb": 8.59375,
      "xla_out_mb": 8.593788146972656,
      "residual_mb": 0.0,
      "rel_err": 0.0,
      "rel_err_grads": {
        "q": 0.0,
        "k": 0.0,
        "v": 0.0,
        "log_alpha": 0.0,
        "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.499837004004803,
      "inner": 2,
      "median_ms": 22.208802998648025,
      "q1_ms": 22.20620249863714,
      "q3_ms": 22.20935399964219,
      "xla_temp_mb": 2316.100685119629,
      "xla_arg_mb": 8.59375,
      "xla_out_mb": 8.593788146972656,
      "residual_mb": 1566.1875,
      "rel_err": 5.26019025298885e-06,
      "rel_err_grads": {
        "q": 2.499045952456974e-06,
        "k": 2.1277970052343548e-06,
        "v": 3.5251926656944142e-06,
        "log_alpha": 5.26019025298885e-06,
        "h0": 3.8170695379986984e-07
      }
    },
    {
      "impl": "pallas",
      "mode": "fwd_bwd",
      "config": "g1h24_dk128_dv64",
      "batch": 2,
      "groups": 1,
      "heads": 24,
      "seq": 512,
      "dk": 128,
      "dv": 64,
      "compile_s": 3.5778912640016642,
      "inner": 37,
      "median_ms": 1.1992973784406393,
      "q1_ms": 1.1933374323259185,
      "q3_ms": 1.20041543237135,
      "xla_temp_mb": 36.000518798828125,
      "xla_arg_mb": 8.59375,
      "xla_out_mb": 8.593788146972656,
      "residual_mb": 6.0,
      "rel_err": 0.002952096324064406,
      "rel_err_grads": {
        "q": 0.00046906462516032747,
        "k": 0.0005133454913897119,
        "v": 0.0003763297784342195,
        "log_alpha": 0.002952096324064406,
        "h0": 0.00041300692401145914
      }
    },
    {
      "impl": "chunked",
      "mode": "fwd",
      "config": "g1h24_dk128_dv64",
      "batch": 2,
      "groups": 1,
      "heads": 24,
      "seq": 2048,
      "dk": 128,
      "dv": 64,
      "compile_s": 0.3318625649990281,
      "inner": 42,
      "median_ms": 0.9425275001419886,
      "q1_ms": 0.942123880945257,
      "q3_ms": 0.9452686428544222,
      "xla_temp_mb": 124.37577056884766,
      "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.21756057200400392,
      "inner": 2,
      "median_ms": 18.78984949871665,
      "q1_ms": 18.762772997433785,
      "q3_ms": 18.80919099858147,
      "xla_temp_mb": 121.94780731201172,
      "xla_arg_mb": 29.875,
      "xla_out_mb": 24.0,
      "rel_err": 2.959777345078615e-06
    },
    {
      "impl": "pallas",
      "mode": "fwd",
      "config": "g1h24_dk128_dv64",
      "batch": 2,
      "groups": 1,
      "heads": 24,
      "seq": 2048,
      "dk": 128,
      "dv": 64,
      "compile_s": 1.0756795139968744,
      "inner": 36,
      "median_ms": 1.2227948611527812,
      "q1_ms": 1.2225596943204033,
      "q3_ms": 1.222992694440311,
      "xla_temp_mb": 48.76564025878906,
      "xla_arg_mb": 29.875,
      "xla_out_mb": 24.0,
      "rel_err": 0.00047079822816110305
    },
    {
      "impl": "chunked",
      "mode": "fwd_bwd",
      "config": "g1h24_dk128_dv64",
      "batch": 2,
      "groups": 1,
      "heads": 24,
      "seq": 2048,
      "dk": 128,
      "dv": 64,
      "compile_s": 0.8669668520014966,
      "inner": 17,
      "median_ms": 2.63207676471211,
      "q1_ms": 2.630459000048337,
      "q3_ms": 2.632117470698979,
      "xla_temp_mb": 221.50494384765625,
      "xla_arg_mb": 29.875,
      "xla_out_mb": 29.875038146972656,
      "residual_mb": 0.0,
      "rel_err": 1.1326413649687618e-07,
      "rel_err_grads": {
        "q": 0.0,
        "k": 0.0,
        "v": 0.0,
        "log_alpha": 1.1326413649687618e-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,
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 9.05GiB with allocator GPU_0_bfc on device 0. [executable_name='jit_loss'] [tf-allocator-allocation-error='']"
    },
    {
      "impl": "pallas",
      "mode": "fwd_bwd",
      "config": "g1h24_dk128_dv64",
      "batch": 2,
      "groups": 1,
      "heads": 24,
      "seq": 2048,
      "dk": 128,
      "dv": 64,
      "compile_s": 3.1432279660002678,
      "inner": 10,
      "median_ms": 4.4786664002458565,
      "q1_ms": 4.4246372002817225,
      "q3_ms": 4.57682600026601,
      "xla_temp_mb": 144.00027465820312,
      "xla_arg_mb": 29.875,
      "xla_out_mb": 29.875038146972656,
      "residual_mb": 24.0,
      "rel_err": 0.006695439532808089,
      "rel_err_grads": {
        "q": 0.0005107696544685715,
        "k": 0.0004202110012905677,
        "v": 0.0005120651574347953,
        "log_alpha": 0.006695439532808089,
        "h0": 0.0005424099125370212
      }
    },
    {
      "impl": "chunked",
      "mode": "fwd",
      "config": "g1h24_dk128_dv64",
      "batch": 2,
      "groups": 1,
      "heads": 24,
      "seq": 8192,
      "dk": 128,
      "dv": 64,
      "compile_s": 0.3011560910017579,
      "inner": 14,
      "median_ms": 3.4607668572529553,
      "q1_ms": 3.4562890713069856,
      "q3_ms": 3.4636165712851965,
      "xla_temp_mb": 499.00101470947266,
      "xla_arg_mb": 115.0,
      "xla_out_mb": 96.0,
      "rel_err": 0.0
    },
    {
      "impl": "scan",
      "mode": "fwd",
      "config": "g1h24_dk128_dv64",
      "batch": 2,
      "groups": 1,
      "heads": 24,
      "seq": 8192,
      "dk": 128,
      "dv": 64,
      "compile_s": 0.31423923300462775,
      "inner": 1,
      "median_ms": 82.7596570015885,
      "q1_ms": 82.7430950012058,
      "q3_ms": 82.77912400080822,
      "xla_temp_mb": 483.0730514526367,
      "xla_arg_mb": 115.0,
      "xla_out_mb": 96.0,
      "rel_err": 6.825061912213799e-06
    },
    {
      "impl": "pallas",
      "mode": "fwd",
      "config": "g1h24_dk128_dv64",
      "batch": 2,
      "groups": 1,
      "heads": 24,
      "seq": 8192,
      "dk": 128,
      "dv": 64,
      "compile_s": 1.0669034899983671,
      "inner": 10,
      "median_ms": 4.831674299930455,
      "q1_ms": 4.828473200177541,
      "q3_ms": 4.835764000017662,
      "xla_temp_mb": 195.01564025878906,
      "xla_arg_mb": 115.0,
      "xla_out_mb": 96.0,
      "rel_err": 0.0005050545815038211
    },
    {
      "impl": "chunked",
      "mode": "fwd_bwd",
      "config": "g1h24_dk128_dv64",
      "batch": 2,
      "groups": 1,
      "heads": 24,
      "seq": 8192,
      "dk": 128,
      "dv": 64,
      "compile_s": 2.3273530879960163,
      "inner": 4,
      "median_ms": 10.695072249291115,
      "q1_ms": 10.679590250219917,
      "q3_ms": 10.708317000535317,
      "xla_temp_mb": 887.5095825195312,
      "xla_arg_mb": 115.0,
      "xla_out_mb": 115.00003814697266,
      "residual_mb": 0.0,
      "rel_err": 6.901977154179541e-08,
      "rel_err_grads": {
        "q": 0.0,
        "k": 0.0,
        "v": 0.0,
        "log_alpha": 6.901977154179541e-08,
        "h0": 0.0
      }
    },
    {
      "impl": "scan",
      "mode": "fwd_bwd",
      "config": "g1h24_dk128_dv64",
      "batch": 2,
      "groups": 1,
      "heads": 24,
      "seq": 8192,
      "dk": 128,
      "dv": 64,
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 12.02GiB with allocator GPU_0_bfc on device 0. [tf-allocator-allocation-error='']"
    },
    {
      "impl": "pallas",
      "mode": "fwd_bwd",
      "config": "g1h24_dk128_dv64",
      "batch": 2,
      "groups": 1,
      "heads": 24,
      "seq": 8192,
      "dk": 128,
      "dv": 64,
      "compile_s": 4.112501876996248,
      "inner": 2,
      "median_ms": 18.353548999584746,
      "q1_ms": 18.132730499928584,
      "q3_ms": 18.35601300263079,
      "xla_temp_mb": 576.0002746582031,
      "xla_arg_mb": 115.0,
      "xla_out_mb": 115.00003814697266,
      "residual_mb": 96.0,
      "rel_err": 0.008256032914700812,
      "rel_err_grads": {
        "q": 0.0003377103959070551,
        "k": 0.0005014725507644826,
        "v": 0.00044730719220852213,
        "log_alpha": 0.008256032914700812,
        "h0": 0.00036265720918122205
      }
    },
    {
      "impl": "chunked",
      "mode": "fwd",
      "config": "g1h1_d128",
      "batch": 4,
      "groups": 1,
      "heads": 1,
      "seq": 512,
      "dk": 128,
      "dv": 128,
      "compile_s": 0.2855615820008097,
      "inner": 200,
      "median_ms": 0.08902542500436539,
      "q1_ms": 0.08875426497979788,
      "q3_ms": 0.08977348497865023,
      "xla_temp_mb": 8.508583068847656,
      "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.1398782600008417,
      "inner": 11,
      "median_ms": 4.339630909354045,
      "q1_ms": 4.308238727389835,
      "q3_ms": 4.353620182056065,
      "xla_temp_mb": 3.2610397338867188,
      "xla_arg_mb": 3.2578125,
      "xla_out_mb": 1.0,
      "rel_err": 2.75443728829395e-06
    },
    {
      "impl": "pallas",
      "mode": "fwd",
      "config": "g1h1_d128",
      "batch": 4,
      "groups": 1,
      "heads": 1,
      "seq": 512,
      "dk": 128,
      "dv": 128,
      "compile_s": 1.002362967999943,
      "inner": 200,
      "median_ms": 0.10316061503544915,
      "q1_ms": 0.10306347998266574,
      "q3_ms": 0.10319428001821507,
      "xla_temp_mb": 2.0315093994140625,
      "xla_arg_mb": 3.2578125,
      "xla_out_mb": 1.0,
      "rel_err": 0.0005205151958265354
    },
    {
      "impl": "chunked",
      "mode": "fwd_bwd",
      "config": "g1h1_d128",
      "batch": 4,
      "groups": 1,
      "heads": 1,
      "seq": 512,
      "dk": 128,
      "dv": 128,
      "compile_s": 4.24866419799946,
      "inner": 137,
      "median_ms": 0.25188338686414674,
      "q1_ms": 0.23096739417528325,
      "q3_ms": 0.25251311680002914,
      "xla_temp_mb": 10.51934814453125,
      "xla_arg_mb": 3.2578125,
      "xla_out_mb": 3.2578506469726562,
      "residual_mb": 0.0,
      "rel_err": 4.557407382088478e-08,
      "rel_err_grads": {
        "q": 0.0,
        "k": 0.0,
        "v": 0.0,
        "log_alpha": 4.557407382088478e-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.28990367599908495,
      "inner": 3,
      "median_ms": 12.457390999770723,
      "q1_ms": 12.4400983331725,
      "q3_ms": 12.457925666240044,
      "xla_temp_mb": 384.0196304321289,
      "xla_arg_mb": 3.2578125,
      "xla_out_mb": 3.2578506469726562,
      "residual_mb": 259.015625,
      "rel_err": 4.530455169394323e-06,
      "rel_err_grads": {
        "q": 4.530455169394323e-06,
        "k": 3.755990133550489e-06,
        "v": 2.9337191127399615e-06,
        "log_alpha": 3.007888872178395e-06,
        "h0": 2.5231927673856815e-07
      }
    },
    {
      "impl": "pallas",
      "mode": "fwd_bwd",
      "config": "g1h1_d128",
      "batch": 4,
      "groups": 1,
      "heads": 1,
      "seq": 512,
      "dk": 128,
      "dv": 128,
      "compile_s": 2.680049692004104,
      "inner": 93,
      "median_ms": 0.37115367748449124,
      "q1_ms": 0.3711307311483148,
      "q3_ms": 0.3712898494565599,
      "xla_temp_mb": 6.008575439453125,
      "xla_arg_mb": 3.2578125,
      "xla_out_mb": 3.2578506469726562,
      "residual_mb": 1.0,
      "rel_err": 0.002034973544250147,
      "rel_err_grads": {
        "q": 0.000555433803767744,
        "k": 0.0005625042366674423,
        "v": 0.0004298597004714696,
        "log_alpha": 0.002034973544250147,
        "h0": 0.0005006855514748987
      }
    },
    {
      "impl": "chunked",
      "mode": "fwd",
      "config": "g1h1_d128",
      "batch": 4,
      "groups": 1,
      "heads": 1,
      "seq": 2048,
      "dk": 128,
      "dv": 128,
      "compile_s": 0.564304663996154,
      "inner": 117,
      "median_ms": 0.2836962905861676,
      "q1_ms": 0.28367787177931186,
      "q3_ms": 0.28377600852748314,
      "xla_temp_mb": 24.032020568847656,
      "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.16024735500104725,
      "inner": 2,
      "median_ms": 17.4577474972466,
      "q1_ms": 17.411334498319775,
      "q3_ms": 17.4800345012045,
      "xla_temp_mb": 12.284477233886719,
      "xla_arg_mb": 12.28125,
      "xla_out_mb": 4.0,
      "rel_err": 3.296607296600703e-06
    },
    {
      "impl": "pallas",
      "mode": "fwd",
      "config": "g1h1_d128",
      "batch": 4,
      "groups": 1,
      "heads": 1,
      "seq": 2048,
      "dk": 128,
      "dv": 128,
      "compile_s": 1.006203163000464,
      "inner": 89,
      "median_ms": 0.3885450673857285,
      "q1_ms": 0.3884585056450019,
      "q3_ms": 0.3885598202481908,
      "xla_temp_mb": 8.078384399414062,
      "xla_arg_mb": 12.28125,
      "xla_out_mb": 4.0,
      "rel_err": 0.000539166991290913
    },
    {
      "impl": "chunked",
      "mode": "fwd_bwd",
      "config": "g1h1_d128",
      "batch": 4,
      "groups": 1,
      "heads": 1,
      "seq": 2048,
      "dk": 128,
      "dv": 128,
      "compile_s": 4.387380149004457,
      "inner": 61,
      "median_ms": 0.6560785902261383,
      "q1_ms": 0.6428238360610119,
      "q3_ms": 0.6751246064744859,
      "xla_temp_mb": 32.09722900390625,
      "xla_arg_mb": 12.28125,
      "xla_out_mb": 12.281288146972656,
      "residual_mb": 0.0,
      "rel_err": 9.114834703384138e-08,
      "rel_err_grads": {
        "q": 0.0,
        "k": 0.0,
        "v": 0.0,
        "log_alpha": 9.114834703384138e-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.4511605919979047,
      "inner": 1,
      "median_ms": 49.58043900114717,
      "q1_ms": 49.5738559984602,
      "q3_ms": 49.619213001278695,
      "xla_temp_mb": 1536.066505432129,
      "xla_arg_mb": 12.28125,
      "xla_out_mb": 12.281288146972656,
      "residual_mb": 1036.0625,
      "rel_err": 4.587634023832845e-06,
      "rel_err_grads": {
        "q": 4.559414824877832e-06,
        "k": 4.587634023832845e-06,
        "v": 3.248046160533529e-06,
        "log_alpha": 3.5547855343198135e-06,
        "h0": 2.5231927673856815e-07
      }
    },
    {
      "impl": "pallas",
      "mode": "fwd_bwd",
      "config": "g1h1_d128",
      "batch": 4,
      "groups": 1,
      "heads": 1,
      "seq": 2048,
      "dk": 128,
      "dv": 128,
      "compile_s": 2.7009687579993624,
      "inner": 32,
      "median_ms": 1.3750241873822233,
      "q1_ms": 1.3235414376140398,
      "q3_ms": 1.3761478437572805,
      "xla_temp_mb": 24.016387939453125,
      "xla_arg_mb": 12.28125,
      "xla_out_mb": 12.281288146972656,
      "residual_mb": 4.0,
      "rel_err": 0.0071477109551016,
      "rel_err_grads": {
        "q": 0.0005025268078724046,
        "k": 0.0005445435027157062,
        "v": 0.00046705506779370844,
        "log_alpha": 0.0071477109551016,
        "h0": 0.0005006855514748987
      }
    },
    {
      "impl": "chunked",
      "mode": "fwd",
      "config": "g1h1_d128",
      "batch": 4,
      "groups": 1,
      "heads": 1,
      "seq": 8192,
      "dk": 128,
      "dv": 128,
      "compile_s": 0.29895854699861957,
      "inner": 48,
      "median_ms": 0.9092048750668861,
      "q1_ms": 0.905471791763072,
      "q3_ms": 0.9197888542378981,
      "xla_temp_mb": 96.12577056884766,
      "xla_arg_mb": 48.375,
      "xla_out_mb": 16.0,
      "rel_err": 0.0
    },
    {
      "impl": "scan",
      "mode": "fwd",
      "config": "g1h1_d128",
      "batch": 4,
      "groups": 1,
      "heads": 1,
      "seq": 8192,
      "dk": 128,
      "dv": 128,
      "compile_s": 0.263448015000904,
      "inner": 1,
      "median_ms": 67.89502799801994,
      "q1_ms": 67.79234299756354,
      "q3_ms": 68.35342399426736,
      "xla_temp_mb": 48.37822723388672,
      "xla_arg_mb": 48.375,
      "xla_out_mb": 16.0,
      "rel_err": 3.552041562980165e-06
    },
    {
      "impl": "pallas",
      "mode": "fwd",
      "config": "g1h1_d128",
      "batch": 4,
      "groups": 1,
      "heads": 1,
      "seq": 8192,
      "dk": 128,
      "dv": 128,
      "compile_s": 0.999691935001465,
      "inner": 27,
      "median_ms": 1.528030074108823,
      "q1_ms": 1.5279625556474827,
      "q3_ms": 1.5283332594898875,
      "xla_temp_mb": 32.26588439941406,
      "xla_arg_mb": 48.375,
      "xla_out_mb": 16.0,
      "rel_err": 0.0004726322422063693
    },
    {
      "impl": "chunked",
      "mode": "fwd_bwd",
      "config": "g1h1_d128",
      "batch": 4,
      "groups": 1,
      "heads": 1,
      "seq": 8192,
      "dk": 128,
      "dv": 128,
      "compile_s": 4.405083441000897,
      "inner": 19,
      "median_ms": 2.4386077367173704,
      "q1_ms": 2.4170522102442766,
      "q3_ms": 2.441102947494456,
      "xla_temp_mb": 128.37872314453125,
      "xla_arg_mb": 48.375,
      "xla_out_mb": 48.375038146972656,
      "residual_mb": 0.0,
      "rel_err": 7.925038016407364e-08,
      "rel_err_grads": {
        "q": 0.0,
        "k": 0.0,
        "v": 0.0,
        "log_alpha": 7.925038016407364e-08,
        "h0": 0.0
      }
    },
    {
      "impl": "scan",
      "mode": "fwd_bwd",
      "config": "g1h1_d128",
      "batch": 4,
      "groups": 1,
      "heads": 1,
      "seq": 8192,
      "dk": 128,
      "dv": 128,
      "compile_s": 1.0976248890001443,
      "inner": 1,
      "median_ms": 196.8485499965027,
      "q1_ms": 196.83704700582894,
      "q3_ms": 196.87502000306267,
      "xla_temp_mb": 6144.254005432129,
      "xla_arg_mb": 48.375,
      "xla_out_mb": 48.375038146972656,
      "residual_mb": 4144.25,
      "rel_err": 5.434472041864538e-06,
      "rel_err_grads": {
        "q": 5.434472041864538e-06,
        "k": 3.925617958481908e-06,
        "v": 3.5085408680238063e-06,
        "log_alpha": 3.9030812230806275e-06,
        "h0": 2.5231927673856815e-07
      }
    },
    {
      "impl": "pallas",
      "mode": "fwd_bwd",
      "config": "g1h1_d128",
      "batch": 4,
      "groups": 1,
      "heads": 1,
      "seq": 8192,
      "dk": 128,
      "dv": 128,
      "compile_s": 2.764372388999618,
      "inner": 9,
      "median_ms": 5.369288999645505,
      "q1_ms": 4.974668111471046,
      "q3_ms": 5.3762833332358335,
      "xla_temp_mb": 96.00076293945312,
      "xla_arg_mb": 48.375,
      "xla_out_mb": 48.375038146972656,
      "residual_mb": 16.0,
      "rel_err": 0.007790381714211083,
      "rel_err_grads": {
        "q": 0.0004968424894031912,
        "k": 0.00035685286140657823,
        "v": 0.0004938945991141204,
        "log_alpha": 0.007790381714211083,
        "h0": 0.0005006855514748987
      }
    },
    {
      "impl": "chunked",
      "mode": "fwd",
      "config": "g2h4_dk48_dv40",
      "batch": 4,
      "groups": 2,
      "heads": 4,
      "seq": 512,
      "dk": 48,
      "dv": 40,
      "compile_s": 0.5911743780015968,
      "inner": 200,
      "median_ms": 0.11194227998203132,
      "q1_ms": 0.11189434000698384,
      "q3_ms": 0.11199487998965196,
      "xla_temp_mb": 17.281776428222656,
      "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": "pallas",
      "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.8213768479981809,
      "inner": 114,
      "median_ms": 0.29574836841212554,
      "q1_ms": 0.29563069294549915,
      "q3_ms": 0.2957815964978899,
      "xla_temp_mb": 21.37823486328125,
      "xla_arg_mb": 2.8984375,
      "xla_out_mb": 2.8984756469726562,
      "residual_mb": 0.0,
      "rel_err": 0.0,
      "rel_err_grads": {
        "q": 0.0,
        "k": 0.0,
        "v": 0.0,
        "log_alpha": 0.0,
        "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": "pallas",
      "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.33340897200105246,
      "inner": 104,
      "median_ms": 0.37085747114290896,
      "q1_ms": 0.36486623076388447,
      "q3_ms": 0.3735193365103512,
      "xla_temp_mb": 69.12577056884766,
      "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": "pallas",
      "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.8505750769982114,
      "inner": 45,
      "median_ms": 0.9682896445156075,
      "q1_ms": 0.9680636666922106,
      "q3_ms": 0.9686839333476706,
      "xla_temp_mb": 85.50347900390625,
      "xla_arg_mb": 11.2421875,
      "xla_out_mb": 11.242225646972656,
      "residual_mb": 0.0,
      "rel_err": 7.245602372210217e-08,
      "rel_err_grads": {
        "q": 0.0,
        "k": 0.0,
        "v": 0.0,
        "log_alpha": 7.245602372210217e-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": "pallas",
      "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.593568807002157,
      "inner": 31,
      "median_ms": 1.4480533871157546,
      "q1_ms": 1.4274282258420794,
      "q3_ms": 1.4487711934282657,
      "xla_temp_mb": 276.50052642822266,
      "xla_arg_mb": 44.6171875,
      "xla_out_mb": 20.0,
      "rel_err": 0.0
    },
    {
      "impl": "scan",
      "mode": "fwd",
      "config": "g2h4_dk48_dv40",
      "batch": 4,
      "groups": 2,
      "heads": 4,
      "seq": 8192,
      "dk": 48,
      "dv": 40,
      "skipped": "unsupported shape"
    },
    {
      "impl": "pallas",
      "mode": "fwd",
      "config": "g2h4_dk48_dv40",
      "batch": 4,
      "groups": 2,
      "heads": 4,
      "seq": 8192,
      "dk": 48,
      "dv": 40,
      "skipped": "unsupported shape"
    },
    {
      "impl": "chunked",
      "mode": "fwd_bwd",
      "config": "g2h4_dk48_dv40",
      "batch": 4,
      "groups": 2,
      "heads": 4,
      "seq": 8192,
      "dk": 48,
      "dv": 40,
      "compile_s": 1.0067422129941406,
      "inner": 12,
      "median_ms": 3.735824332882961,
      "q1_ms": 3.733850583254631,
      "q3_ms": 3.737344666660647,
      "xla_temp_mb": 342.50518798828125,
      "xla_arg_mb": 44.6171875,
      "xla_out_mb": 44.617225646972656,
      "residual_mb": 0.0,
      "rel_err": 0.0,
      "rel_err_grads": {
        "q": 0.0,
        "k": 0.0,
        "v": 0.0,
        "log_alpha": 0.0,
        "h0": 0.0
      }
    },
    {
      "impl": "scan",
      "mode": "fwd_bwd",
      "config": "g2h4_dk48_dv40",
      "batch": 4,
      "groups": 2,
      "heads": 4,
      "seq": 8192,
      "dk": 48,
      "dv": 40,
      "skipped": "unsupported shape"
    },
    {
      "impl": "pallas",
      "mode": "fwd_bwd",
      "config": "g2h4_dk48_dv40",
      "batch": 4,
      "groups": 2,
      "heads": 4,
      "seq": 8192,
      "dk": 48,
      "dv": 40,
      "skipped": "unsupported shape"
    },
    {
      "kind": "peak",
      "impl": "chunked",
      "mode": "fwd",
      "config": "g2h8_d128",
      "seq": 8192,
      "peak_mb": 1128.0
    },
    {
      "kind": "peak",
      "impl": "scan",
      "mode": "fwd",
      "config": "g2h8_d128",
      "seq": 8192,
      "peak_mb": 963.99951171875
    },
    {
      "kind": "peak",
      "impl": "pallas",
      "mode": "fwd",
      "config": "g2h8_d128",
      "seq": 8192,
      "peak_mb": 963.99951171875
    },
    {
      "kind": "peak",
      "impl": "chunked",
      "mode": "fwd_bwd",
      "config": "g2h8_d128",
      "seq": 8192,
      "peak_mb": 1323.007568359375
    },
    {
      "kind": "peak",
      "impl": "scan",
      "mode": "fwd_bwd",
      "config": "g2h8_d128",
      "seq": 8192,
      "error": "EXECUTION FAILED: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 16.02GiB with allocator GPU_0_bfc on device 0. [tf-allocator-allocation-error='']"
    },
    {
      "kind": "peak",
      "impl": "pallas",
      "mode": "fwd_bwd",
      "config": "g2h8_d128",
      "seq": 8192,
      "peak_mb": 1287.000732421875
    },
    {
      "kind": "peak",
      "impl": "chunked",
      "mode": "fwd",
      "config": "g1h1_d128",
      "seq": 8192,
      "peak_mb": 320.375732421875
    },
    {
      "kind": "peak",
      "impl": "scan",
      "mode": "fwd",
      "config": "g1h1_d128",
      "seq": 8192,
      "peak_mb": 144.375732421875
    },
    {
      "kind": "peak",
      "impl": "pallas",
      "mode": "fwd",
      "config": "g1h1_d128",
      "seq": 8192,
      "peak_mb": 176.375732421875
    },
    {
      "kind": "peak",
      "impl": "chunked",
      "mode": "fwd_bwd",
      "config": "g1h1_d128",
      "seq": 8192,
      "peak_mb": 368.751220703125
    },
    {
      "kind": "peak",
      "impl": "scan",
      "mode": "fwd_bwd",
      "config": "g1h1_d128",
      "seq": 8192,
      "peak_mb": 6257.00537109375
    },
    {
      "kind": "peak",
      "impl": "pallas",
      "mode": "fwd_bwd",
      "config": "g1h1_d128",
      "seq": 8192,
      "peak_mb": 240.751220703125
    },
    {
      "kind": "max_len",
      "impl": "chunked",
      "mode": "fwd_bwd",
      "config": "g2h8_d128",
      "largest_ok": {
        "seq": 65536,
        "peak_mb": 12076.037109375
      },
      "first_failure_error": "EXECUTION FAILED: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 4.02GiB with allocator GPU_0_bfc on device 0. [tf-allocator-allocation-error='']"
    },
    {
      "kind": "max_len",
      "impl": "scan",
      "mode": "fwd_bwd",
      "config": "g2h8_d128",
      "largest_ok": {
        "seq": 1024,
        "peak_mb": 6228.383056640625
      },
      "first_failure_error": "jax.errors.JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 12.03GiB with allocator GPU_0_bfc on device 0. [tf-allocator-allocation-error=''] [executable_name='jit_loss']"
    },
    {
      "kind": "max_len",
      "impl": "pallas",
      "mode": "fwd_bwd",
      "config": "g2h8_d128",
      "largest_ok": {
        "seq": 65536,
        "peak_mb": 10260.001220703125
      },
      "first_failure_error": "jax.errors.JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 4.02GiB with allocator GPU_0_bfc on device 0. [tf-allocator-allocation-error='']"
    }
  ]
}