{
  "device": {
    "jax": "0.11.2",
    "backend": "cpu",
    "devices": [
      "cpu:0"
    ],
    "python": "3.14.8",
    "device_kind": "cpu",
    "gpu": "NVIDIA GeForce RTX 3090",
    "driver": "615.78.08",
    "gpu_memory_gb": 24.0,
    "torch": "2.14.1+cu130",
    "torch_cuda": "13.0"
  },
  "config": {
    "device": "cpu",
    "batch": 2,
    "heads": 8,
    "seq_lens": [
      512,
      2048
    ],
    "dims": [
      128
    ],
    "dtypes": [
      "float32"
    ],
    "impls": [
      "recurrent",
      "chunked",
      "chunked_fat",
      "ssd"
    ],
    "chunk_sizes": [
      64
    ],
    "repeats": 3
  },
  "results": [
    {
      "impl": "recurrent",
      "mode": "fwd",
      "batch": 2,
      "heads": 8,
      "seq": 512,
      "dk": 64,
      "dv": 64,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 0.10149578101118095,
      "inner": 4,
      "median_ms": 8.85094025215949,
      "q1_ms": 8.444520126431598,
      "q3_ms": 9.923574625645415,
      "xla_temp_mb": 8.312564849853516,
      "xla_arg_mb": 6.0625,
      "xla_out_mb": 2.0,
      "rel_err": 4.110092105109029e-07
    },
    {
      "impl": "chunked",
      "mode": "fwd",
      "batch": 2,
      "heads": 8,
      "seq": 512,
      "dk": 64,
      "dv": 64,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 0.3579625929996837,
      "inner": 5,
      "median_ms": 7.64616179803852,
      "q1_ms": 7.530311598384287,
      "q3_ms": 7.724216498900205,
      "xla_temp_mb": 15.858039855957031,
      "xla_arg_mb": 6.0625,
      "xla_out_mb": 2.0,
      "rel_err": 0.0
    },
    {
      "impl": "chunked_fat",
      "mode": "fwd",
      "batch": 2,
      "heads": 8,
      "seq": 512,
      "dk": 64,
      "dv": 64,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 0.3511552039999515,
      "inner": 6,
      "median_ms": 7.546704335254617,
      "q1_ms": 7.541573834411489,
      "q3_ms": 7.655609584617196,
      "xla_temp_mb": 15.858039855957031,
      "xla_arg_mb": 6.0625,
      "xla_out_mb": 2.0,
      "rel_err": 0.0
    },
    {
      "impl": "ssd",
      "mode": "fwd",
      "batch": 2,
      "heads": 8,
      "seq": 512,
      "dk": 64,
      "dv": 64,
      "dtype": "float32",
      "chunk_size": 64,
      "error": "RuntimeError: ssd requires an accelerator backend."
    },
    {
      "impl": "recurrent",
      "mode": "fwd_bwd",
      "batch": 2,
      "heads": 8,
      "seq": 512,
      "dk": 64,
      "dv": 64,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 0.35004277899861336,
      "inner": 1,
      "median_ms": 175.00102600024547,
      "q1_ms": 172.86996449547587,
      "q3_ms": 175.8514550019754,
      "xla_temp_mb": 394.91455841064453,
      "xla_arg_mb": 6.0625,
      "xla_out_mb": 6.062538146972656,
      "residual_mb": 392.09375,
      "rel_err": 5.684113627136641e-07,
      "rel_err_grads": {
        "q": 4.1552297430680173e-07,
        "k": 5.684113627136641e-07,
        "v": 5.304520498223198e-07,
        "log_alpha": 4.0527970023892254e-07,
        "beta": 4.109619159537676e-07
      }
    },
    {
      "impl": "chunked",
      "mode": "fwd_bwd",
      "batch": 2,
      "heads": 8,
      "seq": 512,
      "dk": 64,
      "dv": 64,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 0.583031544010737,
      "inner": 2,
      "median_ms": 22.920143004739657,
      "q1_ms": 22.272555503150215,
      "q3_ms": 23.408429253322538,
      "xla_temp_mb": 29.74341583251953,
      "xla_arg_mb": 6.0625,
      "xla_out_mb": 6.062538146972656,
      "residual_mb": 2.0,
      "rel_err": 0.0,
      "rel_err_grads": {
        "q": 0.0,
        "k": 0.0,
        "v": 0.0,
        "log_alpha": 0.0,
        "beta": 0.0
      }
    },
    {
      "impl": "chunked_fat",
      "mode": "fwd_bwd",
      "batch": 2,
      "heads": 8,
      "seq": 512,
      "dk": 64,
      "dv": 64,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 0.4981553760007955,
      "inner": 2,
      "median_ms": 21.27086750260787,
      "q1_ms": 21.201168252446223,
      "q3_ms": 21.43326475197682,
      "xla_temp_mb": 28.56372833251953,
      "xla_arg_mb": 6.0625,
      "xla_out_mb": 6.062538146972656,
      "residual_mb": 14.0,
      "rel_err": 0.0,
      "rel_err_grads": {
        "q": 0.0,
        "k": 0.0,
        "v": 0.0,
        "log_alpha": 0.0,
        "beta": 0.0
      }
    },
    {
      "impl": "ssd",
      "mode": "fwd_bwd",
      "batch": 2,
      "heads": 8,
      "seq": 512,
      "dk": 64,
      "dv": 64,
      "dtype": "float32",
      "chunk_size": 64,
      "error": "RuntimeError: ssd requires an accelerator backend."
    },
    {
      "impl": "recurrent",
      "mode": "fwd",
      "batch": 2,
      "heads": 8,
      "seq": 2048,
      "dk": 64,
      "dv": 64,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 0.13443950800865423,
      "inner": 1,
      "median_ms": 39.794737996999174,
      "q1_ms": 39.010178501484916,
      "q3_ms": 41.412788996240124,
      "xla_temp_mb": 32.500064849853516,
      "xla_arg_mb": 24.25,
      "xla_out_mb": 8.0,
      "rel_err": 4.942062960399496e-07
    },
    {
      "impl": "chunked",
      "mode": "fwd",
      "batch": 2,
      "heads": 8,
      "seq": 2048,
      "dk": 64,
      "dv": 64,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 0.38728979699953925,
      "inner": 1,
      "median_ms": 40.734016001806594,
      "q1_ms": 40.695898002013564,
      "q3_ms": 41.375364002306014,
      "xla_temp_mb": 62.68177032470703,
      "xla_arg_mb": 24.25,
      "xla_out_mb": 8.0,
      "rel_err": 0.0
    },
    {
      "impl": "chunked_fat",
      "mode": "fwd",
      "batch": 2,
      "heads": 8,
      "seq": 2048,
      "dk": 64,
      "dv": 64,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 0.37774297599389683,
      "inner": 1,
      "median_ms": 42.33516500971746,
      "q1_ms": 40.89740950439591,
      "q3_ms": 42.819210000743624,
      "xla_temp_mb": 62.68177032470703,
      "xla_arg_mb": 24.25,
      "xla_out_mb": 8.0,
      "rel_err": 0.0
    },
    {
      "impl": "ssd",
      "mode": "fwd",
      "batch": 2,
      "heads": 8,
      "seq": 2048,
      "dk": 64,
      "dv": 64,
      "dtype": "float32",
      "chunk_size": 64,
      "error": "RuntimeError: ssd requires an accelerator backend."
    },
    {
      "impl": "recurrent",
      "mode": "fwd_bwd",
      "batch": 2,
      "heads": 8,
      "seq": 2048,
      "dk": 64,
      "dv": 64,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 0.8268223549966933,
      "inner": 1,
      "median_ms": 755.5524680064991,
      "q1_ms": 746.5287410013843,
      "q3_ms": 758.3149495039834,
      "xla_temp_mb": 1577.3833084106445,
      "xla_arg_mb": 24.25,
      "xla_out_mb": 24.250038146972656,
      "residual_mb": 1568.375,
      "rel_err": 5.889679625641493e-07,
      "rel_err_grads": {
        "q": 4.98627569168162e-07,
        "k": 4.1287042144397063e-07,
        "v": 5.889679625641493e-07,
        "log_alpha": 4.3694217682008835e-07,
        "beta": 3.9958733284179654e-07
      }
    },
    {
      "impl": "chunked",
      "mode": "fwd_bwd",
      "batch": 2,
      "heads": 8,
      "seq": 2048,
      "dk": 64,
      "dv": 64,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 0.7264187500113621,
      "inner": 1,
      "median_ms": 124.57269099832047,
      "q1_ms": 121.94399249710841,
      "q3_ms": 125.14269450184656,
      "xla_temp_mb": 117.47290802001953,
      "xla_arg_mb": 24.25,
      "xla_out_mb": 24.250038146972656,
      "residual_mb": 8.0,
      "rel_err": 0.0,
      "rel_err_grads": {
        "q": 0.0,
        "k": 0.0,
        "v": 0.0,
        "log_alpha": 0.0,
        "beta": 0.0
      }
    },
    {
      "impl": "chunked_fat",
      "mode": "fwd_bwd",
      "batch": 2,
      "heads": 8,
      "seq": 2048,
      "dk": 64,
      "dv": 64,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 0.641810629007523,
      "inner": 1,
      "median_ms": 112.35994500748347,
      "q1_ms": 110.5135465040803,
      "q3_ms": 113.19896200438961,
      "xla_temp_mb": 112.75415802001953,
      "xla_arg_mb": 24.25,
      "xla_out_mb": 24.250038146972656,
      "residual_mb": 56.0,
      "rel_err": 0.0,
      "rel_err_grads": {
        "q": 0.0,
        "k": 0.0,
        "v": 0.0,
        "log_alpha": 0.0,
        "beta": 0.0
      }
    },
    {
      "impl": "ssd",
      "mode": "fwd_bwd",
      "batch": 2,
      "heads": 8,
      "seq": 2048,
      "dk": 64,
      "dv": 64,
      "dtype": "float32",
      "chunk_size": 64,
      "error": "RuntimeError: ssd requires an accelerator backend."
    },
    {
      "impl": "chunked",
      "mode": "fwd",
      "batch": 2,
      "heads": 8,
      "seq": 8192,
      "dk": 64,
      "dv": 64,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 0.48901010300323833,
      "inner": 1,
      "median_ms": 179.07036500400864,
      "q1_ms": 177.08122200565413,
      "q3_ms": 179.42523099918617,
      "xla_temp_mb": 249.97669219970703,
      "xla_arg_mb": 97.0,
      "xla_out_mb": 32.0,
      "rel_err": 0.0
    },
    {
      "impl": "chunked_fat",
      "mode": "fwd",
      "batch": 2,
      "heads": 8,
      "seq": 8192,
      "dk": 64,
      "dv": 64,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 0.49725457100430503,
      "inner": 1,
      "median_ms": 179.24015699827578,
      "q1_ms": 178.78247600310715,
      "q3_ms": 181.70115750399418,
      "xla_temp_mb": 249.97669219970703,
      "xla_arg_mb": 97.0,
      "xla_out_mb": 32.0,
      "rel_err": 0.0
    },
    {
      "impl": "ssd",
      "mode": "fwd",
      "batch": 2,
      "heads": 8,
      "seq": 8192,
      "dk": 64,
      "dv": 64,
      "dtype": "float32",
      "chunk_size": 64,
      "error": "RuntimeError: ssd requires an accelerator backend."
    },
    {
      "impl": "chunked",
      "mode": "fwd_bwd",
      "batch": 2,
      "heads": 8,
      "seq": 8192,
      "dk": 64,
      "dv": 64,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 1.678737681999337,
      "inner": 1,
      "median_ms": 534.8212539975066,
      "q1_ms": 534.7246314995573,
      "q3_ms": 535.2296220007702,
      "xla_temp_mb": 468.39087677001953,
      "xla_arg_mb": 97.0,
      "xla_out_mb": 97.00003814697266,
      "residual_mb": 32.0,
      "rel_err": 0.0,
      "rel_err_grads": {
        "q": 0.0,
        "k": 0.0,
        "v": 0.0,
        "log_alpha": 0.0,
        "beta": 0.0
      }
    },
    {
      "impl": "chunked_fat",
      "mode": "fwd_bwd",
      "batch": 2,
      "heads": 8,
      "seq": 8192,
      "dk": 64,
      "dv": 64,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 1.3872163699998055,
      "inner": 1,
      "median_ms": 483.3967140002642,
      "q1_ms": 482.28388350253226,
      "q3_ms": 483.97503799787955,
      "xla_temp_mb": 449.51587677001953,
      "xla_arg_mb": 97.0,
      "xla_out_mb": 97.00003814697266,
      "residual_mb": 224.0,
      "rel_err": 0.0,
      "rel_err_grads": {
        "q": 0.0,
        "k": 0.0,
        "v": 0.0,
        "log_alpha": 0.0,
        "beta": 0.0
      }
    },
    {
      "impl": "ssd",
      "mode": "fwd_bwd",
      "batch": 2,
      "heads": 8,
      "seq": 8192,
      "dk": 64,
      "dv": 64,
      "dtype": "float32",
      "chunk_size": 64,
      "error": "RuntimeError: ssd requires an accelerator backend."
    },
    {
      "impl": "recurrent",
      "mode": "fwd",
      "batch": 2,
      "heads": 8,
      "seq": 512,
      "dk": 128,
      "dv": 128,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 0.15532219599117525,
      "inner": 1,
      "median_ms": 48.21182299929205,
      "q1_ms": 43.66324250440812,
      "q3_ms": 69.74907949916087,
      "xla_temp_mb": 17.062564849853516,
      "xla_arg_mb": 12.0625,
      "xla_out_mb": 4.0,
      "rel_err": 4.993866105603432e-07
    },
    {
      "impl": "chunked",
      "mode": "fwd",
      "batch": 2,
      "heads": 8,
      "seq": 512,
      "dk": 128,
      "dv": 128,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 0.4055430319858715,
      "inner": 3,
      "median_ms": 15.80632399903455,
      "q1_ms": 15.077660498112284,
      "q3_ms": 17.093328999180812,
      "xla_temp_mb": 29.00061798095703,
      "xla_arg_mb": 12.0625,
      "xla_out_mb": 4.0,
      "rel_err": 0.0
    },
    {
      "impl": "chunked_fat",
      "mode": "fwd",
      "batch": 2,
      "heads": 8,
      "seq": 512,
      "dk": 128,
      "dv": 128,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 0.3637767629988957,
      "inner": 3,
      "median_ms": 17.355288000544533,
      "q1_ms": 16.528490499088853,
      "q3_ms": 17.391272835084237,
      "xla_temp_mb": 29.00061798095703,
      "xla_arg_mb": 12.0625,
      "xla_out_mb": 4.0,
      "rel_err": 0.0
    },
    {
      "impl": "ssd",
      "mode": "fwd",
      "batch": 2,
      "heads": 8,
      "seq": 512,
      "dk": 128,
      "dv": 128,
      "dtype": "float32",
      "chunk_size": 64,
      "error": "RuntimeError: ssd requires an accelerator backend."
    },
    {
      "impl": "recurrent",
      "mode": "fwd_bwd",
      "batch": 2,
      "heads": 8,
      "seq": 512,
      "dk": 128,
      "dv": 128,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 0.9543200589978369,
      "inner": 1,
      "median_ms": 800.9884589991998,
      "q1_ms": 774.5471890011686,
      "q3_ms": 806.4735765001387,
      "xla_temp_mb": 1559.1723709106445,
      "xla_arg_mb": 12.0625,
      "xla_out_mb": 12.062538146972656,
      "residual_mb": 1552.09375,
      "rel_err": 6.60262324526051e-07,
      "rel_err_grads": {
        "q": 3.433744610094007e-07,
        "k": 4.1408719102711707e-07,
        "v": 5.881581992634887e-07,
        "log_alpha": 6.60262324526051e-07,
        "beta": 4.104130659760345e-07
      }
    },
    {
      "impl": "chunked",
      "mode": "fwd_bwd",
      "batch": 2,
      "heads": 8,
      "seq": 512,
      "dk": 128,
      "dv": 128,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 0.6557113669987302,
      "inner": 1,
      "median_ms": 59.29503099469002,
      "q1_ms": 58.77287299517775,
      "q3_ms": 61.98362749273656,
      "xla_temp_mb": 62.56372833251953,
      "xla_arg_mb": 12.0625,
      "xla_out_mb": 12.062538146972656,
      "residual_mb": 8.0,
      "rel_err": 0.0,
      "rel_err_grads": {
        "q": 0.0,
        "k": 0.0,
        "v": 0.0,
        "log_alpha": 0.0,
        "beta": 0.0
      }
    },
    {
      "impl": "chunked_fat",
      "mode": "fwd_bwd",
      "batch": 2,
      "heads": 8,
      "seq": 512,
      "dk": 128,
      "dv": 128,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 0.5582565520016942,
      "inner": 1,
      "median_ms": 56.376759996055625,
      "q1_ms": 54.62180449831067,
      "q3_ms": 56.43424349545967,
      "xla_temp_mb": 56.06372833251953,
      "xla_arg_mb": 12.0625,
      "xla_out_mb": 12.062538146972656,
      "residual_mb": 28.0,
      "rel_err": 0.0,
      "rel_err_grads": {
        "q": 0.0,
        "k": 0.0,
        "v": 0.0,
        "log_alpha": 0.0,
        "beta": 0.0
      }
    },
    {
      "impl": "ssd",
      "mode": "fwd_bwd",
      "batch": 2,
      "heads": 8,
      "seq": 512,
      "dk": 128,
      "dv": 128,
      "dtype": "float32",
      "chunk_size": 64,
      "error": "RuntimeError: ssd requires an accelerator backend."
    },
    {
      "impl": "recurrent",
      "mode": "fwd",
      "batch": 2,
      "heads": 8,
      "seq": 2048,
      "dk": 128,
      "dv": 128,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 0.256415374999051,
      "inner": 1,
      "median_ms": 162.73475899652112,
      "q1_ms": 162.6155089979875,
      "q3_ms": 214.64553449914092,
      "xla_temp_mb": 65.25006484985352,
      "xla_arg_mb": 48.25,
      "xla_out_mb": 16.0,
      "rel_err": 5.675590618513569e-07
    },
    {
      "impl": "chunked",
      "mode": "fwd",
      "batch": 2,
      "heads": 8,
      "seq": 2048,
      "dk": 128,
      "dv": 128,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 0.4561459179967642,
      "inner": 1,
      "median_ms": 86.19231800548732,
      "q1_ms": 84.81863400083967,
      "q3_ms": 88.07997449912364,
      "xla_temp_mb": 113.00208282470703,
      "xla_arg_mb": 48.25,
      "xla_out_mb": 16.0,
      "rel_err": 0.0
    },
    {
      "impl": "chunked_fat",
      "mode": "fwd",
      "batch": 2,
      "heads": 8,
      "seq": 2048,
      "dk": 128,
      "dv": 128,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 0.4469447800074704,
      "inner": 1,
      "median_ms": 84.43067599728238,
      "q1_ms": 84.0296419992228,
      "q3_ms": 85.49050099827582,
      "xla_temp_mb": 113.00208282470703,
      "xla_arg_mb": 48.25,
      "xla_out_mb": 16.0,
      "rel_err": 0.0
    },
    {
      "impl": "ssd",
      "mode": "fwd",
      "batch": 2,
      "heads": 8,
      "seq": 2048,
      "dk": 128,
      "dv": 128,
      "dtype": "float32",
      "chunk_size": 64,
      "error": "RuntimeError: ssd requires an accelerator backend."
    },
    {
      "impl": "recurrent",
      "mode": "fwd_bwd",
      "batch": 2,
      "heads": 8,
      "seq": 2048,
      "dk": 128,
      "dv": 128,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 3.137445005006157,
      "inner": 1,
      "median_ms": 3273.8520670100115,
      "q1_ms": 3273.24619051069,
      "q3_ms": 3361.608156505099,
      "xla_temp_mb": 6227.6411209106445,
      "xla_arg_mb": 48.25,
      "xla_out_mb": 48.250038146972656,
      "residual_mb": 6208.375,
      "rel_err": 7.674112060092812e-07,
      "rel_err_grads": {
        "q": 3.940906578760279e-07,
        "k": 3.3533579100934565e-07,
        "v": 7.674112060092812e-07,
        "log_alpha": 4.974823879529196e-07,
        "beta": 5.222107485189131e-07
      }
    },
    {
      "impl": "chunked",
      "mode": "fwd_bwd",
      "batch": 2,
      "heads": 8,
      "seq": 2048,
      "dk": 128,
      "dv": 128,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 0.9827960019902093,
      "inner": 1,
      "median_ms": 259.1154719993938,
      "q1_ms": 258.68527850252576,
      "q3_ms": 259.468004500377,
      "xla_temp_mb": 244.25415802001953,
      "xla_arg_mb": 48.25,
      "xla_out_mb": 48.250038146972656,
      "residual_mb": 32.0,
      "rel_err": 0.0,
      "rel_err_grads": {
        "q": 0.0,
        "k": 0.0,
        "v": 0.0,
        "log_alpha": 0.0,
        "beta": 0.0
      }
    },
    {
      "impl": "chunked_fat",
      "mode": "fwd_bwd",
      "batch": 2,
      "heads": 8,
      "seq": 2048,
      "dk": 128,
      "dv": 128,
      "dtype": "float32",
      "chunk_size": 64,
      "compile_s": 0.9201118960045278,
      "inner": 1,
      "median_ms": 248.0533179914346,
      "q1_ms": 240.173728998343,
      "q3_ms": 248.388216496096,
      "xla_temp_mb": 218.25415802001953,
      "xla_arg_mb": 48.25,
      "xla_out_mb": 48.250038146972656,
      "residual_mb": 112.0,
      "rel_err": 0.0,
      "rel_err_grads": {
        "q": 0.0,
        "k": 0.0,
        "v": 0.0,
        "log_alpha": 0.0,
        "beta": 0.0
      }
    },
    {
      "impl": "ssd",
      "mode": "fwd_bwd",
      "batch": 2,
      "heads": 8,
      "seq": 2048,
      "dk": 128,
      "dv": 128,
      "dtype": "float32",
      "chunk_size": 64,
      "error": "RuntimeError: ssd requires an accelerator backend."
    }
  ]
}