{
  "peaks": [
    {
      "impl": "probjax_mha",
      "mode": "fwdbwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 4096,
      "heads": 8,
      "dim": 64,
      "device": "cpu",
      "xla_temp_mb": 40.12550354003906,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 24.000022888183594,
      "baseline_rss_mb": 875.59375,
      "peak_rss_mb": 919.09375,
      "peak_delta_mb": 43.5
    },
    {
      "impl": "xla_dense",
      "mode": "fwdbwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 4096,
      "heads": 8,
      "dim": 64,
      "device": "cpu",
      "xla_temp_mb": 2568.125,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 24.000022888183594,
      "baseline_rss_mb": 847.9140625,
      "peak_rss_mb": 3437.6484375,
      "peak_delta_mb": 2589.734375
    },
    {
      "impl": "jax_xla",
      "mode": "fwdbwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 4096,
      "heads": 8,
      "dim": 64,
      "device": "cpu",
      "xla_temp_mb": 1544.25,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 24.000022888183594,
      "baseline_rss_mb": 871.5,
      "peak_rss_mb": 2539.125,
      "peak_delta_mb": 1667.625
    },
    {
      "impl": "torch_sdpa",
      "mode": "fwdbwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 4096,
      "heads": 8,
      "dim": 64,
      "device": "cpu",
      "baseline_rss_mb": 995.84375,
      "peak_rss_mb": 1192.64453125,
      "peak_delta_mb": 196.80078125
    },
    {
      "impl": "probjax_mha",
      "mode": "fwdbwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 16384,
      "heads": 8,
      "dim": 64,
      "device": "cpu",
      "xla_temp_mb": 160.50086975097656,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 96.0000228881836,
      "baseline_rss_mb": 1113.10546875,
      "peak_rss_mb": 1301.9375,
      "peak_delta_mb": 188.83203125
    },
    {
      "impl": "torch_sdpa",
      "mode": "fwdbwd",
      "variant": "plain",
      "dtype": "float32",
      "batch": 1,
      "seq": 16384,
      "heads": 8,
      "dim": 64,
      "device": "cpu",
      "baseline_rss_mb": 1059.24609375,
      "peak_rss_mb": 1523.05859375,
      "peak_delta_mb": 463.8125
    }
  ],
  "config": {
    "device": "cpu",
    "append": false,
    "seqs": "128,256,512,1024,2048,4096",
    "dims": "64,128",
    "dtypes": "float32",
    "variants": "plain,causal,masked_bias",
    "modes": "fwd,fwdbwd",
    "batch": 1,
    "heads": 8,
    "impls": "probjax_mha,xla_dense,xla_tiled,mha_interpret,jax_xla,torch_sdpa",
    "interpret_max_seq": 1024,
    "no_ref": false,
    "peaks": true,
    "peak_seqs": "4096",
    "peak_impls": "probjax_mha,xla_dense,jax_xla,torch_sdpa",
    "peak_timeout": 120,
    "max_fit": false,
    "fit_seqs": "4096,8192,16384,32768,65536,131072",
    "fit_impls": "xla_tiled,xla_dense,jax_xla"
  },
  "configs": [
    {
      "device": "cpu",
      "append": true,
      "seqs": "128,256,512,1024,2048,4096",
      "dims": "64,128",
      "dtypes": "float32",
      "variants": "plain,causal,masked_bias",
      "modes": "fwd,fwdbwd",
      "batch": 1,
      "heads": 8,
      "impls": "probjax_mha,xla_dense,xla_tiled,mha_interpret,jax_xla,torch_sdpa",
      "interpret_max_seq": 1024,
      "no_ref": false,
      "peaks": true,
      "peak_seqs": "16384",
      "peak_impls": "probjax_mha,torch_sdpa",
      "peak_timeout": 280,
      "max_fit": false,
      "fit_seqs": "4096,8192,16384,32768,65536,131072",
      "fit_impls": "xla_tiled,xla_dense,jax_xla"
    },
    {
      "device": "cpu",
      "append": true,
      "seqs": "128,256,512,1024,2048,4096",
      "dims": "64,128",
      "dtypes": "float32",
      "variants": "plain,causal,masked_bias",
      "modes": "fwd,fwdbwd",
      "batch": 1,
      "heads": 8,
      "impls": "probjax_mha,xla_dense,xla_tiled,mha_interpret,jax_xla,torch_sdpa",
      "interpret_max_seq": 1024,
      "no_ref": false,
      "peaks": false,
      "peak_seqs": "4096,8192,16384",
      "peak_impls": "xla_tiled,xla_dense,jax_xla,torch_sdpa",
      "peak_timeout": 300,
      "max_fit": true,
      "fit_seqs": "4096,8192,16384,32768,65536,131072",
      "fit_impls": "xla_tiled,xla_dense,jax_xla"
    }
  ],
  "max_fit": [
    {
      "impl": "xla_tiled",
      "mode": "fwdbwd",
      "seq": 4096,
      "batch": 1,
      "heads": 8,
      "head_dim": 64,
      "dtype": "float32",
      "xla_temp_mb": 40.12550354003906,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 24.000022888183594,
      "total_mb": 88.12552642822266,
      "fits_host_ram": true,
      "host_ram_mb": 63681.82421875
    },
    {
      "impl": "xla_tiled",
      "mode": "fwdbwd",
      "seq": 8192,
      "batch": 1,
      "heads": 8,
      "head_dim": 64,
      "dtype": "float32",
      "xla_temp_mb": 80.25062561035156,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 48.000022888183594,
      "total_mb": 176.25064849853516,
      "fits_host_ram": true,
      "host_ram_mb": 63681.82421875
    },
    {
      "impl": "xla_tiled",
      "mode": "fwdbwd",
      "seq": 16384,
      "batch": 1,
      "heads": 8,
      "head_dim": 64,
      "dtype": "float32",
      "xla_temp_mb": 160.50086975097656,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 96.0000228881836,
      "total_mb": 352.50089263916016,
      "fits_host_ram": true,
      "host_ram_mb": 63681.82421875
    },
    {
      "impl": "xla_tiled",
      "mode": "fwdbwd",
      "seq": 32768,
      "batch": 1,
      "heads": 8,
      "head_dim": 64,
      "dtype": "float32",
      "xla_temp_mb": 321.00135803222656,
      "xla_arg_mb": 192.0,
      "xla_out_mb": 192.0000228881836,
      "total_mb": 705.0013809204102,
      "fits_host_ram": true,
      "host_ram_mb": 63681.82421875
    },
    {
      "impl": "xla_tiled",
      "mode": "fwdbwd",
      "seq": 65536,
      "batch": 1,
      "heads": 8,
      "head_dim": 64,
      "dtype": "float32",
      "xla_temp_mb": 642.0023345947266,
      "xla_arg_mb": 384.0,
      "xla_out_mb": 384.0000228881836,
      "total_mb": 1410.0023574829102,
      "fits_host_ram": true,
      "host_ram_mb": 63681.82421875
    },
    {
      "impl": "xla_tiled",
      "mode": "fwdbwd",
      "seq": 131072,
      "batch": 1,
      "heads": 8,
      "head_dim": 64,
      "dtype": "float32",
      "xla_temp_mb": 1284.0042877197266,
      "xla_arg_mb": 768.0,
      "xla_out_mb": 768.0000228881836,
      "total_mb": 2820.00431060791,
      "fits_host_ram": true,
      "host_ram_mb": 63681.82421875
    },
    {
      "impl": "xla_dense",
      "mode": "fwdbwd",
      "seq": 4096,
      "batch": 1,
      "heads": 8,
      "head_dim": 64,
      "dtype": "float32",
      "xla_temp_mb": 2576.125,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 24.000022888183594,
      "total_mb": 2624.1250228881836,
      "fits_host_ram": true,
      "host_ram_mb": 63681.82421875
    },
    {
      "impl": "xla_dense",
      "mode": "fwdbwd",
      "seq": 8192,
      "batch": 1,
      "heads": 8,
      "head_dim": 64,
      "dtype": "float32",
      "xla_temp_mb": 10272.25,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 48.000022888183594,
      "total_mb": 10368.250022888184,
      "fits_host_ram": true,
      "host_ram_mb": 63681.82421875
    },
    {
      "impl": "xla_dense",
      "mode": "fwdbwd",
      "seq": 16384,
      "batch": 1,
      "heads": 8,
      "head_dim": 64,
      "dtype": "float32",
      "xla_temp_mb": 41024.5,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 96.0000228881836,
      "total_mb": 41216.50002288818,
      "fits_host_ram": true,
      "host_ram_mb": 63681.82421875
    },
    {
      "impl": "xla_dense",
      "mode": "fwdbwd",
      "seq": 32768,
      "batch": 1,
      "heads": 8,
      "head_dim": 64,
      "dtype": "float32",
      "xla_temp_mb": 163969.0,
      "xla_arg_mb": 192.0,
      "xla_out_mb": 192.0000228881836,
      "total_mb": 164353.00002288818,
      "fits_host_ram": false,
      "host_ram_mb": 63681.82421875
    },
    {
      "impl": "xla_dense",
      "mode": "fwdbwd",
      "seq": 65536,
      "batch": 1,
      "heads": 8,
      "head_dim": 64,
      "dtype": "float32",
      "xla_temp_mb": 655618.0,
      "xla_arg_mb": 384.0,
      "xla_out_mb": 384.0000228881836,
      "total_mb": 656386.0000228882,
      "fits_host_ram": false,
      "host_ram_mb": 63681.82421875
    },
    {
      "impl": "xla_dense",
      "mode": "fwdbwd",
      "seq": 131072,
      "batch": 1,
      "heads": 8,
      "head_dim": 64,
      "dtype": "float32",
      "xla_temp_mb": 2621956.0,
      "xla_arg_mb": 768.0,
      "xla_out_mb": 768.0000228881836,
      "total_mb": 2623492.000022888,
      "fits_host_ram": false,
      "host_ram_mb": 63681.82421875
    },
    {
      "impl": "jax_xla",
      "mode": "fwdbwd",
      "seq": 4096,
      "batch": 1,
      "heads": 8,
      "head_dim": 64,
      "dtype": "float32",
      "xla_temp_mb": 1552.25,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 24.000022888183594,
      "total_mb": 1600.2500228881836,
      "fits_host_ram": true,
      "host_ram_mb": 63681.82421875
    },
    {
      "impl": "jax_xla",
      "mode": "fwdbwd",
      "seq": 8192,
      "batch": 1,
      "heads": 8,
      "head_dim": 64,
      "dtype": "float32",
      "xla_temp_mb": 6176.5,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 48.000022888183594,
      "total_mb": 6272.500022888184,
      "fits_host_ram": true,
      "host_ram_mb": 63681.82421875
    },
    {
      "impl": "jax_xla",
      "mode": "fwdbwd",
      "seq": 16384,
      "batch": 1,
      "heads": 8,
      "head_dim": 64,
      "dtype": "float32",
      "xla_temp_mb": 24641.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 96.0000228881836,
      "total_mb": 24833.000022888184,
      "fits_host_ram": true,
      "host_ram_mb": 63681.82421875
    },
    {
      "impl": "jax_xla",
      "mode": "fwdbwd",
      "seq": 32768,
      "batch": 1,
      "heads": 8,
      "head_dim": 64,
      "dtype": "float32",
      "xla_temp_mb": 98434.0,
      "xla_arg_mb": 192.0,
      "xla_out_mb": 192.0000228881836,
      "total_mb": 98818.00002288818,
      "fits_host_ram": false,
      "host_ram_mb": 63681.82421875
    },
    {
      "impl": "jax_xla",
      "mode": "fwdbwd",
      "seq": 65536,
      "batch": 1,
      "heads": 8,
      "head_dim": 64,
      "dtype": "float32",
      "xla_temp_mb": 393476.0,
      "xla_arg_mb": 384.0,
      "xla_out_mb": 384.0000228881836,
      "total_mb": 394244.0000228882,
      "fits_host_ram": false,
      "host_ram_mb": 63681.82421875
    },
    {
      "impl": "jax_xla",
      "mode": "fwdbwd",
      "seq": 131072,
      "batch": 1,
      "heads": 8,
      "head_dim": 64,
      "dtype": "float32",
      "xla_temp_mb": 1573384.0,
      "xla_arg_mb": 768.0,
      "xla_out_mb": 768.0000228881836,
      "total_mb": 1574920.0000228882,
      "fits_host_ram": false,
      "host_ram_mb": 63681.82421875
    }
  ]
}