{
  "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",
    "batch": 4,
    "heads": 8,
    "seq_lens": [
      512,
      2048,
      8192
    ],
    "dims": [
      64,
      128
    ],
    "dtypes": [
      "bfloat16",
      "float32"
    ],
    "impls": [
      "recurrent",
      "chunked",
      "chunked_fat",
      "pallas",
      "pallas_unfused",
      "ssd"
    ],
    "chunk_sizes": [
      64
    ],
    "repeats": 5
  },
  "results": [
    {
      "impl": "recurrent",
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "seq": 512,
      "dk": 64,
      "dv": 64,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 0.18810881199897267,
      "inner": 11,
      "median_ms": 4.383482728206383,
      "q1_ms": 4.282633182645606,
      "q3_ms": 4.525865090809847,
      "xla_temp_mb": 16.14166259765625,
      "xla_arg_mb": 6.125,
      "xla_out_mb": 2.0,
      "rel_err": 0.0029469328250949786
    },
    {
      "impl": "chunked",
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "seq": 512,
      "dk": 64,
      "dv": 64,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 1.3929519950033864,
      "inner": 73,
      "median_ms": 0.478813383442809,
      "q1_ms": 0.46895916433367013,
      "q3_ms": 0.4886866713417949,
      "xla_temp_mb": 28.00666046142578,
      "xla_arg_mb": 6.125,
      "xla_out_mb": 2.0,
      "rel_err": 0.00308954090667691
    },
    {
      "impl": "chunked_fat",
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "seq": 512,
      "dk": 64,
      "dv": 64,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 0.5182805200020084,
      "inner": 71,
      "median_ms": 0.47820800006173186,
      "q1_ms": 0.47436271835951477,
      "q3_ms": 0.4787617182673913,
      "xla_temp_mb": 28.00666046142578,
      "xla_arg_mb": 6.125,
      "xla_out_mb": 2.0,
      "rel_err": 0.00308954090667691
    },
    {
      "impl": "pallas",
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "seq": 512,
      "dk": 64,
      "dv": 64,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 4.009437854998396,
      "inner": 187,
      "median_ms": 0.15716751336112597,
      "q1_ms": 0.155123010670572,
      "q3_ms": 0.16263381814282826,
      "xla_temp_mb": 12.562522888183594,
      "xla_arg_mb": 6.125,
      "xla_out_mb": 2.0,
      "rel_err": 0.00308954090667691
    },
    {
      "impl": "pallas_unfused",
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "seq": 512,
      "dk": 64,
      "dv": 64,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 4.465855225993437,
      "inner": 51,
      "median_ms": 0.28452833323492943,
      "q1_ms": 0.2743785686713333,
      "q3_ms": 0.286277509967833,
      "xla_temp_mb": 32.06303405761719,
      "xla_arg_mb": 6.125,
      "xla_out_mb": 2.0,
      "rel_err": 0.00308954090667691
    },
    {
      "impl": "ssd",
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "seq": 512,
      "dk": 64,
      "dv": 64,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 1.4207836870045867,
      "inner": 193,
      "median_ms": 0.15238296369348842,
      "q1_ms": 0.1390484092911522,
      "q3_ms": 0.1533450829462327,
      "xla_temp_mb": 13.140884399414062,
      "xla_arg_mb": 6.0625,
      "xla_out_mb": 4.0
    },
    {
      "impl": "recurrent",
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "seq": 512,
      "dk": 64,
      "dv": 64,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 0.5786340580089018,
      "inner": 1,
      "median_ms": 25.060240004677325,
      "q1_ms": 24.81355499185156,
      "q3_ms": 25.431911999476142,
      "xla_temp_mb": 1052.130744934082,
      "xla_arg_mb": 6.125,
      "xla_out_mb": 6.125038146972656,
      "residual_mb": 784.1875,
      "rel_err": 0.005434782608695652,
      "rel_err_grads": {
        "q": 0.005434782608695652,
        "k": 0.004464285714285714,
        "v": 0.004629629629629629,
        "log_alpha": 0.0017262102924015948,
        "beta": 0.001581245204573256
      }
    },
    {
      "impl": "chunked",
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "seq": 512,
      "dk": 64,
      "dv": 64,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 2.538857239007484,
      "inner": 27,
      "median_ms": 1.7379009255415037,
      "q1_ms": 1.6892609263002802,
      "q3_ms": 1.7829240737396672,
      "xla_temp_mb": 64.07013702392578,
      "xla_arg_mb": 6.125,
      "xla_out_mb": 6.125038146972656,
      "residual_mb": 4.0,
      "rel_err": 0.004629629629629629,
      "rel_err_grads": {
        "q": 0.002717391304347826,
        "k": 0.004464285714285714,
        "v": 0.004629629629629629,
        "log_alpha": 0.001752602156777603,
        "beta": 0.001573163649906778
      }
    },
    {
      "impl": "chunked_fat",
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "seq": 512,
      "dk": 64,
      "dv": 64,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 0.9784406770049827,
      "inner": 32,
      "median_ms": 1.406895718901069,
      "q1_ms": 1.3865920936950715,
      "q3_ms": 1.409326531302213,
      "xla_temp_mb": 60.13263702392578,
      "xla_arg_mb": 6.125,
      "xla_out_mb": 6.125038146972656,
      "residual_mb": 28.0,
      "rel_err": 0.004629629629629629,
      "rel_err_grads": {
        "q": 0.002717391304347826,
        "k": 0.004464285714285714,
        "v": 0.004629629629629629,
        "log_alpha": 0.001752602156777603,
        "beta": 0.001573163649906778
      }
    },
    {
      "impl": "pallas",
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "seq": 512,
      "dk": 64,
      "dv": 64,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 4.573872296998161,
      "inner": 46,
      "median_ms": 0.8700203043585846,
      "q1_ms": 0.8659360654469904,
      "q3_ms": 0.8763524782135274,
      "xla_temp_mb": 68.06400299072266,
      "xla_arg_mb": 6.125,
      "xla_out_mb": 6.125038146972656,
      "residual_mb": 4.0,
      "rel_err": 0.005434782608695652,
      "rel_err_grads": {
        "q": 0.005434782608695652,
        "k": 0.004464285714285714,
        "v": 0.004629629629629629,
        "log_alpha": 0.0017827679057908716,
        "beta": 0.0015818958382116588
      }
    },
    {
      "impl": "pallas_unfused",
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "seq": 512,
      "dk": 64,
      "dv": 64,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 3.7625171249965206,
      "inner": 51,
      "median_ms": 0.8657196469321919,
      "q1_ms": 0.8542499411965301,
      "q3_ms": 0.8674531372711427,
      "xla_temp_mb": 68.06450653076172,
      "xla_arg_mb": 6.125,
      "xla_out_mb": 6.125038146972656,
      "residual_mb": 4.0,
      "rel_err": 0.005434782608695652,
      "rel_err_grads": {
        "q": 0.005434782608695652,
        "k": 0.004464285714285714,
        "v": 0.004629629629629629,
        "log_alpha": 0.001782995858554196,
        "beta": 0.0015924771958046322
      }
    },
    {
      "impl": "ssd",
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "seq": 512,
      "dk": 64,
      "dv": 64,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 3.545908528001746,
      "inner": 90,
      "median_ms": 0.49675157772273654,
      "q1_ms": 0.48862022215164164,
      "q3_ms": 0.5028916112173142,
      "xla_temp_mb": 24.078887939453125,
      "xla_arg_mb": 6.0625,
      "xla_out_mb": 6.062530517578125,
      "residual_mb": 17.0
    },
    {
      "impl": "recurrent",
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "seq": 2048,
      "dk": 64,
      "dv": 64,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 0.20071625300624873,
      "inner": 2,
      "median_ms": 17.294388999289367,
      "q1_ms": 17.103489000874106,
      "q3_ms": 17.3059309963719,
      "xla_temp_mb": 64.51666259765625,
      "xla_arg_mb": 24.5,
      "xla_out_mb": 8.0,
      "rel_err": 0.002686979629155663
    },
    {
      "impl": "chunked",
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "seq": 2048,
      "dk": 64,
      "dv": 64,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 0.9180588839953998,
      "inner": 27,
      "median_ms": 1.7202264447135989,
      "q1_ms": 1.6943631484604407,
      "q3_ms": 1.727051518779869,
      "xla_temp_mb": 128.01862335205078,
      "xla_arg_mb": 24.5,
      "xla_out_mb": 8.0,
      "rel_err": 0.0028170080461255237
    },
    {
      "impl": "chunked_fat",
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "seq": 2048,
      "dk": 64,
      "dv": 64,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 0.3952932270040037,
      "inner": 28,
      "median_ms": 1.6728892499356465,
      "q1_ms": 1.6696828216871447,
      "q3_ms": 1.6878797141544055,
      "xla_temp_mb": 128.01862335205078,
      "xla_arg_mb": 24.5,
      "xla_out_mb": 8.0,
      "rel_err": 0.0028170080461255237
    },
    {
      "impl": "pallas",
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "seq": 2048,
      "dk": 64,
      "dv": 64,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 3.0461733099946287,
      "inner": 75,
      "median_ms": 0.5363931600004435,
      "q1_ms": 0.5363489468193924,
      "q3_ms": 0.5562612134963274,
      "xla_temp_mb": 48.750022888183594,
      "xla_arg_mb": 24.5,
      "xla_out_mb": 8.0,
      "rel_err": 0.0028170080461255237
    },
    {
      "impl": "pallas_unfused",
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "seq": 2048,
      "dk": 64,
      "dv": 64,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 3.288244705996476,
      "inner": 51,
      "median_ms": 0.7791189022028965,
      "q1_ms": 0.7715153333265334,
      "q3_ms": 0.8100250980693519,
      "xla_temp_mb": 128.25076293945312,
      "xla_arg_mb": 24.5,
      "xla_out_mb": 8.0,
      "rel_err": 0.0028170080461255237
    },
    {
      "impl": "ssd",
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "seq": 2048,
      "dk": 64,
      "dv": 64,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 1.1205344690097263,
      "inner": 74,
      "median_ms": 0.5197288783902705,
      "q1_ms": 0.5117882162291633,
      "q3_ms": 0.5427700810754245,
      "xla_temp_mb": 49.51564025878906,
      "xla_arg_mb": 24.25,
      "xla_out_mb": 16.0
    },
    {
      "impl": "recurrent",
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "seq": 2048,
      "dk": 64,
      "dv": 64,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 1.027117700999952,
      "inner": 1,
      "median_ms": 89.07616299984511,
      "q1_ms": 88.8462389993947,
      "q3_ms": 89.69952000188641,
      "xla_temp_mb": 4208.507453918457,
      "xla_arg_mb": 24.5,
      "xla_out_mb": 24.500038146972656,
      "residual_mb": 3136.75,
      "rel_err": 0.005154639175257732,
      "rel_err_grads": {
        "q": 0.005154639175257732,
        "k": 0.0038461538461538464,
        "v": 0.003875968992248062,
        "log_alpha": 0.0017439652798560878,
        "beta": 0.0016426578992973116
      }
    },
    {
      "impl": "chunked",
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "seq": 2048,
      "dk": 64,
      "dv": 64,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 2.1101839060138445,
      "inner": 8,
      "median_ms": 5.540912874494097,
      "q1_ms": 5.494884750078199,
      "q3_ms": 5.570338626057492,
      "xla_temp_mb": 240.26227569580078,
      "xla_arg_mb": 24.5,
      "xla_out_mb": 24.500038146972656,
      "residual_mb": 16.0,
      "rel_err": 0.005154639175257732,
      "rel_err_grads": {
        "q": 0.005154639175257732,
        "k": 0.0038461538461538464,
        "v": 0.003875968992248062,
        "log_alpha": 0.0019420562311850322,
        "beta": 0.0016279104490763813
      }
    },
    {
      "impl": "chunked_fat",
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "seq": 2048,
      "dk": 64,
      "dv": 64,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 1.1257316499977605,
      "inner": 11,
      "median_ms": 4.66963572729252,
      "q1_ms": 4.578346000264652,
      "q3_ms": 5.042865726863965,
      "xla_temp_mb": 272.5069046020508,
      "xla_arg_mb": 24.5,
      "xla_out_mb": 24.500038146972656,
      "residual_mb": 112.0,
      "rel_err": 0.005154639175257732,
      "rel_err_grads": {
        "q": 0.005154639175257732,
        "k": 0.0038461538461538464,
        "v": 0.003875968992248062,
        "log_alpha": 0.0019420055750154046,
        "beta": 0.0016279104490763813
      }
    },
    {
      "impl": "pallas",
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "seq": 2048,
      "dk": 64,
      "dv": 64,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 3.903889845998492,
      "inner": 14,
      "median_ms": 3.3175301432493143,
      "q1_ms": 3.2923354291207425,
      "q3_ms": 3.4297577855925607,
      "xla_temp_mb": 272.25174713134766,
      "xla_arg_mb": 24.5,
      "xla_out_mb": 24.500038146972656,
      "residual_mb": 16.0,
      "rel_err": 0.005154639175257732,
      "rel_err_grads": {
        "q": 0.005154639175257732,
        "k": 0.0038461538461538464,
        "v": 0.003875968992248062,
        "log_alpha": 0.001835070400931425,
        "beta": 0.0016401795933473212
      }
    },
    {
      "impl": "pallas_unfused",
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "seq": 2048,
      "dk": 64,
      "dv": 64,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 3.5916401940048672,
      "inner": 14,
      "median_ms": 3.273628000377877,
      "q1_ms": 3.2658847849234007,
      "q3_ms": 3.3008830713307753,
      "xla_temp_mb": 272.50272369384766,
      "xla_arg_mb": 24.5,
      "xla_out_mb": 24.500038146972656,
      "residual_mb": 16.0,
      "rel_err": 0.005154639175257732,
      "rel_err_grads": {
        "q": 0.005154639175257732,
        "k": 0.0038461538461538464,
        "v": 0.003875968992248062,
        "log_alpha": 0.0018401360178941907,
        "beta": 0.0016396900514312738
      }
    },
    {
      "impl": "ssd",
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "seq": 2048,
      "dk": 64,
      "dv": 64,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 4.240862473001471,
      "inner": 25,
      "median_ms": 1.839933120063506,
      "q1_ms": 1.8242058396572247,
      "q3_ms": 1.9710987602593377,
      "xla_temp_mb": 96.26614379882812,
      "xla_arg_mb": 24.25,
      "xla_out_mb": 24.250030517578125,
      "residual_mb": 65.0
    },
    {
      "impl": "recurrent",
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "seq": 8192,
      "dk": 64,
      "dv": 64,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 0.3660411939927144,
      "inner": 1,
      "median_ms": 72.0824620075291,
      "q1_ms": 72.07997700606938,
      "q3_ms": 75.73365399730392,
      "xla_temp_mb": 257.51666259765625,
      "xla_arg_mb": 98.0,
      "xla_out_mb": 32.0,
      "rel_err": 0.0024147740373423763,
      "peak_mb": 514.516357421875,
      "inputs_mb": 194.99951171875
    },
    {
      "impl": "chunked",
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "seq": 8192,
      "dk": 64,
      "dv": 64,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 0.7701338619954186,
      "inner": 6,
      "median_ms": 7.228217499990326,
      "q1_ms": 7.199306332040578,
      "q3_ms": 7.515003999287728,
      "xla_temp_mb": 387.0186233520508,
      "xla_arg_mb": 98.0,
      "xla_out_mb": 32.0,
      "rel_err": 0.002780443193082685,
      "peak_mb": 768.99951171875,
      "inputs_mb": 194.999755859375
    },
    {
      "impl": "chunked_fat",
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "seq": 8192,
      "dk": 64,
      "dv": 64,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 0.32748644798994064,
      "inner": 6,
      "median_ms": 7.2088843323096325,
      "q1_ms": 7.191321165009867,
      "q3_ms": 7.2709281666902825,
      "xla_temp_mb": 387.0186233520508,
      "xla_arg_mb": 98.0,
      "xla_out_mb": 32.0,
      "rel_err": 0.002780443193082685,
      "peak_mb": 768.99951171875,
      "inputs_mb": 194.999755859375
    },
    {
      "impl": "pallas",
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "seq": 8192,
      "dk": 64,
      "dv": 64,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 2.886095166002633,
      "inner": 23,
      "median_ms": 2.1012830003873324,
      "q1_ms": 2.0326930870119804,
      "q3_ms": 2.126484477902641,
      "xla_temp_mb": 193.0000228881836,
      "xla_arg_mb": 98.0,
      "xla_out_mb": 32.0,
      "rel_err": 0.002780443193082685,
      "peak_mb": 660.999755859375,
      "inputs_mb": 194.999755859375
    },
    {
      "impl": "pallas_unfused",
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "seq": 8192,
      "dk": 64,
      "dv": 64,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 3.619571562987403,
      "inner": 15,
      "median_ms": 2.8746208661080646,
      "q1_ms": 2.7945277329611904,
      "q3_ms": 2.875469799619168,
      "xla_temp_mb": 449.5007629394531,
      "xla_arg_mb": 98.0,
      "xla_out_mb": 32.0,
      "rel_err": 0.002780443193082685,
      "peak_mb": 768.99951171875,
      "inputs_mb": 194.99951171875
    },
    {
      "impl": "ssd",
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "seq": 8192,
      "dk": 64,
      "dv": 64,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 1.3875893509975867,
      "inner": 24,
      "median_ms": 2.03361791682255,
      "q1_ms": 1.9486020003872302,
      "q3_ms": 2.1126767921183878,
      "xla_temp_mb": 195.01564025878906,
      "xla_arg_mb": 97.0,
      "xla_out_mb": 64.0,
      "peak_mb": 454.015380859375,
      "inputs_mb": 194.99951171875
    },
    {
      "impl": "recurrent",
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "seq": 8192,
      "dk": 64,
      "dv": 64,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "error": "JaxRuntimeError: NOT_FOUND: All configs failed during profiling or were excluded from selection."
    },
    {
      "impl": "recurrent",
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "seq": 512,
      "dk": 128,
      "dv": 128,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 0.16068139299750328,
      "inner": 6,
      "median_ms": 7.392794500143888,
      "q1_ms": 7.385191832630274,
      "q3_ms": 7.394933500715221,
      "xla_temp_mb": 32.18951416015625,
      "xla_arg_mb": 12.125,
      "xla_out_mb": 4.0,
      "rel_err": 0.0030254911251562868
    },
    {
      "impl": "chunked",
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "seq": 512,
      "dk": 128,
      "dv": 128,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 1.453663410997251,
      "inner": 57,
      "median_ms": 0.6969865437560273,
      "q1_ms": 0.6962581402356655,
      "q3_ms": 0.69813871947158,
      "xla_temp_mb": 60.06501007080078,
      "xla_arg_mb": 12.125,
      "xla_out_mb": 4.0,
      "rel_err": 0.0030254911251562868
    },
    {
      "impl": "chunked_fat",
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "seq": 512,
      "dk": 128,
      "dv": 128,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 0.43101623900292907,
      "inner": 62,
      "median_ms": 0.6970398871119194,
      "q1_ms": 0.6958990160191822,
      "q3_ms": 0.6977568872013851,
      "xla_temp_mb": 60.06501007080078,
      "xla_arg_mb": 12.125,
      "xla_out_mb": 4.0,
      "rel_err": 0.0030254911251562868
    },
    {
      "impl": "pallas",
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "seq": 512,
      "dk": 128,
      "dv": 128,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 2.8591300010011764,
      "inner": 95,
      "median_ms": 0.4063537684391792,
      "q1_ms": 0.40360901039093733,
      "q3_ms": 0.4066469578231734,
      "xla_temp_mb": 60.125518798828125,
      "xla_arg_mb": 12.125,
      "xla_out_mb": 4.0,
      "rel_err": 0.0030254911251562868
    },
    {
      "impl": "pallas_unfused",
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "seq": 512,
      "dk": 128,
      "dv": 128,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 2.4103034870058764,
      "inner": 95,
      "median_ms": 0.4067742526163592,
      "q1_ms": 0.4066883999927852,
      "q3_ms": 0.40707091577228555,
      "xla_temp_mb": 60.125518798828125,
      "xla_arg_mb": 12.125,
      "xla_out_mb": 4.0,
      "rel_err": 0.0030254911251562868
    },
    {
      "impl": "ssd",
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "seq": 512,
      "dk": 128,
      "dv": 128,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 1.0900448789907387,
      "inner": 94,
      "median_ms": 0.4285587659823966,
      "q1_ms": 0.4283117021723809,
      "q3_ms": 0.4286315639508254,
      "xla_temp_mb": 36.14088439941406,
      "xla_arg_mb": 12.0625,
      "xla_out_mb": 8.0
    },
    {
      "impl": "recurrent",
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "seq": 512,
      "dk": 128,
      "dv": 128,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 0.7369675190129783,
      "inner": 1,
      "median_ms": 38.78517099656165,
      "q1_ms": 38.68738700111862,
      "q3_ms": 40.56080599548295,
      "xla_temp_mb": 4152.132942199707,
      "xla_arg_mb": 12.125,
      "xla_out_mb": 12.125038146972656,
      "residual_mb": 3104.1875,
      "rel_err": 0.006097560975609756,
      "rel_err_grads": {
        "q": 0.005649717514124294,
        "k": 0.004310344827586207,
        "v": 0.006097560975609756,
        "log_alpha": 0.0018838770241178578,
        "beta": 0.0018779195333843166
      }
    },
    {
      "impl": "chunked",
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "seq": 512,
      "dk": 128,
      "dv": 128,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 2.6955230720050167,
      "inner": 19,
      "median_ms": 2.3418689471375393,
      "q1_ms": 2.3356081058553078,
      "q3_ms": 2.343139683944173,
      "xla_temp_mb": 140.07135772705078,
      "xla_arg_mb": 12.125,
      "xla_out_mb": 12.125038146972656,
      "residual_mb": 16.0,
      "rel_err": 0.006097560975609756,
      "rel_err_grads": {
        "q": 0.005649717514124294,
        "k": 0.004310344827586207,
        "v": 0.006097560975609756,
        "log_alpha": 0.0018410609729794184,
        "beta": 0.0019041905722173781
      }
    },
    {
      "impl": "chunked_fat",
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "seq": 512,
      "dk": 128,
      "dv": 128,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 1.5454645179997897,
      "inner": 22,
      "median_ms": 2.083605818122371,
      "q1_ms": 2.0816480001402935,
      "q3_ms": 2.0870181820514104,
      "xla_temp_mb": 128.13166046142578,
      "xla_arg_mb": 12.125,
      "xla_out_mb": 12.125038146972656,
      "residual_mb": 56.0,
      "rel_err": 0.006097560975609756,
      "rel_err_grads": {
        "q": 0.005649717514124294,
        "k": 0.004310344827586207,
        "v": 0.006097560975609756,
        "log_alpha": 0.0018410609729794184,
        "beta": 0.0019041905722173781
      }
    },
    {
      "impl": "pallas",
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "seq": 512,
      "dk": 128,
      "dv": 128,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 3.52356327099551,
      "inner": 30,
      "median_ms": 1.4336243334885996,
      "q1_ms": 1.4257514664980893,
      "q3_ms": 1.4403637663538877,
      "xla_temp_mb": 110.25247955322266,
      "xla_arg_mb": 12.125,
      "xla_out_mb": 12.125038146972656,
      "residual_mb": 16.0,
      "rel_err": 0.006097560975609756,
      "rel_err_grads": {
        "q": 0.005649717514124294,
        "k": 0.004310344827586207,
        "v": 0.006097560975609756,
        "log_alpha": 0.001853701206506637,
        "beta": 0.001867805183433588
      }
    },
    {
      "impl": "pallas_unfused",
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "seq": 512,
      "dk": 128,
      "dv": 128,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 3.031590546001098,
      "inner": 30,
      "median_ms": 1.453051700082142,
      "q1_ms": 1.4467400998303979,
      "q3_ms": 1.4652410662771822,
      "xla_temp_mb": 110.25247955322266,
      "xla_arg_mb": 12.125,
      "xla_out_mb": 12.125038146972656,
      "residual_mb": 16.0,
      "rel_err": 0.006097560975609756,
      "rel_err_grads": {
        "q": 0.005649717514124294,
        "k": 0.004310344827586207,
        "v": 0.006097560975609756,
        "log_alpha": 0.001853701206506637,
        "beta": 0.001867805183433588
      }
    },
    {
      "impl": "ssd",
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "seq": 512,
      "dk": 128,
      "dv": 128,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 3.139777778997086,
      "inner": 27,
      "median_ms": 1.669225296441086,
      "q1_ms": 1.6667108887944508,
      "q3_ms": 1.673481851833424,
      "xla_temp_mb": 72.07888793945312,
      "xla_arg_mb": 12.0625,
      "xla_out_mb": 12.062530517578125,
      "residual_mb": 36.0
    },
    {
      "impl": "recurrent",
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "seq": 2048,
      "dk": 128,
      "dv": 128,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 0.25079948001075536,
      "inner": 1,
      "median_ms": 27.572459002840333,
      "q1_ms": 27.559083988307975,
      "q3_ms": 30.94476999831386,
      "xla_temp_mb": 128.56475830078125,
      "xla_arg_mb": 48.5,
      "xla_out_mb": 16.0,
      "rel_err": 0.003008762022728728
    },
    {
      "impl": "chunked",
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "seq": 2048,
      "dk": 128,
      "dv": 128,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 0.895533645991236,
      "inner": 18,
      "median_ms": 2.652286722473541,
      "q1_ms": 2.650977555377823,
      "q3_ms": 2.6529707780961567,
      "xla_temp_mb": 208.50275421142578,
      "xla_arg_mb": 48.5,
      "xla_out_mb": 16.0,
      "rel_err": 0.003008762022728728
    },
    {
      "impl": "chunked_fat",
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "seq": 2048,
      "dk": 128,
      "dv": 128,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 0.333755227009533,
      "inner": 17,
      "median_ms": 2.6634292940691333,
      "q1_ms": 2.659765293869628,
      "q3_ms": 2.66647147038705,
      "xla_temp_mb": 208.50275421142578,
      "xla_arg_mb": 48.5,
      "xla_out_mb": 16.0,
      "rel_err": 0.003008762022728728
    },
    {
      "impl": "pallas",
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "seq": 2048,
      "dk": 128,
      "dv": 128,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 3.2402618399937637,
      "inner": 30,
      "median_ms": 1.5096513001481071,
      "q1_ms": 1.5088658333600808,
      "q3_ms": 1.5138582331322445,
      "xla_temp_mb": 228.5007553100586,
      "xla_arg_mb": 48.5,
      "xla_out_mb": 16.0,
      "rel_err": 0.003008762022728728
    },
    {
      "impl": "pallas_unfused",
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "seq": 2048,
      "dk": 128,
      "dv": 128,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 2.8757183389971033,
      "inner": 28,
      "median_ms": 1.5108803927432746,
      "q1_ms": 1.5100695709406864,
      "q3_ms": 1.5109730359524423,
      "xla_temp_mb": 228.5007553100586,
      "xla_arg_mb": 48.5,
      "xla_out_mb": 16.0,
      "rel_err": 0.003008762022728728
    },
    {
      "impl": "ssd",
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "seq": 2048,
      "dk": 128,
      "dv": 128,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 1.2245948950003367,
      "inner": 6,
      "median_ms": 1.8604180004331283,
      "q1_ms": 1.813793166851004,
      "q3_ms": 1.9917226672987454,
      "xla_temp_mb": 132.51564025878906,
      "xla_arg_mb": 48.25,
      "xla_out_mb": 32.0
    },
    {
      "impl": "recurrent",
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "seq": 2048,
      "dk": 128,
      "dv": 128,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "error": "JaxRuntimeError: NOT_FOUND: All configs failed during profiling or were excluded from selection."
    },
    {
      "impl": "chunked",
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "seq": 2048,
      "dk": 128,
      "dv": 128,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 3.919557948000147,
      "inner": 5,
      "median_ms": 8.683903201017529,
      "q1_ms": 8.66865240095649,
      "q3_ms": 8.717893401626498,
      "xla_temp_mb": 498.7620315551758,
      "xla_arg_mb": 48.5,
      "xla_out_mb": 48.500038146972656,
      "residual_mb": 64.0,
      "rel_err": 0.0049261083743842365,
      "rel_err_grads": {
        "q": 0.0049261083743842365,
        "k": 0.0027472527472527475,
        "v": 0.004464285714285714,
        "log_alpha": 0.0018224315645679208,
        "beta": 0.0016595985363016428
      }
    },
    {
      "impl": "chunked_fat",
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "seq": 2048,
      "dk": 128,
      "dv": 128,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 2.7097517820075154,
      "inner": 6,
      "median_ms": 7.939155000106743,
      "q1_ms": 7.828213664955305,
      "q3_ms": 7.951197832881007,
      "xla_temp_mb": 482.5108108520508,
      "xla_arg_mb": 48.5,
      "xla_out_mb": 48.500038146972656,
      "residual_mb": 224.0,
      "rel_err": 0.0049261083743842365,
      "rel_err_grads": {
        "q": 0.0049261083743842365,
        "k": 0.0027472527472527475,
        "v": 0.004464285714285714,
        "log_alpha": 0.0018224315645679208,
        "beta": 0.0016595985363016428
      }
    },
    {
      "impl": "pallas",
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "seq": 2048,
      "dk": 128,
      "dv": 128,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 5.111209887007135,
      "inner": 9,
      "median_ms": 5.403605000335826,
      "q1_ms": 5.374695889056763,
      "q3_ms": 5.413606111283621,
      "xla_temp_mb": 496.75296783447266,
      "xla_arg_mb": 48.5,
      "xla_out_mb": 48.500038146972656,
      "residual_mb": 64.0,
      "rel_err": 0.0049261083743842365,
      "rel_err_grads": {
        "q": 0.0049261083743842365,
        "k": 0.0027472527472527475,
        "v": 0.004464285714285714,
        "log_alpha": 0.0019196672228953319,
        "beta": 0.0016280552871294383
      }
    },
    {
      "impl": "pallas_unfused",
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "seq": 2048,
      "dk": 128,
      "dv": 128,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 3.9020145849935943,
      "inner": 9,
      "median_ms": 5.3698166658351605,
      "q1_ms": 5.363892110633767,
      "q3_ms": 5.373229777129988,
      "xla_temp_mb": 496.75296783447266,
      "xla_arg_mb": 48.5,
      "xla_out_mb": 48.500038146972656,
      "residual_mb": 64.0,
      "rel_err": 0.0049261083743842365,
      "rel_err_grads": {
        "q": 0.0049261083743842365,
        "k": 0.0027472527472527475,
        "v": 0.004464285714285714,
        "log_alpha": 0.0019196672228953319,
        "beta": 0.0016280552871294383
      }
    },
    {
      "impl": "ssd",
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "seq": 2048,
      "dk": 128,
      "dv": 128,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 3.4160278490016935,
      "inner": 7,
      "median_ms": 6.167502855532803,
      "q1_ms": 6.13487399823498,
      "q3_ms": 6.2018520007508675,
      "xla_temp_mb": 288.2661437988281,
      "xla_arg_mb": 48.25,
      "xla_out_mb": 48.250030517578125,
      "residual_mb": 132.0
    },
    {
      "impl": "recurrent",
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "seq": 8192,
      "dk": 128,
      "dv": 128,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 0.36304995899263304,
      "inner": 1,
      "median_ms": 109.22125600336585,
      "q1_ms": 109.20726000040304,
      "q3_ms": 109.22157700406387,
      "xla_temp_mb": 514.0647583007812,
      "xla_arg_mb": 194.0,
      "xla_out_mb": 64.0,
      "rel_err": 0.003008762022728728,
      "peak_mb": 901.064697265625,
      "inputs_mb": 322.999755859375
    },
    {
      "impl": "chunked",
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "seq": 8192,
      "dk": 128,
      "dv": 128,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 0.9249058540008264,
      "inner": 4,
      "median_ms": 10.722173501562793,
      "q1_ms": 10.720430000219494,
      "q3_ms": 10.740550249465741,
      "xla_temp_mb": 710.5186233520508,
      "xla_arg_mb": 194.0,
      "xla_out_mb": 64.0,
      "rel_err": 0.003008762022728728,
      "peak_mb": 1223.518310546875,
      "inputs_mb": 386.99951171875
    },
    {
      "impl": "chunked_fat",
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "seq": 8192,
      "dk": 128,
      "dv": 128,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 0.288999544995022,
      "inner": 4,
      "median_ms": 10.722957249527099,
      "q1_ms": 10.72201325223432,
      "q3_ms": 10.730804500781232,
      "xla_temp_mb": 710.5186233520508,
      "xla_arg_mb": 194.0,
      "xla_out_mb": 64.0,
      "rel_err": 0.003008762022728728,
      "peak_mb": 1223.518310546875,
      "inputs_mb": 386.99951171875
    },
    {
      "impl": "pallas",
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "seq": 8192,
      "dk": 128,
      "dv": 128,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 2.7255404590105172,
      "inner": 8,
      "median_ms": 5.862349875314976,
      "q1_ms": 5.856791874975897,
      "q3_ms": 5.865047500265064,
      "xla_temp_mb": 902.0007553100586,
      "xla_arg_mb": 194.0,
      "xla_out_mb": 64.0,
      "rel_err": 0.003008762022728728,
      "peak_mb": 1410.999755859375,
      "inputs_mb": 322.999755859375
    },
    {
      "impl": "pallas_unfused",
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "seq": 8192,
      "dk": 128,
      "dv": 128,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 2.409292827011086,
      "inner": 8,
      "median_ms": 5.851708625414176,
      "q1_ms": 5.8225171251251595,
      "q3_ms": 5.854366125277011,
      "xla_temp_mb": 902.0007553100586,
      "xla_arg_mb": 194.0,
      "xla_out_mb": 64.0,
      "rel_err": 0.003008762022728728,
      "peak_mb": 1410.999755859375,
      "inputs_mb": 322.999755859375
    },
    {
      "impl": "ssd",
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "seq": 8192,
      "dk": 128,
      "dv": 128,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 1.1730863760021748,
      "inner": 6,
      "median_ms": 7.0886866669752635,
      "q1_ms": 7.0814779998424155,
      "q3_ms": 7.109073500032537,
      "xla_temp_mb": 518.0156402587891,
      "xla_arg_mb": 193.0,
      "xla_out_mb": 128.0,
      "peak_mb": 1042.99951171875,
      "inputs_mb": 386.99951171875
    },
    {
      "impl": "recurrent",
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "seq": 8192,
      "dk": 128,
      "dv": 128,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "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": "chunked",
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "seq": 8192,
      "dk": 128,
      "dv": 128,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 4.988440075001563,
      "inner": 1,
      "median_ms": 33.054815998184495,
      "q1_ms": 33.0524820019491,
      "q3_ms": 33.101774999522604,
      "xla_temp_mb": 1798.0078811645508,
      "xla_arg_mb": 194.0,
      "xla_out_mb": 194.00003814697266,
      "residual_mb": 256.0,
      "rel_err": 0.004464285714285714,
      "rel_err_grads": {
        "q": 0.004405286343612335,
        "k": 0.0027472527472527475,
        "v": 0.004464285714285714,
        "log_alpha": 0.0021078481114667935,
        "beta": 0.0018873565299928812
      },
      "peak_mb": 2315.00830078125,
      "inputs_mb": 322.999755859375
    },
    {
      "impl": "chunked_fat",
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "seq": 8192,
      "dk": 128,
      "dv": 128,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 3.8479564909939654,
      "inner": 1,
      "median_ms": 30.355226001120172,
      "q1_ms": 30.348412998137064,
      "q3_ms": 30.564871995011345,
      "xla_temp_mb": 1545.0230178833008,
      "xla_arg_mb": 194.0,
      "xla_out_mb": 194.00003814697266,
      "residual_mb": 896.0,
      "rel_err": 0.004464285714285714,
      "rel_err_grads": {
        "q": 0.004405286343612335,
        "k": 0.0027472527472527475,
        "v": 0.004464285714285714,
        "log_alpha": 0.0021078481114667935,
        "beta": 0.0018873565299928812
      },
      "peak_mb": 2188.02294921875,
      "inputs_mb": 386.99951171875
    },
    {
      "impl": "pallas",
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "seq": 8192,
      "dk": 128,
      "dv": 128,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 5.119308908004314,
      "inner": 2,
      "median_ms": 20.533691502350848,
      "q1_ms": 20.46710050490219,
      "q3_ms": 20.553859998472035,
      "xla_temp_mb": 1802.0185928344727,
      "xla_arg_mb": 194.0,
      "xla_out_mb": 194.00003814697266,
      "residual_mb": 256.0,
      "rel_err": 0.004464285714285714,
      "rel_err_grads": {
        "q": 0.004405286343612335,
        "k": 0.0027472527472527475,
        "v": 0.004464285714285714,
        "log_alpha": 0.0020855293935204927,
        "beta": 0.001895734720470261
      },
      "peak_mb": 2319.01904296875,
      "inputs_mb": 322.999755859375
    },
    {
      "impl": "pallas_unfused",
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "seq": 8192,
      "dk": 128,
      "dv": 128,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 4.7286836550047155,
      "inner": 2,
      "median_ms": 20.570556000166107,
      "q1_ms": 20.472666001296602,
      "q3_ms": 20.572640001773834,
      "xla_temp_mb": 1802.0185928344727,
      "xla_arg_mb": 194.0,
      "xla_out_mb": 194.00003814697266,
      "residual_mb": 256.0,
      "rel_err": 0.004464285714285714,
      "rel_err_grads": {
        "q": 0.004405286343612335,
        "k": 0.0027472527472527475,
        "v": 0.004464285714285714,
        "log_alpha": 0.0020855293935204927,
        "beta": 0.001895734720470261
      },
      "peak_mb": 2319.01904296875,
      "inputs_mb": 322.999755859375
    },
    {
      "impl": "ssd",
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "seq": 8192,
      "dk": 128,
      "dv": 128,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 4.561912848002976,
      "inner": 2,
      "median_ms": 24.151932004315313,
      "q1_ms": 24.03991000028327,
      "q3_ms": 24.265691499749664,
      "xla_temp_mb": 1153.0161437988281,
      "xla_arg_mb": 193.0,
      "xla_out_mb": 193.00003051757812,
      "residual_mb": 516.0,
      "peak_mb": 1669.0166015625,
      "inputs_mb": 322.999755859375
    },
    {
      "impl": "recurrent",
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "seq": 2048,
      "dk": 64,
      "dv": 64,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 0.2056194280012278,
      "inner": 2,
      "median_ms": 16.979221501969732,
      "q1_ms": 16.96941800037166,
      "q3_ms": 16.99564999944414,
      "xla_temp_mb": 49.01666259765625,
      "xla_arg_mb": 48.5,
      "xla_out_mb": 16.0,
      "rel_err": 4.118384782667069e-07
    },
    {
      "impl": "chunked",
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "seq": 2048,
      "dk": 64,
      "dv": 64,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 1.1225365999998758,
      "inner": 26,
      "median_ms": 1.7338536154756394,
      "q1_ms": 1.7297919036466807,
      "q3_ms": 1.7340069806288319,
      "xla_temp_mb": 112.26862335205078,
      "xla_arg_mb": 48.5,
      "xla_out_mb": 16.0,
      "rel_err": 0.00075012260431498
    },
    {
      "impl": "chunked_fat",
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "seq": 2048,
      "dk": 64,
      "dv": 64,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 0.5011378189956304,
      "inner": 26,
      "median_ms": 1.7318124228040688,
      "q1_ms": 1.7278046538054836,
      "q3_ms": 1.7431619612929912,
      "xla_temp_mb": 112.26862335205078,
      "xla_arg_mb": 48.5,
      "xla_out_mb": 16.0,
      "rel_err": 0.00075012260431498
    },
    {
      "impl": "pallas_unfused",
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "seq": 2048,
      "dk": 64,
      "dv": 64,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 3.2273307309951633,
      "inner": 55,
      "median_ms": 0.7884728362445127,
      "q1_ms": 0.7845643543987535,
      "q3_ms": 0.806423372705467,
      "xla_temp_mb": 112.25076293945312,
      "xla_arg_mb": 48.5,
      "xla_out_mb": 16.0,
      "rel_err": 0.0007593066023803275
    },
    {
      "impl": "ssd",
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "seq": 2048,
      "dk": 64,
      "dv": 64,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 1.1152649029972963,
      "inner": 100,
      "median_ms": 0.3864710099878721,
      "q1_ms": 0.3859355949680321,
      "q3_ms": 0.3875377250369638,
      "xla_temp_mb": 1.5156402587890625,
      "xla_arg_mb": 48.25,
      "xla_out_mb": 16.0
    },
    {
      "impl": "recurrent",
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "seq": 2048,
      "dk": 64,
      "dv": 64,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 1.0138614060124382,
      "inner": 1,
      "median_ms": 94.17849800956901,
      "q1_ms": 91.72184200724587,
      "q3_ms": 94.85872300137999,
      "xla_temp_mb": 4161.257453918457,
      "xla_arg_mb": 48.5,
      "xla_out_mb": 48.500038146972656,
      "residual_mb": 3136.75,
      "rel_err": 5.304251888313672e-07,
      "rel_err_grads": {
        "q": 3.547636599973464e-07,
        "k": 4.71851854560631e-07,
        "v": 5.304251888313672e-07,
        "log_alpha": 5.065995226414018e-07,
        "beta": 3.745361179089841e-07
      }
    },
    {
      "impl": "chunked",
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "seq": 2048,
      "dk": 64,
      "dv": 64,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 2.0593828609999036,
      "inner": 8,
      "median_ms": 5.57560324887163,
      "q1_ms": 5.522693249986332,
      "q3_ms": 5.57887562445103,
      "xla_temp_mb": 240.25861358642578,
      "xla_arg_mb": 48.5,
      "xla_out_mb": 48.500038146972656,
      "residual_mb": 16.0,
      "rel_err": 0.0005990775671821856,
      "rel_err_grads": {
        "q": 0.0005990775671821856,
        "k": 0.00048679628751679353,
        "v": 0.0005034324403330599,
        "log_alpha": 0.0004555596207352806,
        "beta": 0.0005359688283219993
      }
    },
    {
      "impl": "chunked_fat",
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "seq": 2048,
      "dk": 64,
      "dv": 64,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 1.0179480149963638,
      "inner": 10,
      "median_ms": 4.685077800240833,
      "q1_ms": 4.674416699708672,
      "q3_ms": 4.695974400237901,
      "xla_temp_mb": 240.50690460205078,
      "xla_arg_mb": 48.5,
      "xla_out_mb": 48.500038146972656,
      "residual_mb": 112.0,
      "rel_err": 0.0005990775671821856,
      "rel_err_grads": {
        "q": 0.0005990775671821856,
        "k": 0.00048679628751679353,
        "v": 0.0005034324403330599,
        "log_alpha": 0.0004555849507114126,
        "beta": 0.0005359688283219993
      }
    },
    {
      "impl": "pallas_unfused",
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "seq": 2048,
      "dk": 64,
      "dv": 64,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 3.545153337996453,
      "inner": 14,
      "median_ms": 3.4281584999657104,
      "q1_ms": 3.3749651426790352,
      "q3_ms": 3.4916216074114867,
      "xla_temp_mb": 257.25296783447266,
      "xla_arg_mb": 48.5,
      "xla_out_mb": 48.500038146972656,
      "residual_mb": 16.0,
      "rel_err": 0.0005966336397466483,
      "rel_err_grads": {
        "q": 0.0005966336397466483,
        "k": 0.0005882223181916466,
        "v": 0.0005102690316557753,
        "log_alpha": 0.000433674521357172,
        "beta": 0.0005288144241105134
      }
    },
    {
      "impl": "ssd",
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "seq": 2048,
      "dk": 64,
      "dv": 64,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 3.5962700259988196,
      "inner": 26,
      "median_ms": 1.6343544999612136,
      "q1_ms": 1.6165843847445034,
      "q3_ms": 1.6405082309435013,
      "xla_temp_mb": 4.765899658203125,
      "xla_arg_mb": 48.25,
      "xla_out_mb": 48.250030517578125,
      "residual_mb": 65.0
    },
    {
      "impl": "recurrent",
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "seq": 8192,
      "dk": 64,
      "dv": 64,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 0.3703455920040142,
      "inner": 1,
      "median_ms": 73.82928500010166,
      "q1_ms": 71.12137950025499,
      "q3_ms": 77.41359150531935,
      "xla_temp_mb": 194.51666259765625,
      "xla_arg_mb": 194.0,
      "xla_out_mb": 64.0,
      "rel_err": 4.4253968806557493e-07,
      "peak_mb": 530.99951171875,
      "inputs_mb": 258.99951171875
    },
    {
      "impl": "chunked",
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "seq": 8192,
      "dk": 64,
      "dv": 64,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 0.8380801159946714,
      "inner": 7,
      "median_ms": 6.9970084296073765,
      "q1_ms": 6.9668428581540605,
      "q3_ms": 7.192116357535789,
      "xla_temp_mb": 324.0186233520508,
      "xla_arg_mb": 194.0,
      "xla_out_mb": 64.0,
      "rel_err": 0.0006716646115615263,
      "peak_mb": 678.999755859375,
      "inputs_mb": 258.99951171875
    },
    {
      "impl": "chunked_fat",
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "seq": 8192,
      "dk": 64,
      "dv": 64,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 0.5477012250048574,
      "inner": 6,
      "median_ms": 7.222724833506315,
      "q1_ms": 7.130380416735231,
      "q3_ms": 7.295377916307189,
      "xla_temp_mb": 324.0186233520508,
      "xla_arg_mb": 194.0,
      "xla_out_mb": 64.0,
      "rel_err": 0.0006716646115615263,
      "peak_mb": 740.999755859375,
      "inputs_mb": 322.99951171875
    },
    {
      "impl": "pallas_unfused",
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "seq": 8192,
      "dk": 64,
      "dv": 64,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 2.8316427629906684,
      "inner": 16,
      "median_ms": 2.764918374850822,
      "q1_ms": 2.7646723124234995,
      "q3_ms": 2.77204431267819,
      "xla_temp_mb": 386.0007629394531,
      "xla_arg_mb": 194.0,
      "xla_out_mb": 64.0,
      "rel_err": 0.0006798884740980783,
      "peak_mb": 834.99951171875,
      "inputs_mb": 258.99951171875
    },
    {
      "impl": "ssd",
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "seq": 8192,
      "dk": 64,
      "dv": 64,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 2.7375307509937556,
      "inner": 25,
      "median_ms": 1.5502803202252835,
      "q1_ms": 1.5395723000983708,
      "q3_ms": 1.550437820260413,
      "xla_temp_mb": 3.0156402587890625,
      "xla_arg_mb": 193.0,
      "xla_out_mb": 64.0,
      "peak_mb": 326.015380859375,
      "inputs_mb": 258.99951171875
    },
    {
      "impl": "recurrent",
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "seq": 8192,
      "dk": 64,
      "dv": 64,
      "dtype": "float32",
      "chunk_size": 64,
      "error": "JaxRuntimeError: NOT_FOUND: All configs failed during profiling or were excluded from selection."
    },
    {
      "impl": "chunked",
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "seq": 8192,
      "dk": 64,
      "dv": 64,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 3.696694533995469,
      "inner": 1,
      "median_ms": 25.9293250128394,
      "q1_ms": 25.78519300732296,
      "q3_ms": 26.215871002932545,
      "xla_temp_mb": 774.0081253051758,
      "xla_arg_mb": 194.0,
      "xla_out_mb": 194.00003814697266,
      "residual_mb": 64.0,
      "rel_err": 0.000694621334717311,
      "rel_err_grads": {
        "q": 0.000592237339302774,
        "k": 0.0004392150695031213,
        "v": 0.000694621334717311,
        "log_alpha": 0.00043262193182458034,
        "beta": 0.0004491625678382584
      },
      "peak_mb": 1227.008056640625,
      "inputs_mb": 258.99951171875
    },
    {
      "impl": "chunked_fat",
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "seq": 8192,
      "dk": 64,
      "dv": 64,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 4.47666955999739,
      "inner": 2,
      "median_ms": 17.96387149806833,
      "q1_ms": 17.913686748215696,
      "q3_ms": 18.45923749715439,
      "xla_temp_mb": 709.0232620239258,
      "xla_arg_mb": 194.0,
      "xla_out_mb": 194.00003814697266,
      "residual_mb": 448.0,
      "rel_err": 0.000694621334717311,
      "rel_err_grads": {
        "q": 0.000592237339302774,
        "k": 0.0004392150695031213,
        "v": 0.000694621334717311,
        "log_alpha": 0.00043262193182458034,
        "beta": 0.0004491625678382584
      },
      "peak_mb": 1162.023193359375,
      "inputs_mb": 258.99951171875
    },
    {
      "impl": "pallas_unfused",
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "seq": 8192,
      "dk": 64,
      "dv": 64,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 3.4127626400004374,
      "inner": 4,
      "median_ms": 11.67107350192964,
      "q1_ms": 11.659420377327478,
      "q3_ms": 11.685431874866481,
      "xla_temp_mb": 901.0183486938477,
      "xla_arg_mb": 194.0,
      "xla_out_mb": 194.00003814697266,
      "residual_mb": 64.0,
      "rel_err": 0.000694621334717311,
      "rel_err_grads": {
        "q": 0.0005298857122862197,
        "k": 0.00043282303055203116,
        "v": 0.000694621334717311,
        "log_alpha": 0.00045774419210714226,
        "beta": 0.00044326804595061725
      },
      "peak_mb": 1602.999755859375,
      "inputs_mb": 322.99951171875
    },
    {
      "impl": "ssd",
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "seq": 8192,
      "dk": 64,
      "dv": 64,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 3.5387873120052973,
      "inner": 8,
      "median_ms": 5.852231999597279,
      "q1_ms": 5.8479420622461475,
      "q3_ms": 5.855794999661157,
      "xla_temp_mb": 17.515899658203125,
      "xla_arg_mb": 193.0,
      "xla_out_mb": 193.00003051757812,
      "residual_mb": 257.0,
      "peak_mb": 628.99951171875,
      "inputs_mb": 258.99951171875
    },
    {
      "impl": "recurrent",
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "seq": 2048,
      "dk": 128,
      "dv": 128,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 0.21786299798986875,
      "inner": 1,
      "median_ms": 31.043545997818,
      "q1_ms": 31.042869995872024,
      "q3_ms": 31.054932995175477,
      "xla_temp_mb": 98.56475830078125,
      "xla_arg_mb": 96.5,
      "xla_out_mb": 32.0,
      "rel_err": 5.202623749372019e-07
    },
    {
      "impl": "chunked",
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "seq": 2048,
      "dk": 128,
      "dv": 128,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 1.2498875059973216,
      "inner": 17,
      "median_ms": 2.792395294390564,
      "q1_ms": 2.7850505297572132,
      "q3_ms": 2.7942367061554894,
      "xla_temp_mb": 182.50690460205078,
      "xla_arg_mb": 96.5,
      "xla_out_mb": 32.0,
      "rel_err": 0.0006296120668330938
    },
    {
      "impl": "chunked_fat",
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "seq": 2048,
      "dk": 128,
      "dv": 128,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 0.35296564899908844,
      "inner": 17,
      "median_ms": 2.788965294123901,
      "q1_ms": 2.7850802939878228,
      "q3_ms": 2.791124088295927,
      "xla_temp_mb": 182.50690460205078,
      "xla_arg_mb": 96.5,
      "xla_out_mb": 32.0,
      "rel_err": 0.0006296120668330938
    },
    {
      "impl": "pallas",
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "seq": 2048,
      "dk": 128,
      "dv": 128,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 2.7455739860015456,
      "inner": 28,
      "median_ms": 1.6129072859517433,
      "q1_ms": 1.6127317856963992,
      "q3_ms": 1.6136819823649213,
      "xla_temp_mb": 212.5007553100586,
      "xla_arg_mb": 96.5,
      "xla_out_mb": 32.0,
      "rel_err": 0.000630889074480667
    },
    {
      "impl": "pallas_unfused",
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "seq": 2048,
      "dk": 128,
      "dv": 128,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 2.282383123994805,
      "inner": 28,
      "median_ms": 1.613263678466735,
      "q1_ms": 1.6128126608236926,
      "q3_ms": 1.61327280329715,
      "xla_temp_mb": 212.5007553100586,
      "xla_arg_mb": 96.5,
      "xla_out_mb": 32.0,
      "rel_err": 0.000630889074480667
    },
    {
      "impl": "ssd",
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "seq": 2048,
      "dk": 128,
      "dv": 128,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 1.0612177449947922,
      "inner": 28,
      "median_ms": 1.5800251785549335,
      "q1_ms": 1.5794598213168294,
      "q3_ms": 1.5811625180504052,
      "xla_temp_mb": 66.51564025878906,
      "xla_arg_mb": 96.25,
      "xla_out_mb": 32.0
    },
    {
      "impl": "recurrent",
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "seq": 2048,
      "dk": 128,
      "dv": 128,
      "dtype": "float32",
      "chunk_size": 64,
      "error": "JaxRuntimeError: NOT_FOUND: All configs failed during profiling or were excluded from selection."
    },
    {
      "impl": "chunked",
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "seq": 2048,
      "dk": 128,
      "dv": 128,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 3.168335644993931,
      "inner": 5,
      "median_ms": 8.865429801517166,
      "q1_ms": 8.831807901151478,
      "q3_ms": 8.88892050134018,
      "xla_temp_mb": 434.7620315551758,
      "xla_arg_mb": 96.5,
      "xla_out_mb": 96.50003814697266,
      "residual_mb": 64.0,
      "rel_err": 0.0005915486972441561,
      "rel_err_grads": {
        "q": 0.0005915486972441561,
        "k": 0.00034321609577858453,
        "v": 0.00046780018082676577,
        "log_alpha": 0.0005155089417596588,
        "beta": 0.00046064108966604637
      }
    },
    {
      "impl": "chunked_fat",
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "seq": 2048,
      "dk": 128,
      "dv": 128,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 1.971819358994253,
      "inner": 6,
      "median_ms": 7.987525998032652,
      "q1_ms": 7.982408082170878,
      "q3_ms": 8.007868081525277,
      "xla_temp_mb": 418.5066604614258,
      "xla_arg_mb": 96.5,
      "xla_out_mb": 96.50003814697266,
      "residual_mb": 224.0,
      "rel_err": 0.0005915486972441561,
      "rel_err_grads": {
        "q": 0.0005915486972441561,
        "k": 0.00034321609577858453,
        "v": 0.00046780018082676577,
        "log_alpha": 0.0005155089417596588,
        "beta": 0.00046064108966604637
      }
    },
    {
      "impl": "pallas",
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "seq": 2048,
      "dk": 128,
      "dv": 128,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 3.7066287159977946,
      "inner": 8,
      "median_ms": 5.799031625429052,
      "q1_ms": 5.795323437268962,
      "q3_ms": 5.802423687782721,
      "xla_temp_mb": 421.25296783447266,
      "xla_arg_mb": 96.5,
      "xla_out_mb": 96.50003814697266,
      "residual_mb": 64.0,
      "rel_err": 0.0005988573488341043,
      "rel_err_grads": {
        "q": 0.0005988573488341043,
        "k": 0.0003500276023203554,
        "v": 0.00046780018082676577,
        "log_alpha": 0.0005589088377562575,
        "beta": 0.00043639832820543115
      }
    },
    {
      "impl": "pallas_unfused",
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "seq": 2048,
      "dk": 128,
      "dv": 128,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 2.878156188991852,
      "inner": 8,
      "median_ms": 5.80758062460518,
      "q1_ms": 5.800697499580565,
      "q3_ms": 5.816002187202685,
      "xla_temp_mb": 421.25296783447266,
      "xla_arg_mb": 96.5,
      "xla_out_mb": 96.50003814697266,
      "residual_mb": 64.0,
      "rel_err": 0.0005988573488341043,
      "rel_err_grads": {
        "q": 0.0005988573488341043,
        "k": 0.0003500276023203554,
        "v": 0.00046780018082676577,
        "log_alpha": 0.0005589088377562575,
        "beta": 0.00043639832820543115
      }
    },
    {
      "impl": "ssd",
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "seq": 2048,
      "dk": 128,
      "dv": 128,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 3.2293930089945206,
      "inner": 7,
      "median_ms": 6.108969572648805,
      "q1_ms": 6.097124714870006,
      "q3_ms": 6.109396071613966,
      "xla_temp_mb": 192.26614379882812,
      "xla_arg_mb": 96.25,
      "xla_out_mb": 96.25003051757812,
      "residual_mb": 132.0
    },
    {
      "impl": "recurrent",
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "seq": 8192,
      "dk": 128,
      "dv": 128,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 0.3611300500051584,
      "inner": 1,
      "median_ms": 108.90974999347236,
      "q1_ms": 108.88129749946529,
      "q3_ms": 109.00232199492166,
      "xla_temp_mb": 388.06475830078125,
      "xla_arg_mb": 386.0,
      "xla_out_mb": 128.0,
      "rel_err": 5.202624241504682e-07,
      "peak_mb": 1042.99951171875,
      "inputs_mb": 514.99951171875
    },
    {
      "impl": "chunked",
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "seq": 8192,
      "dk": 128,
      "dv": 128,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 0.8086644430004526,
      "inner": 4,
      "median_ms": 11.269664497376652,
      "q1_ms": 11.260525872785365,
      "q3_ms": 11.2946719982574,
      "xla_temp_mb": 585.5186233520508,
      "xla_arg_mb": 386.0,
      "xla_out_mb": 128.0,
      "rel_err": 0.0006810708097969765,
      "peak_mb": 1252.999755859375,
      "inputs_mb": 514.99951171875
    },
    {
      "impl": "chunked_fat",
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "seq": 8192,
      "dk": 128,
      "dv": 128,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 0.27161346400680486,
      "inner": 4,
      "median_ms": 11.326434749207692,
      "q1_ms": 11.291787624941207,
      "q3_ms": 11.34773899866559,
      "xla_temp_mb": 585.5186233520508,
      "xla_arg_mb": 386.0,
      "xla_out_mb": 128.0,
      "rel_err": 0.0006810708097969765,
      "peak_mb": 1252.999755859375,
      "inputs_mb": 514.99951171875
    },
    {
      "impl": "pallas",
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "seq": 8192,
      "dk": 128,
      "dv": 128,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 2.5788371140079107,
      "inner": 8,
      "median_ms": 6.102772749727592,
      "q1_ms": 6.098464000388049,
      "q3_ms": 6.103097750383313,
      "xla_temp_mb": 838.0007553100586,
      "xla_arg_mb": 386.0,
      "xla_out_mb": 128.0,
      "rel_err": 0.0006848545365180709,
      "peak_mb": 1481.00048828125,
      "inputs_mb": 514.99951171875
    },
    {
      "impl": "pallas_unfused",
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "seq": 8192,
      "dk": 128,
      "dv": 128,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 2.264276143992902,
      "inner": 8,
      "median_ms": 6.105817124989699,
      "q1_ms": 6.084496000767103,
      "q3_ms": 6.125978000454779,
      "xla_temp_mb": 838.0007553100586,
      "xla_arg_mb": 386.0,
      "xla_out_mb": 128.0,
      "rel_err": 0.0006848545365180709,
      "peak_mb": 1481.00048828125,
      "inputs_mb": 514.99951171875
    },
    {
      "impl": "ssd",
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "seq": 8192,
      "dk": 128,
      "dv": 128,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 1.0726006780023454,
      "inner": 7,
      "median_ms": 6.304382286283986,
      "q1_ms": 6.300462000321464,
      "q3_ms": 6.30800271444189,
      "xla_temp_mb": 260.01564025878906,
      "xla_arg_mb": 385.0,
      "xla_out_mb": 128.0,
      "peak_mb": 1170.99951171875,
      "inputs_mb": 514.99951171875
    },
    {
      "impl": "recurrent",
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "seq": 8192,
      "dk": 128,
      "dv": 128,
      "dtype": "float32",
      "chunk_size": 64,
      "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": "chunked",
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "seq": 8192,
      "dk": 128,
      "dv": 128,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 4.156336589003331,
      "inner": 1,
      "median_ms": 34.25731101015117,
      "q1_ms": 34.243821501149796,
      "q3_ms": 34.27818100317381,
      "xla_temp_mb": 1606.0078811645508,
      "xla_arg_mb": 386.0,
      "xla_out_mb": 386.00003814697266,
      "residual_mb": 256.0,
      "rel_err": 0.0006009233952694608,
      "rel_err_grads": {
        "q": 0.0006009233952694608,
        "k": 0.0005195608476382679,
        "v": 0.0005217574352140946,
        "log_alpha": 0.00051454467210003,
        "beta": 0.0004192740065433316
      },
      "peak_mb": 2761.0078125,
      "inputs_mb": 642.99951171875
    },
    {
      "impl": "chunked_fat",
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "seq": 8192,
      "dk": 128,
      "dv": 128,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 3.1619864199892618,
      "inner": 1,
      "median_ms": 31.12076100660488,
      "q1_ms": 31.114971498027444,
      "q3_ms": 31.127358502999414,
      "xla_temp_mb": 1417.0230178833008,
      "xla_arg_mb": 386.0,
      "xla_out_mb": 386.00003814697266,
      "residual_mb": 896.0,
      "rel_err": 0.0006009233952694608,
      "rel_err_grads": {
        "q": 0.0006009233952694608,
        "k": 0.0005195608476382679,
        "v": 0.0005217574352140946,
        "log_alpha": 0.00051454467210003,
        "beta": 0.0004192740065433316
      },
      "peak_mb": 2318.02294921875,
      "inputs_mb": 514.99951171875
    },
    {
      "impl": "pallas",
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "seq": 8192,
      "dk": 128,
      "dv": 128,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 4.3199507639947115,
      "inner": 2,
      "median_ms": 22.072126994316932,
      "q1_ms": 22.04513974720612,
      "q3_ms": 22.072469495469704,
      "xla_temp_mb": 1674.0185928344727,
      "xla_arg_mb": 386.0,
      "xla_out_mb": 386.00003814697266,
      "residual_mb": 256.0,
      "rel_err": 0.0005433846115488683,
      "rel_err_grads": {
        "q": 0.0005433846115488683,
        "k": 0.0004468768210303812,
        "v": 0.0005247607485512767,
        "log_alpha": 0.0005198616356360746,
        "beta": 0.0003972096956677676
      },
      "peak_mb": 2575.0185546875,
      "inputs_mb": 514.99951171875
    },
    {
      "impl": "pallas_unfused",
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "seq": 8192,
      "dk": 128,
      "dv": 128,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 4.384038342992426,
      "inner": 2,
      "median_ms": 22.13559149822686,
      "q1_ms": 22.094012001616647,
      "q3_ms": 22.175968249939615,
      "xla_temp_mb": 1674.0185928344727,
      "xla_arg_mb": 386.0,
      "xla_out_mb": 386.00003814697266,
      "residual_mb": 256.0,
      "rel_err": 0.0005433846115488683,
      "rel_err_grads": {
        "q": 0.0005433846115488683,
        "k": 0.0004468768210303812,
        "v": 0.0005247607485512767,
        "log_alpha": 0.0005198616356360746,
        "beta": 0.0003972096956677676
      },
      "peak_mb": 2829.0185546875,
      "inputs_mb": 642.99951171875
    },
    {
      "impl": "ssd",
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "seq": 8192,
      "dk": 128,
      "dv": 128,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 4.483348562003812,
      "inner": 2,
      "median_ms": 24.11245849361876,
      "q1_ms": 24.093083997286158,
      "q3_ms": 24.14131824480137,
      "xla_temp_mb": 769.0161437988281,
      "xla_arg_mb": 385.0,
      "xla_out_mb": 385.0000305175781,
      "residual_mb": 516.0,
      "peak_mb": 1669.01611328125,
      "inputs_mb": 514.99951171875
    },
    {
      "impl": "chunked",
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "seq": 8192,
      "dk": 64,
      "dv": 64,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 2.9630880920012714,
      "inner": 2,
      "median_ms": 19.32677099830471,
      "q1_ms": 18.952538499434013,
      "q3_ms": 19.60665600199718,
      "xla_temp_mb": 900.0081253051758,
      "xla_arg_mb": 98.0,
      "xla_out_mb": 98.00003814697266,
      "residual_mb": 64.0,
      "rel_err": 0.0045045045045045045,
      "rel_err_grads": {
        "q": 0.0045045045045045045,
        "k": 0.003289473684210526,
        "v": 0.003875968992248062,
        "log_alpha": 0.0016641917127602447,
        "beta": 0.0016420397444260053
      },
      "peak_mb": 1346.999755859375,
      "inputs_mb": 194.99951171875
    },
    {
      "impl": "chunked_fat",
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "seq": 8192,
      "dk": 64,
      "dv": 64,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 2.3510819829971297,
      "inner": 2,
      "median_ms": 17.306565001490526,
      "q1_ms": 17.054628500773106,
      "q3_ms": 18.11594099854119,
      "xla_temp_mb": 835.5232620239258,
      "xla_arg_mb": 98.0,
      "xla_out_mb": 98.00003814697266,
      "residual_mb": 448.0,
      "rel_err": 0.0045045045045045045,
      "rel_err_grads": {
        "q": 0.0045045045045045045,
        "k": 0.003289473684210526,
        "v": 0.003875968992248062,
        "log_alpha": 0.0016641917127602447,
        "beta": 0.0016420397444260053
      },
      "peak_mb": 1158.523193359375,
      "inputs_mb": 194.99951171875
    },
    {
      "impl": "pallas",
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "seq": 8192,
      "dk": 64,
      "dv": 64,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 3.986390100995777,
      "inner": 4,
      "median_ms": 10.832887750439113,
      "q1_ms": 10.744067250925582,
      "q3_ms": 10.879558500164421,
      "xla_temp_mb": 705.0017471313477,
      "xla_arg_mb": 98.0,
      "xla_out_mb": 98.00003814697266,
      "residual_mb": 64.0,
      "rel_err": 0.0045045045045045045,
      "rel_err_grads": {
        "q": 0.0045045045045045045,
        "k": 0.003289473684210526,
        "v": 0.003875968992248062,
        "log_alpha": 0.0016237241618838171,
        "beta": 0.0016714525102892136
      },
      "peak_mb": 1028.001708984375,
      "inputs_mb": 194.99951171875
    },
    {
      "impl": "pallas_unfused",
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "seq": 8192,
      "dk": 64,
      "dv": 64,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 5.3058095409942325,
      "inner": 4,
      "median_ms": 11.085317499237135,
      "q1_ms": 10.968270999001106,
      "q3_ms": 11.356753500876948,
      "xla_temp_mb": 1091.5183486938477,
      "xla_arg_mb": 98.0,
      "xla_out_mb": 98.00003814697266,
      "residual_mb": 64.0,
      "rel_err": 0.0045045045045045045,
      "rel_err_grads": {
        "q": 0.0045045045045045045,
        "k": 0.003289473684210526,
        "v": 0.003875968992248062,
        "log_alpha": 0.001623641151523045,
        "beta": 0.0016709911335697906
      },
      "peak_mb": 1414.518310546875,
      "inputs_mb": 194.99951171875
    },
    {
      "impl": "ssd",
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "seq": 8192,
      "dk": 64,
      "dv": 64,
      "dtype": "bfloat16",
      "chunk_size": 64,
      "compile_s": 4.47720797501097,
      "inner": 7,
      "median_ms": 7.209459571250981,
      "q1_ms": 7.1883509995781685,
      "q3_ms": 7.28273085717644,
      "xla_temp_mb": 384.0161437988281,
      "xla_arg_mb": 97.0,
      "xla_out_mb": 97.00003051757812,
      "residual_mb": 257.0,
      "peak_mb": 833.999755859375,
      "inputs_mb": 194.99951171875
    },
    {
      "impl": "pallas",
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "seq": 2048,
      "dk": 64,
      "dv": 64,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 2.568447887999355,
      "inner": 54,
      "median_ms": 0.7712878704532395,
      "q1_ms": 0.7672494074439682,
      "q3_ms": 0.7715121944245227,
      "xla_temp_mb": 112.25076293945312,
      "xla_arg_mb": 48.5,
      "xla_out_mb": 16.0,
      "rel_err": 0.0007593066023803275
    },
    {
      "impl": "pallas",
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "seq": 2048,
      "dk": 64,
      "dv": 64,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 3.3412853740010178,
      "inner": 14,
      "median_ms": 3.334809785883408,
      "q1_ms": 3.33191503575238,
      "q3_ms": 3.3349468571291903,
      "xla_temp_mb": 257.25296783447266,
      "xla_arg_mb": 48.5,
      "xla_out_mb": 48.500038146972656,
      "residual_mb": 16.0,
      "rel_err": 0.0005966336397466483,
      "rel_err_grads": {
        "q": 0.0005966336397466483,
        "k": 0.0005882223181916466,
        "v": 0.0005102690316557753,
        "log_alpha": 0.000433674521357172,
        "beta": 0.0005288144241105134
      }
    },
    {
      "impl": "pallas",
      "mode": "fwd",
      "batch": 4,
      "heads": 8,
      "seq": 8192,
      "dk": 64,
      "dv": 64,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 2.6029828140017344,
      "inner": 17,
      "median_ms": 2.767113294205902,
      "q1_ms": 2.7633332061026628,
      "q3_ms": 2.767928058935521,
      "xla_temp_mb": 386.0007629394531,
      "xla_arg_mb": 194.0,
      "xla_out_mb": 64.0,
      "rel_err": 0.0006799253524054171,
      "peak_mb": 898.99951171875,
      "inputs_mb": 322.99951171875
    },
    {
      "impl": "pallas",
      "mode": "fwd_bwd",
      "batch": 4,
      "heads": 8,
      "seq": 8192,
      "dk": 64,
      "dv": 64,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 3.734180957006174,
      "inner": 4,
      "median_ms": 11.64855000024545,
      "q1_ms": 11.638263125860249,
      "q3_ms": 11.655474500003038,
      "xla_temp_mb": 901.0183486938477,
      "xla_arg_mb": 194.0,
      "xla_out_mb": 194.00003814697266,
      "residual_mb": 64.0,
      "rel_err": 0.0006946212528407215,
      "rel_err_grads": {
        "q": 0.0005298513398347221,
        "k": 0.00043286076488601864,
        "v": 0.0006946212528407215,
        "log_alpha": 0.00045872001544043184,
        "beta": 0.00044329367430665043
      },
      "peak_mb": 1602.999755859375,
      "inputs_mb": 322.99951171875
    }
  ]
}