{
  "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": {
    "note": "RTX 3090. gqa: flex_native (no K/V repeat) vs flex_repeat (previous behaviour). decode: probjax_splitkv = default flex_attention (auto split-KV), probjax_nosplit = split_kv=False, flex_repeat = previous GQA path. splitsweep: num_splits=0 is the regular forward.",
    "runs": {
      "gqa_1k_bfloat16": {
        "device": "gpu",
        "suite": "gqa",
        "output": "/tmp/claude-1003/-home-manuel-src/80785a0e-56be-4da7-a2a8-e1d6effd6199/scratchpad/attn-ext/gqa_1k_bfloat16.json",
        "append": false,
        "tag": null,
        "repeats": 3,
        "batch": 1,
        "batches": [
          1,
          8
        ],
        "heads": 32,
        "kv_heads": [
          8,
          4,
          1
        ],
        "head_dims": [
          64,
          128
        ],
        "seq_lens": [
          1024
        ],
        "q_lens": [
          1,
          4,
          16
        ],
        "kv_lens": [
          4096,
          16384,
          65536
        ],
        "splits": [
          2,
          4,
          8,
          16,
          32,
          64
        ],
        "dtypes": [
          "bfloat16"
        ],
        "modes": [
          "fwd",
          "fwd_bwd"
        ],
        "impls": [
          "flex_native",
          "flex_repeat",
          "jax_xla",
          "jax_cudnn",
          "torch_sdpa"
        ],
        "baseline_impl": "flex_repeat",
        "ref_max_seq": 1024,
        "max_kv_gb": 4.5,
        "child": null
      },
      "gqa_4k_bfloat16": {
        "device": "gpu",
        "suite": "gqa",
        "output": "/tmp/claude-1003/-home-manuel-src/80785a0e-56be-4da7-a2a8-e1d6effd6199/scratchpad/attn-ext/gqa_4k_bfloat16.json",
        "append": false,
        "tag": null,
        "repeats": 3,
        "batch": 1,
        "batches": [
          1,
          8
        ],
        "heads": 32,
        "kv_heads": [
          8,
          4,
          1
        ],
        "head_dims": [
          64,
          128
        ],
        "seq_lens": [
          4096
        ],
        "q_lens": [
          1,
          4,
          16
        ],
        "kv_lens": [
          4096,
          16384,
          65536
        ],
        "splits": [
          2,
          4,
          8,
          16,
          32,
          64
        ],
        "dtypes": [
          "bfloat16"
        ],
        "modes": [
          "fwd",
          "fwd_bwd"
        ],
        "impls": [
          "flex_native",
          "flex_repeat",
          "jax_cudnn",
          "torch_sdpa"
        ],
        "baseline_impl": "flex_repeat",
        "ref_max_seq": 1024,
        "max_kv_gb": 4.5,
        "child": null
      },
      "gqa_1k_float32": {
        "device": "gpu",
        "suite": "gqa",
        "output": "/tmp/claude-1003/-home-manuel-src/80785a0e-56be-4da7-a2a8-e1d6effd6199/scratchpad/attn-ext/gqa_1k_float32.json",
        "append": false,
        "tag": null,
        "repeats": 3,
        "batch": 1,
        "batches": [
          1,
          8
        ],
        "heads": 32,
        "kv_heads": [
          8,
          4,
          1
        ],
        "head_dims": [
          64,
          128
        ],
        "seq_lens": [
          1024
        ],
        "q_lens": [
          1,
          4,
          16
        ],
        "kv_lens": [
          4096,
          16384,
          65536
        ],
        "splits": [
          2,
          4,
          8,
          16,
          32,
          64
        ],
        "dtypes": [
          "float32"
        ],
        "modes": [
          "fwd",
          "fwd_bwd"
        ],
        "impls": [
          "flex_native",
          "flex_repeat",
          "jax_xla",
          "jax_cudnn",
          "torch_sdpa"
        ],
        "baseline_impl": "flex_repeat",
        "ref_max_seq": 1024,
        "max_kv_gb": 4.5,
        "child": null
      },
      "gqa_4k_float32": {
        "device": "gpu",
        "suite": "gqa",
        "output": "/tmp/claude-1003/-home-manuel-src/80785a0e-56be-4da7-a2a8-e1d6effd6199/scratchpad/attn-ext/gqa_4k_float32.json",
        "append": false,
        "tag": null,
        "repeats": 3,
        "batch": 1,
        "batches": [
          1,
          8
        ],
        "heads": 32,
        "kv_heads": [
          8,
          4,
          1
        ],
        "head_dims": [
          64,
          128
        ],
        "seq_lens": [
          4096
        ],
        "q_lens": [
          1,
          4,
          16
        ],
        "kv_lens": [
          4096,
          16384,
          65536
        ],
        "splits": [
          2,
          4,
          8,
          16,
          32,
          64
        ],
        "dtypes": [
          "float32"
        ],
        "modes": [
          "fwd",
          "fwd_bwd"
        ],
        "impls": [
          "flex_native",
          "flex_repeat",
          "jax_cudnn",
          "torch_sdpa"
        ],
        "baseline_impl": "flex_repeat",
        "ref_max_seq": 1024,
        "max_kv_gb": 4.5,
        "child": null
      },
      "decode2_h8": {
        "device": "gpu",
        "suite": "decode",
        "output": "/tmp/claude-1003/-home-manuel-src/80785a0e-56be-4da7-a2a8-e1d6effd6199/scratchpad/attn-ext/decode2_h8.json",
        "append": false,
        "tag": null,
        "repeats": 3,
        "batch": 1,
        "batches": [
          1,
          8
        ],
        "heads": 32,
        "kv_heads": [
          8
        ],
        "head_dims": [
          128
        ],
        "seq_lens": [
          1024,
          4096
        ],
        "q_lens": [
          1,
          4,
          16
        ],
        "kv_lens": [
          4096,
          16384,
          65536
        ],
        "splits": [
          2,
          4,
          8,
          16,
          32,
          64
        ],
        "block_ks": [
          32,
          64
        ],
        "dtypes": [
          "bfloat16"
        ],
        "modes": [
          "fwd"
        ],
        "impls": [
          "probjax_splitkv",
          "probjax_nosplit",
          "flex_repeat",
          "jax_xla",
          "jax_cudnn",
          "torch_sdpa"
        ],
        "baseline_impl": "probjax_nosplit",
        "ref_max_seq": 1024,
        "max_kv_gb": 4.5,
        "child": null
      },
      "decode2_h32": {
        "device": "gpu",
        "suite": "decode",
        "output": "/tmp/claude-1003/-home-manuel-src/80785a0e-56be-4da7-a2a8-e1d6effd6199/scratchpad/attn-ext/decode2_h32.json",
        "append": false,
        "tag": null,
        "repeats": 3,
        "batch": 1,
        "batches": [
          1,
          8
        ],
        "heads": 32,
        "kv_heads": [
          32
        ],
        "head_dims": [
          128
        ],
        "seq_lens": [
          1024,
          4096
        ],
        "q_lens": [
          1,
          4,
          16
        ],
        "kv_lens": [
          4096,
          16384,
          65536
        ],
        "splits": [
          2,
          4,
          8,
          16,
          32,
          64
        ],
        "block_ks": [
          32,
          64
        ],
        "dtypes": [
          "bfloat16"
        ],
        "modes": [
          "fwd"
        ],
        "impls": [
          "probjax_splitkv",
          "probjax_nosplit",
          "flex_repeat",
          "jax_xla",
          "jax_cudnn",
          "torch_sdpa"
        ],
        "baseline_impl": "probjax_nosplit",
        "ref_max_seq": 1024,
        "max_kv_gb": 4.5,
        "child": null
      },
      "decode2_f32": {
        "device": "gpu",
        "suite": "decode",
        "output": "/tmp/claude-1003/-home-manuel-src/80785a0e-56be-4da7-a2a8-e1d6effd6199/scratchpad/attn-ext/decode2_f32.json",
        "append": false,
        "tag": null,
        "repeats": 3,
        "batch": 1,
        "batches": [
          1
        ],
        "heads": 32,
        "kv_heads": [
          8
        ],
        "head_dims": [
          64
        ],
        "seq_lens": [
          1024,
          4096
        ],
        "q_lens": [
          1,
          4,
          16
        ],
        "kv_lens": [
          4096,
          16384,
          65536
        ],
        "splits": [
          2,
          4,
          8,
          16,
          32,
          64
        ],
        "block_ks": [
          32,
          64
        ],
        "dtypes": [
          "float32"
        ],
        "modes": [
          "fwd"
        ],
        "impls": [
          "probjax_splitkv",
          "probjax_nosplit",
          "jax_xla",
          "torch_sdpa"
        ],
        "baseline_impl": "probjax_nosplit",
        "ref_max_seq": 1024,
        "max_kv_gb": 4.5,
        "child": null
      },
      "sweep_h8": {
        "device": "gpu",
        "suite": "splitsweep",
        "output": "/tmp/claude-1003/-home-manuel-src/80785a0e-56be-4da7-a2a8-e1d6effd6199/scratchpad/attn-ext/sweep_h8.json",
        "append": false,
        "tag": null,
        "repeats": 2,
        "batch": 1,
        "batches": [
          1,
          8
        ],
        "heads": 32,
        "kv_heads": [
          8
        ],
        "head_dims": [
          128
        ],
        "seq_lens": [
          1024,
          4096
        ],
        "q_lens": [
          1,
          4,
          16
        ],
        "kv_lens": [
          4096,
          16384,
          65536
        ],
        "splits": [
          2,
          4,
          8,
          16,
          32,
          64
        ],
        "dtypes": [
          "bfloat16"
        ],
        "modes": [
          "fwd",
          "fwd_bwd"
        ],
        "impls": [
          "probjax_splitkv",
          "probjax_nosplit",
          "jax_xla",
          "jax_cudnn",
          "torch_sdpa"
        ],
        "baseline_impl": "probjax_nosplit",
        "ref_max_seq": 1024,
        "max_kv_gb": 4.5,
        "child": null
      },
      "sweep2": {
        "device": "gpu",
        "suite": "splitsweep",
        "output": "/tmp/claude-1003/-home-manuel-src/80785a0e-56be-4da7-a2a8-e1d6effd6199/scratchpad/attn-ext/sweep2.json",
        "append": false,
        "tag": null,
        "repeats": 2,
        "batch": 1,
        "batches": [
          1,
          8
        ],
        "heads": 32,
        "kv_heads": [
          32,
          8
        ],
        "head_dims": [
          128
        ],
        "seq_lens": [
          1024,
          4096
        ],
        "q_lens": [
          1,
          32,
          64
        ],
        "kv_lens": [
          1024,
          2048,
          16384
        ],
        "splits": [
          2,
          4,
          8,
          16
        ],
        "dtypes": [
          "bfloat16"
        ],
        "modes": [
          "fwd",
          "fwd_bwd"
        ],
        "impls": [
          "probjax_splitkv",
          "probjax_nosplit",
          "jax_xla",
          "jax_cudnn",
          "torch_sdpa"
        ],
        "baseline_impl": "probjax_nosplit",
        "ref_max_seq": 1024,
        "max_kv_gb": 4.5,
        "child": null
      },
      "blockk": {
        "device": "gpu",
        "suite": "blockk",
        "output": "/tmp/claude-1003/-home-manuel-src/80785a0e-56be-4da7-a2a8-e1d6effd6199/scratchpad/attn-ext/blockk.json",
        "append": false,
        "tag": null,
        "repeats": 3,
        "batch": 1,
        "batches": [
          1,
          8
        ],
        "heads": 32,
        "kv_heads": [
          32,
          8
        ],
        "head_dims": [
          128
        ],
        "seq_lens": [
          1024,
          4096
        ],
        "q_lens": [
          1,
          16
        ],
        "kv_lens": [
          16384,
          65536
        ],
        "splits": [
          2,
          4,
          8,
          16,
          32,
          64
        ],
        "block_ks": [
          32,
          64
        ],
        "dtypes": [
          "bfloat16"
        ],
        "modes": [
          "fwd",
          "fwd_bwd"
        ],
        "impls": [
          "probjax_splitkv",
          "probjax_nosplit",
          "jax_xla",
          "jax_cudnn",
          "torch_sdpa"
        ],
        "baseline_impl": "probjax_nosplit",
        "ref_max_seq": 1024,
        "max_kv_gb": 4.5,
        "child": null
      },
      "peak": {
        "device": "gpu",
        "suite": "peak",
        "output": "/tmp/claude-1003/-home-manuel-src/80785a0e-56be-4da7-a2a8-e1d6effd6199/scratchpad/attn-ext/peak.json",
        "append": false,
        "tag": null,
        "repeats": 3,
        "batch": 1,
        "batches": [
          1,
          8
        ],
        "heads": 32,
        "kv_heads": [
          32,
          8
        ],
        "head_dims": [
          64,
          128
        ],
        "seq_lens": [
          1024,
          4096
        ],
        "q_lens": [
          1,
          4,
          16
        ],
        "kv_lens": [
          4096,
          16384,
          65536
        ],
        "splits": [
          2,
          4,
          8,
          16,
          32,
          64
        ],
        "block_ks": [
          32,
          64
        ],
        "dtypes": [
          "bfloat16"
        ],
        "modes": [
          "fwd",
          "fwd_bwd"
        ],
        "impls": [
          "probjax_splitkv",
          "probjax_nosplit",
          "jax_xla",
          "jax_cudnn",
          "torch_sdpa"
        ],
        "baseline_impl": "probjax_nosplit",
        "ref_max_seq": 1024,
        "max_kv_gb": 4.5,
        "child": null
      }
    }
  },
  "results": [
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 4.0,
      "rel_err": 0.002210708102211356,
      "rel_err_parts": [
        0.002210708102211356
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.19102628900145646,
      "inner": 181,
      "median_ms": 0.17685209946367295,
      "q1_ms": 0.15715906634966673,
      "q3_ms": 0.18647585632742603,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 8.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 4.0,
      "rel_err": 0.002210708102211356,
      "rel_err_parts": [
        0.002210708102211356
      ],
      "compile_s": 0.19738474900077563,
      "inner": 200,
      "median_ms": 0.17382645499310456,
      "q1_ms": 0.1728600300339167,
      "q3_ms": 0.17414433750673197,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 192.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 4.0,
      "rel_err": 0.002334752818569541,
      "rel_err_parts": [
        0.002334752818569541
      ],
      "rel_vs_flex_repeat": 0.0030383500270545483,
      "compile_s": 0.4909398799936753,
      "inner": 62,
      "median_ms": 0.7000998387323512,
      "q1_ms": 0.6911743871422649,
      "q3_ms": 0.7479468063871947,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 4.0,
      "rel_err": 0.002210714155808091,
      "rel_err_parts": [
        0.002210714155808091
      ],
      "rel_vs_flex_repeat": 1.4772143913432956e-05,
      "compile_s": 0.05410140300227795,
      "inner": 158,
      "median_ms": 0.17381440501813364,
      "q1_ms": 0.17312972469673835,
      "q3_ms": 0.1746035253100799,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_err": 0.0022156843915581703,
      "rel_err_parts": [
        0.0022156843915581703
      ],
      "rel_vs_flex_repeat": 0.002827050630003214,
      "compile_s": 0.0003838459961116314,
      "inner": 200,
      "median_ms": 0.18097823500283994,
      "q1_ms": 0.16246929499175167,
      "q3_ms": 0.18188169749919325,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 4.0,
      "rel_err": 0.002210708102211356,
      "rel_err_parts": [
        0.002210708102211356
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.1856462910072878,
      "inner": 190,
      "median_ms": 0.156780078970386,
      "q1_ms": 0.156197613174108,
      "q3_ms": 0.15913881051298967,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 8.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 4.0,
      "rel_err": 0.002210708102211356,
      "rel_err_parts": [
        0.002210708102211356
      ],
      "compile_s": 0.1916508709982736,
      "inner": 200,
      "median_ms": 0.17326584493275732,
      "q1_ms": 0.1703889825148508,
      "q3_ms": 0.178544337504718,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 192.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 4.0,
      "rel_err": 0.002334752818569541,
      "rel_err_parts": [
        0.002334752818569541
      ],
      "rel_vs_flex_repeat": 0.0030383500270545483,
      "compile_s": 0.5036439550021896,
      "inner": 64,
      "median_ms": 0.698243968827228,
      "q1_ms": 0.6955018282042147,
      "q3_ms": 0.6999492968589038,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 4.0,
      "rel_err": 0.002210714155808091,
      "rel_err_parts": [
        0.002210714155808091
      ],
      "rel_vs_flex_repeat": 1.4772143913432956e-05,
      "compile_s": 0.052898507987265475,
      "inner": 149,
      "median_ms": 0.18119192617748484,
      "q1_ms": 0.17566482213032025,
      "q3_ms": 0.18502040606010828,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_err": 0.0022156843915581703,
      "rel_err_parts": [
        0.0022156843915581703
      ],
      "rel_vs_flex_repeat": 0.002827050630003214,
      "compile_s": 0.0003779149992624298,
      "inner": 200,
      "median_ms": 0.17762790004780982,
      "q1_ms": 0.17046472501533572,
      "q3_ms": 0.18500545746064745,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 4.0,
      "rel_err": 0.002210708102211356,
      "rel_err_parts": [
        0.002210708102211356
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.35218679599347524,
      "inner": 41,
      "median_ms": 0.17208507300138745,
      "q1_ms": 0.17132448787068402,
      "q3_ms": 0.1732713415815563,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 8.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 4.0,
      "rel_err": 0.002210708102211356,
      "rel_err_parts": [
        0.002210708102211356
      ],
      "compile_s": 0.1972200169984717,
      "inner": 185,
      "median_ms": 0.20322047026097975,
      "q1_ms": 0.19548583785197823,
      "q3_ms": 0.2074199783729037,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 192.0,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 4.0,
      "rel_err": 0.002334752818569541,
      "rel_err_parts": [
        0.002334752818569541
      ],
      "rel_vs_flex_repeat": 0.0030383500270545483,
      "compile_s": 0.6686502070078859,
      "inner": 11,
      "median_ms": 0.8027788188139147,
      "q1_ms": 0.7230246813957241,
      "q3_ms": 0.8568920457037166,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 6.0,
      "xla_out_mb": 4.0,
      "rel_err": 0.002210714155808091,
      "rel_err_parts": [
        0.002210714155808091
      ],
      "rel_vs_flex_repeat": 1.4772143913432956e-05,
      "compile_s": 0.053623409010469913,
      "inner": 163,
      "median_ms": 0.17356842945892637,
      "q1_ms": 0.17260338344569803,
      "q3_ms": 0.17389469018328843,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_err": 0.0022156843915581703,
      "rel_err_parts": [
        0.0022156843915581703
      ],
      "rel_vs_flex_repeat": 0.002827050630003214,
      "compile_s": 0.0003916409914381802,
      "inner": 164,
      "median_ms": 0.16179823779100666,
      "q1_ms": 0.1610154755858608,
      "q3_ms": 0.16263100006631803,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 16.000267028808594,
      "xla_arg_mb": 14.0,
      "xla_out_mb": 6.000022888183594,
      "residual_mb": 4.125011444091797,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 1024, 32, 64]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.003367821453139186,
      "rel_err_parts": [
        0.0029164892621338367,
        0.003337046131491661,
        0.003367821453139186
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.45790790699538775,
      "inner": 61,
      "median_ms": 0.7124846559265995,
      "q1_ms": 0.7060859097095329,
      "q3_ms": 0.7136328935243388,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 24.000267028808594,
      "xla_arg_mb": 14.0,
      "xla_out_mb": 6.000022888183594,
      "residual_mb": 12.125011444091797,
      "residual_shapes": [
        "bfloat16[1, 1024, 32, 64]",
        "bfloat16[1, 1024, 32, 64]",
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 1024, 32, 64]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.003367821453139186,
      "rel_err_parts": [
        0.0029164892621338367,
        0.003337046131491661,
        0.003367821453139186
      ],
      "compile_s": 0.4673764909966849,
      "inner": 60,
      "median_ms": 0.7165625000197906,
      "q1_ms": 0.7125360250938684,
      "q3_ms": 0.7181740250719789,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "jax_xla",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 388.12501525878906,
      "xla_arg_mb": 14.0,
      "xla_out_mb": 6.000022888183594,
      "rel_err": 0.0033365730196237564,
      "rel_err_parts": [
        0.0033365730196237564,
        0.003317640395835042,
        0.0029108314774930477
      ],
      "rel_vs_flex_repeat": 0.003988709766417742,
      "compile_s": 0.19992530399758834,
      "inner": 24,
      "median_ms": 1.9419376670460526,
      "q1_ms": 1.925308083324732,
      "q3_ms": 1.9572070001837951,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "jax_cudnn",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 32.250732421875,
      "xla_arg_mb": 14.0,
      "xla_out_mb": 6.000022888183594,
      "rel_err": 0.0033677832689136267,
      "rel_err_parts": [
        0.002916590543463826,
        0.003337414702400565,
        0.0033677832689136267
      ],
      "rel_vs_flex_repeat": 0.00023062218679115176,
      "compile_s": 0.11843102700368036,
      "inner": 52,
      "median_ms": 0.7597938075867625,
      "q1_ms": 0.7576896058218643,
      "q3_ms": 0.7650721538364171,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 0,
      "rel_err": 0.0033680007327347994,
      "rel_err_parts": [
        0.0029168121982365847,
        0.0033360689412802458,
        0.0033680007327347994
      ],
      "rel_vs_flex_repeat": 0.002138735493645072,
      "compile_s": 0.0011416490015108138,
      "inner": 65,
      "median_ms": 0.7644593384331809,
      "q1_ms": 0.7556785770038215,
      "q3_ms": 0.7664279922699699,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 16.000267028808594,
      "xla_arg_mb": 14.0,
      "xla_out_mb": 6.000022888183594,
      "residual_mb": 4.125011444091797,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 1024, 32, 64]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.003367821453139186,
      "rel_err_parts": [
        0.0029164892621338367,
        0.003337046131491661,
        0.003367821453139186
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.47311603899288457,
      "inner": 67,
      "median_ms": 0.6810157015092616,
      "q1_ms": 0.6785677984198764,
      "q3_ms": 0.6869798283997937,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 24.000267028808594,
      "xla_arg_mb": 14.0,
      "xla_out_mb": 6.000022888183594,
      "residual_mb": 12.125011444091797,
      "residual_shapes": [
        "bfloat16[1, 1024, 32, 64]",
        "bfloat16[1, 1024, 32, 64]",
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 1024, 32, 64]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.003367821453139186,
      "rel_err_parts": [
        0.0029164892621338367,
        0.003337046131491661,
        0.003367821453139186
      ],
      "compile_s": 0.5979960480035516,
      "inner": 64,
      "median_ms": 0.7213966719064047,
      "q1_ms": 0.7142067421455067,
      "q3_ms": 0.7459709765953448,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "jax_xla",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 388.12501525878906,
      "xla_arg_mb": 14.0,
      "xla_out_mb": 6.000022888183594,
      "rel_err": 0.0033365730196237564,
      "rel_err_parts": [
        0.0033365730196237564,
        0.003317640395835042,
        0.0029108314774930477
      ],
      "rel_vs_flex_repeat": 0.003988709766417742,
      "compile_s": 0.20314364398655016,
      "inner": 23,
      "median_ms": 2.236228477989283,
      "q1_ms": 1.9815275217811374,
      "q3_ms": 2.237551652322746,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "jax_cudnn",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 32.250732421875,
      "xla_arg_mb": 14.0,
      "xla_out_mb": 6.000022888183594,
      "rel_err": 0.0033677832689136267,
      "rel_err_parts": [
        0.002916590543463826,
        0.003337414702400565,
        0.0033677832689136267
      ],
      "rel_vs_flex_repeat": 0.00023062218679115176,
      "compile_s": 0.1303524110116996,
      "inner": 51,
      "median_ms": 0.7619697842564361,
      "q1_ms": 0.7605377743625557,
      "q3_ms": 0.8063034705735523,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 1,
      "rel_err": 0.0033680007327347994,
      "rel_err_parts": [
        0.0029168121982365847,
        0.0033360689412802458,
        0.0033680007327347994
      ],
      "rel_vs_flex_repeat": 0.002138735493645072,
      "compile_s": 0.001028736005537212,
      "inner": 65,
      "median_ms": 0.7618681384393803,
      "q1_ms": 0.7600873000158642,
      "q3_ms": 0.7640078692929819,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 16.000267028808594,
      "xla_arg_mb": 14.0,
      "xla_out_mb": 6.000022888183594,
      "residual_mb": 4.125011444091797,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 1024, 32, 64]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.003367821453139186,
      "rel_err_parts": [
        0.0029164892621338367,
        0.003337046131491661,
        0.003367821453139186
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.4553493170096772,
      "inner": 67,
      "median_ms": 0.6760365971270254,
      "q1_ms": 0.6717718058654956,
      "q3_ms": 0.6825328431415385,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 24.000267028808594,
      "xla_arg_mb": 14.0,
      "xla_out_mb": 6.000022888183594,
      "residual_mb": 12.125011444091797,
      "residual_shapes": [
        "bfloat16[1, 1024, 32, 64]",
        "bfloat16[1, 1024, 32, 64]",
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 1024, 32, 64]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.003367821453139186,
      "rel_err_parts": [
        0.0029164892621338367,
        0.003337046131491661,
        0.003367821453139186
      ],
      "compile_s": 0.45528518500213977,
      "inner": 62,
      "median_ms": 0.7116707420209423,
      "q1_ms": 0.7099874032240722,
      "q3_ms": 0.7148733387904,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "jax_xla",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 388.12501525878906,
      "xla_arg_mb": 14.0,
      "xla_out_mb": 6.000022888183594,
      "rel_err": 0.0033365730196237564,
      "rel_err_parts": [
        0.0033365730196237564,
        0.003317640395835042,
        0.0029108314774930477
      ],
      "rel_vs_flex_repeat": 0.003988709766417742,
      "compile_s": 0.19636294498923235,
      "inner": 24,
      "median_ms": 1.952371166529095,
      "q1_ms": 1.9215405415404045,
      "q3_ms": 1.9525915622580214,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "jax_cudnn",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 32.250732421875,
      "xla_arg_mb": 14.0,
      "xla_out_mb": 6.000022888183594,
      "rel_err": 0.0033677832689136267,
      "rel_err_parts": [
        0.002916590543463826,
        0.003337414702400565,
        0.0033677832689136267
      ],
      "rel_vs_flex_repeat": 0.00023062218679115176,
      "compile_s": 0.11912242499238346,
      "inner": 55,
      "median_ms": 0.7592629818563266,
      "q1_ms": 0.7562434454236857,
      "q3_ms": 0.7652617726390334,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 2,
      "rel_err": 0.0033680007327347994,
      "rel_err_parts": [
        0.0029168121982365847,
        0.0033360689412802458,
        0.0033680007327347994
      ],
      "rel_vs_flex_repeat": 0.002138735493645072,
      "compile_s": 0.001032954009133391,
      "inner": 55,
      "median_ms": 0.762521690921858,
      "q1_ms": 0.758875527323901,
      "q3_ms": 0.7634674818043343,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 5.0,
      "xla_out_mb": 4.0,
      "rel_err": 0.0022335934918373823,
      "rel_err_parts": [
        0.0022335934918373823
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.20347174401103985,
      "inner": 187,
      "median_ms": 0.18187901068959425,
      "q1_ms": 0.1727685427702814,
      "q3_ms": 0.18656555350613407,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 8.0,
      "xla_arg_mb": 5.0,
      "xla_out_mb": 4.0,
      "rel_err": 0.0022335934918373823,
      "rel_err_parts": [
        0.0022335934918373823
      ],
      "compile_s": 0.19630694000807125,
      "inner": 187,
      "median_ms": 0.17426857216791675,
      "q1_ms": 0.17114753478031136,
      "q3_ms": 0.17895752675040688,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 192.0,
      "xla_arg_mb": 5.0,
      "xla_out_mb": 4.0,
      "rel_err": 0.0023604859597980976,
      "rel_err_parts": [
        0.0023604859597980976
      ],
      "rel_vs_flex_repeat": 0.0030836209189146757,
      "compile_s": 0.5058503390027909,
      "inner": 63,
      "median_ms": 0.6923792539127467,
      "q1_ms": 0.6887123255716963,
      "q3_ms": 0.6986507301224352,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 5.0,
      "xla_out_mb": 4.0,
      "rel_err": 0.002233583014458418,
      "rel_err_parts": [
        0.002233583014458418
      ],
      "rel_vs_flex_repeat": 1.96528581000166e-05,
      "compile_s": 0.049051529000280425,
      "inner": 170,
      "median_ms": 0.1737706353785196,
      "q1_ms": 0.17355042355655528,
      "q3_ms": 0.17441049413344656,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_err": 0.00223964499309659,
      "rel_err_parts": [
        0.00223964499309659
      ],
      "rel_vs_flex_repeat": 0.0028679301030933857,
      "compile_s": 0.00038241299625951797,
      "inner": 200,
      "median_ms": 0.16308278500218876,
      "q1_ms": 0.1624166225155932,
      "q3_ms": 0.1634772300167242,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 5.0,
      "xla_out_mb": 4.0,
      "rel_err": 0.0022335934918373823,
      "rel_err_parts": [
        0.0022335934918373823
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.1973475990089355,
      "inner": 189,
      "median_ms": 0.1588958518354163,
      "q1_ms": 0.15430714281845012,
      "q3_ms": 0.158969293619203,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 8.0,
      "xla_arg_mb": 5.0,
      "xla_out_mb": 4.0,
      "rel_err": 0.0022335934918373823,
      "rel_err_parts": [
        0.0022335934918373823
      ],
      "compile_s": 0.18740051999338903,
      "inner": 162,
      "median_ms": 0.17280327777786433,
      "q1_ms": 0.17006111729755008,
      "q3_ms": 0.17293906479787918,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 192.0,
      "xla_arg_mb": 5.0,
      "xla_out_mb": 4.0,
      "rel_err": 0.0023604859597980976,
      "rel_err_parts": [
        0.0023604859597980976
      ],
      "rel_vs_flex_repeat": 0.0030836209189146757,
      "compile_s": 0.5025332529912703,
      "inner": 63,
      "median_ms": 0.7212440317426438,
      "q1_ms": 0.706686015912546,
      "q3_ms": 0.7756452063663997,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 5.0,
      "xla_out_mb": 4.0,
      "rel_err": 0.002233583014458418,
      "rel_err_parts": [
        0.002233583014458418
      ],
      "rel_vs_flex_repeat": 1.96528581000166e-05,
      "compile_s": 0.048612979997415096,
      "inner": 168,
      "median_ms": 0.17686276190770078,
      "q1_ms": 0.17390081544969388,
      "q3_ms": 0.1842945237599668,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_err": 0.00223964499309659,
      "rel_err_parts": [
        0.00223964499309659
      ],
      "rel_vs_flex_repeat": 0.0028679301030933857,
      "compile_s": 0.00042795999615918845,
      "inner": 200,
      "median_ms": 0.16223622995312326,
      "q1_ms": 0.16106598999613198,
      "q3_ms": 0.1681094925152138,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 5.0,
      "xla_out_mb": 4.0,
      "rel_err": 0.0022335934918373823,
      "rel_err_parts": [
        0.0022335934918373823
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.38643905700882897,
      "inner": 43,
      "median_ms": 0.1713802792987418,
      "q1_ms": 0.16074154649855743,
      "q3_ms": 0.18493988369131287,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 8.0,
      "xla_arg_mb": 5.0,
      "xla_out_mb": 4.0,
      "rel_err": 0.0022335934918373823,
      "rel_err_parts": [
        0.0022335934918373823
      ],
      "compile_s": 0.32075002099736594,
      "inner": 172,
      "median_ms": 0.19605626161941217,
      "q1_ms": 0.19191269479220907,
      "q3_ms": 0.20013345352327994,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 192.0,
      "xla_arg_mb": 5.0,
      "xla_out_mb": 4.0,
      "rel_err": 0.0023604859597980976,
      "rel_err_parts": [
        0.0023604859597980976
      ],
      "rel_vs_flex_repeat": 0.0030836209189146757,
      "compile_s": 0.5756928030023118,
      "inner": 62,
      "median_ms": 0.6944812905144758,
      "q1_ms": 0.6938259434865622,
      "q3_ms": 0.7030472581625932,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 5.0,
      "xla_out_mb": 4.0,
      "rel_err": 0.002233583014458418,
      "rel_err_parts": [
        0.002233583014458418
      ],
      "rel_vs_flex_repeat": 1.96528581000166e-05,
      "compile_s": 0.04655437500332482,
      "inner": 167,
      "median_ms": 0.1871533533119078,
      "q1_ms": 0.17531427543806916,
      "q3_ms": 0.20054794609557813,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_err": 0.00223964499309659,
      "rel_err_parts": [
        0.00223964499309659
      ],
      "rel_vs_flex_repeat": 0.0028679301030933857,
      "compile_s": 0.00041416399471927434,
      "inner": 200,
      "median_ms": 0.16638417502690572,
      "q1_ms": 0.16108192252431763,
      "q3_ms": 0.16941994752414757,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 16.000267028808594,
      "xla_arg_mb": 13.0,
      "xla_out_mb": 5.000022888183594,
      "residual_mb": 4.125011444091797,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 1024, 32, 64]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.0033154087141156197,
      "rel_err_parts": [
        0.0029097700025886297,
        0.003308901796117425,
        0.0033154087141156197
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.4698809189867461,
      "inner": 61,
      "median_ms": 0.7191704262666343,
      "q1_ms": 0.7063519918373176,
      "q3_ms": 0.7354277376940505,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 24.000267028808594,
      "xla_arg_mb": 13.0,
      "xla_out_mb": 5.000022888183594,
      "residual_mb": 12.125011444091797,
      "residual_shapes": [
        "bfloat16[1, 1024, 32, 64]",
        "bfloat16[1, 1024, 32, 64]",
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 1024, 32, 64]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.0033154087141156197,
      "rel_err_parts": [
        0.0029097700025886297,
        0.003308901796117425,
        0.0033154087141156197
      ],
      "compile_s": 0.48255562799749896,
      "inner": 60,
      "median_ms": 0.7186470999537656,
      "q1_ms": 0.7075602000744159,
      "q3_ms": 0.7348499249928864,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "jax_xla",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 388.12501525878906,
      "xla_arg_mb": 13.0,
      "xla_out_mb": 5.000022888183594,
      "rel_err": 0.0033371674362570047,
      "rel_err_parts": [
        0.0033371674362570047,
        0.003313656197860837,
        0.0028673475608229637
      ],
      "rel_vs_flex_repeat": 0.003959808498620987,
      "compile_s": 0.19858455899520777,
      "inner": 24,
      "median_ms": 2.042906625016864,
      "q1_ms": 1.9412008544653268,
      "q3_ms": 2.1132190001177755,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "jax_cudnn",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 32.250732421875,
      "xla_arg_mb": 13.0,
      "xla_out_mb": 5.000022888183594,
      "rel_err": 0.00331525644287467,
      "rel_err_parts": [
        0.0029097539372742176,
        0.0033090035431087017,
        0.00331525644287467
      ],
      "rel_vs_flex_repeat": 0.00025606012786738575,
      "compile_s": 0.12115861699567176,
      "inner": 54,
      "median_ms": 0.756214740653781,
      "q1_ms": 0.7553302962656339,
      "q3_ms": 0.7598378056341231,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 0,
      "rel_err": 0.0033152687828987837,
      "rel_err_parts": [
        0.0029092549812048674,
        0.003308732993900776,
        0.0033152687828987837
      ],
      "rel_vs_flex_repeat": 0.0022868006490170956,
      "compile_s": 0.0010683010041248053,
      "inner": 64,
      "median_ms": 0.7587745781165722,
      "q1_ms": 0.7542711327914731,
      "q3_ms": 0.8034825235654353,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 16.000267028808594,
      "xla_arg_mb": 13.0,
      "xla_out_mb": 5.000022888183594,
      "residual_mb": 4.125011444091797,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 1024, 32, 64]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.0033154087141156197,
      "rel_err_parts": [
        0.0029097700025886297,
        0.003308901796117425,
        0.0033154087141156197
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.48043033698922954,
      "inner": 65,
      "median_ms": 0.6770243845844212,
      "q1_ms": 0.6761919692507945,
      "q3_ms": 0.7329756998492835,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 24.000267028808594,
      "xla_arg_mb": 13.0,
      "xla_out_mb": 5.000022888183594,
      "residual_mb": 12.125011444091797,
      "residual_shapes": [
        "bfloat16[1, 1024, 32, 64]",
        "bfloat16[1, 1024, 32, 64]",
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 1024, 32, 64]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.0033154087141156197,
      "rel_err_parts": [
        0.0029097700025886297,
        0.003308901796117425,
        0.0033154087141156197
      ],
      "compile_s": 0.48989042500033975,
      "inner": 65,
      "median_ms": 0.7487679538406575,
      "q1_ms": 0.7103778768886024,
      "q3_ms": 0.7870809076568828,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "jax_xla",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 388.12501525878906,
      "xla_arg_mb": 13.0,
      "xla_out_mb": 5.000022888183594,
      "rel_err": 0.0033371674362570047,
      "rel_err_parts": [
        0.0033371674362570047,
        0.003313656197860837,
        0.0028673475608229637
      ],
      "rel_vs_flex_repeat": 0.003959808498620987,
      "compile_s": 0.19927931299025659,
      "inner": 24,
      "median_ms": 1.9481978330683585,
      "q1_ms": 1.9415289583169701,
      "q3_ms": 1.967039749918816,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "jax_cudnn",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 32.250732421875,
      "xla_arg_mb": 13.0,
      "xla_out_mb": 5.000022888183594,
      "rel_err": 0.00331525644287467,
      "rel_err_parts": [
        0.0029097539372742176,
        0.0033090035431087017,
        0.00331525644287467
      ],
      "rel_vs_flex_repeat": 0.00025606012786738575,
      "compile_s": 0.11944873200263828,
      "inner": 53,
      "median_ms": 0.7559216228744261,
      "q1_ms": 0.752881075418115,
      "q3_ms": 0.764042367902804,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 1,
      "rel_err": 0.0033152687828987837,
      "rel_err_parts": [
        0.0029092549812048674,
        0.003308732993900776,
        0.0033152687828987837
      ],
      "rel_vs_flex_repeat": 0.0022868006490170956,
      "compile_s": 0.0010099200007971376,
      "inner": 65,
      "median_ms": 0.7663754307521651,
      "q1_ms": 0.7603798538464337,
      "q3_ms": 0.7760863385136383,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 16.000267028808594,
      "xla_arg_mb": 13.0,
      "xla_out_mb": 5.000022888183594,
      "residual_mb": 4.125011444091797,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 1024, 32, 64]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.0033154087141156197,
      "rel_err_parts": [
        0.0029097700025886297,
        0.003308901796117425,
        0.0033154087141156197
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.4627806960052112,
      "inner": 61,
      "median_ms": 0.6927831149011179,
      "q1_ms": 0.6765284345076862,
      "q3_ms": 0.7461962214008844,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 24.000267028808594,
      "xla_arg_mb": 13.0,
      "xla_out_mb": 5.000022888183594,
      "residual_mb": 12.125011444091797,
      "residual_shapes": [
        "bfloat16[1, 1024, 32, 64]",
        "bfloat16[1, 1024, 32, 64]",
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 1024, 32, 64]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.0033154087141156197,
      "rel_err_parts": [
        0.0029097700025886297,
        0.003308901796117425,
        0.0033154087141156197
      ],
      "compile_s": 0.4758190129941795,
      "inner": 55,
      "median_ms": 0.7025848544964736,
      "q1_ms": 0.6997416364769875,
      "q3_ms": 0.7091528999312273,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "jax_xla",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 388.12501525878906,
      "xla_arg_mb": 13.0,
      "xla_out_mb": 5.000022888183594,
      "rel_err": 0.0033371674362570047,
      "rel_err_parts": [
        0.0033371674362570047,
        0.003313656197860837,
        0.0028673475608229637
      ],
      "rel_vs_flex_repeat": 0.003959808498620987,
      "compile_s": 0.20021081500453874,
      "inner": 24,
      "median_ms": 1.9431633336353116,
      "q1_ms": 1.913383854116546,
      "q3_ms": 1.9937715621078191,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "jax_cudnn",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 32.250732421875,
      "xla_arg_mb": 13.0,
      "xla_out_mb": 5.000022888183594,
      "rel_err": 0.00331525644287467,
      "rel_err_parts": [
        0.0029097539372742176,
        0.0033090035431087017,
        0.00331525644287467
      ],
      "rel_vs_flex_repeat": 0.00025606012786738575,
      "compile_s": 0.15221886000654195,
      "inner": 51,
      "median_ms": 0.8313235488386971,
      "q1_ms": 0.7787764805149944,
      "q3_ms": 0.8433282256352387,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 2,
      "rel_err": 0.0033152687828987837,
      "rel_err_parts": [
        0.0029092549812048674,
        0.003308732993900776,
        0.0033152687828987837
      ],
      "rel_vs_flex_repeat": 0.0022868006490170956,
      "compile_s": 0.0011737309978343546,
      "inner": 64,
      "median_ms": 0.7588522187234048,
      "q1_ms": 0.7552241796702219,
      "q3_ms": 0.7631321093413135,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 4.25,
      "xla_out_mb": 4.0,
      "rel_err": 0.002221020171418786,
      "rel_err_parts": [
        0.002221020171418786
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.19421677399077453,
      "inner": 183,
      "median_ms": 0.17157366667566262,
      "q1_ms": 0.17067822677819974,
      "q3_ms": 0.17748014482000302,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 8.0,
      "xla_arg_mb": 4.25,
      "xla_out_mb": 4.0,
      "rel_err": 0.002221020171418786,
      "rel_err_parts": [
        0.002221020171418786
      ],
      "compile_s": 0.3861973890016088,
      "inner": 62,
      "median_ms": 0.17693893553974527,
      "q1_ms": 0.1732645726024789,
      "q3_ms": 0.17725115321809426,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 192.0,
      "xla_arg_mb": 4.25,
      "xla_out_mb": 4.0,
      "rel_err": 0.002346956869587302,
      "rel_err_parts": [
        0.002346956869587302
      ],
      "rel_vs_flex_repeat": 0.0030613942071795464,
      "compile_s": 0.39798779600823764,
      "inner": 69,
      "median_ms": 0.5470938694602607,
      "q1_ms": 0.5403469058101241,
      "q3_ms": 0.5536255797040581,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 4.25,
      "xla_out_mb": 4.0,
      "rel_err": 0.002221018308773637,
      "rel_err_parts": [
        0.002221018308773637
      ],
      "rel_vs_flex_repeat": 1.5433326552738436e-05,
      "compile_s": 0.051954170994577,
      "inner": 113,
      "median_ms": 0.19789642481281694,
      "q1_ms": 0.19611283191283999,
      "q3_ms": 0.198486265448227,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_err": 0.002228848170489073,
      "rel_err_parts": [
        0.002228848170489073
      ],
      "rel_vs_flex_repeat": 0.0028533251024782658,
      "compile_s": 0.00038777399458922446,
      "inner": 200,
      "median_ms": 0.16398408501117956,
      "q1_ms": 0.1635465625076904,
      "q3_ms": 0.16532026249478804,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 4.25,
      "xla_out_mb": 4.0,
      "rel_err": 0.002221020171418786,
      "rel_err_parts": [
        0.002221020171418786
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.19950187399808783,
      "inner": 167,
      "median_ms": 0.1547620120258935,
      "q1_ms": 0.15336458084506713,
      "q3_ms": 0.15477998207273064,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 8.0,
      "xla_arg_mb": 4.25,
      "xla_out_mb": 4.0,
      "rel_err": 0.002221020171418786,
      "rel_err_parts": [
        0.002221020171418786
      ],
      "compile_s": 0.19721646000107285,
      "inner": 197,
      "median_ms": 0.16994735531551472,
      "q1_ms": 0.16933530962790794,
      "q3_ms": 0.17559769288441734,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 192.0,
      "xla_arg_mb": 4.25,
      "xla_out_mb": 4.0,
      "rel_err": 0.002346956869587302,
      "rel_err_parts": [
        0.002346956869587302
      ],
      "rel_vs_flex_repeat": 0.0030613942071795464,
      "compile_s": 0.20323098900553305,
      "inner": 79,
      "median_ms": 0.5344640505220152,
      "q1_ms": 0.5330780126475669,
      "q3_ms": 0.5356447722009807,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 4.25,
      "xla_out_mb": 4.0,
      "rel_err": 0.002221018308773637,
      "rel_err_parts": [
        0.002221018308773637
      ],
      "rel_vs_flex_repeat": 1.5433326552738436e-05,
      "compile_s": 0.04694469399692025,
      "inner": 175,
      "median_ms": 0.17351512572661573,
      "q1_ms": 0.17309464572463185,
      "q3_ms": 0.18069296285310493,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_err": 0.002228848170489073,
      "rel_err_parts": [
        0.002228848170489073
      ],
      "rel_vs_flex_repeat": 0.0028533251024782658,
      "compile_s": 0.00037939700996503234,
      "inner": 200,
      "median_ms": 0.18303071999980602,
      "q1_ms": 0.1763599249534309,
      "q3_ms": 0.18432193999615265,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 4.25,
      "xla_out_mb": 4.0,
      "rel_err": 0.002221020171418786,
      "rel_err_parts": [
        0.002221020171418786
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.19251853199966718,
      "inner": 195,
      "median_ms": 0.16798935382841873,
      "q1_ms": 0.15304649486773625,
      "q3_ms": 0.17226904101568297,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 8.0,
      "xla_arg_mb": 4.25,
      "xla_out_mb": 4.0,
      "rel_err": 0.002221020171418786,
      "rel_err_parts": [
        0.002221020171418786
      ],
      "compile_s": 0.20734547299798578,
      "inner": 178,
      "median_ms": 0.17176957861808118,
      "q1_ms": 0.16934704215375684,
      "q3_ms": 0.1725685870859594,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 192.0,
      "xla_arg_mb": 4.25,
      "xla_out_mb": 4.0,
      "rel_err": 0.002346956869587302,
      "rel_err_parts": [
        0.002346956869587302
      ],
      "rel_vs_flex_repeat": 0.0030613942071795464,
      "compile_s": 0.20815009500074666,
      "inner": 79,
      "median_ms": 0.5984123797329803,
      "q1_ms": 0.5747958608089556,
      "q3_ms": 0.6062621013128753,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 4.25,
      "xla_out_mb": 4.0,
      "rel_err": 0.002221018308773637,
      "rel_err_parts": [
        0.002221018308773637
      ],
      "rel_vs_flex_repeat": 1.5433326552738436e-05,
      "compile_s": 0.0493059610016644,
      "inner": 166,
      "median_ms": 0.19390974693485616,
      "q1_ms": 0.1774051445655536,
      "q3_ms": 0.19429870477017766,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_err": 0.002228848170489073,
      "rel_err_parts": [
        0.002228848170489073
      ],
      "rel_vs_flex_repeat": 0.0028533251024782658,
      "compile_s": 0.0003991550038335845,
      "inner": 200,
      "median_ms": 0.1645533650298603,
      "q1_ms": 0.16306477249599993,
      "q3_ms": 0.16885973753232975,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 16.000267028808594,
      "xla_arg_mb": 12.25,
      "xla_out_mb": 4.250022888183594,
      "residual_mb": 4.125011444091797,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 1024, 32, 64]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.003306942293420434,
      "rel_err_parts": [
        0.002911139512434602,
        0.003306942293420434,
        0.0032303633634001017
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.9799894629977643,
      "inner": 38,
      "median_ms": 0.7690966051182171,
      "q1_ms": 0.7204797760718257,
      "q3_ms": 0.7895761447894314,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 24.000267028808594,
      "xla_arg_mb": 12.25,
      "xla_out_mb": 4.250022888183594,
      "residual_mb": 12.125011444091797,
      "residual_shapes": [
        "bfloat16[1, 1024, 32, 64]",
        "bfloat16[1, 1024, 32, 64]",
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 1024, 32, 64]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.003306942293420434,
      "rel_err_parts": [
        0.002911139512434602,
        0.003306942293420434,
        0.0032303633634001017
      ],
      "compile_s": 1.0097769379935926,
      "inner": 9,
      "median_ms": 0.8884942217264324,
      "q1_ms": 0.8322862777276896,
      "q3_ms": 0.8964966109488159,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "jax_xla",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 384.12501525878906,
      "xla_arg_mb": 12.25,
      "xla_out_mb": 4.250022888183594,
      "rel_err": 0.0033396189101040363,
      "rel_err_parts": [
        0.0033396189101040363,
        0.003291150089353323,
        0.002781081013381481
      ],
      "rel_vs_flex_repeat": 0.003921579569578171,
      "compile_s": 0.4265160270006163,
      "inner": 22,
      "median_ms": 2.006065681846601,
      "q1_ms": 1.958520068001235,
      "q3_ms": 2.2336150452081873,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "jax_cudnn",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 32.250732421875,
      "xla_arg_mb": 12.25,
      "xla_out_mb": 4.250022888183594,
      "rel_err": 0.0033070126082748175,
      "rel_err_parts": [
        0.0029113164637237787,
        0.0033070126082748175,
        0.0032301251776516438
      ],
      "rel_vs_flex_repeat": 0.00033537071431055665,
      "compile_s": 0.1283872829953907,
      "inner": 55,
      "median_ms": 0.78492861807304,
      "q1_ms": 0.7524669363523241,
      "q3_ms": 0.8612300818309779,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 0,
      "rel_err": 0.003308331361040473,
      "rel_err_parts": [
        0.0029122005216777325,
        0.003308331361040473,
        0.0032280059531331062
      ],
      "rel_vs_flex_repeat": 0.0023965779691934586,
      "compile_s": 0.0011311800044495612,
      "inner": 65,
      "median_ms": 0.8352216767469565,
      "q1_ms": 0.7726785152604302,
      "q3_ms": 0.850243561525041,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 16.000267028808594,
      "xla_arg_mb": 12.25,
      "xla_out_mb": 4.250022888183594,
      "residual_mb": 4.125011444091797,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 1024, 32, 64]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.003306942293420434,
      "rel_err_parts": [
        0.002911139512434602,
        0.003306942293420434,
        0.0032303633634001017
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.4754147680068854,
      "inner": 62,
      "median_ms": 0.6745610322292535,
      "q1_ms": 0.6737372016012217,
      "q3_ms": 0.6799477500129386,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 24.000267028808594,
      "xla_arg_mb": 12.25,
      "xla_out_mb": 4.250022888183594,
      "residual_mb": 12.125011444091797,
      "residual_shapes": [
        "bfloat16[1, 1024, 32, 64]",
        "bfloat16[1, 1024, 32, 64]",
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 1024, 32, 64]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.003306942293420434,
      "rel_err_parts": [
        0.002911139512434602,
        0.003306942293420434,
        0.0032303633634001017
      ],
      "compile_s": 0.5701458889961941,
      "inner": 62,
      "median_ms": 0.7120731289990456,
      "q1_ms": 0.7082321452241647,
      "q3_ms": 0.714518806372296,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "jax_xla",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 384.12501525878906,
      "xla_arg_mb": 12.25,
      "xla_out_mb": 4.250022888183594,
      "rel_err": 0.0033396189101040363,
      "rel_err_parts": [
        0.0033396189101040363,
        0.003291150089353323,
        0.002781081013381481
      ],
      "rel_vs_flex_repeat": 0.003921579569578171,
      "compile_s": 0.18350553000345826,
      "inner": 25,
      "median_ms": 1.989659399841912,
      "q1_ms": 1.956703239993658,
      "q3_ms": 2.019664359977469,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "jax_cudnn",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 32.250732421875,
      "xla_arg_mb": 12.25,
      "xla_out_mb": 4.250022888183594,
      "rel_err": 0.0033070126082748175,
      "rel_err_parts": [
        0.0029113164637237787,
        0.0033070126082748175,
        0.0032301251776516438
      ],
      "rel_vs_flex_repeat": 0.00033537071431055665,
      "compile_s": 0.21623052199720405,
      "inner": 47,
      "median_ms": 0.7625347448567404,
      "q1_ms": 0.7586642340298723,
      "q3_ms": 0.76823094670854,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 1,
      "rel_err": 0.003308331361040473,
      "rel_err_parts": [
        0.0029122005216777325,
        0.003308331361040473,
        0.0032280059531331062
      ],
      "rel_vs_flex_repeat": 0.0023965779691934586,
      "compile_s": 0.001411220000591129,
      "inner": 63,
      "median_ms": 0.817537793496047,
      "q1_ms": 0.772649936340121,
      "q3_ms": 0.8445739443968272,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 16.000267028808594,
      "xla_arg_mb": 12.25,
      "xla_out_mb": 4.250022888183594,
      "residual_mb": 4.125011444091797,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 1024, 32, 64]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.003306942293420434,
      "rel_err_parts": [
        0.002911139512434602,
        0.003306942293420434,
        0.0032303633634001017
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.886082331999205,
      "inner": 63,
      "median_ms": 0.7173212063554232,
      "q1_ms": 0.7126047539280815,
      "q3_ms": 0.735560722306319,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 24.000267028808594,
      "xla_arg_mb": 12.25,
      "xla_out_mb": 4.250022888183594,
      "residual_mb": 12.125011444091797,
      "residual_shapes": [
        "bfloat16[1, 1024, 32, 64]",
        "bfloat16[1, 1024, 32, 64]",
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 1024, 32, 64]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.003306942293420434,
      "rel_err_parts": [
        0.002911139512434602,
        0.003306942293420434,
        0.0032303633634001017
      ],
      "compile_s": 0.5024895399983507,
      "inner": 63,
      "median_ms": 0.7846469998762514,
      "q1_ms": 0.7222752618771564,
      "q3_ms": 0.7922082064425691,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "jax_xla",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 384.12501525878906,
      "xla_arg_mb": 12.25,
      "xla_out_mb": 4.250022888183594,
      "rel_err": 0.0033396189101040363,
      "rel_err_parts": [
        0.0033396189101040363,
        0.003291150089353323,
        0.002781081013381481
      ],
      "rel_vs_flex_repeat": 0.003921579569578171,
      "compile_s": 0.19475538400001824,
      "inner": 24,
      "median_ms": 1.9413565829988026,
      "q1_ms": 1.9304133544816673,
      "q3_ms": 1.9489649169675736,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "jax_cudnn",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 32.250732421875,
      "xla_arg_mb": 12.25,
      "xla_out_mb": 4.250022888183594,
      "rel_err": 0.0033070126082748175,
      "rel_err_parts": [
        0.0029113164637237787,
        0.0033070126082748175,
        0.0032301251776516438
      ],
      "rel_vs_flex_repeat": 0.00033537071431055665,
      "compile_s": 0.12291950000508223,
      "inner": 55,
      "median_ms": 0.7608341272729873,
      "q1_ms": 0.7546853181503883,
      "q3_ms": 0.7837950273162939,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 2,
      "rel_err": 0.003308331361040473,
      "rel_err_parts": [
        0.0029122005216777325,
        0.003308331361040473,
        0.0032280059531331062
      ],
      "rel_vs_flex_repeat": 0.0023965779691934586,
      "compile_s": 0.0011520400003064424,
      "inner": 63,
      "median_ms": 0.8445318888998516,
      "q1_ms": 0.7813327062490678,
      "q3_ms": 0.8469578967797422,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 8.0,
      "rel_err": 0.0022291461937129498,
      "rel_err_parts": [
        0.0022291461937129498
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.28060713499144185,
      "inner": 99,
      "median_ms": 0.383884343348039,
      "q1_ms": 0.3822577828858894,
      "q3_ms": 0.38843982827419304,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 16.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 8.0,
      "rel_err": 0.0022291461937129498,
      "rel_err_parts": [
        0.0022291461937129498
      ],
      "compile_s": 0.2790545580064645,
      "inner": 101,
      "median_ms": 0.37657411879684666,
      "q1_ms": 0.37532705949944456,
      "q3_ms": 0.37790926732065266,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 192.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 8.0,
      "rel_err": 0.00235233292914927,
      "rel_err_parts": [
        0.00235233292914927
      ],
      "rel_vs_flex_repeat": 0.003067564219236374,
      "compile_s": 0.5116554609994637,
      "inner": 55,
      "median_ms": 0.7723739637929776,
      "q1_ms": 0.7719387818889862,
      "q3_ms": 0.8019944637040184,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 8.0,
      "rel_err": 0.002229159465059638,
      "rel_err_parts": [
        0.002229159465059638
      ],
      "rel_vs_flex_repeat": 2.1710608052671887e-05,
      "compile_s": 0.045031072993879206,
      "inner": 107,
      "median_ms": 0.3339171588581891,
      "q1_ms": 0.33329265423621723,
      "q3_ms": 0.3350427429731413,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_err": 0.002223329385742545,
      "rel_err_parts": [
        0.002223329385742545
      ],
      "rel_vs_flex_repeat": 0.0028966718818992376,
      "compile_s": 0.0005399110086727887,
      "inner": 149,
      "median_ms": 0.35499207380121484,
      "q1_ms": 0.34508107046654524,
      "q3_ms": 0.35807219795506107,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 8.0,
      "rel_err": 0.0022291461937129498,
      "rel_err_parts": [
        0.0022291461937129498
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.2610280539956875,
      "inner": 109,
      "median_ms": 0.3457291284309051,
      "q1_ms": 0.3452968394246012,
      "q3_ms": 0.3465557798135539,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 16.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 8.0,
      "rel_err": 0.0022291461937129498,
      "rel_err_parts": [
        0.0022291461937129498
      ],
      "compile_s": 0.2671836899971822,
      "inner": 106,
      "median_ms": 0.3923387641475958,
      "q1_ms": 0.37914347174008517,
      "q3_ms": 0.398293004729807,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 192.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 8.0,
      "rel_err": 0.00235233292914927,
      "rel_err_parts": [
        0.00235233292914927
      ],
      "rel_vs_flex_repeat": 0.003067564219236374,
      "compile_s": 0.6375653280119877,
      "inner": 59,
      "median_ms": 0.7869335253016582,
      "q1_ms": 0.7751735762921872,
      "q3_ms": 0.8687996864170862,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 8.0,
      "rel_err": 0.002229159465059638,
      "rel_err_parts": [
        0.002229159465059638
      ],
      "rel_vs_flex_repeat": 2.1710608052671887e-05,
      "compile_s": 0.08782671800872777,
      "inner": 103,
      "median_ms": 0.3371358835262708,
      "q1_ms": 0.33052774275894936,
      "q3_ms": 0.3435133349713282,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_err": 0.002223329385742545,
      "rel_err_parts": [
        0.002223329385742545
      ],
      "rel_vs_flex_repeat": 0.0028966718818992376,
      "compile_s": 0.0005661720060743392,
      "inner": 150,
      "median_ms": 0.3149541199672967,
      "q1_ms": 0.3149213300396999,
      "q3_ms": 0.3184103967214469,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 8.0,
      "rel_err": 0.0022291461937129498,
      "rel_err_parts": [
        0.0022291461937129498
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.27348812800482847,
      "inner": 101,
      "median_ms": 0.35085505949988005,
      "q1_ms": 0.34802857422475053,
      "q3_ms": 0.35899064845364026,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 16.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 8.0,
      "rel_err": 0.0022291461937129498,
      "rel_err_parts": [
        0.0022291461937129498
      ],
      "compile_s": 0.26890031700895634,
      "inner": 101,
      "median_ms": 0.3771716831102137,
      "q1_ms": 0.37669394553424607,
      "q3_ms": 0.3800045197920662,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 192.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 8.0,
      "rel_err": 0.00235233292914927,
      "rel_err_parts": [
        0.00235233292914927
      ],
      "rel_vs_flex_repeat": 0.003067564219236374,
      "compile_s": 0.5418317510047927,
      "inner": 58,
      "median_ms": 0.7607886379038723,
      "q1_ms": 0.7592869052163662,
      "q3_ms": 0.7840885775984698,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 8.0,
      "rel_err": 0.002229159465059638,
      "rel_err_parts": [
        0.002229159465059638
      ],
      "rel_vs_flex_repeat": 2.1710608052671887e-05,
      "compile_s": 0.04454798900405876,
      "inner": 118,
      "median_ms": 0.33247416951965886,
      "q1_ms": 0.33019426275316976,
      "q3_ms": 0.3360653347538003,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_err": 0.002223329385742545,
      "rel_err_parts": [
        0.002223329385742545
      ],
      "rel_vs_flex_repeat": 0.0028966718818992376,
      "compile_s": 0.0005296919989632443,
      "inner": 148,
      "median_ms": 0.3147174323424445,
      "q1_ms": 0.31443453716181813,
      "q3_ms": 0.3194996621164751,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 32.000267028808594,
      "xla_arg_mb": 28.0,
      "xla_out_mb": 12.000022888183594,
      "residual_mb": 8.125011444091797,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 1024, 32, 128]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.003344526281580329,
      "rel_err_parts": [
        0.0028945563826709986,
        0.0033384172711521387,
        0.003344526281580329
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.5114037649909733,
      "inner": 28,
      "median_ms": 1.6570553213698855,
      "q1_ms": 1.641537732211873,
      "q3_ms": 1.6847121070376099,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 48.000267028808594,
      "xla_arg_mb": 28.0,
      "xla_out_mb": 12.000022888183594,
      "residual_mb": 24.125011444091797,
      "residual_shapes": [
        "bfloat16[1, 1024, 32, 128]",
        "bfloat16[1, 1024, 32, 128]",
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 1024, 32, 128]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.003344526281580329,
      "rel_err_parts": [
        0.0028945563826709986,
        0.0033384172711521387,
        0.003344526281580329
      ],
      "compile_s": 0.5248940570018021,
      "inner": 29,
      "median_ms": 1.5658973104015377,
      "q1_ms": 1.5590712932392088,
      "q3_ms": 1.6085740862763487,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "jax_xla",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 392.12501525878906,
      "xla_arg_mb": 28.0,
      "xla_out_mb": 12.000022888183594,
      "rel_err": 0.0033211391419172287,
      "rel_err_parts": [
        0.0033211391419172287,
        0.003312610788270831,
        0.0028861421160399914
      ],
      "rel_vs_flex_repeat": 0.0039905039593577385,
      "compile_s": 0.3439257369900588,
      "inner": 23,
      "median_ms": 2.063271739525194,
      "q1_ms": 2.06270565199109,
      "q3_ms": 2.0841549131032284,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "jax_cudnn",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 64.250732421875,
      "xla_arg_mb": 28.0,
      "xla_out_mb": 12.000022888183594,
      "rel_err": 0.0033443919382989407,
      "rel_err_parts": [
        0.0028946080710738897,
        0.0033382740803062916,
        0.0033443919382989407
      ],
      "rel_vs_flex_repeat": 0.00022865409846417606,
      "compile_s": 0.12008083899854682,
      "inner": 32,
      "median_ms": 1.3733651248912793,
      "q1_ms": 1.367475062579615,
      "q3_ms": 1.4508266717712104,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 0,
      "rel_err": 0.003344413358718157,
      "rel_err_parts": [
        0.002894316567108035,
        0.0033402168191969395,
        0.003344413358718157
      ],
      "rel_vs_flex_repeat": 0.003851044923067093,
      "compile_s": 0.0017526659939903766,
      "inner": 36,
      "median_ms": 1.4269151110460774,
      "q1_ms": 1.383155249843387,
      "q3_ms": 1.5494544307633382,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 32.000267028808594,
      "xla_arg_mb": 28.0,
      "xla_out_mb": 12.000022888183594,
      "residual_mb": 8.125011444091797,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 1024, 32, 128]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.003344526281580329,
      "rel_err_parts": [
        0.0028945563826709986,
        0.0033384172711521387,
        0.003344526281580329
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.6195968029933283,
      "inner": 31,
      "median_ms": 1.5372531291382807,
      "q1_ms": 1.509975419462388,
      "q3_ms": 1.5597518389357343,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 48.000267028808594,
      "xla_arg_mb": 28.0,
      "xla_out_mb": 12.000022888183594,
      "residual_mb": 24.125011444091797,
      "residual_shapes": [
        "bfloat16[1, 1024, 32, 128]",
        "bfloat16[1, 1024, 32, 128]",
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 1024, 32, 128]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.003344526281580329,
      "rel_err_parts": [
        0.0028945563826709986,
        0.0033384172711521387,
        0.003344526281580329
      ],
      "compile_s": 0.5404538739967393,
      "inner": 30,
      "median_ms": 1.5659365662334799,
      "q1_ms": 1.5609452998129805,
      "q3_ms": 1.5858821333191977,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "jax_xla",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 392.12501525878906,
      "xla_arg_mb": 28.0,
      "xla_out_mb": 12.000022888183594,
      "rel_err": 0.0033211391419172287,
      "rel_err_parts": [
        0.0033211391419172287,
        0.003312610788270831,
        0.0028861421160399914
      ],
      "rel_vs_flex_repeat": 0.0039905039593577385,
      "compile_s": 0.34236094600055367,
      "inner": 23,
      "median_ms": 2.103750478177655,
      "q1_ms": 2.100350999883275,
      "q3_ms": 2.1214014562682006,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "jax_cudnn",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 64.250732421875,
      "xla_arg_mb": 28.0,
      "xla_out_mb": 12.000022888183594,
      "rel_err": 0.0033443919382989407,
      "rel_err_parts": [
        0.0028946080710738897,
        0.0033382740803062916,
        0.0033443919382989407
      ],
      "rel_vs_flex_repeat": 0.00022865409846417606,
      "compile_s": 0.11566527500690427,
      "inner": 32,
      "median_ms": 1.3687213122466346,
      "q1_ms": 1.3630543592171307,
      "q3_ms": 1.3877484218483005,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 1,
      "rel_err": 0.003344413358718157,
      "rel_err_parts": [
        0.002894316567108035,
        0.0033402168191969395,
        0.003344413358718157
      ],
      "rel_vs_flex_repeat": 0.003851044923067093,
      "compile_s": 0.0015459639980690554,
      "inner": 36,
      "median_ms": 1.4063403331318922,
      "q1_ms": 1.3933616387981198,
      "q3_ms": 1.4264204166667898,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 32.000267028808594,
      "xla_arg_mb": 28.0,
      "xla_out_mb": 12.000022888183594,
      "residual_mb": 8.125011444091797,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 1024, 32, 128]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.003344526281580329,
      "rel_err_parts": [
        0.0028945563826709986,
        0.0033384172711521387,
        0.003344526281580329
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.5267950240086066,
      "inner": 31,
      "median_ms": 1.5164221291460336,
      "q1_ms": 1.5135143551179358,
      "q3_ms": 1.5337361774366018,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 48.000267028808594,
      "xla_arg_mb": 28.0,
      "xla_out_mb": 12.000022888183594,
      "residual_mb": 24.125011444091797,
      "residual_shapes": [
        "bfloat16[1, 1024, 32, 128]",
        "bfloat16[1, 1024, 32, 128]",
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 1024, 32, 128]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.003344526281580329,
      "rel_err_parts": [
        0.0028945563826709986,
        0.0033384172711521387,
        0.003344526281580329
      ],
      "compile_s": 0.5281770690053236,
      "inner": 30,
      "median_ms": 1.5723583666840568,
      "q1_ms": 1.5616452668230827,
      "q3_ms": 1.581774533163601,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "jax_xla",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 392.12501525878906,
      "xla_arg_mb": 28.0,
      "xla_out_mb": 12.000022888183594,
      "rel_err": 0.0033211391419172287,
      "rel_err_parts": [
        0.0033211391419172287,
        0.003312610788270831,
        0.0028861421160399914
      ],
      "rel_vs_flex_repeat": 0.0039905039593577385,
      "compile_s": 0.48537045299599413,
      "inner": 23,
      "median_ms": 2.0811668699657866,
      "q1_ms": 2.0766888473704253,
      "q3_ms": 2.1132414568337086,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "jax_cudnn",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 64.250732421875,
      "xla_arg_mb": 28.0,
      "xla_out_mb": 12.000022888183594,
      "rel_err": 0.0033443919382989407,
      "rel_err_parts": [
        0.0028946080710738897,
        0.0033382740803062916,
        0.0033443919382989407
      ],
      "rel_vs_flex_repeat": 0.00022865409846417606,
      "compile_s": 0.11894128199492116,
      "inner": 32,
      "median_ms": 1.3547329685934528,
      "q1_ms": 1.3534419531424646,
      "q3_ms": 1.372012406363865,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 2,
      "rel_err": 0.003344413358718157,
      "rel_err_parts": [
        0.002894316567108035,
        0.0033402168191969395,
        0.003344413358718157
      ],
      "rel_vs_flex_repeat": 0.003851044923067093,
      "compile_s": 0.0015605520020471886,
      "inner": 36,
      "median_ms": 1.4467648890988334,
      "q1_ms": 1.4030285417700523,
      "q3_ms": 1.5306903331091033,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 8.0,
      "rel_err": 0.0022144638933241367,
      "rel_err_parts": [
        0.0022144638933241367
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.27051878000202123,
      "inner": 104,
      "median_ms": 0.38343224038656515,
      "q1_ms": 0.37727433174530656,
      "q3_ms": 0.38383930768759456,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 16.0,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 8.0,
      "rel_err": 0.0022144638933241367,
      "rel_err_parts": [
        0.0022144638933241367
      ],
      "compile_s": 0.27686615800485015,
      "inner": 104,
      "median_ms": 0.3757826634682715,
      "q1_ms": 0.3711857019040438,
      "q3_ms": 0.3864832019294576,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 192.0,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 8.0,
      "rel_err": 0.0023377228062599897,
      "rel_err_parts": [
        0.0023377228062599897
      ],
      "rel_vs_flex_repeat": 0.0030424734577536583,
      "compile_s": 0.508667237998452,
      "inner": 58,
      "median_ms": 0.7731808622464024,
      "q1_ms": 0.7710481292853164,
      "q3_ms": 0.8528912413598375,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 8.0,
      "rel_err": 0.002214468317106366,
      "rel_err_parts": [
        0.002214468317106366
      ],
      "rel_vs_flex_repeat": 1.7807384210755117e-05,
      "compile_s": 0.045962003991007805,
      "inner": 106,
      "median_ms": 0.33468648119728556,
      "q1_ms": 0.33349224532633703,
      "q3_ms": 0.3395390235967187,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_err": 0.0022089353296905756,
      "rel_err_parts": [
        0.0022089353296905756
      ],
      "rel_vs_flex_repeat": 0.0028709552716463804,
      "compile_s": 0.0005363550008041784,
      "inner": 150,
      "median_ms": 0.3599450399633497,
      "q1_ms": 0.34260133000013104,
      "q3_ms": 0.3605568999773823,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 8.0,
      "rel_err": 0.0022144638933241367,
      "rel_err_parts": [
        0.0022144638933241367
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.30855232500471175,
      "inner": 116,
      "median_ms": 0.3425185258775259,
      "q1_ms": 0.34223514219398754,
      "q3_ms": 0.3438692241114292,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 16.0,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 8.0,
      "rel_err": 0.0022144638933241367,
      "rel_err_parts": [
        0.0022144638933241367
      ],
      "compile_s": 0.29599410500668455,
      "inner": 104,
      "median_ms": 0.3766364903780274,
      "q1_ms": 0.37574702883620587,
      "q3_ms": 0.3768226586963955,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 192.0,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 8.0,
      "rel_err": 0.0023377228062599897,
      "rel_err_parts": [
        0.0023377228062599897
      ],
      "rel_vs_flex_repeat": 0.0030424734577536583,
      "compile_s": 0.5553113829955691,
      "inner": 58,
      "median_ms": 0.8805755863030409,
      "q1_ms": 0.8471763793405162,
      "q3_ms": 0.8831064224068541,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 8.0,
      "rel_err": 0.002214468317106366,
      "rel_err_parts": [
        0.002214468317106366
      ],
      "rel_vs_flex_repeat": 1.7807384210755117e-05,
      "compile_s": 0.047701785995741375,
      "inner": 106,
      "median_ms": 0.35427128311294953,
      "q1_ms": 0.34976435844722764,
      "q3_ms": 0.37435701411801725,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_err": 0.0022089353296905756,
      "rel_err_parts": [
        0.0022089353296905756
      ],
      "rel_vs_flex_repeat": 0.0028709552716463804,
      "compile_s": 0.0005364449898479506,
      "inner": 150,
      "median_ms": 0.31957819999661297,
      "q1_ms": 0.31631012335613684,
      "q3_ms": 0.3228330633040362,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 8.0,
      "rel_err": 0.0022144638933241367,
      "rel_err_parts": [
        0.0022144638933241367
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.3157808810065035,
      "inner": 101,
      "median_ms": 0.33792703960547316,
      "q1_ms": 0.3370228069338782,
      "q3_ms": 0.3416790841763959,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 16.0,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 8.0,
      "rel_err": 0.0022144638933241367,
      "rel_err_parts": [
        0.0022144638933241367
      ],
      "compile_s": 0.2717469329945743,
      "inner": 107,
      "median_ms": 0.37935348586658035,
      "q1_ms": 0.37847991589952923,
      "q3_ms": 0.4011850654355372,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 192.0,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 8.0,
      "rel_err": 0.0023377228062599897,
      "rel_err_parts": [
        0.0023377228062599897
      ],
      "rel_vs_flex_repeat": 0.0030424734577536583,
      "compile_s": 0.5883989010035293,
      "inner": 57,
      "median_ms": 0.8763809825738956,
      "q1_ms": 0.7767133070605301,
      "q3_ms": 0.8801499648701825,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 8.0,
      "rel_err": 0.002214468317106366,
      "rel_err_parts": [
        0.002214468317106366
      ],
      "rel_vs_flex_repeat": 1.7807384210755117e-05,
      "compile_s": 0.046194725000532344,
      "inner": 110,
      "median_ms": 0.3351260272955353,
      "q1_ms": 0.33351274098756467,
      "q3_ms": 0.3366401090963998,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_err": 0.0022089353296905756,
      "rel_err_parts": [
        0.0022089353296905756
      ],
      "rel_vs_flex_repeat": 0.0028709552716463804,
      "compile_s": 0.0010616179934004322,
      "inner": 129,
      "median_ms": 0.357080674337678,
      "q1_ms": 0.35440148839857,
      "q3_ms": 0.3574903605550044,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 32.000267028808594,
      "xla_arg_mb": 26.0,
      "xla_out_mb": 10.000022888183594,
      "residual_mb": 8.125011444091797,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 1024, 32, 128]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.003335731104016304,
      "rel_err_parts": [
        0.002897482132539153,
        0.003335731104016304,
        0.0033338258508592844
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.6370077729952754,
      "inner": 28,
      "median_ms": 1.544237607179509,
      "q1_ms": 1.5211882681407067,
      "q3_ms": 1.643806500006966,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 48.000267028808594,
      "xla_arg_mb": 26.0,
      "xla_out_mb": 10.000022888183594,
      "residual_mb": 24.125011444091797,
      "residual_shapes": [
        "bfloat16[1, 1024, 32, 128]",
        "bfloat16[1, 1024, 32, 128]",
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 1024, 32, 128]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.003335731104016304,
      "rel_err_parts": [
        0.002897482132539153,
        0.003335731104016304,
        0.0033338258508592844
      ],
      "compile_s": 0.548741794002126,
      "inner": 30,
      "median_ms": 1.750930866789228,
      "q1_ms": 1.6354507664800622,
      "q3_ms": 1.7571270999421056,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "jax_xla",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 392.12501525878906,
      "xla_arg_mb": 26.0,
      "xla_out_mb": 10.000022888183594,
      "rel_err": 0.0033327690325677395,
      "rel_err_parts": [
        0.0033327690325677395,
        0.0033268060069531202,
        0.0028921631164848804
      ],
      "rel_vs_flex_repeat": 0.003990147262811661,
      "compile_s": 0.30328845599433407,
      "inner": 22,
      "median_ms": 2.2326411357922056,
      "q1_ms": 2.1248123865172435,
      "q3_ms": 2.3940997952277856,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "jax_cudnn",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 64.250732421875,
      "xla_arg_mb": 26.0,
      "xla_out_mb": 10.000022888183594,
      "rel_err": 0.0033355483319610357,
      "rel_err_parts": [
        0.0028974306769669056,
        0.0033355483319610357,
        0.00333409383893013
      ],
      "rel_vs_flex_repeat": 0.00026004452956840396,
      "compile_s": 0.12583198898937553,
      "inner": 33,
      "median_ms": 1.3766374847336644,
      "q1_ms": 1.3621916666372935,
      "q3_ms": 1.3791616365779191,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 0,
      "rel_err": 0.0033337268978357315,
      "rel_err_parts": [
        0.002895575715228915,
        0.0033317178022116423,
        0.0033337268978357315
      ],
      "rel_vs_flex_repeat": 0.0038867106195539236,
      "compile_s": 0.0015474480023840442,
      "inner": 36,
      "median_ms": 1.384809360994647,
      "q1_ms": 1.3824005972209206,
      "q3_ms": 1.4053502917603409,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 32.000267028808594,
      "xla_arg_mb": 26.0,
      "xla_out_mb": 10.000022888183594,
      "residual_mb": 8.125011444091797,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 1024, 32, 128]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.003335731104016304,
      "rel_err_parts": [
        0.002897482132539153,
        0.003335731104016304,
        0.0033338258508592844
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.5177162479958497,
      "inner": 30,
      "median_ms": 1.5065849331828456,
      "q1_ms": 1.5030783166973074,
      "q3_ms": 1.5267225501399178,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 48.000267028808594,
      "xla_arg_mb": 26.0,
      "xla_out_mb": 10.000022888183594,
      "residual_mb": 24.125011444091797,
      "residual_shapes": [
        "bfloat16[1, 1024, 32, 128]",
        "bfloat16[1, 1024, 32, 128]",
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 1024, 32, 128]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.003335731104016304,
      "rel_err_parts": [
        0.002897482132539153,
        0.003335731104016304,
        0.0033338258508592844
      ],
      "compile_s": 0.6575766570022097,
      "inner": 28,
      "median_ms": 1.5767845356354624,
      "q1_ms": 1.5709372143776688,
      "q3_ms": 1.5870743573032087,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "jax_xla",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 392.12501525878906,
      "xla_arg_mb": 26.0,
      "xla_out_mb": 10.000022888183594,
      "rel_err": 0.0033327690325677395,
      "rel_err_parts": [
        0.0033327690325677395,
        0.0033268060069531202,
        0.0028921631164848804
      ],
      "rel_vs_flex_repeat": 0.003990147262811661,
      "compile_s": 0.3795563659950858,
      "inner": 23,
      "median_ms": 2.157843260489323,
      "q1_ms": 2.1286053262251103,
      "q3_ms": 2.1718094563845587,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "jax_cudnn",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 64.250732421875,
      "xla_arg_mb": 26.0,
      "xla_out_mb": 10.000022888183594,
      "rel_err": 0.0033355483319610357,
      "rel_err_parts": [
        0.0028974306769669056,
        0.0033355483319610357,
        0.00333409383893013
      ],
      "rel_vs_flex_repeat": 0.00026004452956840396,
      "compile_s": 0.13736531799077056,
      "inner": 32,
      "median_ms": 1.3641808750435302,
      "q1_ms": 1.3502551405508711,
      "q3_ms": 1.4552384375292604,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 1,
      "rel_err": 0.0033337268978357315,
      "rel_err_parts": [
        0.002895575715228915,
        0.0033317178022116423,
        0.0033337268978357315
      ],
      "rel_vs_flex_repeat": 0.0038867106195539236,
      "compile_s": 0.001562015007948503,
      "inner": 36,
      "median_ms": 1.559519583275283,
      "q1_ms": 1.5029186112466655,
      "q3_ms": 1.5683243056830383,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 32.000267028808594,
      "xla_arg_mb": 26.0,
      "xla_out_mb": 10.000022888183594,
      "residual_mb": 8.125011444091797,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 1024, 32, 128]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.003335731104016304,
      "rel_err_parts": [
        0.002897482132539153,
        0.003335731104016304,
        0.0033338258508592844
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.5059396979922894,
      "inner": 31,
      "median_ms": 1.5838254513519427,
      "q1_ms": 1.5290888226328176,
      "q3_ms": 1.6359626292063285,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 48.000267028808594,
      "xla_arg_mb": 26.0,
      "xla_out_mb": 10.000022888183594,
      "residual_mb": 24.125011444091797,
      "residual_shapes": [
        "bfloat16[1, 1024, 32, 128]",
        "bfloat16[1, 1024, 32, 128]",
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 1024, 32, 128]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.003335731104016304,
      "rel_err_parts": [
        0.002897482132539153,
        0.003335731104016304,
        0.0033338258508592844
      ],
      "compile_s": 0.5253508010064252,
      "inner": 30,
      "median_ms": 1.5592756336748912,
      "q1_ms": 1.5560653667004467,
      "q3_ms": 1.562763899952794,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "jax_xla",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 392.12501525878906,
      "xla_arg_mb": 26.0,
      "xla_out_mb": 10.000022888183594,
      "rel_err": 0.0033327690325677395,
      "rel_err_parts": [
        0.0033327690325677395,
        0.0033268060069531202,
        0.0028921631164848804
      ],
      "rel_vs_flex_repeat": 0.003990147262811661,
      "compile_s": 0.2535245080071036,
      "inner": 23,
      "median_ms": 2.1387959131971,
      "q1_ms": 2.1304743043045793,
      "q3_ms": 2.248079630510841,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "jax_cudnn",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 64.250732421875,
      "xla_arg_mb": 26.0,
      "xla_out_mb": 10.000022888183594,
      "rel_err": 0.0033355483319610357,
      "rel_err_parts": [
        0.0028974306769669056,
        0.0033355483319610357,
        0.00333409383893013
      ],
      "rel_vs_flex_repeat": 0.00026004452956840396,
      "compile_s": 0.11776411598839331,
      "inner": 33,
      "median_ms": 1.3590778786844264,
      "q1_ms": 1.345370560584618,
      "q3_ms": 1.369685515149313,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 2,
      "rel_err": 0.0033337268978357315,
      "rel_err_parts": [
        0.002895575715228915,
        0.0033317178022116423,
        0.0033337268978357315
      ],
      "rel_vs_flex_repeat": 0.0038867106195539236,
      "compile_s": 0.0015684770041843876,
      "inner": 36,
      "median_ms": 1.3660496943985991,
      "q1_ms": 1.364504944326149,
      "q3_ms": 1.4344123471043226,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 8.5,
      "xla_out_mb": 8.0,
      "rel_err": 0.002219589427113533,
      "rel_err_parts": [
        0.002219589427113533
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.2770718370011309,
      "inner": 104,
      "median_ms": 0.3809439711823451,
      "q1_ms": 0.3693400961380505,
      "q3_ms": 0.3859388029573333,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 16.0,
      "xla_arg_mb": 8.5,
      "xla_out_mb": 8.0,
      "rel_err": 0.002219589427113533,
      "rel_err_parts": [
        0.002219589427113533
      ],
      "compile_s": 0.27028911499655806,
      "inner": 100,
      "median_ms": 0.36674404007499106,
      "q1_ms": 0.3666830199654214,
      "q3_ms": 0.3668346149788704,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 192.0,
      "xla_arg_mb": 8.5,
      "xla_out_mb": 8.0,
      "rel_err": 0.002344554290175438,
      "rel_err_parts": [
        0.002344554290175438
      ],
      "rel_vs_flex_repeat": 0.003055271925404668,
      "compile_s": 0.213386511997669,
      "inner": 72,
      "median_ms": 0.5858700555513173,
      "q1_ms": 0.5839829583338644,
      "q3_ms": 0.6748657293024329,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 8.5,
      "xla_out_mb": 8.0,
      "rel_err": 0.002219591522589326,
      "rel_err_parts": [
        0.002219591522589326
      ],
      "rel_vs_flex_repeat": 1.7334348740405403e-05,
      "compile_s": 0.048267275997204706,
      "inner": 109,
      "median_ms": 0.33448933017089827,
      "q1_ms": 0.3341849449564865,
      "q3_ms": 0.3355239449098809,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_err": 0.002215860178694129,
      "rel_err_parts": [
        0.002215860178694129
      ],
      "rel_vs_flex_repeat": 0.0028834594413638115,
      "compile_s": 0.0005402920069172978,
      "inner": 150,
      "median_ms": 0.3544241266596752,
      "q1_ms": 0.3125512399613702,
      "q3_ms": 0.3562223300104961,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 8.5,
      "xla_out_mb": 8.0,
      "rel_err": 0.002219589427113533,
      "rel_err_parts": [
        0.002219589427113533
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.2587935760093387,
      "inner": 117,
      "median_ms": 0.33902960683171374,
      "q1_ms": 0.3375940683544574,
      "q3_ms": 0.33909893165836036,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 16.0,
      "xla_arg_mb": 8.5,
      "xla_out_mb": 8.0,
      "rel_err": 0.002219589427113533,
      "rel_err_parts": [
        0.002219589427113533
      ],
      "compile_s": 0.26396129399654455,
      "inner": 108,
      "median_ms": 0.3673393148346804,
      "q1_ms": 0.36694723611516467,
      "q3_ms": 0.36754071758574963,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 192.0,
      "xla_arg_mb": 8.5,
      "xla_out_mb": 8.0,
      "rel_err": 0.002344554290175438,
      "rel_err_parts": [
        0.002344554290175438
      ],
      "rel_vs_flex_repeat": 0.003055271925404668,
      "compile_s": 0.15728627699718345,
      "inner": 75,
      "median_ms": 0.5897363733189801,
      "q1_ms": 0.587215799993525,
      "q3_ms": 0.5920273799953671,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 8.5,
      "xla_out_mb": 8.0,
      "rel_err": 0.002219591522589326,
      "rel_err_parts": [
        0.002219591522589326
      ],
      "rel_vs_flex_repeat": 1.7334348740405403e-05,
      "compile_s": 0.11050305899698287,
      "inner": 94,
      "median_ms": 0.3317109681603974,
      "q1_ms": 0.3297133563969413,
      "q3_ms": 0.3335214095522075,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_err": 0.002215860178694129,
      "rel_err_parts": [
        0.002215860178694129
      ],
      "rel_vs_flex_repeat": 0.0028834594413638115,
      "compile_s": 0.0005744969967054203,
      "inner": 147,
      "median_ms": 0.3161303945310211,
      "q1_ms": 0.3156119286000323,
      "q3_ms": 0.31861015309407026,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 8.5,
      "xla_out_mb": 8.0,
      "rel_err": 0.002219589427113533,
      "rel_err_parts": [
        0.002219589427113533
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.27300397100043483,
      "inner": 111,
      "median_ms": 0.3431487477411776,
      "q1_ms": 0.3418896487166241,
      "q3_ms": 0.3444555000274081,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 16.0,
      "xla_arg_mb": 8.5,
      "xla_out_mb": 8.0,
      "rel_err": 0.002219589427113533,
      "rel_err_parts": [
        0.002219589427113533
      ],
      "compile_s": 0.27199049398768693,
      "inner": 98,
      "median_ms": 0.37312631632500726,
      "q1_ms": 0.3716323674093381,
      "q3_ms": 0.38663207654893517,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 192.0,
      "xla_arg_mb": 8.5,
      "xla_out_mb": 8.0,
      "rel_err": 0.002344554290175438,
      "rel_err_parts": [
        0.002344554290175438
      ],
      "rel_vs_flex_repeat": 0.003055271925404668,
      "compile_s": 0.15669092099415138,
      "inner": 76,
      "median_ms": 0.5968975525128802,
      "q1_ms": 0.590241269731594,
      "q3_ms": 0.6006308617782595,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 8.5,
      "xla_out_mb": 8.0,
      "rel_err": 0.002219591522589326,
      "rel_err_parts": [
        0.002219591522589326
      ],
      "rel_vs_flex_repeat": 1.7334348740405403e-05,
      "compile_s": 0.04572217998793349,
      "inner": 111,
      "median_ms": 0.33109016212538256,
      "q1_ms": 0.329399720801309,
      "q3_ms": 0.33200744144631095,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_err": 0.002215860178694129,
      "rel_err_parts": [
        0.002215860178694129
      ],
      "rel_vs_flex_repeat": 0.0028834594413638115,
      "compile_s": 0.000471791994641535,
      "inner": 148,
      "median_ms": 0.31456756079097503,
      "q1_ms": 0.3143719864916371,
      "q3_ms": 0.31779636825573304,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 32.000267028808594,
      "xla_arg_mb": 24.5,
      "xla_out_mb": 8.500022888183594,
      "residual_mb": 8.125011444091797,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 1024, 32, 128]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.0033434818033128977,
      "rel_err_parts": [
        0.002896428108215332,
        0.0033434818033128977,
        0.003335095476359129
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.4950896800000919,
      "inner": 31,
      "median_ms": 1.5336647094602909,
      "q1_ms": 1.5275882418087172,
      "q3_ms": 1.5906559354099896,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 48.000267028808594,
      "xla_arg_mb": 24.5,
      "xla_out_mb": 8.500022888183594,
      "residual_mb": 24.125011444091797,
      "residual_shapes": [
        "bfloat16[1, 1024, 32, 128]",
        "bfloat16[1, 1024, 32, 128]",
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 1024, 32, 128]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.0033434818033128977,
      "rel_err_parts": [
        0.002896428108215332,
        0.0033434818033128977,
        0.003335095476359129
      ],
      "compile_s": 0.5953040449967375,
      "inner": 30,
      "median_ms": 1.58080503315432,
      "q1_ms": 1.5676307836353467,
      "q3_ms": 1.6071149667065279,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "jax_xla",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 384.12525939941406,
      "xla_arg_mb": 24.5,
      "xla_out_mb": 8.500022888183594,
      "rel_err": 0.0033315212931483984,
      "rel_err_parts": [
        0.0033315212931483984,
        0.003330466104671359,
        0.00289628142490983
      ],
      "rel_vs_flex_repeat": 0.003974394407123327,
      "compile_s": 0.29048244800651446,
      "inner": 23,
      "median_ms": 2.066301825992844,
      "q1_ms": 2.0640950432330696,
      "q3_ms": 2.0777266520029944,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "jax_cudnn",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 64.250732421875,
      "xla_arg_mb": 24.5,
      "xla_out_mb": 8.500022888183594,
      "rel_err": 0.0033441714476794004,
      "rel_err_parts": [
        0.0028964574448764324,
        0.0033441714476794004,
        0.0033359418157488108
      ],
      "rel_vs_flex_repeat": 0.0003665741242002696,
      "compile_s": 0.11789827000757214,
      "inner": 33,
      "median_ms": 1.361129333610844,
      "q1_ms": 1.343070621122465,
      "q3_ms": 1.3775498029479587,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 0,
      "rel_err": 0.0033472594805061817,
      "rel_err_parts": [
        0.0028994225431233644,
        0.0033472594805061817,
        0.003335174871608615
      ],
      "rel_vs_flex_repeat": 0.003893638961017132,
      "compile_s": 0.0015587990055792034,
      "inner": 36,
      "median_ms": 1.364687110860056,
      "q1_ms": 1.3639756388228206,
      "q3_ms": 1.3674662501822845,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 32.000267028808594,
      "xla_arg_mb": 24.5,
      "xla_out_mb": 8.500022888183594,
      "residual_mb": 8.125011444091797,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 1024, 32, 128]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.0033434818033128977,
      "rel_err_parts": [
        0.002896428108215332,
        0.0033434818033128977,
        0.003335095476359129
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.5086798620031914,
      "inner": 31,
      "median_ms": 1.527153386925948,
      "q1_ms": 1.49797641948211,
      "q3_ms": 1.5365886612280812,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 48.000267028808594,
      "xla_arg_mb": 24.5,
      "xla_out_mb": 8.500022888183594,
      "residual_mb": 24.125011444091797,
      "residual_shapes": [
        "bfloat16[1, 1024, 32, 128]",
        "bfloat16[1, 1024, 32, 128]",
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 1024, 32, 128]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.0033434818033128977,
      "rel_err_parts": [
        0.002896428108215332,
        0.0033434818033128977,
        0.003335095476359129
      ],
      "compile_s": 0.525314772996353,
      "inner": 29,
      "median_ms": 1.5611956553365072,
      "q1_ms": 1.5605976551686864,
      "q3_ms": 1.575039258600084,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "jax_xla",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 384.12525939941406,
      "xla_arg_mb": 24.5,
      "xla_out_mb": 8.500022888183594,
      "rel_err": 0.0033315212931483984,
      "rel_err_parts": [
        0.0033315212931483984,
        0.003330466104671359,
        0.00289628142490983
      ],
      "rel_vs_flex_repeat": 0.003974394407123327,
      "compile_s": 0.28969382599461824,
      "inner": 22,
      "median_ms": 2.1135103181057437,
      "q1_ms": 2.110884658825605,
      "q3_ms": 2.146584909165871,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "jax_cudnn",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 64.250732421875,
      "xla_arg_mb": 24.5,
      "xla_out_mb": 8.500022888183594,
      "rel_err": 0.0033441714476794004,
      "rel_err_parts": [
        0.0028964574448764324,
        0.0033441714476794004,
        0.0033359418157488108
      ],
      "rel_vs_flex_repeat": 0.0003665741242002696,
      "compile_s": 0.11635058200045023,
      "inner": 33,
      "median_ms": 1.3669226971229143,
      "q1_ms": 1.3493944242134495,
      "q3_ms": 1.374634287926848,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 1,
      "rel_err": 0.0033472594805061817,
      "rel_err_parts": [
        0.0028994225431233644,
        0.0033472594805061817,
        0.003335174871608615
      ],
      "rel_vs_flex_repeat": 0.003893638961017132,
      "compile_s": 0.001555723007186316,
      "inner": 36,
      "median_ms": 1.4212240554722182,
      "q1_ms": 1.4065347362096268,
      "q3_ms": 1.5141034026681963,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 32.000267028808594,
      "xla_arg_mb": 24.5,
      "xla_out_mb": 8.500022888183594,
      "residual_mb": 8.125011444091797,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 1024, 32, 128]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.0033434818033128977,
      "rel_err_parts": [
        0.002896428108215332,
        0.0033434818033128977,
        0.003335095476359129
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.49348631799512077,
      "inner": 30,
      "median_ms": 1.5094881334031622,
      "q1_ms": 1.4960576998419128,
      "q3_ms": 1.5211999333890465,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 48.000267028808594,
      "xla_arg_mb": 24.5,
      "xla_out_mb": 8.500022888183594,
      "residual_mb": 24.125011444091797,
      "residual_shapes": [
        "bfloat16[1, 1024, 32, 128]",
        "bfloat16[1, 1024, 32, 128]",
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 1024, 32, 128]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.0033434818033128977,
      "rel_err_parts": [
        0.002896428108215332,
        0.0033434818033128977,
        0.003335095476359129
      ],
      "compile_s": 0.5120955140009755,
      "inner": 30,
      "median_ms": 1.5519748335160937,
      "q1_ms": 1.551185333179698,
      "q3_ms": 1.5525734668093112,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "jax_xla",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 384.12525939941406,
      "xla_arg_mb": 24.5,
      "xla_out_mb": 8.500022888183594,
      "rel_err": 0.0033315212931483984,
      "rel_err_parts": [
        0.0033315212931483984,
        0.003330466104671359,
        0.00289628142490983
      ],
      "rel_vs_flex_repeat": 0.003974394407123327,
      "compile_s": 0.28732405199843924,
      "inner": 22,
      "median_ms": 2.064256227061957,
      "q1_ms": 2.063448567985853,
      "q3_ms": 2.0945637952536345,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "jax_cudnn",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 64.250732421875,
      "xla_arg_mb": 24.5,
      "xla_out_mb": 8.500022888183594,
      "rel_err": 0.0033441714476794004,
      "rel_err_parts": [
        0.0028964574448764324,
        0.0033441714476794004,
        0.0033359418157488108
      ],
      "rel_vs_flex_repeat": 0.0003665741242002696,
      "compile_s": 0.11655608000000939,
      "inner": 30,
      "median_ms": 1.3397032331946928,
      "q1_ms": 1.3390541666012723,
      "q3_ms": 1.3400583833572455,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 2,
      "rel_err": 0.0033472594805061817,
      "rel_err_parts": [
        0.0028994225431233644,
        0.0033472594805061817,
        0.003335174871608615
      ],
      "rel_vs_flex_repeat": 0.003893638961017132,
      "compile_s": 0.0014903390110703185,
      "inner": 36,
      "median_ms": 1.3610176663658218,
      "q1_ms": 1.360886999868348,
      "q3_ms": 1.3616676387755433,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 16.0,
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.18771574499260169,
      "inner": 21,
      "median_ms": 2.13298438064244,
      "q1_ms": 2.1253580475569747,
      "q3_ms": 2.139890333291121,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 32.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.19235639499675017,
      "inner": 23,
      "median_ms": 2.1725571739649316,
      "q1_ms": 2.171294782348686,
      "q3_ms": 2.1759172172017354,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 16.0,
      "rel_vs_flex_repeat": 2.3770026018610224e-05,
      "compile_s": 0.04744924799888395,
      "inner": 21,
      "median_ms": 2.1977383335719685,
      "q1_ms": 2.186699333140582,
      "q3_ms": 2.2124993573968474,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_vs_flex_repeat": 0.003037688322365284,
      "compile_s": 0.002150439002434723,
      "inner": 25,
      "median_ms": 2.048674639663659,
      "q1_ms": 2.0480426400899887,
      "q3_ms": 2.049227499810513,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 16.0,
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.1868698259931989,
      "inner": 22,
      "median_ms": 2.1132407273398712,
      "q1_ms": 2.1127400227256166,
      "q3_ms": 2.114524272738808,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 32.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.1916672009974718,
      "inner": 22,
      "median_ms": 2.171068318024151,
      "q1_ms": 2.1658994317642115,
      "q3_ms": 2.1760944089857563,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 16.0,
      "rel_vs_flex_repeat": 2.3770026018610224e-05,
      "compile_s": 0.04630487300164532,
      "inner": 21,
      "median_ms": 2.195840904494703,
      "q1_ms": 2.1943452382055555,
      "q3_ms": 2.2041540474414694,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_vs_flex_repeat": 0.003037688322365284,
      "compile_s": 0.0021629719994962215,
      "inner": 25,
      "median_ms": 2.0493182801874354,
      "q1_ms": 2.048958179948386,
      "q3_ms": 2.0507337598246522,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 16.0,
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.20053452700085472,
      "inner": 24,
      "median_ms": 2.118182791794728,
      "q1_ms": 2.1149288747134656,
      "q3_ms": 2.134493208359345,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 32.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.2019031969975913,
      "inner": 23,
      "median_ms": 2.2796929128083363,
      "q1_ms": 2.188570978393292,
      "q3_ms": 2.328546869609019,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 16.0,
      "rel_vs_flex_repeat": 2.3770026018610224e-05,
      "compile_s": 0.050794216003851034,
      "inner": 21,
      "median_ms": 2.2133536189184744,
      "q1_ms": 2.203620928325801,
      "q3_ms": 2.237512167041478,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_vs_flex_repeat": 0.003037688322365284,
      "compile_s": 0.0021411199995782226,
      "inner": 25,
      "median_ms": 2.0546799997100607,
      "q1_ms": 2.0492323199869134,
      "q3_ms": 2.070924559957348,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 64.0002670288086,
      "xla_arg_mb": 56.0,
      "xla_out_mb": 24.000022888183594,
      "residual_mb": 16.500011444091797,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 4096, 32, 64]",
        "float32[1, 32, 4096]"
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.7420745570125291,
      "inner": 5,
      "median_ms": 9.372499599703588,
      "q1_ms": 9.206496599654201,
      "q3_ms": 9.444475900090765,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 96.0002670288086,
      "xla_arg_mb": 56.0,
      "xla_out_mb": 24.000022888183594,
      "residual_mb": 48.5000114440918,
      "residual_shapes": [
        "bfloat16[1, 4096, 32, 64]",
        "bfloat16[1, 4096, 32, 64]",
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 4096, 32, 64]",
        "float32[1, 32, 4096]"
      ],
      "compile_s": 0.46807437200914137,
      "inner": 5,
      "median_ms": 10.15527659910731,
      "q1_ms": 10.029283199401107,
      "q3_ms": 10.304286499740556,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "jax_cudnn",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 129.000732421875,
      "xla_arg_mb": 56.0,
      "xla_out_mb": 24.000022888183594,
      "rel_vs_flex_repeat": 0.0003526470100041479,
      "compile_s": 0.12155833400902338,
      "inner": 5,
      "median_ms": 8.753027798957191,
      "q1_ms": 8.641655999235809,
      "q3_ms": 8.929289101797622,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 0,
      "rel_vs_flex_repeat": 0.001878076116554439,
      "compile_s": 0.008471135995932855,
      "inner": 5,
      "median_ms": 8.24372399947606,
      "q1_ms": 8.126468300179113,
      "q3_ms": 8.656691400392447,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 64.0002670288086,
      "xla_arg_mb": 56.0,
      "xla_out_mb": 24.000022888183594,
      "residual_mb": 16.500011444091797,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 4096, 32, 64]",
        "float32[1, 32, 4096]"
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.5384749299992109,
      "inner": 5,
      "median_ms": 10.156917598214932,
      "q1_ms": 9.845722001045942,
      "q3_ms": 10.315857500245329,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 96.0002670288086,
      "xla_arg_mb": 56.0,
      "xla_out_mb": 24.000022888183594,
      "residual_mb": 48.5000114440918,
      "residual_shapes": [
        "bfloat16[1, 4096, 32, 64]",
        "bfloat16[1, 4096, 32, 64]",
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 4096, 32, 64]",
        "float32[1, 32, 4096]"
      ],
      "compile_s": 0.5902053790050559,
      "inner": 4,
      "median_ms": 10.015362997364718,
      "q1_ms": 9.875230249235756,
      "q3_ms": 10.290553500453825,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "jax_cudnn",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 129.000732421875,
      "xla_arg_mb": 56.0,
      "xla_out_mb": 24.000022888183594,
      "rel_vs_flex_repeat": 0.0003526470100041479,
      "compile_s": 0.12930765500641428,
      "inner": 5,
      "median_ms": 9.235516597982496,
      "q1_ms": 9.219774800294545,
      "q3_ms": 9.25751510221744,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 1,
      "rel_vs_flex_repeat": 0.001878076116554439,
      "compile_s": 0.00832378699851688,
      "inner": 6,
      "median_ms": 8.262257833848707,
      "q1_ms": 7.901495334711702,
      "q3_ms": 8.779796917224303,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 64.0002670288086,
      "xla_arg_mb": 56.0,
      "xla_out_mb": 24.000022888183594,
      "residual_mb": 16.500011444091797,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 4096, 32, 64]",
        "float32[1, 32, 4096]"
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.5421168110042345,
      "inner": 5,
      "median_ms": 9.727080402080901,
      "q1_ms": 9.529901400674134,
      "q3_ms": 10.063212398381438,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 96.0002670288086,
      "xla_arg_mb": 56.0,
      "xla_out_mb": 24.000022888183594,
      "residual_mb": 48.5000114440918,
      "residual_shapes": [
        "bfloat16[1, 4096, 32, 64]",
        "bfloat16[1, 4096, 32, 64]",
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 4096, 32, 64]",
        "float32[1, 32, 4096]"
      ],
      "compile_s": 0.5325215269986074,
      "inner": 5,
      "median_ms": 9.865823001018725,
      "q1_ms": 9.788629900140222,
      "q3_ms": 9.921167700667866,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "jax_cudnn",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 129.000732421875,
      "xla_arg_mb": 56.0,
      "xla_out_mb": 24.000022888183594,
      "rel_vs_flex_repeat": 0.0003526470100041479,
      "compile_s": 0.14081593698938377,
      "inner": 5,
      "median_ms": 8.944694200181402,
      "q1_ms": 8.48263100051554,
      "q3_ms": 9.02228109916905,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 2,
      "rel_vs_flex_repeat": 0.001878076116554439,
      "compile_s": 0.007686551994993351,
      "inner": 6,
      "median_ms": 8.267417666502297,
      "q1_ms": 7.885720500780735,
      "q3_ms": 8.654626749679057,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 16.0,
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.19373183799325489,
      "inner": 21,
      "median_ms": 2.1465834766526575,
      "q1_ms": 2.131352024055862,
      "q3_ms": 2.2375834998369792,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 32.0,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.19436747100553475,
      "inner": 23,
      "median_ms": 2.212447434430942,
      "q1_ms": 2.1855876738023334,
      "q3_ms": 2.2377041307703145,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 16.0,
      "rel_vs_flex_repeat": 2.5218849259545095e-05,
      "compile_s": 0.051510331002646126,
      "inner": 22,
      "median_ms": 2.235024772579147,
      "q1_ms": 2.217454022518888,
      "q3_ms": 2.2825979546971875,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_vs_flex_repeat": 0.0029643187299370766,
      "compile_s": 0.002162181001040153,
      "inner": 25,
      "median_ms": 2.059929920360446,
      "q1_ms": 2.0578966799075715,
      "q3_ms": 2.1002181599033065,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 16.0,
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.19124954100698233,
      "inner": 24,
      "median_ms": 2.125463208358269,
      "q1_ms": 2.1015210416711243,
      "q3_ms": 2.147129395780212,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 32.0,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.1957934380043298,
      "inner": 23,
      "median_ms": 2.17961095673356,
      "q1_ms": 2.175743217736928,
      "q3_ms": 2.1993955651060277,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 16.0,
      "rel_vs_flex_repeat": 2.5218849259545095e-05,
      "compile_s": 0.04850805099704303,
      "inner": 21,
      "median_ms": 2.2299755239925747,
      "q1_ms": 2.2112097857253894,
      "q3_ms": 2.2551340478544475,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_vs_flex_repeat": 0.0029643187299370766,
      "compile_s": 0.0021477130067069083,
      "inner": 24,
      "median_ms": 2.0794323754671495,
      "q1_ms": 2.0647579376600333,
      "q3_ms": 2.0878843955263924,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 16.0,
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.18956135700864252,
      "inner": 23,
      "median_ms": 2.1632543914323756,
      "q1_ms": 2.1478713476782385,
      "q3_ms": 2.1688786740845027,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 32.0,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.19378562898782548,
      "inner": 23,
      "median_ms": 2.203595913116537,
      "q1_ms": 2.1852498694383744,
      "q3_ms": 2.223249804481621,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 16.0,
      "rel_vs_flex_repeat": 2.5218849259545095e-05,
      "compile_s": 0.04911038000136614,
      "inner": 21,
      "median_ms": 2.234720238318135,
      "q1_ms": 2.2300415955777724,
      "q3_ms": 2.247550642918906,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_vs_flex_repeat": 0.0029643187299370766,
      "compile_s": 0.002150458996766247,
      "inner": 24,
      "median_ms": 2.0721105829579756,
      "q1_ms": 2.067720874947554,
      "q3_ms": 2.0852483751999293,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 64.0002670288086,
      "xla_arg_mb": 52.0,
      "xla_out_mb": 20.000022888183594,
      "residual_mb": 16.500011444091797,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 4096, 32, 64]",
        "float32[1, 32, 4096]"
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.5686019269924145,
      "inner": 4,
      "median_ms": 10.423675248603104,
      "q1_ms": 10.306564998245449,
      "q3_ms": 10.457804874022258,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 96.0002670288086,
      "xla_arg_mb": 52.0,
      "xla_out_mb": 20.000022888183594,
      "residual_mb": 48.5000114440918,
      "residual_shapes": [
        "bfloat16[1, 4096, 32, 64]",
        "bfloat16[1, 4096, 32, 64]",
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 4096, 32, 64]",
        "float32[1, 32, 4096]"
      ],
      "compile_s": 0.906609447003575,
      "inner": 5,
      "median_ms": 10.797451398684643,
      "q1_ms": 10.782240699336398,
      "q3_ms": 10.839883799781092,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "jax_cudnn",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 129.000732421875,
      "xla_arg_mb": 52.0,
      "xla_out_mb": 20.000022888183594,
      "rel_vs_flex_repeat": 0.0004151734465267509,
      "compile_s": 0.1835849710041657,
      "inner": 5,
      "median_ms": 8.311829000012949,
      "q1_ms": 8.276232601201627,
      "q3_ms": 8.32504499994684,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 0,
      "rel_vs_flex_repeat": 0.002029506489634514,
      "compile_s": 0.0076202369964448735,
      "inner": 6,
      "median_ms": 7.832177499949466,
      "q1_ms": 7.812775917651985,
      "q3_ms": 7.843961499020224,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 64.0002670288086,
      "xla_arg_mb": 52.0,
      "xla_out_mb": 20.000022888183594,
      "residual_mb": 16.500011444091797,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 4096, 32, 64]",
        "float32[1, 32, 4096]"
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.4760349920106819,
      "inner": 5,
      "median_ms": 9.11288240167778,
      "q1_ms": 9.104853299504612,
      "q3_ms": 9.140837298764382,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 96.0002670288086,
      "xla_arg_mb": 52.0,
      "xla_out_mb": 20.000022888183594,
      "residual_mb": 48.5000114440918,
      "residual_shapes": [
        "bfloat16[1, 4096, 32, 64]",
        "bfloat16[1, 4096, 32, 64]",
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 4096, 32, 64]",
        "float32[1, 32, 4096]"
      ],
      "compile_s": 0.4977103090059245,
      "inner": 5,
      "median_ms": 9.596387200872414,
      "q1_ms": 9.411983400059398,
      "q3_ms": 10.592330498911906,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "jax_cudnn",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 129.000732421875,
      "xla_arg_mb": 52.0,
      "xla_out_mb": 20.000022888183594,
      "rel_vs_flex_repeat": 0.0004151734465267509,
      "compile_s": 0.22406817000592127,
      "inner": 5,
      "median_ms": 8.389203401748091,
      "q1_ms": 8.383798001159448,
      "q3_ms": 8.403921200078912,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 1,
      "rel_vs_flex_repeat": 0.002029506489634514,
      "compile_s": 0.0076783760014222935,
      "inner": 6,
      "median_ms": 8.93054549912146,
      "q1_ms": 8.910171916795662,
      "q3_ms": 8.957504082597248,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 64.0002670288086,
      "xla_arg_mb": 52.0,
      "xla_out_mb": 20.000022888183594,
      "residual_mb": 16.500011444091797,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 4096, 32, 64]",
        "float32[1, 32, 4096]"
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.5205203229998006,
      "inner": 5,
      "median_ms": 9.09630520036444,
      "q1_ms": 9.08685039903503,
      "q3_ms": 9.124417298880871,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 96.0002670288086,
      "xla_arg_mb": 52.0,
      "xla_out_mb": 20.000022888183594,
      "residual_mb": 48.5000114440918,
      "residual_shapes": [
        "bfloat16[1, 4096, 32, 64]",
        "bfloat16[1, 4096, 32, 64]",
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 4096, 32, 64]",
        "float32[1, 32, 4096]"
      ],
      "compile_s": 0.456064369995147,
      "inner": 5,
      "median_ms": 9.641903400188312,
      "q1_ms": 9.606169698236044,
      "q3_ms": 9.707000500930008,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "jax_cudnn",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 129.000732421875,
      "xla_arg_mb": 52.0,
      "xla_out_mb": 20.000022888183594,
      "rel_vs_flex_repeat": 0.0004151734465267509,
      "compile_s": 0.11612128700653557,
      "inner": 6,
      "median_ms": 8.412413832653934,
      "q1_ms": 8.310416498716222,
      "q3_ms": 8.631311000499409,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 2,
      "rel_vs_flex_repeat": 0.002029506489634514,
      "compile_s": 0.007704325995291583,
      "inner": 6,
      "median_ms": 7.829465833007513,
      "q1_ms": 7.799362916557584,
      "q3_ms": 7.892482499300968,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 17.0,
      "xla_out_mb": 16.0,
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.19180043299274985,
      "inner": 22,
      "median_ms": 2.1021055449223653,
      "q1_ms": 2.101843704374253,
      "q3_ms": 2.119862363708142,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 32.0,
      "xla_arg_mb": 17.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.18655921799654607,
      "inner": 23,
      "median_ms": 2.2844928696387163,
      "q1_ms": 2.1425351303111277,
      "q3_ms": 2.371598304307554,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 17.0,
      "xla_out_mb": 16.0,
      "rel_vs_flex_repeat": 2.428273182886187e-05,
      "compile_s": 0.04686671600211412,
      "inner": 21,
      "median_ms": 2.2124647616042887,
      "q1_ms": 2.2026440475158218,
      "q3_ms": 2.225644213987315,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_vs_flex_repeat": 0.003032179083675146,
      "compile_s": 0.0020885819976683706,
      "inner": 25,
      "median_ms": 2.0486578001873568,
      "q1_ms": 2.047923039935995,
      "q3_ms": 2.0505081201554276,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 17.0,
      "xla_out_mb": 16.0,
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.18870053900172934,
      "inner": 24,
      "median_ms": 2.116229457897134,
      "q1_ms": 2.1137314164055474,
      "q3_ms": 2.119801228824751,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 32.0,
      "xla_arg_mb": 17.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.1909011429961538,
      "inner": 23,
      "median_ms": 2.1634713041788456,
      "q1_ms": 2.152206282556543,
      "q3_ms": 2.1662391306523676,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 17.0,
      "xla_out_mb": 16.0,
      "rel_vs_flex_repeat": 2.428273182886187e-05,
      "compile_s": 0.04906278099224437,
      "inner": 21,
      "median_ms": 2.236750714170436,
      "q1_ms": 2.1887642138781738,
      "q3_ms": 2.3273105952761224,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_vs_flex_repeat": 0.003032179083675146,
      "compile_s": 0.002146911996533163,
      "inner": 25,
      "median_ms": 2.0522509998409078,
      "q1_ms": 2.0501870999578387,
      "q3_ms": 2.0812051801476628,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 17.0,
      "xla_out_mb": 16.0,
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.18453993801085744,
      "inner": 23,
      "median_ms": 2.146356739103794,
      "q1_ms": 2.1367476956979576,
      "q3_ms": 2.156200369506183,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 32.0,
      "xla_arg_mb": 17.0,
      "xla_out_mb": 16.0,
      "compile_s": 0.18495939100102987,
      "inner": 23,
      "median_ms": 2.165293434898005,
      "q1_ms": 2.1599287176049193,
      "q3_ms": 2.1656528042512946,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 17.0,
      "xla_out_mb": 16.0,
      "rel_vs_flex_repeat": 2.428273182886187e-05,
      "compile_s": 0.04753087300923653,
      "inner": 21,
      "median_ms": 2.1922708096792034,
      "q1_ms": 2.1896053093812604,
      "q3_ms": 2.196502928633154,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_vs_flex_repeat": 0.003032179083675146,
      "compile_s": 0.0021039210114395246,
      "inner": 25,
      "median_ms": 2.0533069997327402,
      "q1_ms": 2.0521173599991016,
      "q3_ms": 2.0535614801337942,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 64.0002670288086,
      "xla_arg_mb": 49.0,
      "xla_out_mb": 17.000022888183594,
      "residual_mb": 16.500011444091797,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 4096, 32, 64]",
        "float32[1, 32, 4096]"
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.43804618100693915,
      "inner": 5,
      "median_ms": 9.169193200068548,
      "q1_ms": 9.127204198739491,
      "q3_ms": 9.205274099076632,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 96.0002670288086,
      "xla_arg_mb": 49.0,
      "xla_out_mb": 17.000022888183594,
      "residual_mb": 48.5000114440918,
      "residual_shapes": [
        "bfloat16[1, 4096, 32, 64]",
        "bfloat16[1, 4096, 32, 64]",
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 4096, 32, 64]",
        "float32[1, 32, 4096]"
      ],
      "compile_s": 0.4414741569926264,
      "inner": 5,
      "median_ms": 9.5475690002786,
      "q1_ms": 9.540233998268377,
      "q3_ms": 9.549414398497902,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "jax_cudnn",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 129.000732421875,
      "xla_arg_mb": 49.0,
      "xla_out_mb": 17.000022888183594,
      "rel_vs_flex_repeat": 0.0005571952206082642,
      "compile_s": 0.11296414500975516,
      "inner": 6,
      "median_ms": 8.285782333890287,
      "q1_ms": 8.2723052497992,
      "q3_ms": 8.297052000367936,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 0,
      "rel_vs_flex_repeat": 0.0021972053218632936,
      "compile_s": 0.007633030996657908,
      "inner": 6,
      "median_ms": 7.811641999675582,
      "q1_ms": 7.810241998716568,
      "q3_ms": 7.812319166987436,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 64.0002670288086,
      "xla_arg_mb": 49.0,
      "xla_out_mb": 17.000022888183594,
      "residual_mb": 16.500011444091797,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 4096, 32, 64]",
        "float32[1, 32, 4096]"
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.4377589370124042,
      "inner": 5,
      "median_ms": 9.034177401917987,
      "q1_ms": 9.022311998705845,
      "q3_ms": 9.077859300305136,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 96.0002670288086,
      "xla_arg_mb": 49.0,
      "xla_out_mb": 17.000022888183594,
      "residual_mb": 48.5000114440918,
      "residual_shapes": [
        "bfloat16[1, 4096, 32, 64]",
        "bfloat16[1, 4096, 32, 64]",
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 4096, 32, 64]",
        "float32[1, 32, 4096]"
      ],
      "compile_s": 0.44254698599979747,
      "inner": 5,
      "median_ms": 9.583752800244838,
      "q1_ms": 9.555374800402205,
      "q3_ms": 9.629523899639025,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "jax_cudnn",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 129.000732421875,
      "xla_arg_mb": 49.0,
      "xla_out_mb": 17.000022888183594,
      "rel_vs_flex_repeat": 0.0005571952206082642,
      "compile_s": 0.12223854099283926,
      "inner": 6,
      "median_ms": 8.335631665734885,
      "q1_ms": 8.276176749253258,
      "q3_ms": 8.342345999456786,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 1,
      "rel_vs_flex_repeat": 0.0021972053218632936,
      "compile_s": 0.007650093000847846,
      "inner": 6,
      "median_ms": 7.861404497816693,
      "q1_ms": 7.830435917033659,
      "q3_ms": 8.0686299164275,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 64.0002670288086,
      "xla_arg_mb": 49.0,
      "xla_out_mb": 17.000022888183594,
      "residual_mb": 16.500011444091797,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 4096, 32, 64]",
        "float32[1, 32, 4096]"
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.4455315230006818,
      "inner": 5,
      "median_ms": 8.993478599586524,
      "q1_ms": 8.972026100673247,
      "q3_ms": 9.180423300131224,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 96.0002670288086,
      "xla_arg_mb": 49.0,
      "xla_out_mb": 17.000022888183594,
      "residual_mb": 48.5000114440918,
      "residual_shapes": [
        "bfloat16[1, 4096, 32, 64]",
        "bfloat16[1, 4096, 32, 64]",
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 4096, 32, 64]",
        "float32[1, 32, 4096]"
      ],
      "compile_s": 0.45802937699772883,
      "inner": 5,
      "median_ms": 9.550163801759481,
      "q1_ms": 9.53214869950898,
      "q3_ms": 9.606100599921774,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "jax_cudnn",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 129.000732421875,
      "xla_arg_mb": 49.0,
      "xla_out_mb": 17.000022888183594,
      "rel_vs_flex_repeat": 0.0005571952206082642,
      "compile_s": 0.12135078100254759,
      "inner": 6,
      "median_ms": 8.378188833982373,
      "q1_ms": 8.315032666966243,
      "q3_ms": 8.60693908301376,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 2,
      "rel_vs_flex_repeat": 0.0021972053218632936,
      "compile_s": 0.007696952001424506,
      "inner": 6,
      "median_ms": 8.035058000435432,
      "q1_ms": 7.93331650129403,
      "q3_ms": 8.82493658355088,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 32.0,
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.2648582909896504,
      "inner": 9,
      "median_ms": 5.481530444942312,
      "q1_ms": 5.015262778518567,
      "q3_ms": 5.625189944112208,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 64.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.27338947099633515,
      "inner": 11,
      "median_ms": 5.041519727621397,
      "q1_ms": 4.874570773443503,
      "q3_ms": 5.058733272810721,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 32.0,
      "rel_vs_flex_repeat": 2.4349483283003792e-05,
      "compile_s": 0.04910656300489791,
      "inner": 12,
      "median_ms": 4.212242334081869,
      "q1_ms": 4.166293416346889,
      "q3_ms": 4.518109833346292,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_vs_flex_repeat": 0.003019785974174738,
      "compile_s": 0.0041277390118921176,
      "inner": 12,
      "median_ms": 4.593185665726196,
      "q1_ms": 4.332813000776999,
      "q3_ms": 4.6067245420999825,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 32.0,
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.2909378990007099,
      "inner": 11,
      "median_ms": 4.944041727205992,
      "q1_ms": 4.5962660905471155,
      "q3_ms": 4.965628499286884,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 64.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.29432405599800404,
      "inner": 10,
      "median_ms": 4.560911700536963,
      "q1_ms": 4.51648265006952,
      "q3_ms": 4.571206199761946,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 32.0,
      "rel_vs_flex_repeat": 2.4349483283003792e-05,
      "compile_s": 0.04893172299489379,
      "inner": 12,
      "median_ms": 4.228325333921627,
      "q1_ms": 4.207257958720826,
      "q3_ms": 4.307677583104427,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_vs_flex_repeat": 0.003019785974174738,
      "compile_s": 0.004091130002052523,
      "inner": 12,
      "median_ms": 4.188857249876794,
      "q1_ms": 4.150855833358946,
      "q3_ms": 4.329894166706557,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 32.0,
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.27385015100298915,
      "inner": 11,
      "median_ms": 4.5386876365889535,
      "q1_ms": 4.4489810454913155,
      "q3_ms": 5.02643350046128,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 64.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.27095023600850254,
      "inner": 8,
      "median_ms": 4.554102124529891,
      "q1_ms": 4.5332864374358905,
      "q3_ms": 4.8778708132886095,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 32.0,
      "rel_vs_flex_repeat": 2.4349483283003792e-05,
      "compile_s": 0.04705332899175119,
      "inner": 12,
      "median_ms": 4.702475666514753,
      "q1_ms": 4.294674208722427,
      "q3_ms": 4.712125666628708,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_vs_flex_repeat": 0.003019785974174738,
      "compile_s": 0.004093995987204835,
      "inner": 12,
      "median_ms": 4.658369750056106,
      "q1_ms": 4.503347667196067,
      "q3_ms": 4.691192708681531,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 128.0002670288086,
      "xla_arg_mb": 112.0,
      "xla_out_mb": 48.000022888183594,
      "residual_mb": 32.5000114440918,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 4096, 32, 128]",
        "float32[1, 32, 4096]"
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.5170011449954472,
      "inner": 2,
      "median_ms": 21.524176001548767,
      "q1_ms": 20.80021149595268,
      "q3_ms": 22.479313251096755,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 192.0002670288086,
      "xla_arg_mb": 112.0,
      "xla_out_mb": 48.000022888183594,
      "residual_mb": 96.5000114440918,
      "residual_shapes": [
        "bfloat16[1, 4096, 32, 128]",
        "bfloat16[1, 4096, 32, 128]",
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 4096, 32, 128]",
        "float32[1, 32, 4096]"
      ],
      "compile_s": 0.5635778919968288,
      "inner": 2,
      "median_ms": 21.200654002313968,
      "q1_ms": 20.984818249417003,
      "q3_ms": 21.458574999996927,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "jax_cudnn",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 257.000732421875,
      "xla_arg_mb": 112.0,
      "xla_out_mb": 48.000022888183594,
      "rel_vs_flex_repeat": 0.0003562509664334357,
      "compile_s": 0.13521383798797615,
      "inner": 3,
      "median_ms": 17.78807966426636,
      "q1_ms": 16.29159683458662,
      "q3_ms": 18.167273667737998,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 0,
      "rel_vs_flex_repeat": 0.0038485340774059296,
      "compile_s": 0.014565525998477824,
      "inner": 3,
      "median_ms": 15.662608663357485,
      "q1_ms": 15.471206335253859,
      "q3_ms": 15.97470733395312,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 128.0002670288086,
      "xla_arg_mb": 112.0,
      "xla_out_mb": 48.000022888183594,
      "residual_mb": 32.5000114440918,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 4096, 32, 128]",
        "float32[1, 32, 4096]"
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.548400097002741,
      "inner": 2,
      "median_ms": 20.467602495045867,
      "q1_ms": 20.358788999146782,
      "q3_ms": 20.86228650296107,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 192.0002670288086,
      "xla_arg_mb": 112.0,
      "xla_out_mb": 48.000022888183594,
      "residual_mb": 96.5000114440918,
      "residual_shapes": [
        "bfloat16[1, 4096, 32, 128]",
        "bfloat16[1, 4096, 32, 128]",
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 4096, 32, 128]",
        "float32[1, 32, 4096]"
      ],
      "compile_s": 0.6094337459944654,
      "inner": 2,
      "median_ms": 21.35691000148654,
      "q1_ms": 21.02330599882407,
      "q3_ms": 21.93694700326887,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "jax_cudnn",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 257.000732421875,
      "xla_arg_mb": 112.0,
      "xla_out_mb": 48.000022888183594,
      "rel_vs_flex_repeat": 0.0003562509664334357,
      "compile_s": 0.12411324700224213,
      "inner": 3,
      "median_ms": 16.45421266827422,
      "q1_ms": 16.197033000935335,
      "q3_ms": 16.95159199880436,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 1,
      "rel_vs_flex_repeat": 0.0038485340774059296,
      "compile_s": 0.014570175000699237,
      "inner": 3,
      "median_ms": 15.609207330271602,
      "q1_ms": 15.42524283286184,
      "q3_ms": 15.98374933140197,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 128.0002670288086,
      "xla_arg_mb": 112.0,
      "xla_out_mb": 48.000022888183594,
      "residual_mb": 32.5000114440918,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 4096, 32, 128]",
        "float32[1, 32, 4096]"
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.5376104149909224,
      "inner": 2,
      "median_ms": 22.009093503584154,
      "q1_ms": 20.530032754322747,
      "q3_ms": 22.67689949803753,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 192.0002670288086,
      "xla_arg_mb": 112.0,
      "xla_out_mb": 48.000022888183594,
      "residual_mb": 96.5000114440918,
      "residual_shapes": [
        "bfloat16[1, 4096, 32, 128]",
        "bfloat16[1, 4096, 32, 128]",
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 4096, 32, 128]",
        "float32[1, 32, 4096]"
      ],
      "compile_s": 0.5862645229935879,
      "inner": 2,
      "median_ms": 20.813807495869696,
      "q1_ms": 20.738730250741355,
      "q3_ms": 20.937681751092896,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "jax_cudnn",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 257.000732421875,
      "xla_arg_mb": 112.0,
      "xla_out_mb": 48.000022888183594,
      "rel_vs_flex_repeat": 0.0003562509664334357,
      "compile_s": 0.16331094400084112,
      "inner": 3,
      "median_ms": 18.01163600369667,
      "q1_ms": 16.934747000050265,
      "q3_ms": 18.319878998833396,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 2,
      "rel_vs_flex_repeat": 0.0038485340774059296,
      "compile_s": 0.015178125002421439,
      "inner": 3,
      "median_ms": 15.985936663734417,
      "q1_ms": 15.864338500250597,
      "q3_ms": 16.238917002434995,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 32.0,
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.2631951039948035,
      "inner": 9,
      "median_ms": 4.738337667529574,
      "q1_ms": 4.600572610797826,
      "q3_ms": 4.902387500300797,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 64.0,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.2654379970044829,
      "inner": 11,
      "median_ms": 4.6492449989521205,
      "q1_ms": 4.528385364127726,
      "q3_ms": 4.741678863038858,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 32.0,
      "rel_vs_flex_repeat": 2.5297151296399534e-05,
      "compile_s": 0.048952201002975926,
      "inner": 12,
      "median_ms": 4.205022166085352,
      "q1_ms": 4.16873129142914,
      "q3_ms": 4.280869292415446,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_vs_flex_repeat": 0.0030095174442976713,
      "compile_s": 0.004128360000322573,
      "inner": 12,
      "median_ms": 4.085218667266115,
      "q1_ms": 4.055338458177478,
      "q3_ms": 4.42822683301832,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 32.0,
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.26795627099636476,
      "inner": 10,
      "median_ms": 4.4861442991532385,
      "q1_ms": 4.419424999650801,
      "q3_ms": 4.618262200528989,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 64.0,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.2810744089947548,
      "inner": 11,
      "median_ms": 4.512563636208969,
      "q1_ms": 4.4950071824132465,
      "q3_ms": 4.582227226902349,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 32.0,
      "rel_vs_flex_repeat": 2.5297151296399534e-05,
      "compile_s": 0.0467307289945893,
      "inner": 12,
      "median_ms": 4.24210474981616,
      "q1_ms": 4.187112375196497,
      "q3_ms": 4.280579249704412,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_vs_flex_repeat": 0.0030095174442976713,
      "compile_s": 0.00435355601075571,
      "inner": 11,
      "median_ms": 4.101105182516304,
      "q1_ms": 4.084521954196547,
      "q3_ms": 4.406003817910625,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 32.0,
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.3046490189881297,
      "inner": 11,
      "median_ms": 4.894368636252528,
      "q1_ms": 4.592792181938421,
      "q3_ms": 5.051586499162527,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 64.0,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.2656064460024936,
      "inner": 10,
      "median_ms": 4.534460601280443,
      "q1_ms": 4.508528600126738,
      "q3_ms": 4.583069650107063,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 32.0,
      "rel_vs_flex_repeat": 2.5297151296399534e-05,
      "compile_s": 0.048181895006564446,
      "inner": 12,
      "median_ms": 4.358315415932641,
      "q1_ms": 4.223193166884206,
      "q3_ms": 4.434785875370531,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_vs_flex_repeat": 0.0030095174442976713,
      "compile_s": 0.004106087988475338,
      "inner": 12,
      "median_ms": 4.191382916663618,
      "q1_ms": 4.1092934582896605,
      "q3_ms": 4.491785416272856,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 128.0002670288086,
      "xla_arg_mb": 104.0,
      "xla_out_mb": 40.000022888183594,
      "residual_mb": 32.5000114440918,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 4096, 32, 128]",
        "float32[1, 32, 4096]"
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.59275513200555,
      "inner": 2,
      "median_ms": 21.39293299842393,
      "q1_ms": 21.03063274989836,
      "q3_ms": 23.09921350024524,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 192.0002670288086,
      "xla_arg_mb": 104.0,
      "xla_out_mb": 40.000022888183594,
      "residual_mb": 96.5000114440918,
      "residual_shapes": [
        "bfloat16[1, 4096, 32, 128]",
        "bfloat16[1, 4096, 32, 128]",
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 4096, 32, 128]",
        "float32[1, 32, 4096]"
      ],
      "compile_s": 0.5489335570018739,
      "inner": 2,
      "median_ms": 20.822032500291243,
      "q1_ms": 20.781295748747652,
      "q3_ms": 20.98281450162176,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "jax_cudnn",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 257.000732421875,
      "xla_arg_mb": 104.0,
      "xla_out_mb": 40.000022888183594,
      "rel_vs_flex_repeat": 0.00040300318505614996,
      "compile_s": 0.12904539800365455,
      "inner": 3,
      "median_ms": 16.3076286650418,
      "q1_ms": 16.21020933089312,
      "q3_ms": 16.698234501139574,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 0,
      "rel_vs_flex_repeat": 0.0038634708616882563,
      "compile_s": 0.014541861994075589,
      "inner": 3,
      "median_ms": 15.733061668773491,
      "q1_ms": 15.594833666303504,
      "q3_ms": 16.79658383363858,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 128.0002670288086,
      "xla_arg_mb": 104.0,
      "xla_out_mb": 40.000022888183594,
      "residual_mb": 32.5000114440918,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 4096, 32, 128]",
        "float32[1, 32, 4096]"
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.5151737770065665,
      "inner": 2,
      "median_ms": 20.109319499169942,
      "q1_ms": 20.037125250382815,
      "q3_ms": 20.63535200068145,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 192.0002670288086,
      "xla_arg_mb": 104.0,
      "xla_out_mb": 40.000022888183594,
      "residual_mb": 96.5000114440918,
      "residual_shapes": [
        "bfloat16[1, 4096, 32, 128]",
        "bfloat16[1, 4096, 32, 128]",
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 4096, 32, 128]",
        "float32[1, 32, 4096]"
      ],
      "compile_s": 0.5341661589918658,
      "inner": 2,
      "median_ms": 21.10195749992272,
      "q1_ms": 20.86006975150667,
      "q3_ms": 21.49326274593477,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "jax_cudnn",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 257.000732421875,
      "xla_arg_mb": 104.0,
      "xla_out_mb": 40.000022888183594,
      "rel_vs_flex_repeat": 0.00040300318505614996,
      "compile_s": 0.12515600099868607,
      "inner": 3,
      "median_ms": 16.295045003062114,
      "q1_ms": 16.143758837036636,
      "q3_ms": 16.498372999194544,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 1,
      "rel_vs_flex_repeat": 0.0038634708616882563,
      "compile_s": 0.014825868012849241,
      "inner": 3,
      "median_ms": 15.505611333840838,
      "q1_ms": 15.451836499172108,
      "q3_ms": 15.723156170376265,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 128.0002670288086,
      "xla_arg_mb": 104.0,
      "xla_out_mb": 40.000022888183594,
      "residual_mb": 32.5000114440918,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 4096, 32, 128]",
        "float32[1, 32, 4096]"
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.5386484990012832,
      "inner": 2,
      "median_ms": 20.711844495963305,
      "q1_ms": 20.50876774956123,
      "q3_ms": 21.1153252494114,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 192.0002670288086,
      "xla_arg_mb": 104.0,
      "xla_out_mb": 40.000022888183594,
      "residual_mb": 96.5000114440918,
      "residual_shapes": [
        "bfloat16[1, 4096, 32, 128]",
        "bfloat16[1, 4096, 32, 128]",
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 4096, 32, 128]",
        "float32[1, 32, 4096]"
      ],
      "compile_s": 0.531428970003617,
      "inner": 2,
      "median_ms": 20.789902002434246,
      "q1_ms": 20.71937625078135,
      "q3_ms": 21.218129753833637,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "jax_cudnn",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 257.000732421875,
      "xla_arg_mb": 104.0,
      "xla_out_mb": 40.000022888183594,
      "rel_vs_flex_repeat": 0.00040300318505614996,
      "compile_s": 0.13052359499852173,
      "inner": 3,
      "median_ms": 16.505660004137706,
      "q1_ms": 16.401780167749774,
      "q3_ms": 17.68134600085129,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 2,
      "rel_vs_flex_repeat": 0.0038634708616882563,
      "compile_s": 0.014573010004824027,
      "inner": 3,
      "median_ms": 15.865128332128126,
      "q1_ms": 15.618419832511185,
      "q3_ms": 16.3316293352788,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 34.0,
      "xla_out_mb": 32.0,
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.2603744569933042,
      "inner": 10,
      "median_ms": 4.745936799736228,
      "q1_ms": 4.491439249977702,
      "q3_ms": 4.805694350216072,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 64.0,
      "xla_arg_mb": 34.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.26073006000660826,
      "inner": 10,
      "median_ms": 4.549482099537272,
      "q1_ms": 4.500519499561051,
      "q3_ms": 4.845254099200247,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 34.0,
      "xla_out_mb": 32.0,
      "rel_vs_flex_repeat": 2.50296088779578e-05,
      "compile_s": 0.05152893700869754,
      "inner": 11,
      "median_ms": 4.196595818608661,
      "q1_ms": 4.174809546384495,
      "q3_ms": 4.203718408958098,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_vs_flex_repeat": 0.002989844884723425,
      "compile_s": 0.004091049006092362,
      "inner": 12,
      "median_ms": 4.142659916396951,
      "q1_ms": 4.105865750413311,
      "q3_ms": 4.449982582931019,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 34.0,
      "xla_out_mb": 32.0,
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.26366858900291845,
      "inner": 11,
      "median_ms": 4.443167727334763,
      "q1_ms": 4.415711317878133,
      "q3_ms": 4.499393272536426,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 64.0,
      "xla_arg_mb": 34.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.26841193399741314,
      "inner": 11,
      "median_ms": 4.513656635853377,
      "q1_ms": 4.491344863543583,
      "q3_ms": 4.62550909155239,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 34.0,
      "xla_out_mb": 32.0,
      "rel_vs_flex_repeat": 2.50296088779578e-05,
      "compile_s": 0.04807053299737163,
      "inner": 12,
      "median_ms": 4.2177210840842845,
      "q1_ms": 4.178961541886869,
      "q3_ms": 4.2259720830770675,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_vs_flex_repeat": 0.002989844884723425,
      "compile_s": 0.004096990000107326,
      "inner": 12,
      "median_ms": 4.13606650060198,
      "q1_ms": 4.084687749733954,
      "q3_ms": 4.184481583555074,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 34.0,
      "xla_out_mb": 32.0,
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.26369623199570924,
      "inner": 11,
      "median_ms": 4.450710182109932,
      "q1_ms": 4.438116818412461,
      "q3_ms": 4.466293727587486,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 64.0,
      "xla_arg_mb": 34.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.27281771900015883,
      "inner": 11,
      "median_ms": 4.535032454772259,
      "q1_ms": 4.50889845458071,
      "q3_ms": 4.545330136765683,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 34.0,
      "xla_out_mb": 32.0,
      "rel_vs_flex_repeat": 2.50296088779578e-05,
      "compile_s": 0.050585751989274286,
      "inner": 12,
      "median_ms": 4.228584166412475,
      "q1_ms": 4.166859416121346,
      "q3_ms": 4.253473291707148,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_vs_flex_repeat": 0.002989844884723425,
      "compile_s": 0.004345460998592898,
      "inner": 12,
      "median_ms": 4.132792916304122,
      "q1_ms": 4.0943418747095475,
      "q3_ms": 4.1800114170958595,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 128.0002670288086,
      "xla_arg_mb": 98.0,
      "xla_out_mb": 34.000022888183594,
      "residual_mb": 32.5000114440918,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 4096, 32, 128]",
        "float32[1, 32, 4096]"
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.5219245789921843,
      "inner": 2,
      "median_ms": 20.54376099840738,
      "q1_ms": 20.42215850451612,
      "q3_ms": 21.841661749931518,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 192.0002670288086,
      "xla_arg_mb": 98.0,
      "xla_out_mb": 34.000022888183594,
      "residual_mb": 96.5000114440918,
      "residual_shapes": [
        "bfloat16[1, 4096, 32, 128]",
        "bfloat16[1, 4096, 32, 128]",
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 4096, 32, 128]",
        "float32[1, 32, 4096]"
      ],
      "compile_s": 0.5309559349989286,
      "inner": 2,
      "median_ms": 20.788799498404842,
      "q1_ms": 20.61454024806153,
      "q3_ms": 20.894547498755855,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "jax_cudnn",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 257.000732421875,
      "xla_arg_mb": 98.0,
      "xla_out_mb": 34.000022888183594,
      "rel_vs_flex_repeat": 0.0005787573172710836,
      "compile_s": 0.12859165099507663,
      "inner": 3,
      "median_ms": 16.304953664075583,
      "q1_ms": 16.07861233545312,
      "q3_ms": 16.49009916824677,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 0,
      "rel_vs_flex_repeat": 0.0038645833265036345,
      "compile_s": 0.01516595199063886,
      "inner": 3,
      "median_ms": 15.476376002576822,
      "q1_ms": 15.344882335436221,
      "q3_ms": 15.76499216753291,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 128.0002670288086,
      "xla_arg_mb": 98.0,
      "xla_out_mb": 34.000022888183594,
      "residual_mb": 32.5000114440918,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 4096, 32, 128]",
        "float32[1, 32, 4096]"
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.5107005140016554,
      "inner": 2,
      "median_ms": 20.52686450042529,
      "q1_ms": 20.3889660006098,
      "q3_ms": 20.570512002450414,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 192.0002670288086,
      "xla_arg_mb": 98.0,
      "xla_out_mb": 34.000022888183594,
      "residual_mb": 96.5000114440918,
      "residual_shapes": [
        "bfloat16[1, 4096, 32, 128]",
        "bfloat16[1, 4096, 32, 128]",
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 4096, 32, 128]",
        "float32[1, 32, 4096]"
      ],
      "compile_s": 0.5435745590075385,
      "inner": 2,
      "median_ms": 21.163619006983936,
      "q1_ms": 20.681629503087606,
      "q3_ms": 21.54504074496799,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "jax_cudnn",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 257.000732421875,
      "xla_arg_mb": 98.0,
      "xla_out_mb": 34.000022888183594,
      "rel_vs_flex_repeat": 0.0005787573172710836,
      "compile_s": 0.13056462300301064,
      "inner": 3,
      "median_ms": 16.191001331511263,
      "q1_ms": 16.09376249931908,
      "q3_ms": 16.211872665735427,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 1,
      "rel_vs_flex_repeat": 0.0038645833265036345,
      "compile_s": 0.015287882997654378,
      "inner": 3,
      "median_ms": 15.54942766476112,
      "q1_ms": 15.49417316709878,
      "q3_ms": 15.766790667839814,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 128.0002670288086,
      "xla_arg_mb": 98.0,
      "xla_out_mb": 34.000022888183594,
      "residual_mb": 32.5000114440918,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 4096, 32, 128]",
        "float32[1, 32, 4096]"
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.5122191779955756,
      "inner": 2,
      "median_ms": 20.69503749953583,
      "q1_ms": 20.514901752903825,
      "q3_ms": 20.752846998220775,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 192.0002670288086,
      "xla_arg_mb": 98.0,
      "xla_out_mb": 34.000022888183594,
      "residual_mb": 96.5000114440918,
      "residual_shapes": [
        "bfloat16[1, 4096, 32, 128]",
        "bfloat16[1, 4096, 32, 128]",
        "key<fry>[]",
        "uint32[]",
        "bfloat16[1, 4096, 32, 128]",
        "float32[1, 32, 4096]"
      ],
      "compile_s": 0.5479741919989465,
      "inner": 2,
      "median_ms": 21.08099750330439,
      "q1_ms": 20.90095475068665,
      "q3_ms": 21.91982949807425,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "jax_cudnn",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 257.000732421875,
      "xla_arg_mb": 98.0,
      "xla_out_mb": 34.000022888183594,
      "rel_vs_flex_repeat": 0.0005787573172710836,
      "compile_s": 0.12706370999512728,
      "inner": 3,
      "median_ms": 16.637650003152277,
      "q1_ms": 16.24676533295618,
      "q3_ms": 16.868621666314237,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 2,
      "rel_vs_flex_repeat": 0.0038645833265036345,
      "compile_s": 0.014595171989640221,
      "inner": 3,
      "median_ms": 15.764180667853605,
      "q1_ms": 15.506446334863238,
      "q3_ms": 15.838962002211096,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_bfloat16"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 8.0,
      "rel_err": 0.0004156196373514831,
      "rel_err_parts": [
        0.0004156196373514831
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.22055891100899316,
      "inner": 117,
      "median_ms": 0.33324670084378016,
      "q1_ms": 0.3319733077038326,
      "q3_ms": 0.3433766324345707,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 16.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 8.0,
      "rel_err": 0.0004156196373514831,
      "rel_err_parts": [
        0.0004156196373514831
      ],
      "compile_s": 0.23790546899545006,
      "inner": 115,
      "median_ms": 0.36511257387009327,
      "q1_ms": 0.3613673565585328,
      "q3_ms": 0.3695490999709127,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 256.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 8.0,
      "rel_err": 0.000427189253969118,
      "rel_err_parts": [
        0.000427189253969118
      ],
      "rel_vs_flex_repeat": 0.0002741989155765623,
      "compile_s": 0.49862341600237414,
      "inner": 44,
      "median_ms": 0.9405965455781287,
      "q1_ms": 0.934846238702779,
      "q3_ms": 0.9539561251585837,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_err": 9.228879775946552e-07,
      "rel_err_parts": [
        9.228879775946552e-07
      ],
      "rel_vs_flex_repeat": 0.00041566276922822,
      "compile_s": 0.0012728480069199577,
      "inner": 47,
      "median_ms": 1.1242782765816144,
      "q1_ms": 1.1177669255127338,
      "q3_ms": 1.1391324470094601,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 8.0,
      "rel_err": 0.0004156196373514831,
      "rel_err_parts": [
        0.0004156196373514831
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.22654241100826766,
      "inner": 117,
      "median_ms": 0.3365743589392489,
      "q1_ms": 0.3286502692195722,
      "q3_ms": 0.3368606324865817,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 16.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 8.0,
      "rel_err": 0.0004156196373514831,
      "rel_err_parts": [
        0.0004156196373514831
      ],
      "compile_s": 0.22438329599390272,
      "inner": 117,
      "median_ms": 0.3702050000558504,
      "q1_ms": 0.3605361111361621,
      "q3_ms": 0.3839305427391082,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 256.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 8.0,
      "rel_err": 0.000427189253969118,
      "rel_err_parts": [
        0.000427189253969118
      ],
      "rel_vs_flex_repeat": 0.0002741989155765623,
      "compile_s": 0.5102083340025274,
      "inner": 49,
      "median_ms": 0.95205055104274,
      "q1_ms": 0.9444605715617975,
      "q3_ms": 0.9704972857880235,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_err": 9.228879775946552e-07,
      "rel_err_parts": [
        9.228879775946552e-07
      ],
      "rel_vs_flex_repeat": 0.00041566276922822,
      "compile_s": 0.0012389039911795408,
      "inner": 47,
      "median_ms": 1.134827808585969,
      "q1_ms": 1.1183188191275706,
      "q3_ms": 1.1551498829887468,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 8.0,
      "rel_err": 0.0004156196373514831,
      "rel_err_parts": [
        0.0004156196373514831
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.22861051400832366,
      "inner": 127,
      "median_ms": 0.3295099290773478,
      "q1_ms": 0.3258309212598517,
      "q3_ms": 0.33321776767637157,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 16.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 8.0,
      "rel_err": 0.0004156196373514831,
      "rel_err_parts": [
        0.0004156196373514831
      ],
      "compile_s": 0.22583582399238367,
      "inner": 114,
      "median_ms": 0.36444361403323056,
      "q1_ms": 0.36055067105559846,
      "q3_ms": 0.3744335175597052,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 256.0,
      "xla_arg_mb": 12.0,
      "xla_out_mb": 8.0,
      "rel_err": 0.000427189253969118,
      "rel_err_parts": [
        0.000427189253969118
      ],
      "rel_vs_flex_repeat": 0.0002741989155765623,
      "compile_s": 0.4955689080088632,
      "inner": 50,
      "median_ms": 0.9505771799013019,
      "q1_ms": 0.9376899900962599,
      "q3_ms": 0.9706643800018355,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_err": 9.228879775946552e-07,
      "rel_err_parts": [
        9.228879775946552e-07
      ],
      "rel_vs_flex_repeat": 0.00041566276922822,
      "compile_s": 0.0012586410011863336,
      "inner": 47,
      "median_ms": 1.121712574534653,
      "q1_ms": 1.1087221490092913,
      "q3_ms": 1.1451926277120639,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 24.000267028808594,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "residual_mb": 8.125011444091797,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "float32[1, 1024, 32, 64]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.0005357284680940211,
      "rel_err_parts": [
        0.0005357284680940211,
        0.0005281672347337008,
        0.00043399867718108
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.38358908500231337,
      "inner": 32,
      "median_ms": 1.590927968663891,
      "q1_ms": 1.575220765744234,
      "q3_ms": 1.643637218649019,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 40.000267028808594,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "residual_mb": 24.125011444091797,
      "residual_shapes": [
        "float32[1, 1024, 32, 64]",
        "float32[1, 1024, 32, 64]",
        "key<fry>[]",
        "uint32[]",
        "float32[1, 1024, 32, 64]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.0005357284680940211,
      "rel_err_parts": [
        0.0005357284680940211,
        0.0005281672347337008,
        0.00043399867718108
      ],
      "compile_s": 0.3745703720051097,
      "inner": 32,
      "median_ms": 1.6266847187580424,
      "q1_ms": 1.60807151564768,
      "q3_ms": 1.6382935159526824,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "jax_xla",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 520.1250152587891,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "rel_err": 0.0005327094695530832,
      "rel_err_parts": [
        0.0005327094695530832,
        0.0005276785232126713,
        0.0004340085433796048
      ],
      "rel_vs_flex_repeat": 8.946176967583597e-05,
      "compile_s": 0.405784214992309,
      "inner": 19,
      "median_ms": 2.5306418944015414,
      "q1_ms": 2.480284263009181,
      "q3_ms": 2.6129837363106687,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 0,
      "rel_err": 1.304025545323384e-06,
      "rel_err_parts": [
        1.304025545323384e-06,
        1.3000939134144573e-06,
        1.0802184533531545e-06
      ],
      "rel_vs_flex_repeat": 0.0005357961053960025,
      "compile_s": 0.004404372986755334,
      "inner": 11,
      "median_ms": 4.449140910186212,
      "q1_ms": 4.404991455058652,
      "q3_ms": 4.4620954550654535,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 24.000267028808594,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "residual_mb": 8.125011444091797,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "float32[1, 1024, 32, 64]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.0005357284680940211,
      "rel_err_parts": [
        0.0005357284680940211,
        0.0005281672347337008,
        0.00043399867718108
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.37723939200805034,
      "inner": 33,
      "median_ms": 1.617307333314481,
      "q1_ms": 1.5583638484134679,
      "q3_ms": 1.633644454572773,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 40.000267028808594,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "residual_mb": 24.125011444091797,
      "residual_shapes": [
        "float32[1, 1024, 32, 64]",
        "float32[1, 1024, 32, 64]",
        "key<fry>[]",
        "uint32[]",
        "float32[1, 1024, 32, 64]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.0005357284680940211,
      "rel_err_parts": [
        0.0005357284680940211,
        0.0005281672347337008,
        0.00043399867718108
      ],
      "compile_s": 0.3752806169941323,
      "inner": 32,
      "median_ms": 1.6555702500227198,
      "q1_ms": 1.6301471719089022,
      "q3_ms": 1.6593931250099558,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "jax_xla",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 520.1250152587891,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "rel_err": 0.0005327094695530832,
      "rel_err_parts": [
        0.0005327094695530832,
        0.0005276785232126713,
        0.0004340085433796048
      ],
      "rel_vs_flex_repeat": 8.946176967583597e-05,
      "compile_s": 0.4200757539947517,
      "inner": 19,
      "median_ms": 2.6483978420545005,
      "q1_ms": 2.5205052895073194,
      "q3_ms": 2.7832316053191835,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 1,
      "rel_err": 1.304025545323384e-06,
      "rel_err_parts": [
        1.304025545323384e-06,
        1.3000939134144573e-06,
        1.0802184533531545e-06
      ],
      "rel_vs_flex_repeat": 0.0005357961053960025,
      "compile_s": 0.004419600998517126,
      "inner": 11,
      "median_ms": 4.529688726920126,
      "q1_ms": 4.431727045316206,
      "q3_ms": 4.635736136564942,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 24.000267028808594,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "residual_mb": 8.125011444091797,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "float32[1, 1024, 32, 64]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.0005357284680940211,
      "rel_err_parts": [
        0.0005357284680940211,
        0.0005281672347337008,
        0.00043399867718108
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.3871588179899845,
      "inner": 33,
      "median_ms": 1.55845630268397,
      "q1_ms": 1.5455644999306197,
      "q3_ms": 1.5632981667295098,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 40.000267028808594,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "residual_mb": 24.125011444091797,
      "residual_shapes": [
        "float32[1, 1024, 32, 64]",
        "float32[1, 1024, 32, 64]",
        "key<fry>[]",
        "uint32[]",
        "float32[1, 1024, 32, 64]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.0005357284680940211,
      "rel_err_parts": [
        0.0005357284680940211,
        0.0005281672347337008,
        0.00043399867718108
      ],
      "compile_s": 0.38829853499191813,
      "inner": 31,
      "median_ms": 1.6535890646346456,
      "q1_ms": 1.646910112955788,
      "q3_ms": 1.7120611129587728,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "jax_xla",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 520.1250152587891,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 12.000022888183594,
      "rel_err": 0.0005327094695530832,
      "rel_err_parts": [
        0.0005327094695530832,
        0.0005276785232126713,
        0.0004340085433796048
      ],
      "rel_vs_flex_repeat": 8.946176967583597e-05,
      "compile_s": 0.4298891590005951,
      "inner": 20,
      "median_ms": 2.461803449841682,
      "q1_ms": 2.436092550124158,
      "q3_ms": 2.531381049993797,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 2,
      "rel_err": 1.304025545323384e-06,
      "rel_err_parts": [
        1.304025545323384e-06,
        1.3000939134144573e-06,
        1.0802184533531545e-06
      ],
      "rel_vs_flex_repeat": 0.0005357961053960025,
      "compile_s": 0.004459326999494806,
      "inner": 10,
      "median_ms": 4.5029649991192855,
      "q1_ms": 4.480109850555891,
      "q3_ms": 4.548023850657046,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 8.0,
      "rel_err": 0.0004223655560053885,
      "rel_err_parts": [
        0.0004223655560053885
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.22513558001082856,
      "inner": 103,
      "median_ms": 0.34702886407403727,
      "q1_ms": 0.34167218448433717,
      "q3_ms": 0.3542775290767274,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 16.0,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 8.0,
      "rel_err": 0.0004223655560053885,
      "rel_err_parts": [
        0.0004223655560053885
      ],
      "compile_s": 0.23171082000772003,
      "inner": 108,
      "median_ms": 0.36271844442968293,
      "q1_ms": 0.35840741198895826,
      "q3_ms": 0.38132090271068997,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 256.0,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 8.0,
      "rel_err": 0.00043358586844988167,
      "rel_err_parts": [
        0.00043358586844988167
      ],
      "rel_vs_flex_repeat": 0.00028111293795518577,
      "compile_s": 0.507760882988805,
      "inner": 47,
      "median_ms": 0.9320505530613059,
      "q1_ms": 0.9313072341102908,
      "q3_ms": 0.973169521262691,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_err": 9.329879162578436e-07,
      "rel_err_parts": [
        9.329879162578436e-07
      ],
      "rel_vs_flex_repeat": 0.00042240365291945636,
      "compile_s": 0.0012473300012061372,
      "inner": 47,
      "median_ms": 1.1226790848787518,
      "q1_ms": 1.1191178723783333,
      "q3_ms": 1.1371740638810963,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 8.0,
      "rel_err": 0.0004223655560053885,
      "rel_err_parts": [
        0.0004223655560053885
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.23246083900448866,
      "inner": 123,
      "median_ms": 0.33330210573058877,
      "q1_ms": 0.33186353663573176,
      "q3_ms": 0.34194716668358394,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 16.0,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 8.0,
      "rel_err": 0.0004223655560053885,
      "rel_err_parts": [
        0.0004223655560053885
      ],
      "compile_s": 0.23037402098998427,
      "inner": 119,
      "median_ms": 0.3588268403237795,
      "q1_ms": 0.35728097477773235,
      "q3_ms": 0.3601510672386884,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 256.0,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 8.0,
      "rel_err": 0.00043358586844988167,
      "rel_err_parts": [
        0.00043358586844988167
      ],
      "rel_vs_flex_repeat": 0.00028111293795518577,
      "compile_s": 0.5064782559929881,
      "inner": 45,
      "median_ms": 0.993389155741574,
      "q1_ms": 0.9743284220652034,
      "q3_ms": 1.0425916777522717,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_err": 9.329879162578436e-07,
      "rel_err_parts": [
        9.329879162578436e-07
      ],
      "rel_vs_flex_repeat": 0.00042240365291945636,
      "compile_s": 0.0012530310050351545,
      "inner": 46,
      "median_ms": 1.1198470870706092,
      "q1_ms": 1.1099252934281654,
      "q3_ms": 1.1258653368871738,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 8.0,
      "rel_err": 0.0004223655560053885,
      "rel_err_parts": [
        0.0004223655560053885
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.22068503001355566,
      "inner": 116,
      "median_ms": 0.33880451720223986,
      "q1_ms": 0.32895648272924827,
      "q3_ms": 0.3483598749630017,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 16.0,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 8.0,
      "rel_err": 0.0004223655560053885,
      "rel_err_parts": [
        0.0004223655560053885
      ],
      "compile_s": 0.22349840198876336,
      "inner": 121,
      "median_ms": 0.3633692726870029,
      "q1_ms": 0.3610415619168784,
      "q3_ms": 0.3756895619983331,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 256.0,
      "xla_arg_mb": 10.0,
      "xla_out_mb": 8.0,
      "rel_err": 0.00043358586844988167,
      "rel_err_parts": [
        0.00043358586844988167
      ],
      "rel_vs_flex_repeat": 0.00028111293795518577,
      "compile_s": 0.49793255899567157,
      "inner": 47,
      "median_ms": 0.9694572979648737,
      "q1_ms": 0.9672815957604511,
      "q3_ms": 0.979297819170863,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_err": 9.329879162578436e-07,
      "rel_err_parts": [
        9.329879162578436e-07
      ],
      "rel_vs_flex_repeat": 0.00042240365291945636,
      "compile_s": 0.0012595029984368011,
      "inner": 47,
      "median_ms": 1.1271758935726348,
      "q1_ms": 1.1089649469408385,
      "q3_ms": 1.1420322871925506,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 24.000267028808594,
      "xla_arg_mb": 18.0,
      "xla_out_mb": 10.000022888183594,
      "residual_mb": 8.125011444091797,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "float32[1, 1024, 32, 64]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.0005337766488082707,
      "rel_err_parts": [
        0.0005337766488082707,
        0.0005261579062789679,
        0.0004293412494007498
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.3845008899952518,
      "inner": 28,
      "median_ms": 1.6207051069484026,
      "q1_ms": 1.5557758392138306,
      "q3_ms": 1.6631432499707444,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 40.000267028808594,
      "xla_arg_mb": 18.0,
      "xla_out_mb": 10.000022888183594,
      "residual_mb": 24.125011444091797,
      "residual_shapes": [
        "float32[1, 1024, 32, 64]",
        "float32[1, 1024, 32, 64]",
        "key<fry>[]",
        "uint32[]",
        "float32[1, 1024, 32, 64]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.0005337766488082707,
      "rel_err_parts": [
        0.0005337766488082707,
        0.0005261579062789679,
        0.0004293412494007498
      ],
      "compile_s": 0.4076893200108316,
      "inner": 32,
      "median_ms": 1.7191718438880343,
      "q1_ms": 1.6270494375021372,
      "q3_ms": 1.78863842188548,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "jax_xla",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 520.1250152587891,
      "xla_arg_mb": 18.0,
      "xla_out_mb": 10.000022888183594,
      "rel_err": 0.0005308594554662704,
      "rel_err_parts": [
        0.0005308594554662704,
        0.0005253087729215622,
        0.00042938892147503793
      ],
      "rel_vs_flex_repeat": 9.13066032808274e-05,
      "compile_s": 0.4143319370050449,
      "inner": 19,
      "median_ms": 2.500095684234494,
      "q1_ms": 2.4461916051988832,
      "q3_ms": 2.526659605408847,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 0,
      "rel_err": 1.2914464377900003e-06,
      "rel_err_parts": [
        1.2914464377900003e-06,
        1.287366444557847e-06,
        1.0649637260939926e-06
      ],
      "rel_vs_flex_repeat": 0.00053384032798931,
      "compile_s": 0.004854073005844839,
      "inner": 11,
      "median_ms": 4.434247181052342,
      "q1_ms": 4.39438913665733,
      "q3_ms": 4.509646727290766,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 24.000267028808594,
      "xla_arg_mb": 18.0,
      "xla_out_mb": 10.000022888183594,
      "residual_mb": 8.125011444091797,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "float32[1, 1024, 32, 64]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.0005337766488082707,
      "rel_err_parts": [
        0.0005337766488082707,
        0.0005261579062789679,
        0.0004293412494007498
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.3850665200006915,
      "inner": 33,
      "median_ms": 1.570190000083213,
      "q1_ms": 1.5558023182110805,
      "q3_ms": 1.581265666490336,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 40.000267028808594,
      "xla_arg_mb": 18.0,
      "xla_out_mb": 10.000022888183594,
      "residual_mb": 24.125011444091797,
      "residual_shapes": [
        "float32[1, 1024, 32, 64]",
        "float32[1, 1024, 32, 64]",
        "key<fry>[]",
        "uint32[]",
        "float32[1, 1024, 32, 64]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.0005337766488082707,
      "rel_err_parts": [
        0.0005337766488082707,
        0.0005261579062789679,
        0.0004293412494007498
      ],
      "compile_s": 0.5124033759930171,
      "inner": 32,
      "median_ms": 1.7993954998019035,
      "q1_ms": 1.7744410936302302,
      "q3_ms": 1.8333130783503293,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "jax_xla",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 520.1250152587891,
      "xla_arg_mb": 18.0,
      "xla_out_mb": 10.000022888183594,
      "rel_err": 0.0005308594554662704,
      "rel_err_parts": [
        0.0005308594554662704,
        0.0005253087729215622,
        0.00042938892147503793
      ],
      "rel_vs_flex_repeat": 9.13066032808274e-05,
      "compile_s": 0.4686812089930754,
      "inner": 20,
      "median_ms": 2.4362741503864527,
      "q1_ms": 2.4185077749280026,
      "q3_ms": 2.450646075158147,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 1,
      "rel_err": 1.2914464377900003e-06,
      "rel_err_parts": [
        1.2914464377900003e-06,
        1.287366444557847e-06,
        1.0649637260939926e-06
      ],
      "rel_vs_flex_repeat": 0.00053384032798931,
      "compile_s": 0.004384674990433268,
      "inner": 11,
      "median_ms": 4.471657091264866,
      "q1_ms": 4.388055772737557,
      "q3_ms": 4.53830090952,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 24.000267028808594,
      "xla_arg_mb": 18.0,
      "xla_out_mb": 10.000022888183594,
      "residual_mb": 8.125011444091797,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "float32[1, 1024, 32, 64]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.0005337766488082707,
      "rel_err_parts": [
        0.0005337766488082707,
        0.0005261579062789679,
        0.0004293412494007498
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.4505780600011349,
      "inner": 28,
      "median_ms": 1.7278538928491929,
      "q1_ms": 1.7209447679696106,
      "q3_ms": 1.746206464304123,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 40.000267028808594,
      "xla_arg_mb": 18.0,
      "xla_out_mb": 10.000022888183594,
      "residual_mb": 24.125011444091797,
      "residual_shapes": [
        "float32[1, 1024, 32, 64]",
        "float32[1, 1024, 32, 64]",
        "key<fry>[]",
        "uint32[]",
        "float32[1, 1024, 32, 64]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.0005337766488082707,
      "rel_err_parts": [
        0.0005337766488082707,
        0.0005261579062789679,
        0.0004293412494007498
      ],
      "compile_s": 0.3894354649964953,
      "inner": 32,
      "median_ms": 1.6709455940144835,
      "q1_ms": 1.6318560158197215,
      "q3_ms": 1.6792754220205097,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "jax_xla",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 520.1250152587891,
      "xla_arg_mb": 18.0,
      "xla_out_mb": 10.000022888183594,
      "rel_err": 0.0005308594554662704,
      "rel_err_parts": [
        0.0005308594554662704,
        0.0005253087729215622,
        0.00042938892147503793
      ],
      "rel_vs_flex_repeat": 9.13066032808274e-05,
      "compile_s": 0.4540201630006777,
      "inner": 20,
      "median_ms": 2.4239092999778222,
      "q1_ms": 2.402027349671698,
      "q3_ms": 2.454273424882558,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 2,
      "rel_err": 1.2914464377900003e-06,
      "rel_err_parts": [
        1.2914464377900003e-06,
        1.287366444557847e-06,
        1.0649637260939926e-06
      ],
      "rel_vs_flex_repeat": 0.00053384032798931,
      "compile_s": 0.004392629998619668,
      "inner": 11,
      "median_ms": 4.391374272844669,
      "q1_ms": 4.368919090748849,
      "q3_ms": 4.5052310908266175,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 8.5,
      "xla_out_mb": 8.0,
      "rel_err": 0.00042329682037234306,
      "rel_err_parts": [
        0.00042329682037234306
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.21648488400387578,
      "inner": 115,
      "median_ms": 0.3643893825305302,
      "q1_ms": 0.3365583782357609,
      "q3_ms": 0.3651229869042847,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 16.0,
      "xla_arg_mb": 8.5,
      "xla_out_mb": 8.0,
      "rel_err": 0.00042329682037234306,
      "rel_err_parts": [
        0.00042329682037234306
      ],
      "compile_s": 0.2391228519991273,
      "inner": 120,
      "median_ms": 0.3548318834267169,
      "q1_ms": 0.3532004666946401,
      "q3_ms": 0.35612032497738255,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 256.0,
      "xla_arg_mb": 8.5,
      "xla_out_mb": 8.0,
      "rel_err": 0.0004341566236689687,
      "rel_err_parts": [
        0.0004341566236689687
      ],
      "rel_vs_flex_repeat": 0.00027838849928230047,
      "compile_s": 0.3800664999871515,
      "inner": 58,
      "median_ms": 0.726284723927604,
      "q1_ms": 0.7213991120806895,
      "q3_ms": 0.7284319827108291,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_err": 9.253066650671826e-07,
      "rel_err_parts": [
        9.253066650671826e-07
      ],
      "rel_vs_flex_repeat": 0.00042332554585300386,
      "compile_s": 0.0012547440128400922,
      "inner": 38,
      "median_ms": 1.2218494999902567,
      "q1_ms": 1.2215025130181427,
      "q3_ms": 1.2222652763547988,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 8.5,
      "xla_out_mb": 8.0,
      "rel_err": 0.00042329682037234306,
      "rel_err_parts": [
        0.00042329682037234306
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.21857177099445835,
      "inner": 115,
      "median_ms": 0.3534798782430184,
      "q1_ms": 0.32878184353214,
      "q3_ms": 0.3678410434234968,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 16.0,
      "xla_arg_mb": 8.5,
      "xla_out_mb": 8.0,
      "rel_err": 0.00042329682037234306,
      "rel_err_parts": [
        0.00042329682037234306
      ],
      "compile_s": 0.24819243999081664,
      "inner": 116,
      "median_ms": 0.3564256119967342,
      "q1_ms": 0.35387431897357474,
      "q3_ms": 0.35691377587650563,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 256.0,
      "xla_arg_mb": 8.5,
      "xla_out_mb": 8.0,
      "rel_err": 0.0004341566236689687,
      "rel_err_parts": [
        0.0004341566236689687
      ],
      "rel_vs_flex_repeat": 0.00027838849928230047,
      "compile_s": 0.34140680098789744,
      "inner": 59,
      "median_ms": 0.7096743218549448,
      "q1_ms": 0.7091197881983766,
      "q3_ms": 0.7234493643440054,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_err": 9.253066650671826e-07,
      "rel_err_parts": [
        9.253066650671826e-07
      ],
      "rel_vs_flex_repeat": 0.00042332554585300386,
      "compile_s": 0.0012526089994935319,
      "inner": 47,
      "median_ms": 1.1136498722678765,
      "q1_ms": 1.0980113509774923,
      "q3_ms": 1.1240317553722676,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 8.5,
      "xla_out_mb": 8.0,
      "rel_err": 0.00042329682037234306,
      "rel_err_parts": [
        0.00042329682037234306
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.23518554300244432,
      "inner": 115,
      "median_ms": 0.3654929477741699,
      "q1_ms": 0.3598493999680099,
      "q3_ms": 0.36800495651297516,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 16.0,
      "xla_arg_mb": 8.5,
      "xla_out_mb": 8.0,
      "rel_err": 0.00042329682037234306,
      "rel_err_parts": [
        0.00042329682037234306
      ],
      "compile_s": 0.2527146869979333,
      "inner": 109,
      "median_ms": 0.3538853761210286,
      "q1_ms": 0.352852233929064,
      "q3_ms": 0.35570492196845216,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 256.0,
      "xla_arg_mb": 8.5,
      "xla_out_mb": 8.0,
      "rel_err": 0.0004341566236689687,
      "rel_err_parts": [
        0.0004341566236689687
      ],
      "rel_vs_flex_repeat": 0.00027838849928230047,
      "compile_s": 0.3870068700052798,
      "inner": 62,
      "median_ms": 0.7237363226463899,
      "q1_ms": 0.717466878995765,
      "q3_ms": 0.7278170242282231,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_err": 9.253066650671826e-07,
      "rel_err_parts": [
        9.253066650671826e-07
      ],
      "rel_vs_flex_repeat": 0.00042332554585300386,
      "compile_s": 0.0012388229952193797,
      "inner": 38,
      "median_ms": 1.2211426315173555,
      "q1_ms": 1.2206117630979096,
      "q3_ms": 1.221236342021975,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 24.000267028808594,
      "xla_arg_mb": 16.5,
      "xla_out_mb": 8.500022888183594,
      "residual_mb": 8.125011444091797,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "float32[1, 1024, 32, 64]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.0005333983572199941,
      "rel_err_parts": [
        0.0005333983572199941,
        0.0005212461692281067,
        0.0004193610802758485
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.3774512230011169,
      "inner": 31,
      "median_ms": 1.494108612889484,
      "q1_ms": 1.4938624998678511,
      "q3_ms": 1.4973675161825433,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 40.000267028808594,
      "xla_arg_mb": 16.5,
      "xla_out_mb": 8.500022888183594,
      "residual_mb": 24.125011444091797,
      "residual_shapes": [
        "float32[1, 1024, 32, 64]",
        "float32[1, 1024, 32, 64]",
        "key<fry>[]",
        "uint32[]",
        "float32[1, 1024, 32, 64]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.0005333983572199941,
      "rel_err_parts": [
        0.0005333983572199941,
        0.0005212461692281067,
        0.0004193610802758485
      ],
      "compile_s": 0.39160030100902077,
      "inner": 32,
      "median_ms": 1.5804769377609773,
      "q1_ms": 1.5781484844410443,
      "q3_ms": 1.5838331562463281,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "jax_xla",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 512.1250152587891,
      "xla_arg_mb": 16.5,
      "xla_out_mb": 8.500022888183594,
      "rel_err": 0.0005304054357111454,
      "rel_err_parts": [
        0.0005304054357111454,
        0.0005210100789554417,
        0.00041946006240323186
      ],
      "rel_vs_flex_repeat": 8.945012814365327e-05,
      "compile_s": 0.3901761669985717,
      "inner": 20,
      "median_ms": 2.3471415501262527,
      "q1_ms": 2.346680374830612,
      "q3_ms": 2.3478979248466203,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 0,
      "rel_err": 1.295729248340649e-06,
      "rel_err_parts": [
        1.295729248340649e-06,
        1.2921125289722113e-06,
        1.0399202210464864e-06
      ],
      "rel_vs_flex_repeat": 0.0005334539455361664,
      "compile_s": 0.004548755998257548,
      "inner": 11,
      "median_ms": 4.339607272413559,
      "q1_ms": 4.3083076360000465,
      "q3_ms": 4.365149681689218,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 24.000267028808594,
      "xla_arg_mb": 16.5,
      "xla_out_mb": 8.500022888183594,
      "residual_mb": 8.125011444091797,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "float32[1, 1024, 32, 64]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.0005333983572199941,
      "rel_err_parts": [
        0.0005333983572199941,
        0.0005212461692281067,
        0.0004193610802758485
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.4450734459969681,
      "inner": 34,
      "median_ms": 1.607751970689552,
      "q1_ms": 1.5670657350496828,
      "q3_ms": 1.7056126913196161,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 40.000267028808594,
      "xla_arg_mb": 16.5,
      "xla_out_mb": 8.500022888183594,
      "residual_mb": 24.125011444091797,
      "residual_shapes": [
        "float32[1, 1024, 32, 64]",
        "float32[1, 1024, 32, 64]",
        "key<fry>[]",
        "uint32[]",
        "float32[1, 1024, 32, 64]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.0005333983572199941,
      "rel_err_parts": [
        0.0005333983572199941,
        0.0005212461692281067,
        0.0004193610802758485
      ],
      "compile_s": 0.39272587100276724,
      "inner": 32,
      "median_ms": 1.594008187566942,
      "q1_ms": 1.587145828125358,
      "q3_ms": 1.6153648282397626,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "jax_xla",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 512.1250152587891,
      "xla_arg_mb": 16.5,
      "xla_out_mb": 8.500022888183594,
      "rel_err": 0.0005304054357111454,
      "rel_err_parts": [
        0.0005304054357111454,
        0.0005210100789554417,
        0.00041946006240323186
      ],
      "rel_vs_flex_repeat": 8.945012814365327e-05,
      "compile_s": 0.45182449799904134,
      "inner": 17,
      "median_ms": 2.7456521176431767,
      "q1_ms": 2.7369017351541998,
      "q3_ms": 2.758705881954965,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 1,
      "rel_err": 1.295729248340649e-06,
      "rel_err_parts": [
        1.295729248340649e-06,
        1.2921125289722113e-06,
        1.0399202210464864e-06
      ],
      "rel_vs_flex_repeat": 0.0005334539455361664,
      "compile_s": 0.004387962006148882,
      "inner": 11,
      "median_ms": 4.497249182102017,
      "q1_ms": 4.423196818441449,
      "q3_ms": 4.6323488184663635,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 24.000267028808594,
      "xla_arg_mb": 16.5,
      "xla_out_mb": 8.500022888183594,
      "residual_mb": 8.125011444091797,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "float32[1, 1024, 32, 64]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.0005333983572199941,
      "rel_err_parts": [
        0.0005333983572199941,
        0.0005212461692281067,
        0.0004193610802758485
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.4054784570034826,
      "inner": 34,
      "median_ms": 1.5521959708969328,
      "q1_ms": 1.5402324707418515,
      "q3_ms": 1.5823309412927997,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 40.000267028808594,
      "xla_arg_mb": 16.5,
      "xla_out_mb": 8.500022888183594,
      "residual_mb": 24.125011444091797,
      "residual_shapes": [
        "float32[1, 1024, 32, 64]",
        "float32[1, 1024, 32, 64]",
        "key<fry>[]",
        "uint32[]",
        "float32[1, 1024, 32, 64]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.0005333983572199941,
      "rel_err_parts": [
        0.0005333983572199941,
        0.0005212461692281067,
        0.0004193610802758485
      ],
      "compile_s": 0.40499069399083965,
      "inner": 28,
      "median_ms": 1.6618354285518373,
      "q1_ms": 1.601826339278237,
      "q3_ms": 1.7266387499798188,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "jax_xla",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 512.1250152587891,
      "xla_arg_mb": 16.5,
      "xla_out_mb": 8.500022888183594,
      "rel_err": 0.0005304054357111454,
      "rel_err_parts": [
        0.0005304054357111454,
        0.0005210100789554417,
        0.00041946006240323186
      ],
      "rel_vs_flex_repeat": 8.945012814365327e-05,
      "compile_s": 0.42397582200646866,
      "inner": 19,
      "median_ms": 2.469486684085017,
      "q1_ms": 2.4156751847170344,
      "q3_ms": 2.713225184523157,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 2,
      "rel_err": 1.295729248340649e-06,
      "rel_err_parts": [
        1.295729248340649e-06,
        1.2921125289722113e-06,
        1.0399202210464864e-06
      ],
      "rel_vs_flex_repeat": 0.0005334539455361664,
      "compile_s": 0.00461512099718675,
      "inner": 10,
      "median_ms": 4.389921198890079,
      "q1_ms": 4.358632450021105,
      "q3_ms": 4.430576099548489,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 16.0,
      "rel_err": 0.00041457480983808637,
      "rel_err_parts": [
        0.00041457480983808637
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.3356036420009332,
      "inner": 57,
      "median_ms": 0.7890422454230455,
      "q1_ms": 0.7784283070737219,
      "q3_ms": 0.8359137280964652,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 32.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 16.0,
      "rel_err": 0.00041457480983808637,
      "rel_err_parts": [
        0.00041457480983808637
      ],
      "compile_s": 0.3316689680068521,
      "inner": 58,
      "median_ms": 0.7680927931154468,
      "q1_ms": 0.7638970173188838,
      "q3_ms": 0.7805062844548454,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 256.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 16.0,
      "rel_err": 0.0004252932558301836,
      "rel_err_parts": [
        0.0004252932558301836
      ],
      "rel_vs_flex_repeat": 0.00027938844868913293,
      "compile_s": 0.49537091300589964,
      "inner": 41,
      "median_ms": 1.2195230729270317,
      "q1_ms": 1.1929886828378273,
      "q3_ms": 1.2315352926124454,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_err": 1.1142960829602089e-06,
      "rel_err_parts": [
        1.1142960829602089e-06
      ],
      "rel_vs_flex_repeat": 0.00041461666114628315,
      "compile_s": 0.0019170469895470887,
      "inner": 29,
      "median_ms": 1.78924296553054,
      "q1_ms": 1.7731908793111555,
      "q3_ms": 1.8381382241355384,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 16.0,
      "rel_err": 0.00041457480983808637,
      "rel_err_parts": [
        0.00041457480983808637
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.32858971299719997,
      "inner": 63,
      "median_ms": 0.7189123176171311,
      "q1_ms": 0.7118442619454101,
      "q3_ms": 0.7198205475726844,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 32.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 16.0,
      "rel_err": 0.00041457480983808637,
      "rel_err_parts": [
        0.00041457480983808637
      ],
      "compile_s": 0.32969232799950987,
      "inner": 58,
      "median_ms": 0.7745524311883792,
      "q1_ms": 0.7679181550790011,
      "q3_ms": 0.7995175086574002,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 256.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 16.0,
      "rel_err": 0.0004252932558301836,
      "rel_err_parts": [
        0.0004252932558301836
      ],
      "rel_vs_flex_repeat": 0.00027938844868913293,
      "compile_s": 0.5041794680000748,
      "inner": 41,
      "median_ms": 1.1979055607963989,
      "q1_ms": 1.1921207194795786,
      "q3_ms": 1.2248691340706812,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_err": 1.1142960829602089e-06,
      "rel_err_parts": [
        1.1142960829602089e-06
      ],
      "rel_vs_flex_repeat": 0.00041461666114628315,
      "compile_s": 0.001929078993271105,
      "inner": 29,
      "median_ms": 1.7550631033874855,
      "q1_ms": 1.7346160000607629,
      "q3_ms": 1.7713401551912376,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 16.0,
      "rel_err": 0.00041457480983808637,
      "rel_err_parts": [
        0.00041457480983808637
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.343190826009959,
      "inner": 61,
      "median_ms": 0.7111276559566041,
      "q1_ms": 0.7108394998854759,
      "q3_ms": 0.7451415163575534,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 32.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 16.0,
      "rel_err": 0.00041457480983808637,
      "rel_err_parts": [
        0.00041457480983808637
      ],
      "compile_s": 0.3683734890073538,
      "inner": 60,
      "median_ms": 0.8032528331871921,
      "q1_ms": 0.7889248917005413,
      "q3_ms": 0.8202504833510222,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 256.0,
      "xla_arg_mb": 24.0,
      "xla_out_mb": 16.0,
      "rel_err": 0.0004252932558301836,
      "rel_err_parts": [
        0.0004252932558301836
      ],
      "rel_vs_flex_repeat": 0.00027938844868913293,
      "compile_s": 0.5100355570029933,
      "inner": 41,
      "median_ms": 1.2031545122182452,
      "q1_ms": 1.1830109268463809,
      "q3_ms": 1.2572774024348048,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_err": 1.1142960829602089e-06,
      "rel_err_parts": [
        1.1142960829602089e-06
      ],
      "rel_vs_flex_repeat": 0.00041461666114628315,
      "compile_s": 0.0019237789965700358,
      "inner": 27,
      "median_ms": 1.9802248889054346,
      "q1_ms": 1.877281037094589,
      "q3_ms": 1.9930066479934827,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 48.000267028808594,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "residual_mb": 16.125011444091797,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "float32[1, 1024, 32, 128]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.0005201493622735143,
      "rel_err_parts": [
        0.0005201493622735143,
        0.000516366446390748,
        0.00042170321103185415
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.960016426994116,
      "inner": 13,
      "median_ms": 3.4569393073280272,
      "q1_ms": 3.4282785770931063,
      "q3_ms": 3.6304210001477397,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 80.0002670288086,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "residual_mb": 48.1250114440918,
      "residual_shapes": [
        "float32[1, 1024, 32, 128]",
        "float32[1, 1024, 32, 128]",
        "key<fry>[]",
        "uint32[]",
        "float32[1, 1024, 32, 128]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.0005201493622735143,
      "rel_err_parts": [
        0.0005201493622735143,
        0.000516366446390748,
        0.00042170321103185415
      ],
      "compile_s": 0.970918202001485,
      "inner": 14,
      "median_ms": 3.588349428485214,
      "q1_ms": 3.546041535238536,
      "q3_ms": 3.6473589638522492,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "jax_xla",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 528.1252593994141,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "rel_err": 0.0005184368928894401,
      "rel_err_parts": [
        0.0005184368928894401,
        0.0005159106804057956,
        0.0004220151458866894
      ],
      "rel_vs_flex_repeat": 8.460850222036242e-05,
      "compile_s": 0.43539327300095465,
      "inner": 15,
      "median_ms": 3.0874316667905077,
      "q1_ms": 3.073090266843792,
      "q3_ms": 3.119438599969726,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 0,
      "rel_err": 1.3581701523435186e-06,
      "rel_err_parts": [
        1.3581701523435186e-06,
        1.3558022828874527e-06,
        1.111863753067155e-06
      ],
      "rel_vs_flex_repeat": 0.0005202081520110369,
      "compile_s": 0.007951192994369194,
      "inner": 6,
      "median_ms": 8.276851998137621,
      "q1_ms": 8.112604832907284,
      "q3_ms": 8.286476166783057,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 48.000267028808594,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "residual_mb": 16.125011444091797,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "float32[1, 1024, 32, 128]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.0005201493622735143,
      "rel_err_parts": [
        0.0005201493622735143,
        0.000516366446390748,
        0.00042170321103185415
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.8761359750060365,
      "inner": 14,
      "median_ms": 3.423218499977208,
      "q1_ms": 3.337335249982841,
      "q3_ms": 3.58741589246035,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 80.0002670288086,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "residual_mb": 48.1250114440918,
      "residual_shapes": [
        "float32[1, 1024, 32, 128]",
        "float32[1, 1024, 32, 128]",
        "key<fry>[]",
        "uint32[]",
        "float32[1, 1024, 32, 128]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.0005201493622735143,
      "rel_err_parts": [
        0.0005201493622735143,
        0.000516366446390748,
        0.00042170321103185415
      ],
      "compile_s": 0.8582475319999503,
      "inner": 14,
      "median_ms": 3.48959907153455,
      "q1_ms": 3.460286107188689,
      "q3_ms": 3.5282507856858762,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "jax_xla",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 528.1252593994141,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "rel_err": 0.0005184368928894401,
      "rel_err_parts": [
        0.0005184368928894401,
        0.0005159106804057956,
        0.0004220151458866894
      ],
      "rel_vs_flex_repeat": 8.460850222036242e-05,
      "compile_s": 0.42134428399731405,
      "inner": 15,
      "median_ms": 3.083774800567577,
      "q1_ms": 3.0529137327296967,
      "q3_ms": 3.1903647662450867,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 1,
      "rel_err": 1.3581701523435186e-06,
      "rel_err_parts": [
        1.3581701523435186e-06,
        1.3558022828874527e-06,
        1.111863753067155e-06
      ],
      "rel_vs_flex_repeat": 0.0005202081520110369,
      "compile_s": 0.007954810003866442,
      "inner": 6,
      "median_ms": 8.249271668319125,
      "q1_ms": 8.089457667665556,
      "q3_ms": 8.30005183403652,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 48.000267028808594,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "residual_mb": 16.125011444091797,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "float32[1, 1024, 32, 128]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.0005201493622735143,
      "rel_err_parts": [
        0.0005201493622735143,
        0.000516366446390748,
        0.00042170321103185415
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.865998776993365,
      "inner": 15,
      "median_ms": 3.364879266397717,
      "q1_ms": 3.3417001002817415,
      "q3_ms": 3.378005432993329,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 80.0002670288086,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "residual_mb": 48.1250114440918,
      "residual_shapes": [
        "float32[1, 1024, 32, 128]",
        "float32[1, 1024, 32, 128]",
        "key<fry>[]",
        "uint32[]",
        "float32[1, 1024, 32, 128]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.0005201493622735143,
      "rel_err_parts": [
        0.0005201493622735143,
        0.000516366446390748,
        0.00042170321103185415
      ],
      "compile_s": 0.8621899300051155,
      "inner": 13,
      "median_ms": 3.468206077321576,
      "q1_ms": 3.4453196533901904,
      "q3_ms": 3.509812461113772,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "jax_xla",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 528.1252593994141,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 24.000022888183594,
      "rel_err": 0.0005184368928894401,
      "rel_err_parts": [
        0.0005184368928894401,
        0.0005159106804057956,
        0.0004220151458866894
      ],
      "rel_vs_flex_repeat": 8.460850222036242e-05,
      "compile_s": 0.5391050320031354,
      "inner": 16,
      "median_ms": 3.1003123131085886,
      "q1_ms": 3.0327131871672464,
      "q3_ms": 3.1599976873621927,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 2,
      "rel_err": 1.3581701523435186e-06,
      "rel_err_parts": [
        1.3581701523435186e-06,
        1.3558022828874527e-06,
        1.111863753067155e-06
      ],
      "rel_vs_flex_repeat": 0.0005202081520110369,
      "compile_s": 0.008025542992982082,
      "inner": 6,
      "median_ms": 8.123323333469065,
      "q1_ms": 8.09931291587418,
      "q3_ms": 8.574343333748402,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 16.0,
      "rel_err": 0.00040910390089266,
      "rel_err_parts": [
        0.00040910390089266
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.334631552992505,
      "inner": 57,
      "median_ms": 0.772304982283937,
      "q1_ms": 0.7695806051375704,
      "q3_ms": 0.784931754567783,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 32.0,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 16.0,
      "rel_err": 0.00040910390089266,
      "rel_err_parts": [
        0.00040910390089266
      ],
      "compile_s": 0.32582887999888044,
      "inner": 59,
      "median_ms": 0.7692903220570662,
      "q1_ms": 0.7624187542745149,
      "q3_ms": 0.7800058474304485,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 256.0,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 16.0,
      "rel_err": 0.0004199399845674634,
      "rel_err_parts": [
        0.0004199399845674634
      ],
      "rel_vs_flex_repeat": 0.0002754633314907551,
      "compile_s": 0.5010764970065793,
      "inner": 41,
      "median_ms": 1.2213980485174079,
      "q1_ms": 1.2141466586177647,
      "q3_ms": 1.2340619022568982,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_err": 1.1000706763297785e-06,
      "rel_err_parts": [
        1.1000706763297785e-06
      ],
      "rel_vs_flex_repeat": 0.000409148313337937,
      "compile_s": 0.001930843005538918,
      "inner": 29,
      "median_ms": 1.7262244484676368,
      "q1_ms": 1.7238216036087672,
      "q3_ms": 1.7480640345983658,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 16.0,
      "rel_err": 0.00040910390089266,
      "rel_err_parts": [
        0.00040910390089266
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.32800424499146175,
      "inner": 62,
      "median_ms": 0.716346838734969,
      "q1_ms": 0.7065230241888428,
      "q3_ms": 0.7233693387177247,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 32.0,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 16.0,
      "rel_err": 0.00040910390089266,
      "rel_err_parts": [
        0.00040910390089266
      ],
      "compile_s": 0.3327476190024754,
      "inner": 58,
      "median_ms": 0.7565710001127345,
      "q1_ms": 0.7529807586126543,
      "q3_ms": 0.7716572069731573,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 256.0,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 16.0,
      "rel_err": 0.0004199399845674634,
      "rel_err_parts": [
        0.0004199399845674634
      ],
      "rel_vs_flex_repeat": 0.0002754633314907551,
      "compile_s": 0.48652111999399494,
      "inner": 40,
      "median_ms": 1.2191301750135608,
      "q1_ms": 1.2055135624905233,
      "q3_ms": 1.2584787626110483,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_err": 1.1000706763297785e-06,
      "rel_err_parts": [
        1.1000706763297785e-06
      ],
      "rel_vs_flex_repeat": 0.000409148313337937,
      "compile_s": 0.0019223470008000731,
      "inner": 29,
      "median_ms": 1.7315859651465997,
      "q1_ms": 1.7249984827277989,
      "q3_ms": 1.7528762413160717,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 16.0,
      "rel_err": 0.00040910390089266,
      "rel_err_parts": [
        0.00040910390089266
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.32177686398790684,
      "inner": 63,
      "median_ms": 0.715460380888544,
      "q1_ms": 0.711134658818732,
      "q3_ms": 0.716062000026882,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 32.0,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 16.0,
      "rel_err": 0.00040910390089266,
      "rel_err_parts": [
        0.00040910390089266
      ],
      "compile_s": 0.32622304500546306,
      "inner": 60,
      "median_ms": 0.7709989833529107,
      "q1_ms": 0.7635985582358747,
      "q3_ms": 0.788746791658923,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 256.0,
      "xla_arg_mb": 20.0,
      "xla_out_mb": 16.0,
      "rel_err": 0.0004199399845674634,
      "rel_err_parts": [
        0.0004199399845674634
      ],
      "rel_vs_flex_repeat": 0.0002754633314907551,
      "compile_s": 0.48853758499899413,
      "inner": 39,
      "median_ms": 1.21709620520377,
      "q1_ms": 1.2145706281537572,
      "q3_ms": 1.221333564339707,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_err": 1.1000706763297785e-06,
      "rel_err_parts": [
        1.1000706763297785e-06
      ],
      "rel_vs_flex_repeat": 0.000409148313337937,
      "compile_s": 0.0019146819977322593,
      "inner": 27,
      "median_ms": 1.727404703463738,
      "q1_ms": 1.7240992035820253,
      "q3_ms": 1.7534486108127534,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 48.000267028808594,
      "xla_arg_mb": 36.0,
      "xla_out_mb": 20.000022888183594,
      "residual_mb": 16.125011444091797,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "float32[1, 1024, 32, 128]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.0005203635664656758,
      "rel_err_parts": [
        0.0005203635664656758,
        0.0005166303017176688,
        0.00042248226236552
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.8525806099933106,
      "inner": 13,
      "median_ms": 3.3996886922977865,
      "q1_ms": 3.3879299612285996,
      "q3_ms": 3.60624719234273,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 80.0002670288086,
      "xla_arg_mb": 36.0,
      "xla_out_mb": 20.000022888183594,
      "residual_mb": 48.1250114440918,
      "residual_shapes": [
        "float32[1, 1024, 32, 128]",
        "float32[1, 1024, 32, 128]",
        "key<fry>[]",
        "uint32[]",
        "float32[1, 1024, 32, 128]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.0005203635664656758,
      "rel_err_parts": [
        0.0005203635664656758,
        0.0005166303017176688,
        0.00042248226236552
      ],
      "compile_s": 0.9065955999976723,
      "inner": 14,
      "median_ms": 3.478534499923366,
      "q1_ms": 3.45287314277292,
      "q3_ms": 3.5394927144807298,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "jax_xla",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 528.1252593994141,
      "xla_arg_mb": 36.0,
      "xla_out_mb": 20.000022888183594,
      "rel_err": 0.0005185804911889136,
      "rel_err_parts": [
        0.0005185804911889136,
        0.0005158261046744883,
        0.0004225274024065584
      ],
      "rel_vs_flex_repeat": 8.455044735455886e-05,
      "compile_s": 0.2530316460033646,
      "inner": 16,
      "median_ms": 3.1154828748185537,
      "q1_ms": 3.086446999986947,
      "q3_ms": 3.1639843123230094,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 0,
      "rel_err": 1.358498707304534e-06,
      "rel_err_parts": [
        1.358498707304534e-06,
        1.3539493011194281e-06,
        1.1108487569799763e-06
      ],
      "rel_vs_flex_repeat": 0.0005204262561164796,
      "compile_s": 0.007960701012052596,
      "inner": 6,
      "median_ms": 8.171050331535904,
      "q1_ms": 8.058772084041266,
      "q3_ms": 8.241188083047746,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 48.000267028808594,
      "xla_arg_mb": 36.0,
      "xla_out_mb": 20.000022888183594,
      "residual_mb": 16.125011444091797,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "float32[1, 1024, 32, 128]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.0005203635664656758,
      "rel_err_parts": [
        0.0005203635664656758,
        0.0005166303017176688,
        0.00042248226236552
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.9363675060012611,
      "inner": 15,
      "median_ms": 3.7443846667883918,
      "q1_ms": 3.596009866547926,
      "q3_ms": 3.866548032965511,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 80.0002670288086,
      "xla_arg_mb": 36.0,
      "xla_out_mb": 20.000022888183594,
      "residual_mb": 48.1250114440918,
      "residual_shapes": [
        "float32[1, 1024, 32, 128]",
        "float32[1, 1024, 32, 128]",
        "key<fry>[]",
        "uint32[]",
        "float32[1, 1024, 32, 128]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.0005203635664656758,
      "rel_err_parts": [
        0.0005203635664656758,
        0.0005166303017176688,
        0.00042248226236552
      ],
      "compile_s": 0.9587862900079926,
      "inner": 14,
      "median_ms": 3.7136249998833852,
      "q1_ms": 3.4946747499426625,
      "q3_ms": 3.8335529996922038,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "jax_xla",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 528.1252593994141,
      "xla_arg_mb": 36.0,
      "xla_out_mb": 20.000022888183594,
      "rel_err": 0.0005185804911889136,
      "rel_err_parts": [
        0.0005185804911889136,
        0.0005158261046744883,
        0.0004225274024065584
      ],
      "rel_vs_flex_repeat": 8.455044735455886e-05,
      "compile_s": 0.24478043599810917,
      "inner": 15,
      "median_ms": 3.525520533245678,
      "q1_ms": 3.3109979999911348,
      "q3_ms": 3.555590433340209,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 1,
      "rel_err": 1.358498707304534e-06,
      "rel_err_parts": [
        1.358498707304534e-06,
        1.3539493011194281e-06,
        1.1108487569799763e-06
      ],
      "rel_vs_flex_repeat": 0.0005204262561164796,
      "compile_s": 0.008716260010260157,
      "inner": 6,
      "median_ms": 8.15042450024824,
      "q1_ms": 8.070588583602028,
      "q3_ms": 8.269379666065408,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 48.000267028808594,
      "xla_arg_mb": 36.0,
      "xla_out_mb": 20.000022888183594,
      "residual_mb": 16.125011444091797,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "float32[1, 1024, 32, 128]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.0005203635664656758,
      "rel_err_parts": [
        0.0005203635664656758,
        0.0005166303017176688,
        0.00042248226236552
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.9488368470047135,
      "inner": 14,
      "median_ms": 3.396356928403423,
      "q1_ms": 3.3573932144333543,
      "q3_ms": 3.414739179009173,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 80.0002670288086,
      "xla_arg_mb": 36.0,
      "xla_out_mb": 20.000022888183594,
      "residual_mb": 48.1250114440918,
      "residual_shapes": [
        "float32[1, 1024, 32, 128]",
        "float32[1, 1024, 32, 128]",
        "key<fry>[]",
        "uint32[]",
        "float32[1, 1024, 32, 128]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.0005203635664656758,
      "rel_err_parts": [
        0.0005203635664656758,
        0.0005166303017176688,
        0.00042248226236552
      ],
      "compile_s": 0.864290797006106,
      "inner": 14,
      "median_ms": 3.901053214836533,
      "q1_ms": 3.6816305000684224,
      "q3_ms": 3.9127406785805112,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "jax_xla",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 528.1252593994141,
      "xla_arg_mb": 36.0,
      "xla_out_mb": 20.000022888183594,
      "rel_err": 0.0005185804911889136,
      "rel_err_parts": [
        0.0005185804911889136,
        0.0005158261046744883,
        0.0004225274024065584
      ],
      "rel_vs_flex_repeat": 8.455044735455886e-05,
      "compile_s": 0.25089083300554194,
      "inner": 15,
      "median_ms": 3.1127455333868665,
      "q1_ms": 3.1005244003608823,
      "q3_ms": 3.1559853329478456,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 2,
      "rel_err": 1.358498707304534e-06,
      "rel_err_parts": [
        1.358498707304534e-06,
        1.3539493011194281e-06,
        1.1108487569799763e-06
      ],
      "rel_vs_flex_repeat": 0.0005204262561164796,
      "compile_s": 0.007964709991938435,
      "inner": 6,
      "median_ms": 8.158292832376901,
      "q1_ms": 8.122778833543027,
      "q3_ms": 8.273867248741833,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 17.0,
      "xla_out_mb": 16.0,
      "rel_err": 0.0004141571116633713,
      "rel_err_parts": [
        0.0004141571116633713
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.33156018500449136,
      "inner": 57,
      "median_ms": 0.7844176842920939,
      "q1_ms": 0.7710441314613267,
      "q3_ms": 0.79155723672062,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 32.0,
      "xla_arg_mb": 17.0,
      "xla_out_mb": 16.0,
      "rel_err": 0.0004141571116633713,
      "rel_err_parts": [
        0.0004141571116633713
      ],
      "compile_s": 0.3384127800090937,
      "inner": 57,
      "median_ms": 0.7663290700548443,
      "q1_ms": 0.7555405000426905,
      "q3_ms": 0.7803028683267992,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 256.0,
      "xla_arg_mb": 17.0,
      "xla_out_mb": 16.0,
      "rel_err": 0.0004246723256073892,
      "rel_err_parts": [
        0.0004246723256073892
      ],
      "rel_vs_flex_repeat": 0.0002774011227302253,
      "compile_s": 0.21298242699413095,
      "inner": 52,
      "median_ms": 0.987343692171495,
      "q1_ms": 0.9720399325138155,
      "q3_ms": 0.9938287499748046,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_err": 1.0994126569130458e-06,
      "rel_err_parts": [
        1.0994126569130458e-06
      ],
      "rel_vs_flex_repeat": 0.0004141971003264189,
      "compile_s": 0.0019037419988308102,
      "inner": 29,
      "median_ms": 1.7721588966227941,
      "q1_ms": 1.7431405689547257,
      "q3_ms": 1.8044526208039566,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 17.0,
      "xla_out_mb": 16.0,
      "rel_err": 0.0004141571116633713,
      "rel_err_parts": [
        0.0004141571116633713
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.33867028799431864,
      "inner": 64,
      "median_ms": 0.7212668749616569,
      "q1_ms": 0.714742882905739,
      "q3_ms": 0.7525134141133094,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 32.0,
      "xla_arg_mb": 17.0,
      "xla_out_mb": 16.0,
      "rel_err": 0.0004141571116633713,
      "rel_err_parts": [
        0.0004141571116633713
      ],
      "compile_s": 0.3343559950008057,
      "inner": 48,
      "median_ms": 0.778421145696484,
      "q1_ms": 0.7615886876616667,
      "q3_ms": 0.842652687424561,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 256.0,
      "xla_arg_mb": 17.0,
      "xla_out_mb": 16.0,
      "rel_err": 0.0004246723256073892,
      "rel_err_parts": [
        0.0004246723256073892
      ],
      "rel_vs_flex_repeat": 0.0002774011227302253,
      "compile_s": 0.21254926799156237,
      "inner": 51,
      "median_ms": 1.0279351177707534,
      "q1_ms": 0.9868605393064939,
      "q3_ms": 1.0523179214443172,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_err": 1.0994126569130458e-06,
      "rel_err_parts": [
        1.0994126569130458e-06
      ],
      "rel_vs_flex_repeat": 0.0004141971003264189,
      "compile_s": 0.0019196619978174567,
      "inner": 29,
      "median_ms": 1.7916685861439028,
      "q1_ms": 1.7717657585684532,
      "q3_ms": 1.8249950863437019,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 17.0,
      "xla_out_mb": 16.0,
      "rel_err": 0.0004141571116633713,
      "rel_err_parts": [
        0.0004141571116633713
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.35092100700421724,
      "inner": 61,
      "median_ms": 0.7261772624304763,
      "q1_ms": 0.7149547951097494,
      "q3_ms": 0.7317121557893278,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 32.0,
      "xla_arg_mb": 17.0,
      "xla_out_mb": 16.0,
      "rel_err": 0.0004141571116633713,
      "rel_err_parts": [
        0.0004141571116633713
      ],
      "compile_s": 0.3364527219964657,
      "inner": 60,
      "median_ms": 0.7777036332602923,
      "q1_ms": 0.7641637583825893,
      "q3_ms": 0.8449685751353779,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 256.0,
      "xla_arg_mb": 17.0,
      "xla_out_mb": 16.0,
      "rel_err": 0.0004246723256073892,
      "rel_err_parts": [
        0.0004246723256073892
      ],
      "rel_vs_flex_repeat": 0.0002774011227302253,
      "compile_s": 0.20986011999775656,
      "inner": 51,
      "median_ms": 1.043113843051662,
      "q1_ms": 0.9832513627342825,
      "q3_ms": 1.1087460587469988,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_err": 1.0994126569130458e-06,
      "rel_err_parts": [
        1.0994126569130458e-06
      ],
      "rel_vs_flex_repeat": 0.0004141971003264189,
      "compile_s": 0.0019325759931234643,
      "inner": 28,
      "median_ms": 1.7580069642073275,
      "q1_ms": 1.7486863211420962,
      "q3_ms": 1.770953428734044,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 48.000267028808594,
      "xla_arg_mb": 33.0,
      "xla_out_mb": 17.000022888183594,
      "residual_mb": 16.125011444091797,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "float32[1, 1024, 32, 128]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.0005214691627770662,
      "rel_err_parts": [
        0.0005214691627770662,
        0.0005211885436438024,
        0.0004156175418756902
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.885848656995222,
      "inner": 13,
      "median_ms": 3.6925990000152242,
      "q1_ms": 3.6746619611781517,
      "q3_ms": 3.7153292690233615,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 80.0002670288086,
      "xla_arg_mb": 33.0,
      "xla_out_mb": 17.000022888183594,
      "residual_mb": 48.1250114440918,
      "residual_shapes": [
        "float32[1, 1024, 32, 128]",
        "float32[1, 1024, 32, 128]",
        "key<fry>[]",
        "uint32[]",
        "float32[1, 1024, 32, 128]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.0005214691627770662,
      "rel_err_parts": [
        0.0005214691627770662,
        0.0005211885436438024,
        0.0004156175418756902
      ],
      "compile_s": 0.8688840389950201,
      "inner": 14,
      "median_ms": 3.5661044281109104,
      "q1_ms": 3.4869779644525676,
      "q3_ms": 3.6989948571967295,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "jax_xla",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 512.1255035400391,
      "xla_arg_mb": 33.0,
      "xla_out_mb": 17.000022888183594,
      "rel_err": 0.0005207607173360884,
      "rel_err_parts": [
        0.0005196597776375711,
        0.0005207607173360884,
        0.000415724643971771
      ],
      "rel_vs_flex_repeat": 8.468042506137863e-05,
      "compile_s": 0.22709568100981414,
      "inner": 17,
      "median_ms": 3.0770927054725368,
      "q1_ms": 3.063906823466395,
      "q3_ms": 3.186425088428865,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 0,
      "rel_err": 1.3585588476416888e-06,
      "rel_err_parts": [
        1.3516627177523333e-06,
        1.3585588476416888e-06,
        1.114829387915961e-06
      ],
      "rel_vs_flex_repeat": 0.0005215327837504447,
      "compile_s": 0.008634317011456005,
      "inner": 6,
      "median_ms": 8.458255334213996,
      "q1_ms": 8.125099081856508,
      "q3_ms": 8.58546591674288,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 48.000267028808594,
      "xla_arg_mb": 33.0,
      "xla_out_mb": 17.000022888183594,
      "residual_mb": 16.125011444091797,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "float32[1, 1024, 32, 128]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.0005214691627770662,
      "rel_err_parts": [
        0.0005214691627770662,
        0.0005211885436438024,
        0.0004156175418756902
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.8663481010007672,
      "inner": 15,
      "median_ms": 3.3817879331763834,
      "q1_ms": 3.364989466111486,
      "q3_ms": 3.5504563663077233,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 80.0002670288086,
      "xla_arg_mb": 33.0,
      "xla_out_mb": 17.000022888183594,
      "residual_mb": 48.1250114440918,
      "residual_shapes": [
        "float32[1, 1024, 32, 128]",
        "float32[1, 1024, 32, 128]",
        "key<fry>[]",
        "uint32[]",
        "float32[1, 1024, 32, 128]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.0005214691627770662,
      "rel_err_parts": [
        0.0005214691627770662,
        0.0005211885436438024,
        0.0004156175418756902
      ],
      "compile_s": 0.8690217689872952,
      "inner": 14,
      "median_ms": 3.485635785701951,
      "q1_ms": 3.4640828567457254,
      "q3_ms": 3.6378141069560246,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "jax_xla",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 512.1255035400391,
      "xla_arg_mb": 33.0,
      "xla_out_mb": 17.000022888183594,
      "rel_err": 0.0005207607173360884,
      "rel_err_parts": [
        0.0005196597776375711,
        0.0005207607173360884,
        0.000415724643971771
      ],
      "rel_vs_flex_repeat": 8.468042506137863e-05,
      "compile_s": 0.22784879599930719,
      "inner": 17,
      "median_ms": 3.0887570593382834,
      "q1_ms": 3.012796587855144,
      "q3_ms": 3.1234627646935063,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 1,
      "rel_err": 1.3585588476416888e-06,
      "rel_err_parts": [
        1.3516627177523333e-06,
        1.3585588476416888e-06,
        1.114829387915961e-06
      ],
      "rel_vs_flex_repeat": 0.0005215327837504447,
      "compile_s": 0.008832952997181565,
      "inner": 5,
      "median_ms": 8.197245601331815,
      "q1_ms": 8.175557901267894,
      "q3_ms": 8.302845500293188,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 48.000267028808594,
      "xla_arg_mb": 33.0,
      "xla_out_mb": 17.000022888183594,
      "residual_mb": 16.125011444091797,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "float32[1, 1024, 32, 128]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.0005214691627770662,
      "rel_err_parts": [
        0.0005214691627770662,
        0.0005211885436438024,
        0.0004156175418756902
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.872332674000063,
      "inner": 14,
      "median_ms": 3.399593714025936,
      "q1_ms": 3.3883678214416637,
      "q3_ms": 3.5165954286640044,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 80.0002670288086,
      "xla_arg_mb": 33.0,
      "xla_out_mb": 17.000022888183594,
      "residual_mb": 48.1250114440918,
      "residual_shapes": [
        "float32[1, 1024, 32, 128]",
        "float32[1, 1024, 32, 128]",
        "key<fry>[]",
        "uint32[]",
        "float32[1, 1024, 32, 128]",
        "float32[1, 32, 1024]"
      ],
      "rel_err": 0.0005214691627770662,
      "rel_err_parts": [
        0.0005214691627770662,
        0.0005211885436438024,
        0.0004156175418756902
      ],
      "compile_s": 0.8718820020003477,
      "inner": 14,
      "median_ms": 3.5530490713426843,
      "q1_ms": 3.478407821343613,
      "q3_ms": 3.7782935001554767,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "jax_xla",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 512.1255035400391,
      "xla_arg_mb": 33.0,
      "xla_out_mb": 17.000022888183594,
      "rel_err": 0.0005207607173360884,
      "rel_err_parts": [
        0.0005196597776375711,
        0.0005207607173360884,
        0.000415724643971771
      ],
      "rel_vs_flex_repeat": 8.468042506137863e-05,
      "compile_s": 0.22826228800113313,
      "inner": 16,
      "median_ms": 3.064981625357177,
      "q1_ms": 3.0452273749688175,
      "q3_ms": 3.098347906870913,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 1024,
      "lkv": 1024,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 2,
      "rel_err": 1.3585588476416888e-06,
      "rel_err_parts": [
        1.3516627177523333e-06,
        1.3585588476416888e-06,
        1.114829387915961e-06
      ],
      "rel_vs_flex_repeat": 0.0005215327837504447,
      "compile_s": 0.008469706008327194,
      "inner": 5,
      "median_ms": 8.352611601003446,
      "q1_ms": 8.086653500504326,
      "q3_ms": 8.608225399802905,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_1k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 32.0,
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.22645324499171693,
      "inner": 11,
      "median_ms": 4.685428817703558,
      "q1_ms": 4.641042090952396,
      "q3_ms": 4.703814090383027,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 64.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.233998851996148,
      "inner": 11,
      "median_ms": 4.790935727545398,
      "q1_ms": 4.7454268635853865,
      "q3_ms": 4.7963008176090876,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_vs_flex_repeat": 0.00043491972610354424,
      "compile_s": 0.015241266999510117,
      "inner": 3,
      "median_ms": 16.529798667761497,
      "q1_ms": 16.430010000476614,
      "q3_ms": 18.300993334075123,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 32.0,
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.22993622500507627,
      "inner": 11,
      "median_ms": 4.630237091491422,
      "q1_ms": 4.62716395453423,
      "q3_ms": 4.687750454732767,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 64.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.24167152900190558,
      "inner": 10,
      "median_ms": 5.069803399965167,
      "q1_ms": 4.7075941001821775,
      "q3_ms": 5.209888450190192,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_vs_flex_repeat": 0.00043491972610354424,
      "compile_s": 0.015972512002917938,
      "inner": 3,
      "median_ms": 16.32816066558007,
      "q1_ms": 16.25889633093417,
      "q3_ms": 17.28194533401014,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 32.0,
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.22926217000349425,
      "inner": 11,
      "median_ms": 4.647930545615964,
      "q1_ms": 4.590504182322713,
      "q3_ms": 4.784591863757338,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 64.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.22983101599675138,
      "inner": 11,
      "median_ms": 4.708825726993382,
      "q1_ms": 4.698654772999527,
      "q3_ms": 4.8593150455748075,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_vs_flex_repeat": 0.00043491972610354424,
      "compile_s": 0.015906105996691622,
      "inner": 3,
      "median_ms": 16.825628335936926,
      "q1_ms": 16.649527167222306,
      "q3_ms": 17.060329167482752,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 96.0002670288086,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 48.000022888183594,
      "residual_mb": 32.5000114440918,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "float32[1, 4096, 32, 64]",
        "float32[1, 32, 4096]"
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.4251675159903243,
      "inner": 1,
      "median_ms": 25.484698999207467,
      "q1_ms": 25.341347507492173,
      "q3_ms": 25.500383992039133,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 160.0002670288086,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 48.000022888183594,
      "residual_mb": 96.5000114440918,
      "residual_shapes": [
        "float32[1, 4096, 32, 64]",
        "float32[1, 4096, 32, 64]",
        "key<fry>[]",
        "uint32[]",
        "float32[1, 4096, 32, 64]",
        "float32[1, 32, 4096]"
      ],
      "compile_s": 0.41759436699794605,
      "inner": 2,
      "median_ms": 24.341260497749317,
      "q1_ms": 23.97676575128571,
      "q3_ms": 25.541128004988423,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 0,
      "rel_vs_flex_repeat": 0.0005369650898501277,
      "compile_s": 0.064140018002945,
      "inner": 1,
      "median_ms": 71.07107200135943,
      "q1_ms": 66.57999900198774,
      "q3_ms": 71.65502600400941,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 96.0002670288086,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 48.000022888183594,
      "residual_mb": 32.5000114440918,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "float32[1, 4096, 32, 64]",
        "float32[1, 32, 4096]"
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.4367005270032678,
      "inner": 2,
      "median_ms": 25.581637499271892,
      "q1_ms": 25.302253750851378,
      "q3_ms": 25.686180000775494,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 160.0002670288086,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 48.000022888183594,
      "residual_mb": 96.5000114440918,
      "residual_shapes": [
        "float32[1, 4096, 32, 64]",
        "float32[1, 4096, 32, 64]",
        "key<fry>[]",
        "uint32[]",
        "float32[1, 4096, 32, 64]",
        "float32[1, 32, 4096]"
      ],
      "compile_s": 0.3906293370091589,
      "inner": 2,
      "median_ms": 24.28725299978396,
      "q1_ms": 24.119014753523516,
      "q3_ms": 24.67510449787369,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 1,
      "rel_vs_flex_repeat": 0.0005369650898501277,
      "compile_s": 0.06424787100695539,
      "inner": 1,
      "median_ms": 65.7667460036464,
      "q1_ms": 64.88538799749222,
      "q3_ms": 66.44904550194042,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 96.0002670288086,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 48.000022888183594,
      "residual_mb": 32.5000114440918,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "float32[1, 4096, 32, 64]",
        "float32[1, 32, 4096]"
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.44293317099800333,
      "inner": 2,
      "median_ms": 25.44248400226934,
      "q1_ms": 25.086923749768175,
      "q3_ms": 25.807682250160724,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 160.0002670288086,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 48.000022888183594,
      "residual_mb": 96.5000114440918,
      "residual_shapes": [
        "float32[1, 4096, 32, 64]",
        "float32[1, 4096, 32, 64]",
        "key<fry>[]",
        "uint32[]",
        "float32[1, 4096, 32, 64]",
        "float32[1, 32, 4096]"
      ],
      "compile_s": 0.47406547400169075,
      "inner": 1,
      "median_ms": 27.942739994614385,
      "q1_ms": 27.84240950131789,
      "q3_ms": 28.11780099727912,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 2,
      "rel_vs_flex_repeat": 0.0005369650898501277,
      "compile_s": 0.06391131399141159,
      "inner": 1,
      "median_ms": 64.98436500260141,
      "q1_ms": 64.46709650481353,
      "q3_ms": 65.52171699877363,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 32.0,
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.23277695999422576,
      "inner": 10,
      "median_ms": 4.72959470062051,
      "q1_ms": 4.583757300133584,
      "q3_ms": 4.755198649218073,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 64.0,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.24106575299811084,
      "inner": 10,
      "median_ms": 4.735977799282409,
      "q1_ms": 4.714378849894274,
      "q3_ms": 4.794716900505591,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_vs_flex_repeat": 0.0004134633345529437,
      "compile_s": 0.018817224001395516,
      "inner": 2,
      "median_ms": 19.42388599854894,
      "q1_ms": 19.390287750866264,
      "q3_ms": 19.461362004221883,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 32.0,
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.2919607299991185,
      "inner": 11,
      "median_ms": 5.3442361822817475,
      "q1_ms": 5.201350408738522,
      "q3_ms": 5.43077904545829,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 64.0,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.28484350399230607,
      "inner": 10,
      "median_ms": 4.6261779993074015,
      "q1_ms": 4.585570199560607,
      "q3_ms": 4.7323204002168495,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_vs_flex_repeat": 0.0004134633345529437,
      "compile_s": 0.015304597007343546,
      "inner": 3,
      "median_ms": 16.75873533046494,
      "q1_ms": 16.617628331005108,
      "q3_ms": 17.290414667513687,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 32.0,
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.235928663998493,
      "inner": 11,
      "median_ms": 4.561681000103073,
      "q1_ms": 4.5062870909036565,
      "q3_ms": 4.727165954880713,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 64.0,
      "xla_arg_mb": 40.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.23178951199224684,
      "inner": 11,
      "median_ms": 5.2153550905958665,
      "q1_ms": 4.904037817058535,
      "q3_ms": 5.275166408947288,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_vs_flex_repeat": 0.0004134633345529437,
      "compile_s": 0.017449666003813036,
      "inner": 2,
      "median_ms": 18.35581049817847,
      "q1_ms": 18.34107250033412,
      "q3_ms": 18.400574997940566,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 96.0002670288086,
      "xla_arg_mb": 72.0,
      "xla_out_mb": 40.000022888183594,
      "residual_mb": 32.5000114440918,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "float32[1, 4096, 32, 64]",
        "float32[1, 32, 4096]"
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.3997422199900029,
      "inner": 2,
      "median_ms": 22.081039001932368,
      "q1_ms": 22.06419449794339,
      "q3_ms": 22.22870575133129,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 160.0002670288086,
      "xla_arg_mb": 72.0,
      "xla_out_mb": 40.000022888183594,
      "residual_mb": 96.5000114440918,
      "residual_shapes": [
        "float32[1, 4096, 32, 64]",
        "float32[1, 4096, 32, 64]",
        "key<fry>[]",
        "uint32[]",
        "float32[1, 4096, 32, 64]",
        "float32[1, 32, 4096]"
      ],
      "compile_s": 0.38762786900042556,
      "inner": 2,
      "median_ms": 23.685099004069343,
      "q1_ms": 23.500404498918215,
      "q3_ms": 23.69606724823825,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 0,
      "rel_vs_flex_repeat": 0.0005370639264583588,
      "compile_s": 0.06147215901000891,
      "inner": 1,
      "median_ms": 62.901573997805826,
      "q1_ms": 62.40239000180736,
      "q3_ms": 63.296842003182974,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 96.0002670288086,
      "xla_arg_mb": 72.0,
      "xla_out_mb": 40.000022888183594,
      "residual_mb": 32.5000114440918,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "float32[1, 4096, 32, 64]",
        "float32[1, 32, 4096]"
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.3752356420009164,
      "inner": 2,
      "median_ms": 22.302707999187987,
      "q1_ms": 22.149047250422882,
      "q3_ms": 22.3343907528033,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 160.0002670288086,
      "xla_arg_mb": 72.0,
      "xla_out_mb": 40.000022888183594,
      "residual_mb": 96.5000114440918,
      "residual_shapes": [
        "float32[1, 4096, 32, 64]",
        "float32[1, 4096, 32, 64]",
        "key<fry>[]",
        "uint32[]",
        "float32[1, 4096, 32, 64]",
        "float32[1, 32, 4096]"
      ],
      "compile_s": 0.37171505199512467,
      "inner": 2,
      "median_ms": 23.42970049357973,
      "q1_ms": 23.383107250992907,
      "q3_ms": 23.559649249364156,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 1,
      "rel_vs_flex_repeat": 0.0005370639264583588,
      "compile_s": 0.06205628500902094,
      "inner": 1,
      "median_ms": 62.852561997715384,
      "q1_ms": 62.356998496397864,
      "q3_ms": 63.08347800222691,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 96.0002670288086,
      "xla_arg_mb": 72.0,
      "xla_out_mb": 40.000022888183594,
      "residual_mb": 32.5000114440918,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "float32[1, 4096, 32, 64]",
        "float32[1, 32, 4096]"
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.36170726499403827,
      "inner": 2,
      "median_ms": 22.07655250094831,
      "q1_ms": 22.010115997545654,
      "q3_ms": 22.232658498978708,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 160.0002670288086,
      "xla_arg_mb": 72.0,
      "xla_out_mb": 40.000022888183594,
      "residual_mb": 96.5000114440918,
      "residual_shapes": [
        "float32[1, 4096, 32, 64]",
        "float32[1, 4096, 32, 64]",
        "key<fry>[]",
        "uint32[]",
        "float32[1, 4096, 32, 64]",
        "float32[1, 32, 4096]"
      ],
      "compile_s": 0.3657063319988083,
      "inner": 2,
      "median_ms": 23.39117249357514,
      "q1_ms": 23.268089502380462,
      "q3_ms": 23.63104424875928,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 2,
      "rel_vs_flex_repeat": 0.0005370639264583588,
      "compile_s": 0.06103202699159738,
      "inner": 1,
      "median_ms": 62.9060329956701,
      "q1_ms": 62.56400050915545,
      "q3_ms": 63.0013734989916,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 34.0,
      "xla_out_mb": 32.0,
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.21057167599792592,
      "inner": 11,
      "median_ms": 4.4847745461579915,
      "q1_ms": 4.481551727539309,
      "q3_ms": 4.488566726980604,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 64.0,
      "xla_arg_mb": 34.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.21437469199008774,
      "inner": 10,
      "median_ms": 4.590000500320457,
      "q1_ms": 4.567684349603951,
      "q3_ms": 4.631075850193156,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_vs_flex_repeat": 0.0004361339670140296,
      "compile_s": 0.015162829993641935,
      "inner": 3,
      "median_ms": 16.067630999411147,
      "q1_ms": 16.05049366480671,
      "q3_ms": 16.08344750032605,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 34.0,
      "xla_out_mb": 32.0,
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.21489174000453204,
      "inner": 11,
      "median_ms": 4.470873818552883,
      "q1_ms": 4.453932226831305,
      "q3_ms": 4.511478726767977,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 64.0,
      "xla_arg_mb": 34.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.21469187099137343,
      "inner": 11,
      "median_ms": 4.6079831822267305,
      "q1_ms": 4.606297727150377,
      "q3_ms": 4.6124976367784924,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_vs_flex_repeat": 0.0004361339670140296,
      "compile_s": 0.01518828800180927,
      "inner": 3,
      "median_ms": 16.08087266504299,
      "q1_ms": 16.074160167287726,
      "q3_ms": 16.105506166544124,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 34.0,
      "xla_out_mb": 32.0,
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.2096132930018939,
      "inner": 11,
      "median_ms": 4.5006910910491245,
      "q1_ms": 4.486879090439867,
      "q3_ms": 4.511098454796327,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 64.0,
      "xla_arg_mb": 34.0,
      "xla_out_mb": 32.0,
      "compile_s": 0.21590051900420804,
      "inner": 11,
      "median_ms": 4.576260455492461,
      "q1_ms": 4.569103226954626,
      "q3_ms": 4.580672363359058,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_vs_flex_repeat": 0.0004361339670140296,
      "compile_s": 0.015386421990115196,
      "inner": 3,
      "median_ms": 16.07437066559214,
      "q1_ms": 16.059692667719595,
      "q3_ms": 16.08037183420189,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 96.0002670288086,
      "xla_arg_mb": 66.0,
      "xla_out_mb": 34.000022888183594,
      "residual_mb": 32.5000114440918,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "float32[1, 4096, 32, 64]",
        "float32[1, 32, 4096]"
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.3647345429926645,
      "inner": 2,
      "median_ms": 21.930874994723126,
      "q1_ms": 21.78583274871926,
      "q3_ms": 21.96477674806374,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 160.0002670288086,
      "xla_arg_mb": 66.0,
      "xla_out_mb": 34.000022888183594,
      "residual_mb": 96.5000114440918,
      "residual_shapes": [
        "float32[1, 4096, 32, 64]",
        "float32[1, 4096, 32, 64]",
        "key<fry>[]",
        "uint32[]",
        "float32[1, 4096, 32, 64]",
        "float32[1, 32, 4096]"
      ],
      "compile_s": 0.3705464989907341,
      "inner": 2,
      "median_ms": 23.585710499901325,
      "q1_ms": 23.484316254325677,
      "q3_ms": 23.75601300082053,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 0,
      "rel_vs_flex_repeat": 0.0005381678929552436,
      "compile_s": 0.06218874600017443,
      "inner": 1,
      "median_ms": 62.768092000624165,
      "q1_ms": 62.541252496885136,
      "q3_ms": 62.91876699833665,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 96.0002670288086,
      "xla_arg_mb": 66.0,
      "xla_out_mb": 34.000022888183594,
      "residual_mb": 32.5000114440918,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "float32[1, 4096, 32, 64]",
        "float32[1, 32, 4096]"
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.35873812799400184,
      "inner": 2,
      "median_ms": 21.91308599867625,
      "q1_ms": 21.88804350225837,
      "q3_ms": 22.005165501468582,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 160.0002670288086,
      "xla_arg_mb": 66.0,
      "xla_out_mb": 34.000022888183594,
      "residual_mb": 96.5000114440918,
      "residual_shapes": [
        "float32[1, 4096, 32, 64]",
        "float32[1, 4096, 32, 64]",
        "key<fry>[]",
        "uint32[]",
        "float32[1, 4096, 32, 64]",
        "float32[1, 32, 4096]"
      ],
      "compile_s": 0.36312593999900855,
      "inner": 2,
      "median_ms": 23.580048997246195,
      "q1_ms": 23.382112998660887,
      "q3_ms": 23.61774125165539,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 1,
      "rel_vs_flex_repeat": 0.0005381678929552436,
      "compile_s": 0.06238825200125575,
      "inner": 1,
      "median_ms": 62.953241998911835,
      "q1_ms": 62.548431000323035,
      "q3_ms": 63.188096493831836,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 96.0002670288086,
      "xla_arg_mb": 66.0,
      "xla_out_mb": 34.000022888183594,
      "residual_mb": 32.5000114440918,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "float32[1, 4096, 32, 64]",
        "float32[1, 32, 4096]"
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.36442782200174406,
      "inner": 2,
      "median_ms": 21.97825450275559,
      "q1_ms": 21.936532997642644,
      "q3_ms": 22.096813499956625,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 160.0002670288086,
      "xla_arg_mb": 66.0,
      "xla_out_mb": 34.000022888183594,
      "residual_mb": 96.5000114440918,
      "residual_shapes": [
        "float32[1, 4096, 32, 64]",
        "float32[1, 4096, 32, 64]",
        "key<fry>[]",
        "uint32[]",
        "float32[1, 4096, 32, 64]",
        "float32[1, 32, 4096]"
      ],
      "compile_s": 0.37062945700017735,
      "inner": 2,
      "median_ms": 23.49861600669101,
      "q1_ms": 23.395080002956092,
      "q3_ms": 23.68528924853308,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 64,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 2,
      "rel_vs_flex_repeat": 0.0005381678929552436,
      "compile_s": 0.061673710006289184,
      "inner": 1,
      "median_ms": 62.69053400319535,
      "q1_ms": 62.29156499466626,
      "q3_ms": 62.91577599768061,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 64.0,
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.31242970799212344,
      "inner": 4,
      "median_ms": 10.259389000566443,
      "q1_ms": 9.99898762529483,
      "q3_ms": 10.322609874492628,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 128.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 64.0,
      "compile_s": 0.31284393300302327,
      "inner": 5,
      "median_ms": 10.10821360105183,
      "q1_ms": 10.038480199000333,
      "q3_ms": 10.210751500562765,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_vs_flex_repeat": 0.0004189115425106138,
      "compile_s": 0.0242238329956308,
      "inner": 2,
      "median_ms": 24.427618998743128,
      "q1_ms": 24.423969498457154,
      "q3_ms": 24.429033499473007,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 64.0,
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.31613764401117805,
      "inner": 5,
      "median_ms": 9.955064000678249,
      "q1_ms": 9.87591790035367,
      "q3_ms": 9.98977539857151,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 128.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 64.0,
      "compile_s": 0.31080022499372717,
      "inner": 5,
      "median_ms": 10.123079200275242,
      "q1_ms": 10.112619899155106,
      "q3_ms": 10.138820399879478,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_vs_flex_repeat": 0.0004189115425106138,
      "compile_s": 0.024240936007117853,
      "inner": 2,
      "median_ms": 24.430729499727022,
      "q1_ms": 24.426822001260007,
      "q3_ms": 24.43854950251989,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 64.0,
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.3092231909977272,
      "inner": 5,
      "median_ms": 9.918003200436942,
      "q1_ms": 9.903959001530893,
      "q3_ms": 9.962921599799301,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 128.0,
      "xla_arg_mb": 96.0,
      "xla_out_mb": 64.0,
      "compile_s": 0.3183356340014143,
      "inner": 4,
      "median_ms": 10.123133251909167,
      "q1_ms": 10.114391125171096,
      "q3_ms": 10.149066250960459,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_vs_flex_repeat": 0.0004189115425106138,
      "compile_s": 0.02421841300383676,
      "inner": 2,
      "median_ms": 24.424107003142126,
      "q1_ms": 24.423207501968136,
      "q3_ms": 24.42997575053596,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 192.0002670288086,
      "xla_arg_mb": 160.0,
      "xla_out_mb": 96.0000228881836,
      "residual_mb": 64.5000114440918,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "float32[1, 4096, 32, 128]",
        "float32[1, 32, 4096]"
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.8087498210079502,
      "inner": 1,
      "median_ms": 45.94629298662767,
      "q1_ms": 45.9025010059122,
      "q3_ms": 46.17605348903453,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 320.0002670288086,
      "xla_arg_mb": 160.0,
      "xla_out_mb": 96.0000228881836,
      "residual_mb": 192.5000114440918,
      "residual_shapes": [
        "float32[1, 4096, 32, 128]",
        "float32[1, 4096, 32, 128]",
        "key<fry>[]",
        "uint32[]",
        "float32[1, 4096, 32, 128]",
        "float32[1, 32, 4096]"
      ],
      "compile_s": 0.8239658630045597,
      "inner": 1,
      "median_ms": 46.57405999023467,
      "q1_ms": 46.53398050140822,
      "q3_ms": 46.63306350266794,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 0,
      "rel_vs_flex_repeat": 0.0005224947817623615,
      "compile_s": 0.11377100300160237,
      "inner": 1,
      "median_ms": 115.88450199633371,
      "q1_ms": 115.70876349287573,
      "q3_ms": 116.0726274974877,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 192.0002670288086,
      "xla_arg_mb": 160.0,
      "xla_out_mb": 96.0000228881836,
      "residual_mb": 64.5000114440918,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "float32[1, 4096, 32, 128]",
        "float32[1, 32, 4096]"
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.8152998639998259,
      "inner": 1,
      "median_ms": 46.03735600539949,
      "q1_ms": 45.99370749929221,
      "q3_ms": 46.326309005962685,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 320.0002670288086,
      "xla_arg_mb": 160.0,
      "xla_out_mb": 96.0000228881836,
      "residual_mb": 192.5000114440918,
      "residual_shapes": [
        "float32[1, 4096, 32, 128]",
        "float32[1, 4096, 32, 128]",
        "key<fry>[]",
        "uint32[]",
        "float32[1, 4096, 32, 128]",
        "float32[1, 32, 4096]"
      ],
      "compile_s": 0.8280819979991065,
      "inner": 1,
      "median_ms": 46.46072699688375,
      "q1_ms": 46.410198003286496,
      "q3_ms": 46.56367250572657,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 1,
      "rel_vs_flex_repeat": 0.0005224947817623615,
      "compile_s": 0.1140331290080212,
      "inner": 1,
      "median_ms": 115.80279699410312,
      "q1_ms": 115.72035100107314,
      "q3_ms": 116.07133500365308,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 192.0002670288086,
      "xla_arg_mb": 160.0,
      "xla_out_mb": 96.0000228881836,
      "residual_mb": 64.5000114440918,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "float32[1, 4096, 32, 128]",
        "float32[1, 32, 4096]"
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.8229991419939324,
      "inner": 1,
      "median_ms": 46.01552500389516,
      "q1_ms": 45.937506009067874,
      "q3_ms": 46.26064199692337,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 320.0002670288086,
      "xla_arg_mb": 160.0,
      "xla_out_mb": 96.0000228881836,
      "residual_mb": 192.5000114440918,
      "residual_shapes": [
        "float32[1, 4096, 32, 128]",
        "float32[1, 4096, 32, 128]",
        "key<fry>[]",
        "uint32[]",
        "float32[1, 4096, 32, 128]",
        "float32[1, 32, 4096]"
      ],
      "compile_s": 0.8106907819892513,
      "inner": 1,
      "median_ms": 46.52391899435315,
      "q1_ms": 46.46764349308796,
      "q3_ms": 46.60230350418715,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 2,
      "rel_vs_flex_repeat": 0.0005224947817623615,
      "compile_s": 0.11396230499667581,
      "inner": 1,
      "median_ms": 116.02612900605891,
      "q1_ms": 115.99592200218467,
      "q3_ms": 116.06232249323511,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 64.0,
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.3118464250001125,
      "inner": 4,
      "median_ms": 9.928620002028765,
      "q1_ms": 9.877287249764777,
      "q3_ms": 10.294601876012166,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 128.0,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 64.0,
      "compile_s": 0.31294345000060275,
      "inner": 5,
      "median_ms": 10.108926802058704,
      "q1_ms": 10.105496100732125,
      "q3_ms": 10.140802299429197,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_vs_flex_repeat": 0.00041708172648213804,
      "compile_s": 0.02410111100471113,
      "inner": 2,
      "median_ms": 24.485573499987368,
      "q1_ms": 24.434419250610517,
      "q3_ms": 24.548104003770277,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 64.0,
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.312036805000389,
      "inner": 5,
      "median_ms": 9.894699600408785,
      "q1_ms": 9.877748601138592,
      "q3_ms": 9.951419600110967,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 128.0,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 64.0,
      "compile_s": 0.3172882210055832,
      "inner": 5,
      "median_ms": 10.116829798789695,
      "q1_ms": 10.100470800534822,
      "q3_ms": 10.143456101650372,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_vs_flex_repeat": 0.00041708172648213804,
      "compile_s": 0.02417384900036268,
      "inner": 2,
      "median_ms": 24.46451850119047,
      "q1_ms": 24.452733749058098,
      "q3_ms": 24.657849498908035,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 64.0,
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.318118682000204,
      "inner": 5,
      "median_ms": 9.89114079857245,
      "q1_ms": 9.877101100573782,
      "q3_ms": 9.899880198645405,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 128.0,
      "xla_arg_mb": 80.0,
      "xla_out_mb": 64.0,
      "compile_s": 0.329088849000982,
      "inner": 5,
      "median_ms": 10.053349600639194,
      "q1_ms": 10.03555599891115,
      "q3_ms": 10.069884899712633,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_vs_flex_repeat": 0.00041708172648213804,
      "compile_s": 0.024254270989331417,
      "inner": 2,
      "median_ms": 24.674711006809957,
      "q1_ms": 24.567273005231982,
      "q3_ms": 24.686701250175247,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 192.0002670288086,
      "xla_arg_mb": 144.0,
      "xla_out_mb": 80.0000228881836,
      "residual_mb": 64.5000114440918,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "float32[1, 4096, 32, 128]",
        "float32[1, 32, 4096]"
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.9532744710013503,
      "inner": 1,
      "median_ms": 45.77156199957244,
      "q1_ms": 45.614371010742616,
      "q3_ms": 47.76822550047655,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 320.0002670288086,
      "xla_arg_mb": 144.0,
      "xla_out_mb": 80.0000228881836,
      "residual_mb": 192.5000114440918,
      "residual_shapes": [
        "float32[1, 4096, 32, 128]",
        "float32[1, 4096, 32, 128]",
        "key<fry>[]",
        "uint32[]",
        "float32[1, 4096, 32, 128]",
        "float32[1, 32, 4096]"
      ],
      "compile_s": 0.9838773740048055,
      "inner": 1,
      "median_ms": 46.45398398861289,
      "q1_ms": 46.32846199820051,
      "q3_ms": 46.54424549516989,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 0,
      "rel_vs_flex_repeat": 0.0005227958899922669,
      "compile_s": 0.1138453730090987,
      "inner": 1,
      "median_ms": 115.68848000024445,
      "q1_ms": 115.51331850205315,
      "q3_ms": 115.84688050061231,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 192.0002670288086,
      "xla_arg_mb": 144.0,
      "xla_out_mb": 80.0000228881836,
      "residual_mb": 64.5000114440918,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "float32[1, 4096, 32, 128]",
        "float32[1, 32, 4096]"
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.8617412739986321,
      "inner": 1,
      "median_ms": 45.797041995683685,
      "q1_ms": 45.547829002316575,
      "q3_ms": 45.83059950527968,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 320.0002670288086,
      "xla_arg_mb": 144.0,
      "xla_out_mb": 80.0000228881836,
      "residual_mb": 192.5000114440918,
      "residual_shapes": [
        "float32[1, 4096, 32, 128]",
        "float32[1, 4096, 32, 128]",
        "key<fry>[]",
        "uint32[]",
        "float32[1, 4096, 32, 128]",
        "float32[1, 32, 4096]"
      ],
      "compile_s": 0.8569469930080231,
      "inner": 1,
      "median_ms": 46.3847729988629,
      "q1_ms": 46.34948149760021,
      "q3_ms": 46.445087500615045,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 1,
      "rel_vs_flex_repeat": 0.0005227958899922669,
      "compile_s": 0.11374968200107105,
      "inner": 1,
      "median_ms": 115.86122699372936,
      "q1_ms": 115.84286750439787,
      "q3_ms": 115.94121900270693,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 192.0002670288086,
      "xla_arg_mb": 144.0,
      "xla_out_mb": 80.0000228881836,
      "residual_mb": 64.5000114440918,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "float32[1, 4096, 32, 128]",
        "float32[1, 32, 4096]"
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.8618170780100627,
      "inner": 1,
      "median_ms": 45.67155399126932,
      "q1_ms": 45.6005095038563,
      "q3_ms": 45.80367400194518,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 320.0002670288086,
      "xla_arg_mb": 144.0,
      "xla_out_mb": 80.0000228881836,
      "residual_mb": 192.5000114440918,
      "residual_shapes": [
        "float32[1, 4096, 32, 128]",
        "float32[1, 4096, 32, 128]",
        "key<fry>[]",
        "uint32[]",
        "float32[1, 4096, 32, 128]",
        "float32[1, 32, 4096]"
      ],
      "compile_s": 0.9012143099971581,
      "inner": 1,
      "median_ms": 46.44923500018194,
      "q1_ms": 46.43523399863625,
      "q3_ms": 46.473640999465715,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 4,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 2,
      "rel_vs_flex_repeat": 0.0005227958899922669,
      "compile_s": 0.11367455100116786,
      "inner": 1,
      "median_ms": 115.79015299503226,
      "q1_ms": 115.71141850436106,
      "q3_ms": 115.93357500532875,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 68.0,
      "xla_out_mb": 64.0,
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.34319055000378285,
      "inner": 4,
      "median_ms": 10.129726000741357,
      "q1_ms": 9.949262250302127,
      "q3_ms": 10.316855374185252,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 128.0,
      "xla_arg_mb": 68.0,
      "xla_out_mb": 64.0,
      "compile_s": 0.3277190279914066,
      "inner": 5,
      "median_ms": 10.08978460158687,
      "q1_ms": 10.07702640054049,
      "q3_ms": 10.1031227997737,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_vs_flex_repeat": 0.0004116460040677339,
      "compile_s": 0.024258197998278774,
      "inner": 2,
      "median_ms": 24.4698789974791,
      "q1_ms": 24.458036001306027,
      "q3_ms": 24.47613525146153,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 68.0,
      "xla_out_mb": 64.0,
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.3261044319951907,
      "inner": 5,
      "median_ms": 9.875960001954809,
      "q1_ms": 9.808300000440795,
      "q3_ms": 9.906051900179591,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 128.0,
      "xla_arg_mb": 68.0,
      "xla_out_mb": 64.0,
      "compile_s": 0.32038377900607884,
      "inner": 5,
      "median_ms": 10.101861399016343,
      "q1_ms": 10.064280399819836,
      "q3_ms": 10.136326598876622,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_vs_flex_repeat": 0.0004116460040677339,
      "compile_s": 0.02421222100383602,
      "inner": 2,
      "median_ms": 24.48450150404824,
      "q1_ms": 24.46902749943547,
      "q3_ms": 24.48729899697355,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 68.0,
      "xla_out_mb": 64.0,
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.3327081059978809,
      "inner": 5,
      "median_ms": 9.820420999312773,
      "q1_ms": 9.803457801172044,
      "q3_ms": 9.856446001504082,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 128.0,
      "xla_arg_mb": 68.0,
      "xla_out_mb": 64.0,
      "compile_s": 0.33747540701006074,
      "inner": 5,
      "median_ms": 10.062146198470145,
      "q1_ms": 10.019970400026068,
      "q3_ms": 10.070863799774088,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_vs_flex_repeat": 0.0004116460040677339,
      "compile_s": 0.02422865199332591,
      "inner": 2,
      "median_ms": 24.506458001269493,
      "q1_ms": 24.501227748260135,
      "q3_ms": 24.513355747330934,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 192.0002670288086,
      "xla_arg_mb": 132.0,
      "xla_out_mb": 68.0000228881836,
      "residual_mb": 64.5000114440918,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "float32[1, 4096, 32, 128]",
        "float32[1, 32, 4096]"
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.8629467950086109,
      "inner": 1,
      "median_ms": 45.85225600749254,
      "q1_ms": 45.65559399634367,
      "q3_ms": 46.117412006424274,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 0,
      "xla_temp_mb": 320.0002670288086,
      "xla_arg_mb": 132.0,
      "xla_out_mb": 68.0000228881836,
      "residual_mb": 192.5000114440918,
      "residual_shapes": [
        "float32[1, 4096, 32, 128]",
        "float32[1, 4096, 32, 128]",
        "key<fry>[]",
        "uint32[]",
        "float32[1, 4096, 32, 128]",
        "float32[1, 32, 4096]"
      ],
      "compile_s": 0.8944387189985719,
      "inner": 1,
      "median_ms": 46.60082199552562,
      "q1_ms": 46.524377503374126,
      "q3_ms": 46.65705349907512,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 0,
      "rel_vs_flex_repeat": 0.0005209777154959738,
      "compile_s": 0.11346220900304615,
      "inner": 1,
      "median_ms": 115.79077299393248,
      "q1_ms": 115.62686949764611,
      "q3_ms": 115.99743050464895,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 192.0002670288086,
      "xla_arg_mb": 132.0,
      "xla_out_mb": 68.0000228881836,
      "residual_mb": 64.5000114440918,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "float32[1, 4096, 32, 128]",
        "float32[1, 32, 4096]"
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.870667102994048,
      "inner": 1,
      "median_ms": 45.67190400848631,
      "q1_ms": 45.6313924951246,
      "q3_ms": 45.86778500379296,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 1,
      "xla_temp_mb": 320.0002670288086,
      "xla_arg_mb": 132.0,
      "xla_out_mb": 68.0000228881836,
      "residual_mb": 192.5000114440918,
      "residual_shapes": [
        "float32[1, 4096, 32, 128]",
        "float32[1, 4096, 32, 128]",
        "key<fry>[]",
        "uint32[]",
        "float32[1, 4096, 32, 128]",
        "float32[1, 32, 4096]"
      ],
      "compile_s": 0.9009351120039355,
      "inner": 1,
      "median_ms": 46.37028599972837,
      "q1_ms": 46.283245996164624,
      "q3_ms": 46.412500509177335,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 1,
      "rel_vs_flex_repeat": 0.0005209777154959738,
      "compile_s": 0.11388895599520765,
      "inner": 1,
      "median_ms": 115.74851399927866,
      "q1_ms": 115.5140004993882,
      "q3_ms": 115.82163749699248,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "flex_native",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 192.0002670288086,
      "xla_arg_mb": 132.0,
      "xla_out_mb": 68.0000228881836,
      "residual_mb": 64.5000114440918,
      "residual_shapes": [
        "key<fry>[]",
        "uint32[]",
        "float32[1, 4096, 32, 128]",
        "float32[1, 32, 4096]"
      ],
      "rel_vs_flex_repeat": 0.0,
      "compile_s": 0.9988017819996458,
      "inner": 1,
      "median_ms": 45.57244600437116,
      "q1_ms": 45.50346051109955,
      "q3_ms": 45.86295101034921,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "flex_repeat",
      "mode": "fwd_bwd",
      "round": 2,
      "xla_temp_mb": 320.0002670288086,
      "xla_arg_mb": 132.0,
      "xla_out_mb": 68.0000228881836,
      "residual_mb": 192.5000114440918,
      "residual_shapes": [
        "float32[1, 4096, 32, 128]",
        "float32[1, 4096, 32, 128]",
        "key<fry>[]",
        "uint32[]",
        "float32[1, 4096, 32, 128]",
        "float32[1, 32, 4096]"
      ],
      "compile_s": 0.8798914960061666,
      "inner": 1,
      "median_ms": 46.43380500783678,
      "q1_ms": 46.377314494748134,
      "q3_ms": 46.48057449958287,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "gqa",
      "dtype": "float32",
      "batch": 1,
      "lq": 4096,
      "lkv": 4096,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "impl": "torch_sdpa",
      "mode": "fwd_bwd",
      "round": 2,
      "rel_vs_flex_repeat": 0.0005209777154959738,
      "compile_s": 0.11370491700654384,
      "inner": 1,
      "median_ms": 115.84975600999314,
      "q1_ms": 115.64671650558012,
      "q3_ms": 115.90196449833456,
      "tag": null,
      "device_arg": "gpu",
      "run": "gqa_4k_float32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 16.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 1.0312652587890625,
      "xla_arg_mb": 16.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.002165907295420766,
      "rel_err_parts": [
        0.002165907295420766
      ],
      "rel_vs_probjax_nosplit": 0.0028034208808094263,
      "compile_s": 0.17337676900206134,
      "inner": 200,
      "median_ms": 0.03875558504660148,
      "q1_ms": 0.038743892473576125,
      "q3_ms": 0.03885074253048515,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 16.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 16.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.002332377713173628,
      "rel_err_parts": [
        0.002332377713173628
      ],
      "compile_s": 0.15122345800045878,
      "inner": 177,
      "median_ms": 0.15211955934899757,
      "q1_ms": 0.15205036158019697,
      "q3_ms": 0.152152895525895,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 16.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 64.25,
      "xla_arg_mb": 16.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.002332377713173628,
      "rel_err_parts": [
        0.002332377713173628
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.1533321180031635,
      "inner": 140,
      "median_ms": 0.23740868574740098,
      "q1_ms": 0.23708342856220302,
      "q3_ms": 0.23769680353455314,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 16.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.75,
      "xla_arg_mb": 16.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.002385112689808011,
      "rel_err_parts": [
        0.002385112689808011
      ],
      "rel_vs_probjax_nosplit": 0.0032128961756825447,
      "compile_s": 0.1444752170064021,
      "inner": 200,
      "median_ms": 0.042139944998780265,
      "q1_ms": 0.041197032514901366,
      "q3_ms": 0.045216489997983444,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 16.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.2542724609375,
      "xla_arg_mb": 16.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.002187218749895692,
      "rel_err_parts": [
        0.002187218749895692
      ],
      "rel_vs_probjax_nosplit": 0.002900633029639721,
      "compile_s": 0.04408994400000665,
      "inner": 200,
      "median_ms": 0.06387201501638629,
      "q1_ms": 0.06382598250638694,
      "q3_ms": 0.06430145498597994,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 16.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_err": 0.002181076444685459,
      "rel_err_parts": [
        0.002181076444685459
      ],
      "rel_vs_probjax_nosplit": 0.0028841914609074593,
      "compile_s": 0.00023949300521053374,
      "inner": 200,
      "median_ms": 0.036007689996040426,
      "q1_ms": 0.036007287526445,
      "q3_ms": 0.03602061751735164,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 16.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 1.0312652587890625,
      "xla_arg_mb": 16.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.002165907295420766,
      "rel_err_parts": [
        0.002165907295420766
      ],
      "rel_vs_probjax_nosplit": 0.0028034208808094263,
      "compile_s": 0.1711403069930384,
      "inner": 200,
      "median_ms": 0.044575450010597706,
      "q1_ms": 0.0427633999788668,
      "q3_ms": 0.044667127513093874,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 16.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 16.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.002332377713173628,
      "rel_err_parts": [
        0.002332377713173628
      ],
      "compile_s": 0.15015623799990863,
      "inner": 181,
      "median_ms": 0.1352919281385736,
      "q1_ms": 0.13525462430388568,
      "q3_ms": 0.13541249168080982,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 16.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 64.25,
      "xla_arg_mb": 16.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.002332377713173628,
      "rel_err_parts": [
        0.002332377713173628
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.15414289200271014,
      "inner": 146,
      "median_ms": 0.23701954800444208,
      "q1_ms": 0.2369623219499473,
      "q3_ms": 0.2370389726185176,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 16.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.75,
      "xla_arg_mb": 16.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.002385112689808011,
      "rel_err_parts": [
        0.002385112689808011
      ],
      "rel_vs_probjax_nosplit": 0.0032128961756825447,
      "compile_s": 0.13494495300983544,
      "inner": 200,
      "median_ms": 0.03919231501640752,
      "q1_ms": 0.03756932001124369,
      "q3_ms": 0.04231302496918943,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 16.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.2542724609375,
      "xla_arg_mb": 16.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.002187218749895692,
      "rel_err_parts": [
        0.002187218749895692
      ],
      "rel_vs_probjax_nosplit": 0.002900633029639721,
      "compile_s": 0.04583981500763912,
      "inner": 200,
      "median_ms": 0.0759406600263901,
      "q1_ms": 0.07534149750426877,
      "q3_ms": 0.07607017498230562,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 16.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_err": 0.002181076444685459,
      "rel_err_parts": [
        0.002181076444685459
      ],
      "rel_vs_probjax_nosplit": 0.0028841914609074593,
      "compile_s": 0.0002267399977426976,
      "inner": 200,
      "median_ms": 0.03602327000407968,
      "q1_ms": 0.03601317497668788,
      "q3_ms": 0.036024597502546385,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 16.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 1.0312652587890625,
      "xla_arg_mb": 16.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.002165907295420766,
      "rel_err_parts": [
        0.002165907295420766
      ],
      "rel_vs_probjax_nosplit": 0.0028034208808094263,
      "compile_s": 0.1705694160045823,
      "inner": 200,
      "median_ms": 0.0368969749979442,
      "q1_ms": 0.03676419750263449,
      "q3_ms": 0.03716505249030888,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 16.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 16.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.002332377713173628,
      "rel_err_parts": [
        0.002332377713173628
      ],
      "compile_s": 0.14984924798773136,
      "inner": 179,
      "median_ms": 0.1356290614651256,
      "q1_ms": 0.13559074858925435,
      "q3_ms": 0.13568183796918645,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 16.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 64.25,
      "xla_arg_mb": 16.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.002332377713173628,
      "rel_err_parts": [
        0.002332377713173628
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.15520497199031524,
      "inner": 135,
      "median_ms": 0.23676834082127446,
      "q1_ms": 0.23673902959907772,
      "q3_ms": 0.23687068519561932,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 16.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.75,
      "xla_arg_mb": 16.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.002385112689808011,
      "rel_err_parts": [
        0.002385112689808011
      ],
      "rel_vs_probjax_nosplit": 0.0032128961756825447,
      "compile_s": 0.13646506999793928,
      "inner": 200,
      "median_ms": 0.03713512000103947,
      "q1_ms": 0.03634742995927809,
      "q3_ms": 0.03726323750015581,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 16.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.2542724609375,
      "xla_arg_mb": 16.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.002187218749895692,
      "rel_err_parts": [
        0.002187218749895692
      ],
      "rel_vs_probjax_nosplit": 0.002900633029639721,
      "compile_s": 0.045356069997069426,
      "inner": 200,
      "median_ms": 0.06277428998146206,
      "q1_ms": 0.06270801750360988,
      "q3_ms": 0.06283776248892536,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 16.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_err": 0.002181076444685459,
      "rel_err_parts": [
        0.002181076444685459
      ],
      "rel_vs_probjax_nosplit": 0.0028841914609074593,
      "compile_s": 0.0002202259929617867,
      "inner": 200,
      "median_ms": 0.03602567499910947,
      "q1_ms": 0.03601440250349697,
      "q3_ms": 0.03602935503295157,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 16.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 1.0156402587890625,
      "xla_arg_mb": 16.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.002125017810612917,
      "rel_err_parts": [
        0.002125017810612917
      ],
      "rel_vs_probjax_nosplit": 0.002715064911171794,
      "compile_s": 0.16956656900583766,
      "inner": 200,
      "median_ms": 0.03602682496421039,
      "q1_ms": 0.03590589501982322,
      "q3_ms": 0.036109482498432044,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 16.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 16.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.002276914892718196,
      "rel_err_parts": [
        0.002276914892718196
      ],
      "compile_s": 0.15478732199699152,
      "inner": 182,
      "median_ms": 0.1522457967659428,
      "q1_ms": 0.14407750551561968,
      "q3_ms": 0.15227774997333904,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 16.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 64.25,
      "xla_arg_mb": 16.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.002276914892718196,
      "rel_err_parts": [
        0.002276914892718196
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.15859937701316085,
      "inner": 146,
      "median_ms": 0.23660054105439238,
      "q1_ms": 0.23657261298878487,
      "q3_ms": 0.23660565754765492,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 16.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 3.0,
      "xla_arg_mb": 16.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.0023294826969504356,
      "rel_err_parts": [
        0.0023294826969504356
      ],
      "rel_vs_probjax_nosplit": 0.0031060753390192986,
      "compile_s": 0.29871200800698716,
      "inner": 200,
      "median_ms": 0.04381060498417355,
      "q1_ms": 0.043016477502533235,
      "q3_ms": 0.045753857484669425,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 16.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 16.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.0022767584305256605,
      "rel_err_parts": [
        0.0022767584305256605
      ],
      "rel_vs_probjax_nosplit": 8.162271842593327e-05,
      "compile_s": 0.045516623998992145,
      "inner": 155,
      "median_ms": 0.1586624387381298,
      "q1_ms": 0.1586081741768266,
      "q3_ms": 0.1586743676941091,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 16.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_err": 0.002216945169493556,
      "rel_err_parts": [
        0.002216945169493556
      ],
      "rel_vs_probjax_nosplit": 0.002913597272709012,
      "compile_s": 0.0002370489964960143,
      "inner": 200,
      "median_ms": 0.08095619501546025,
      "q1_ms": 0.08003748494957108,
      "q3_ms": 0.08097370748146204,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 16.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 1.0156402587890625,
      "xla_arg_mb": 16.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.002125017810612917,
      "rel_err_parts": [
        0.002125017810612917
      ],
      "rel_vs_probjax_nosplit": 0.002715064911171794,
      "compile_s": 0.16754165801103227,
      "inner": 200,
      "median_ms": 0.04099678000784479,
      "q1_ms": 0.04015946749859722,
      "q3_ms": 0.044135372518212534,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 16.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 16.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.002276914892718196,
      "rel_err_parts": [
        0.002276914892718196
      ],
      "compile_s": 0.1509312550042523,
      "inner": 197,
      "median_ms": 0.13525748226808165,
      "q1_ms": 0.13520153554636266,
      "q3_ms": 0.13528938833943233,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 16.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 64.25,
      "xla_arg_mb": 16.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.002276914892718196,
      "rel_err_parts": [
        0.002276914892718196
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.1500306410016492,
      "inner": 137,
      "median_ms": 0.23688053290349478,
      "q1_ms": 0.23679653282258592,
      "q3_ms": 0.23690988687818776,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 16.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 3.0,
      "xla_arg_mb": 16.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.0023294826969504356,
      "rel_err_parts": [
        0.0023294826969504356
      ],
      "rel_vs_probjax_nosplit": 0.0031060753390192986,
      "compile_s": 0.30131007499585394,
      "inner": 200,
      "median_ms": 0.044721375015797094,
      "q1_ms": 0.043692204999388196,
      "q3_ms": 0.04703831251390511,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 16.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 16.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.0022767584305256605,
      "rel_err_parts": [
        0.0022767584305256605
      ],
      "rel_vs_probjax_nosplit": 8.162271842593327e-05,
      "compile_s": 0.04751132900128141,
      "inner": 163,
      "median_ms": 0.16026677916603493,
      "q1_ms": 0.1602545460047589,
      "q3_ms": 0.16027590183311205,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 16.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_err": 0.002216945169493556,
      "rel_err_parts": [
        0.002216945169493556
      ],
      "rel_vs_probjax_nosplit": 0.002913597272709012,
      "compile_s": 0.000233401995501481,
      "inner": 200,
      "median_ms": 0.08159054501447827,
      "q1_ms": 0.08012362253793981,
      "q3_ms": 0.08162701749824919,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 16.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 1.0156402587890625,
      "xla_arg_mb": 16.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.002125017810612917,
      "rel_err_parts": [
        0.002125017810612917
      ],
      "rel_vs_probjax_nosplit": 0.002715064911171794,
      "compile_s": 0.1693366849940503,
      "inner": 200,
      "median_ms": 0.04440012002305593,
      "q1_ms": 0.043827837544085924,
      "q3_ms": 0.044710382499033585,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 16.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 16.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.002276914892718196,
      "rel_err_parts": [
        0.002276914892718196
      ],
      "compile_s": 0.15002307599934284,
      "inner": 200,
      "median_ms": 0.13587051500508096,
      "q1_ms": 0.13577318251918769,
      "q3_ms": 0.13592213501397055,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 16.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 64.25,
      "xla_arg_mb": 16.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.002276914892718196,
      "rel_err_parts": [
        0.002276914892718196
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.15702561801299453,
      "inner": 138,
      "median_ms": 0.23785997104098805,
      "q1_ms": 0.23785579348138225,
      "q3_ms": 0.2378823333276449,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 16.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 3.0,
      "xla_arg_mb": 16.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.0023294826969504356,
      "rel_err_parts": [
        0.0023294826969504356
      ],
      "rel_vs_probjax_nosplit": 0.0031060753390192986,
      "compile_s": 0.30585224999231286,
      "inner": 200,
      "median_ms": 0.04306243994506076,
      "q1_ms": 0.04301675500755664,
      "q3_ms": 0.04653891999623738,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 16.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 16.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.0022767584305256605,
      "rel_err_parts": [
        0.0022767584305256605
      ],
      "rel_vs_probjax_nosplit": 8.162271842593327e-05,
      "compile_s": 0.044975448996410705,
      "inner": 164,
      "median_ms": 0.16028056096529772,
      "q1_ms": 0.16026901526280066,
      "q3_ms": 0.16028710063149382,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 16.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_err": 0.002216945169493556,
      "rel_err_parts": [
        0.002216945169493556
      ],
      "rel_vs_probjax_nosplit": 0.002913597272709012,
      "compile_s": 0.00024128699442371726,
      "inner": 200,
      "median_ms": 0.08157852498698048,
      "q1_ms": 0.0801999924442498,
      "q3_ms": 0.08159463253832655,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 16.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 4.0625152587890625,
      "xla_arg_mb": 16.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.002136567374691367,
      "rel_err_parts": [
        0.002136567374691367
      ],
      "rel_vs_probjax_nosplit": 0.002722310833632946,
      "compile_s": 0.37079009400622454,
      "inner": 200,
      "median_ms": 0.05104615003801882,
      "q1_ms": 0.05103515250084456,
      "q3_ms": 0.051079837467113975,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 16.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 16.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.002255132654681802,
      "rel_err_parts": [
        0.002255132654681802
      ],
      "compile_s": 0.13688792000175454,
      "inner": 177,
      "median_ms": 0.14178044635257922,
      "q1_ms": 0.14177772880296385,
      "q3_ms": 0.14182646324927786,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 16.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 64.0,
      "xla_arg_mb": 16.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.002255132654681802,
      "rel_err_parts": [
        0.002255132654681802
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.13474462399608456,
      "inner": 147,
      "median_ms": 0.23235957142991984,
      "q1_ms": 0.23233084698868156,
      "q3_ms": 0.23237459865821603,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 16.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 12.0,
      "xla_arg_mb": 16.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.0023259553126990795,
      "rel_err_parts": [
        0.0023259553126990795
      ],
      "rel_vs_probjax_nosplit": 0.003081369213759899,
      "compile_s": 0.18314006499713287,
      "inner": 200,
      "median_ms": 0.06668639500276186,
      "q1_ms": 0.06618476996663958,
      "q3_ms": 0.0667138949938817,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 16.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 16.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.0022551235742866993,
      "rel_err_parts": [
        0.0022551235742866993
      ],
      "rel_vs_probjax_nosplit": 4.473354420042597e-05,
      "compile_s": 0.04551972000626847,
      "inner": 165,
      "median_ms": 0.15868713936387477,
      "q1_ms": 0.15859475454569538,
      "q3_ms": 0.15877682422852202,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 16.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_err": 0.0022052726708352566,
      "rel_err_parts": [
        0.0022052726708352566
      ],
      "rel_vs_probjax_nosplit": 0.0029075953643769026,
      "compile_s": 0.00022999499924480915,
      "inner": 200,
      "median_ms": 0.0823595050314907,
      "q1_ms": 0.08069730753049953,
      "q3_ms": 0.08236684003350092,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 16.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 4.0625152587890625,
      "xla_arg_mb": 16.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.002136567374691367,
      "rel_err_parts": [
        0.002136567374691367
      ],
      "rel_vs_probjax_nosplit": 0.002722310833632946,
      "compile_s": 0.3570325720065739,
      "inner": 200,
      "median_ms": 0.04741973498312291,
      "q1_ms": 0.047199570035445504,
      "q3_ms": 0.04755198500788538,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 16.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 16.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.002255132654681802,
      "rel_err_parts": [
        0.002255132654681802
      ],
      "compile_s": 0.13566005701432005,
      "inner": 200,
      "median_ms": 0.12608158998773433,
      "q1_ms": 0.12607007000042358,
      "q3_ms": 0.12610671248694416,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 16.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 64.0,
      "xla_arg_mb": 16.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.002255132654681802,
      "rel_err_parts": [
        0.002255132654681802
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.13374604401178658,
      "inner": 128,
      "median_ms": 0.23269103132861346,
      "q1_ms": 0.232645757876071,
      "q3_ms": 0.23271126957524757,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 16.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 12.0,
      "xla_arg_mb": 16.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.0023259553126990795,
      "rel_err_parts": [
        0.0023259553126990795
      ],
      "rel_vs_probjax_nosplit": 0.003081369213759899,
      "compile_s": 0.18596313800662756,
      "inner": 200,
      "median_ms": 0.06672126000921708,
      "q1_ms": 0.06617106999328826,
      "q3_ms": 0.06696156251564389,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 16.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 16.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.0022551235742866993,
      "rel_err_parts": [
        0.0022551235742866993
      ],
      "rel_vs_probjax_nosplit": 4.473354420042597e-05,
      "compile_s": 0.046976325000287034,
      "inner": 177,
      "median_ms": 0.16023485877256483,
      "q1_ms": 0.1602251779670843,
      "q3_ms": 0.16027793219499814,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 16.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_err": 0.0022052726708352566,
      "rel_err_parts": [
        0.0022052726708352566
      ],
      "rel_vs_probjax_nosplit": 0.0029075953643769026,
      "compile_s": 0.000233401995501481,
      "inner": 200,
      "median_ms": 0.08276376502180938,
      "q1_ms": 0.08139242247125367,
      "q3_ms": 0.08281241000076989,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 16.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 4.0625152587890625,
      "xla_arg_mb": 16.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.002136567374691367,
      "rel_err_parts": [
        0.002136567374691367
      ],
      "rel_vs_probjax_nosplit": 0.002722310833632946,
      "compile_s": 0.3595936779893236,
      "inner": 200,
      "median_ms": 0.047407759993802756,
      "q1_ms": 0.04738045998237794,
      "q3_ms": 0.04817210749024525,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 16.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 16.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.002255132654681802,
      "rel_err_parts": [
        0.002255132654681802
      ],
      "compile_s": 0.13583628000924364,
      "inner": 200,
      "median_ms": 0.12667220995354,
      "q1_ms": 0.1266609600133961,
      "q3_ms": 0.12667401500948472,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 16.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 64.0,
      "xla_arg_mb": 16.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.002255132654681802,
      "rel_err_parts": [
        0.002255132654681802
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.13460704300086945,
      "inner": 125,
      "median_ms": 0.23357496794778854,
      "q1_ms": 0.23350090807070956,
      "q3_ms": 0.23358682804973796,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 16.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 12.0,
      "xla_arg_mb": 16.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.0023259553126990795,
      "rel_err_parts": [
        0.0023259553126990795
      ],
      "rel_vs_probjax_nosplit": 0.003081369213759899,
      "compile_s": 0.18844645698845852,
      "inner": 200,
      "median_ms": 0.06665568500466179,
      "q1_ms": 0.06598491501790704,
      "q3_ms": 0.06676950248220237,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 16.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 16.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.0022551235742866993,
      "rel_err_parts": [
        0.0022551235742866993
      ],
      "rel_vs_probjax_nosplit": 4.473354420042597e-05,
      "compile_s": 0.04273693299910519,
      "inner": 166,
      "median_ms": 0.16056927108957647,
      "q1_ms": 0.16040667767548114,
      "q3_ms": 0.16066883437311646,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 16.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_err": 0.0022052726708352566,
      "rel_err_parts": [
        0.0022052726708352566
      ],
      "rel_vs_probjax_nosplit": 0.0029075953643769026,
      "compile_s": 0.00023048599541652948,
      "inner": 200,
      "median_ms": 0.08287668497359846,
      "q1_ms": 0.08148547498421976,
      "q3_ms": 0.08289486249850597,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 128.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 1.7617340087890625,
      "xla_arg_mb": 128.0625,
      "xla_out_mb": 0.0625,
      "rel_err": 0.0022518285550177097,
      "rel_err_parts": [
        0.0022518285550177097
      ],
      "rel_vs_probjax_nosplit": 0.002416403265669942,
      "compile_s": 0.17272672899707686,
      "inner": 170,
      "median_ms": 0.16595557056408486,
      "q1_ms": 0.1659420205910197,
      "q3_ms": 0.16598236179561354,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 128.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 128.0625,
      "xla_out_mb": 0.0625,
      "rel_err": 0.0022940458729863167,
      "rel_err_parts": [
        0.0022940458729863167
      ],
      "compile_s": 0.15490954300912563,
      "inner": 79,
      "median_ms": 0.5191651013049097,
      "q1_ms": 0.5188623102104786,
      "q3_ms": 0.5203363670940995,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 128.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 514.0,
      "xla_arg_mb": 128.0625,
      "xla_out_mb": 0.0625,
      "rel_err": 0.0022940458729863167,
      "rel_err_parts": [
        0.0022940458729863167
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.15551324599073268,
      "inner": 30,
      "median_ms": 1.5288494333314397,
      "q1_ms": 1.5283042499504518,
      "q3_ms": 1.5291800831619184,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 128.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 70.0,
      "xla_arg_mb": 128.0625,
      "xla_out_mb": 0.0625,
      "rel_err": 0.0023471133317798376,
      "rel_err_parts": [
        0.0023471133317798376
      ],
      "rel_vs_probjax_nosplit": 0.0031042236369103193,
      "compile_s": 0.16470215399749577,
      "inner": 84,
      "median_ms": 0.49701135728918455,
      "q1_ms": 0.49694963092319777,
      "q3_ms": 0.49712764883311356,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 128.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 128.0625,
      "xla_out_mb": 0.0625,
      "rel_err": 0.002289640484377742,
      "rel_err_parts": [
        0.002289640484377742
      ],
      "rel_vs_probjax_nosplit": 0.0011085522128269076,
      "compile_s": 0.04682406800566241,
      "inner": 159,
      "median_ms": 0.16104557230020142,
      "q1_ms": 0.16097134590930884,
      "q3_ms": 0.16106161317837853,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 128.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_err": 0.0022249752655625343,
      "rel_err_parts": [
        0.0022249752655625343
      ],
      "rel_vs_probjax_nosplit": 0.002947867149487138,
      "compile_s": 0.0003362159914104268,
      "inner": 200,
      "median_ms": 0.17959591001272202,
      "q1_ms": 0.17952354752196698,
      "q3_ms": 0.1796363099856535,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 128.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 1.7617340087890625,
      "xla_arg_mb": 128.0625,
      "xla_out_mb": 0.0625,
      "rel_err": 0.0022518285550177097,
      "rel_err_parts": [
        0.0022518285550177097
      ],
      "rel_vs_probjax_nosplit": 0.002416403265669942,
      "compile_s": 0.1725985779921757,
      "inner": 153,
      "median_ms": 0.16527923525560623,
      "q1_ms": 0.1652003627337418,
      "q3_ms": 0.1652993137124177,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 128.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 128.0625,
      "xla_out_mb": 0.0625,
      "rel_err": 0.0022940458729863167,
      "rel_err_parts": [
        0.0022940458729863167
      ],
      "compile_s": 0.15556286000355612,
      "inner": 75,
      "median_ms": 0.5226527734582003,
      "q1_ms": 0.5225882533704862,
      "q3_ms": 0.5227968532320422,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 128.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 514.0,
      "xla_arg_mb": 128.0625,
      "xla_out_mb": 0.0625,
      "rel_err": 0.0022940458729863167,
      "rel_err_parts": [
        0.0022940458729863167
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.1558148460026132,
      "inner": 29,
      "median_ms": 1.529568275898406,
      "q1_ms": 1.529320379506379,
      "q3_ms": 1.5297292760304368,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 128.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 70.0,
      "xla_arg_mb": 128.0625,
      "xla_out_mb": 0.0625,
      "rel_err": 0.0023471133317798376,
      "rel_err_parts": [
        0.0023471133317798376
      ],
      "rel_vs_probjax_nosplit": 0.0031042236369103193,
      "compile_s": 0.16194712099968456,
      "inner": 79,
      "median_ms": 0.49744121534337254,
      "q1_ms": 0.4974131139372534,
      "q3_ms": 0.4975928228143661,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 128.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 128.0625,
      "xla_out_mb": 0.0625,
      "rel_err": 0.002289640484377742,
      "rel_err_parts": [
        0.002289640484377742
      ],
      "rel_vs_probjax_nosplit": 0.0011085522128269076,
      "compile_s": 0.04643243600730784,
      "inner": 182,
      "median_ms": 0.16127745054820858,
      "q1_ms": 0.16125126921406763,
      "q3_ms": 0.16132201103123855,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 128.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_err": 0.0022249752655625343,
      "rel_err_parts": [
        0.0022249752655625343
      ],
      "rel_vs_probjax_nosplit": 0.002947867149487138,
      "compile_s": 0.00033667699608486146,
      "inner": 200,
      "median_ms": 0.18006229001912288,
      "q1_ms": 0.17996781250985805,
      "q3_ms": 0.18009247247391613,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 128.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 1.7617340087890625,
      "xla_arg_mb": 128.0625,
      "xla_out_mb": 0.0625,
      "rel_err": 0.0022518285550177097,
      "rel_err_parts": [
        0.0022518285550177097
      ],
      "rel_vs_probjax_nosplit": 0.002416403265669942,
      "compile_s": 0.17149674199754372,
      "inner": 173,
      "median_ms": 0.16521234681016922,
      "q1_ms": 0.16515400002285235,
      "q3_ms": 0.16523621102587746,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 128.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 128.0625,
      "xla_out_mb": 0.0625,
      "rel_err": 0.0022940458729863167,
      "rel_err_parts": [
        0.0022940458729863167
      ],
      "compile_s": 0.15576094400603324,
      "inner": 79,
      "median_ms": 0.5224667975083701,
      "q1_ms": 0.5223593101549386,
      "q3_ms": 0.5226634240921446,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 128.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 514.0,
      "xla_arg_mb": 128.0625,
      "xla_out_mb": 0.0625,
      "rel_err": 0.0022940458729863167,
      "rel_err_parts": [
        0.0022940458729863167
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.15633518100366928,
      "inner": 29,
      "median_ms": 1.5297486207945723,
      "q1_ms": 1.5291972243897458,
      "q3_ms": 1.5302176034310833,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 128.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 70.0,
      "xla_arg_mb": 128.0625,
      "xla_out_mb": 0.0625,
      "rel_err": 0.0023471133317798376,
      "rel_err_parts": [
        0.0023471133317798376
      ],
      "rel_vs_probjax_nosplit": 0.0031042236369103193,
      "compile_s": 0.16865561599843204,
      "inner": 78,
      "median_ms": 0.4974500640319326,
      "q1_ms": 0.4973702371689097,
      "q3_ms": 0.49747331398136274,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 128.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 128.0625,
      "xla_out_mb": 0.0625,
      "rel_err": 0.002289640484377742,
      "rel_err_parts": [
        0.002289640484377742
      ],
      "rel_vs_probjax_nosplit": 0.0011085522128269076,
      "compile_s": 0.04644646300585009,
      "inner": 167,
      "median_ms": 0.16124327543322667,
      "q1_ms": 0.16122623352845572,
      "q3_ms": 0.16127090119864113,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 128.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_err": 0.0022249752655625343,
      "rel_err_parts": [
        0.0022249752655625343
      ],
      "rel_vs_probjax_nosplit": 0.002947867149487138,
      "compile_s": 0.0003361370036145672,
      "inner": 200,
      "median_ms": 0.18010176499956287,
      "q1_ms": 0.17995679248997476,
      "q3_ms": 0.18014757748460397,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 128.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 1.5234527587890625,
      "xla_arg_mb": 128.25,
      "xla_out_mb": 0.25,
      "rel_err": 0.0022290078923106194,
      "rel_err_parts": [
        0.0022290078923106194
      ],
      "rel_vs_probjax_nosplit": 0.0024155171122401953,
      "compile_s": 0.17948761400475632,
      "inner": 178,
      "median_ms": 0.16831575838582108,
      "q1_ms": 0.16829887640648103,
      "q3_ms": 0.16836382867572203,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 128.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 128.25,
      "xla_out_mb": 0.25,
      "rel_err": 0.002274314407259226,
      "rel_err_parts": [
        0.002274314407259226
      ],
      "compile_s": 0.1522852369962493,
      "inner": 75,
      "median_ms": 0.5231378400155032,
      "q1_ms": 0.5209287199735021,
      "q3_ms": 0.5234025199509535,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 128.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 514.0,
      "xla_arg_mb": 128.25,
      "xla_out_mb": 0.25,
      "rel_err": 0.002274314407259226,
      "rel_err_parts": [
        0.002274314407259226
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.1573531180038117,
      "inner": 30,
      "median_ms": 1.5297614998416975,
      "q1_ms": 1.5296284332483387,
      "q3_ms": 1.5300893001646423,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 128.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 88.0,
      "xla_arg_mb": 128.25,
      "xla_out_mb": 0.25,
      "rel_err": 0.0023433782625943422,
      "rel_err_parts": [
        0.0023433782625943422
      ],
      "rel_vs_probjax_nosplit": 0.0030831191688776016,
      "compile_s": 0.3707437670091167,
      "inner": 68,
      "median_ms": 0.5764336177997073,
      "q1_ms": 0.5763998603523599,
      "q3_ms": 0.5766442279898054,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 128.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 128.25,
      "xla_out_mb": 0.25,
      "rel_err": 0.00227429554797709,
      "rel_err_parts": [
        0.00227429554797709
      ],
      "rel_vs_probjax_nosplit": 2.9924729460617527e-05,
      "compile_s": 0.04900432399881538,
      "inner": 66,
      "median_ms": 0.6060189242072573,
      "q1_ms": 0.605221575714302,
      "q3_ms": 0.6060550605814734,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 128.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_err": 0.0022648810409009457,
      "rel_err_parts": [
        0.0022648810409009457
      ],
      "rel_vs_probjax_nosplit": 0.003012953093275428,
      "compile_s": 0.001353240993921645,
      "inner": 40,
      "median_ms": 1.2247786999068921,
      "q1_ms": 1.2245243500728975,
      "q3_ms": 1.2248353250470245,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 128.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 1.5234527587890625,
      "xla_arg_mb": 128.25,
      "xla_out_mb": 0.25,
      "rel_err": 0.0022290078923106194,
      "rel_err_parts": [
        0.0022290078923106194
      ],
      "rel_vs_probjax_nosplit": 0.0024155171122401953,
      "compile_s": 0.17191955199814402,
      "inner": 170,
      "median_ms": 0.16731002354337013,
      "q1_ms": 0.16697915882961004,
      "q3_ms": 0.1674184029124102,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 128.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 128.25,
      "xla_out_mb": 0.25,
      "rel_err": 0.002274314407259226,
      "rel_err_parts": [
        0.002274314407259226
      ],
      "compile_s": 0.1555070440081181,
      "inner": 75,
      "median_ms": 0.5239342667240029,
      "q1_ms": 0.5237839866701203,
      "q3_ms": 0.5258179666513266,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 128.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 514.0,
      "xla_arg_mb": 128.25,
      "xla_out_mb": 0.25,
      "rel_err": 0.002274314407259226,
      "rel_err_parts": [
        0.002274314407259226
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.15802890599297825,
      "inner": 30,
      "median_ms": 1.5293657663278282,
      "q1_ms": 1.5293198496995803,
      "q3_ms": 1.5298239668481983,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 128.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 88.0,
      "xla_arg_mb": 128.25,
      "xla_out_mb": 0.25,
      "rel_err": 0.0023433782625943422,
      "rel_err_parts": [
        0.0023433782625943422
      ],
      "rel_vs_probjax_nosplit": 0.0030831191688776016,
      "compile_s": 0.3783265459933318,
      "inner": 72,
      "median_ms": 0.5767990138944393,
      "q1_ms": 0.5767276874798375,
      "q3_ms": 0.5769462291532867,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 128.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 128.25,
      "xla_out_mb": 0.25,
      "rel_err": 0.00227429554797709,
      "rel_err_parts": [
        0.00227429554797709
      ],
      "rel_vs_probjax_nosplit": 2.9924729460617527e-05,
      "compile_s": 0.048181525999098085,
      "inner": 66,
      "median_ms": 0.606650272719421,
      "q1_ms": 0.606550310652837,
      "q3_ms": 0.6066703031000075,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 128.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_err": 0.0022648810409009457,
      "rel_err_parts": [
        0.0022648810409009457
      ],
      "rel_vs_probjax_nosplit": 0.003012953093275428,
      "compile_s": 0.0013481210044119507,
      "inner": 40,
      "median_ms": 1.2247321250470122,
      "q1_ms": 1.2243050625329488,
      "q3_ms": 1.229225499992026,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 128.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 1.5234527587890625,
      "xla_arg_mb": 128.25,
      "xla_out_mb": 0.25,
      "rel_err": 0.0022290078923106194,
      "rel_err_parts": [
        0.0022290078923106194
      ],
      "rel_vs_probjax_nosplit": 0.0024155171122401953,
      "compile_s": 0.17223512899363413,
      "inner": 167,
      "median_ms": 0.16818221557040525,
      "q1_ms": 0.1681286975473701,
      "q3_ms": 0.168244514977183,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 128.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 128.25,
      "xla_out_mb": 0.25,
      "rel_err": 0.002274314407259226,
      "rel_err_parts": [
        0.002274314407259226
      ],
      "compile_s": 0.15527891300735064,
      "inner": 74,
      "median_ms": 0.5269460946422171,
      "q1_ms": 0.5238834663130645,
      "q3_ms": 0.5277406958771895,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 128.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 514.0,
      "xla_arg_mb": 128.25,
      "xla_out_mb": 0.25,
      "rel_err": 0.002274314407259226,
      "rel_err_parts": [
        0.002274314407259226
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.15764701300940942,
      "inner": 29,
      "median_ms": 1.5307477588066831,
      "q1_ms": 1.5301077241772498,
      "q3_ms": 1.531680724242349,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 128.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 88.0,
      "xla_arg_mb": 128.25,
      "xla_out_mb": 0.25,
      "rel_err": 0.0023433782625943422,
      "rel_err_parts": [
        0.0023433782625943422
      ],
      "rel_vs_probjax_nosplit": 0.0030831191688776016,
      "compile_s": 0.3683523920044536,
      "inner": 73,
      "median_ms": 0.5764282054956748,
      "q1_ms": 0.5763411918202135,
      "q3_ms": 0.5766124590555181,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 128.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 128.25,
      "xla_out_mb": 0.25,
      "rel_err": 0.00227429554797709,
      "rel_err_parts": [
        0.00227429554797709
      ],
      "rel_vs_probjax_nosplit": 2.9924729460617527e-05,
      "compile_s": 0.0469995400053449,
      "inner": 67,
      "median_ms": 0.6062520149688639,
      "q1_ms": 0.6055914478267043,
      "q3_ms": 0.6063068134635479,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 128.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_err": 0.0022648810409009457,
      "rel_err_parts": [
        0.0022648810409009457
      ],
      "rel_vs_probjax_nosplit": 0.003012953093275428,
      "compile_s": 0.0013498140033334494,
      "inner": 40,
      "median_ms": 1.2247213500813814,
      "q1_ms": 1.2244638624906656,
      "q3_ms": 1.2334958375504357,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 128.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 6.0937652587890625,
      "xla_arg_mb": 129.0,
      "xla_out_mb": 1.0,
      "rel_err": 0.0022283876314759254,
      "rel_err_parts": [
        0.0022283876314759254
      ],
      "rel_vs_probjax_nosplit": 0.002399140503257513,
      "compile_s": 0.3699923849926563,
      "inner": 109,
      "median_ms": 0.28437445871531963,
      "q1_ms": 0.28143744032462603,
      "q3_ms": 0.30444777062549383,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 128.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 129.0,
      "xla_out_mb": 1.0,
      "rel_err": 0.0022733707446604967,
      "rel_err_parts": [
        0.0022733707446604967
      ],
      "compile_s": 0.1346633500070311,
      "inner": 79,
      "median_ms": 0.5208231772876734,
      "q1_ms": 0.5204662974611443,
      "q3_ms": 0.5213020569637501,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 128.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 512.0,
      "xla_arg_mb": 129.0,
      "xla_out_mb": 1.0,
      "rel_err": 0.0022733707446604967,
      "rel_err_parts": [
        0.0022733707446604967
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.13608551099605393,
      "inner": 28,
      "median_ms": 1.5267280354497157,
      "q1_ms": 1.526619964351994,
      "q3_ms": 1.52701339279052,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 128.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 128.0,
      "xla_arg_mb": 129.0,
      "xla_out_mb": 1.0,
      "rel_err": 0.0023436483461409807,
      "rel_err_parts": [
        0.0023436483461409807
      ],
      "rel_vs_probjax_nosplit": 0.003100025001913309,
      "compile_s": 0.32219619800162036,
      "inner": 54,
      "median_ms": 0.8171010556031467,
      "q1_ms": 0.8168531852856245,
      "q3_ms": 0.8171905832973102,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 128.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 129.0,
      "xla_out_mb": 1.0,
      "rel_err": 0.0022733542136847973,
      "rel_err_parts": [
        0.0022733542136847973
      ],
      "rel_vs_probjax_nosplit": 2.6027393687400036e-05,
      "compile_s": 0.04616220499156043,
      "inner": 68,
      "median_ms": 0.6123134264592355,
      "q1_ms": 0.611920919147772,
      "q3_ms": 0.6124005662166061,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 128.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_err": 0.002269912278279662,
      "rel_err_parts": [
        0.002269912278279662
      ],
      "rel_vs_probjax_nosplit": 0.0030159414745867252,
      "compile_s": 0.0013516879989765584,
      "inner": 40,
      "median_ms": 1.2261024749022909,
      "q1_ms": 1.2259507875569398,
      "q3_ms": 1.230277362446941,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 128.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 6.0937652587890625,
      "xla_arg_mb": 129.0,
      "xla_out_mb": 1.0,
      "rel_err": 0.0022283876314759254,
      "rel_err_parts": [
        0.0022283876314759254
      ],
      "rel_vs_probjax_nosplit": 0.002399140503257513,
      "compile_s": 0.36482101000729017,
      "inner": 116,
      "median_ms": 0.28185664657563586,
      "q1_ms": 0.27959642240352095,
      "q3_ms": 0.28190769394157167,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 128.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 129.0,
      "xla_out_mb": 1.0,
      "rel_err": 0.0022733707446604967,
      "rel_err_parts": [
        0.0022733707446604967
      ],
      "compile_s": 0.13436476499191485,
      "inner": 75,
      "median_ms": 0.5213987999983752,
      "q1_ms": 0.5213510400305192,
      "q3_ms": 0.5215955133705089,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 128.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 512.0,
      "xla_arg_mb": 129.0,
      "xla_out_mb": 1.0,
      "rel_err": 0.0022733707446604967,
      "rel_err_parts": [
        0.0022733707446604967
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.13660492400231306,
      "inner": 30,
      "median_ms": 1.5270069668379922,
      "q1_ms": 1.5269471834471915,
      "q3_ms": 1.527368832951955,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 128.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 128.0,
      "xla_arg_mb": 129.0,
      "xla_out_mb": 1.0,
      "rel_err": 0.0023436483461409807,
      "rel_err_parts": [
        0.0023436483461409807
      ],
      "rel_vs_probjax_nosplit": 0.003100025001913309,
      "compile_s": 0.3216973950038664,
      "inner": 54,
      "median_ms": 0.8168771110595789,
      "q1_ms": 0.816799277850401,
      "q3_ms": 0.8172107129077496,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 128.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 129.0,
      "xla_out_mb": 1.0,
      "rel_err": 0.0022733542136847973,
      "rel_err_parts": [
        0.0022733542136847973
      ],
      "rel_vs_probjax_nosplit": 2.6027393687400036e-05,
      "compile_s": 0.047160984002402984,
      "inner": 66,
      "median_ms": 0.6121783030032406,
      "q1_ms": 0.6114167120937385,
      "q3_ms": 0.6123909772599511,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 128.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_err": 0.002269912278279662,
      "rel_err_parts": [
        0.002269912278279662
      ],
      "rel_vs_probjax_nosplit": 0.0030159414745867252,
      "compile_s": 0.001366216005408205,
      "inner": 40,
      "median_ms": 1.235774749875418,
      "q1_ms": 1.2355498498436646,
      "q3_ms": 1.2358537875115871,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 128.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 6.0937652587890625,
      "xla_arg_mb": 129.0,
      "xla_out_mb": 1.0,
      "rel_err": 0.0022283876314759254,
      "rel_err_parts": [
        0.0022283876314759254
      ],
      "rel_vs_probjax_nosplit": 0.002399140503257513,
      "compile_s": 0.36532138599432074,
      "inner": 123,
      "median_ms": 0.281416438969526,
      "q1_ms": 0.28140776833506087,
      "q3_ms": 0.2814734187090724,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 128.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 129.0,
      "xla_out_mb": 1.0,
      "rel_err": 0.0022733707446604967,
      "rel_err_parts": [
        0.0022733707446604967
      ],
      "compile_s": 0.13447744799486827,
      "inner": 76,
      "median_ms": 0.5217023815356171,
      "q1_ms": 0.521090960510743,
      "q3_ms": 0.5219012368079826,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 128.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 512.0,
      "xla_arg_mb": 129.0,
      "xla_out_mb": 1.0,
      "rel_err": 0.0022733707446604967,
      "rel_err_parts": [
        0.0022733707446604967
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.1329431949998252,
      "inner": 28,
      "median_ms": 1.5278286783931045,
      "q1_ms": 1.5277039822519458,
      "q3_ms": 1.5278810894545,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 128.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 128.0,
      "xla_arg_mb": 129.0,
      "xla_out_mb": 1.0,
      "rel_err": 0.0023436483461409807,
      "rel_err_parts": [
        0.0023436483461409807
      ],
      "rel_vs_probjax_nosplit": 0.003100025001913309,
      "compile_s": 0.3186892719968455,
      "inner": 53,
      "median_ms": 0.8172858300517788,
      "q1_ms": 0.817224018732613,
      "q3_ms": 0.8173096508505966,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 128.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 129.0,
      "xla_out_mb": 1.0,
      "rel_err": 0.0022733542136847973,
      "rel_err_parts": [
        0.0022733542136847973
      ],
      "rel_vs_probjax_nosplit": 2.6027393687400036e-05,
      "compile_s": 0.0474214080022648,
      "inner": 69,
      "median_ms": 0.6170755507849643,
      "q1_ms": 0.6170064347764086,
      "q3_ms": 0.6172714999158178,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 128.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_err": 0.002269912278279662,
      "rel_err_parts": [
        0.002269912278279662
      ],
      "rel_vs_probjax_nosplit": 0.0030159414745867252,
      "compile_s": 0.0013613060000352561,
      "inner": 40,
      "median_ms": 1.235716424707789,
      "q1_ms": 1.2356689374428242,
      "q3_ms": 1.235767387515807,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 1.3540191650390625,
      "xla_arg_mb": 64.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.00216302159242332,
      "rel_err_parts": [
        0.00216302159242332
      ],
      "rel_vs_probjax_nosplit": 0.0028572012670338154,
      "compile_s": 0.1742217279970646,
      "inner": 200,
      "median_ms": 0.0898054450226482,
      "q1_ms": 0.0897728050040314,
      "q3_ms": 0.08983342002466088,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 64.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.002295232145115733,
      "rel_err_parts": [
        0.002295232145115733
      ],
      "compile_s": 0.1508372770040296,
      "inner": 82,
      "median_ms": 0.5103573536716128,
      "q1_ms": 0.5103120305781288,
      "q3_ms": 0.5105132073456444,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 256.25,
      "xla_arg_mb": 64.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.002295232145115733,
      "rel_err_parts": [
        0.002295232145115733
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.15432994600269012,
      "inner": 47,
      "median_ms": 0.9242257659392867,
      "q1_ms": 0.9241552019994607,
      "q3_ms": 0.9245263510742818,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 3.0,
      "xla_arg_mb": 64.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.0022983872331678867,
      "rel_err_parts": [
        0.0022983872331678867
      ],
      "rel_vs_probjax_nosplit": 0.0030556386336684227,
      "compile_s": 0.24341980600729585,
      "inner": 200,
      "median_ms": 0.09675114000856411,
      "q1_ms": 0.09659634499257663,
      "q3_ms": 0.096833970019361,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.2542724609375,
      "xla_arg_mb": 64.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.002202511066570878,
      "rel_err_parts": [
        0.002202511066570878
      ],
      "rel_vs_probjax_nosplit": 0.0029324116185307503,
      "compile_s": 0.046276291002868675,
      "inner": 165,
      "median_ms": 0.09093566058612795,
      "q1_ms": 0.09087860909692774,
      "q3_ms": 0.09098280910014485,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_err": 0.0021896923426538706,
      "rel_err_parts": [
        0.0021896923426538706
      ],
      "rel_vs_probjax_nosplit": 0.002866102382540703,
      "compile_s": 0.00025633499899413437,
      "inner": 200,
      "median_ms": 0.09747745498316363,
      "q1_ms": 0.09653492750658188,
      "q3_ms": 0.09749872999236686,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 1.3540191650390625,
      "xla_arg_mb": 64.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.00216302159242332,
      "rel_err_parts": [
        0.00216302159242332
      ],
      "rel_vs_probjax_nosplit": 0.0028572012670338154,
      "compile_s": 0.17245001600531396,
      "inner": 200,
      "median_ms": 0.08894552003766876,
      "q1_ms": 0.08884114497050177,
      "q3_ms": 0.08895736500562634,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 64.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.002295232145115733,
      "rel_err_parts": [
        0.002295232145115733
      ],
      "compile_s": 0.1514693939971039,
      "inner": 77,
      "median_ms": 0.5146202076810061,
      "q1_ms": 0.5144303701984296,
      "q3_ms": 0.5148172791663242,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 256.25,
      "xla_arg_mb": 64.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.002295232145115733,
      "rel_err_parts": [
        0.002295232145115733
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.16081214300356805,
      "inner": 47,
      "median_ms": 0.9243216810183242,
      "q1_ms": 0.923867861584245,
      "q3_ms": 0.9243428935309218,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 3.0,
      "xla_arg_mb": 64.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.0022983872331678867,
      "rel_err_parts": [
        0.0022983872331678867
      ],
      "rel_vs_probjax_nosplit": 0.0030556386336684227,
      "compile_s": 0.2506852949882159,
      "inner": 200,
      "median_ms": 0.09648092498537153,
      "q1_ms": 0.0963554124973598,
      "q3_ms": 0.09649009251006646,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.2542724609375,
      "xla_arg_mb": 64.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.002202511066570878,
      "rel_err_parts": [
        0.002202511066570878
      ],
      "rel_vs_probjax_nosplit": 0.0029324116185307503,
      "compile_s": 0.046028010998270474,
      "inner": 142,
      "median_ms": 0.09056307040532054,
      "q1_ms": 0.09053209860084421,
      "q3_ms": 0.0905791126621205,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_err": 0.0021896923426538706,
      "rel_err_parts": [
        0.0021896923426538706
      ],
      "rel_vs_probjax_nosplit": 0.002866102382540703,
      "compile_s": 0.00025761700817383826,
      "inner": 200,
      "median_ms": 0.09780347500054631,
      "q1_ms": 0.09679632003098959,
      "q3_ms": 0.09781880751688732,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 1.3540191650390625,
      "xla_arg_mb": 64.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.00216302159242332,
      "rel_err_parts": [
        0.00216302159242332
      ],
      "rel_vs_probjax_nosplit": 0.0028572012670338154,
      "compile_s": 0.17368286800046917,
      "inner": 200,
      "median_ms": 0.08866709002177231,
      "q1_ms": 0.08859289748215815,
      "q3_ms": 0.0886878024903126,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 64.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.002295232145115733,
      "rel_err_parts": [
        0.002295232145115733
      ],
      "compile_s": 0.15454527499969117,
      "inner": 76,
      "median_ms": 0.5148296316436148,
      "q1_ms": 0.5147554803491068,
      "q3_ms": 0.515018934253787,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 256.25,
      "xla_arg_mb": 64.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.002295232145115733,
      "rel_err_parts": [
        0.002295232145115733
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.15359838299627881,
      "inner": 47,
      "median_ms": 0.9243110212759967,
      "q1_ms": 0.9241203616182518,
      "q3_ms": 0.9244801809434144,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 3.0,
      "xla_arg_mb": 64.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.0022983872331678867,
      "rel_err_parts": [
        0.0022983872331678867
      ],
      "rel_vs_probjax_nosplit": 0.0030556386336684227,
      "compile_s": 0.2425698370061582,
      "inner": 113,
      "median_ms": 0.09691862832956774,
      "q1_ms": 0.09667994247170461,
      "q3_ms": 0.0969391946286268,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.2542724609375,
      "xla_arg_mb": 64.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.002202511066570878,
      "rel_err_parts": [
        0.002202511066570878
      ],
      "rel_vs_probjax_nosplit": 0.0029324116185307503,
      "compile_s": 0.045176840998465195,
      "inner": 185,
      "median_ms": 0.09119418918507526,
      "q1_ms": 0.09115409185461751,
      "q3_ms": 0.09121523510878654,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_err": 0.0021896923426538706,
      "rel_err_parts": [
        0.0021896923426538706
      ],
      "rel_vs_probjax_nosplit": 0.002866102382540703,
      "compile_s": 0.00025893001293297857,
      "inner": 200,
      "median_ms": 0.09747621006681584,
      "q1_ms": 0.09694800999568542,
      "q3_ms": 0.09749822496814886,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 1.3330230712890625,
      "xla_arg_mb": 64.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.002207160461694002,
      "rel_err_parts": [
        0.002207160461694002
      ],
      "rel_vs_probjax_nosplit": 0.0028840333689004183,
      "compile_s": 0.17644726000435185,
      "inner": 200,
      "median_ms": 0.09044570499099791,
      "q1_ms": 0.09040532750077546,
      "q3_ms": 0.09044938753504539,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 64.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.002288778545334935,
      "rel_err_parts": [
        0.002288778545334935
      ],
      "compile_s": 0.15186369899311103,
      "inner": 76,
      "median_ms": 0.5119093816335264,
      "q1_ms": 0.51134533561325,
      "q3_ms": 0.5124008947811824,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 256.25,
      "xla_arg_mb": 64.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.002288778545334935,
      "rel_err_parts": [
        0.002288778545334935
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.15514735400211066,
      "inner": 48,
      "median_ms": 0.9243694375982159,
      "q1_ms": 0.9243126664841839,
      "q3_ms": 0.9244801769151916,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 12.0,
      "xla_arg_mb": 64.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.002326450077816844,
      "rel_err_parts": [
        0.002326450077816844
      ],
      "rel_vs_probjax_nosplit": 0.003079321002587676,
      "compile_s": 0.3006099699996412,
      "inner": 164,
      "median_ms": 0.16774859764931224,
      "q1_ms": 0.16734921341483683,
      "q3_ms": 0.1678298780373042,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 64.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.002288778545334935,
      "rel_err_parts": [
        0.002288778545334935
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.04599668200535234,
      "inner": 66,
      "median_ms": 0.6186973180526614,
      "q1_ms": 0.6186661894113941,
      "q3_ms": 0.6187582577254171,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_err": 0.00224050204269588,
      "rel_err_parts": [
        0.00224050204269588
      ],
      "rel_vs_probjax_nosplit": 0.002992357360199094,
      "compile_s": 0.00041669800702948123,
      "inner": 171,
      "median_ms": 0.27747994151963085,
      "q1_ms": 0.27730440644667165,
      "q3_ms": 0.27764463745811896,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 1.3330230712890625,
      "xla_arg_mb": 64.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.002207160461694002,
      "rel_err_parts": [
        0.002207160461694002
      ],
      "rel_vs_probjax_nosplit": 0.0028840333689004183,
      "compile_s": 0.17680976699921302,
      "inner": 200,
      "median_ms": 0.08943253502366133,
      "q1_ms": 0.08936498499679146,
      "q3_ms": 0.08944773999246536,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 64.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.002288778545334935,
      "rel_err_parts": [
        0.002288778545334935
      ],
      "compile_s": 0.15636881400132552,
      "inner": 80,
      "median_ms": 0.5143142625456676,
      "q1_ms": 0.5141133811775944,
      "q3_ms": 0.5148043187546136,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 256.25,
      "xla_arg_mb": 64.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.002288778545334935,
      "rel_err_parts": [
        0.002288778545334935
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.15239755000220612,
      "inner": 46,
      "median_ms": 0.9239582391076154,
      "q1_ms": 0.9238134021554445,
      "q3_ms": 0.9241489130110258,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 12.0,
      "xla_arg_mb": 64.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.002326450077816844,
      "rel_err_parts": [
        0.002326450077816844
      ],
      "rel_vs_probjax_nosplit": 0.003079321002587676,
      "compile_s": 0.30267981700308155,
      "inner": 159,
      "median_ms": 0.16780471690614127,
      "q1_ms": 0.16747220125014206,
      "q3_ms": 0.16786545911714518,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 64.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.002288778545334935,
      "rel_err_parts": [
        0.002288778545334935
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.04698098500375636,
      "inner": 64,
      "median_ms": 0.6193417029862758,
      "q1_ms": 0.6191982187147005,
      "q3_ms": 0.6193490625037157,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_err": 0.00224050204269588,
      "rel_err_parts": [
        0.00224050204269588
      ],
      "rel_vs_probjax_nosplit": 0.002992357360199094,
      "compile_s": 0.000416006994782947,
      "inner": 170,
      "median_ms": 0.27754998231744943,
      "q1_ms": 0.2775426176359521,
      "q3_ms": 0.27756191766992044,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 1.3330230712890625,
      "xla_arg_mb": 64.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.002207160461694002,
      "rel_err_parts": [
        0.002207160461694002
      ],
      "rel_vs_probjax_nosplit": 0.0028840333689004183,
      "compile_s": 0.17572968199965544,
      "inner": 200,
      "median_ms": 0.08948408503783867,
      "q1_ms": 0.08932819251640467,
      "q3_ms": 0.08951071002229583,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 64.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.002288778545334935,
      "rel_err_parts": [
        0.002288778545334935
      ],
      "compile_s": 0.15302357499604113,
      "inner": 77,
      "median_ms": 0.5146442725800897,
      "q1_ms": 0.5145077922460078,
      "q3_ms": 0.5148235843978012,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 256.25,
      "xla_arg_mb": 64.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.002288778545334935,
      "rel_err_parts": [
        0.002288778545334935
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.15319358700071461,
      "inner": 47,
      "median_ms": 0.9243877659096046,
      "q1_ms": 0.924216393577034,
      "q3_ms": 0.9245910214884405,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 12.0,
      "xla_arg_mb": 64.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.002326450077816844,
      "rel_err_parts": [
        0.002326450077816844
      ],
      "rel_vs_probjax_nosplit": 0.003079321002587676,
      "compile_s": 0.29978675200254656,
      "inner": 161,
      "median_ms": 0.1677346831767513,
      "q1_ms": 0.16747266148153392,
      "q3_ms": 0.16785917083460908,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 64.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.002288778545334935,
      "rel_err_parts": [
        0.002288778545334935
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.04383984100422822,
      "inner": 65,
      "median_ms": 0.6189054921448518,
      "q1_ms": 0.6188934693524901,
      "q3_ms": 0.6189696921725184,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_err": 0.00224050204269588,
      "rel_err_parts": [
        0.00224050204269588
      ],
      "rel_vs_probjax_nosplit": 0.002992357360199094,
      "compile_s": 0.00041487500129733235,
      "inner": 171,
      "median_ms": 0.2775898596240556,
      "q1_ms": 0.2775748888944951,
      "q3_ms": 0.27770947075070407,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 5.0781402587890625,
      "xla_arg_mb": 64.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.0022052691783756018,
      "rel_err_parts": [
        0.0022052691783756018
      ],
      "rel_vs_probjax_nosplit": 0.002877879189327359,
      "compile_s": 0.36938042800466064,
      "inner": 190,
      "median_ms": 0.13209053684381375,
      "q1_ms": 0.132088036786483,
      "q3_ms": 0.1320986631694005,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 64.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.00229478208348155,
      "rel_err_parts": [
        0.00229478208348155
      ],
      "compile_s": 0.13329092299682088,
      "inner": 80,
      "median_ms": 0.48610338744765613,
      "q1_ms": 0.48605916872475063,
      "q3_ms": 0.48654734373485553,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 256.0,
      "xla_arg_mb": 64.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.00229478208348155,
      "rel_err_parts": [
        0.00229478208348155
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.13608741499774624,
      "inner": 50,
      "median_ms": 0.9040217200526968,
      "q1_ms": 0.9039243399456609,
      "q3_ms": 0.904103279899573,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 48.0,
      "xla_arg_mb": 64.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.002328990027308464,
      "rel_err_parts": [
        0.002328990027308464
      ],
      "rel_vs_probjax_nosplit": 0.0031103475484997034,
      "compile_s": 0.2911726809979882,
      "inner": 117,
      "median_ms": 0.26493476921171105,
      "q1_ms": 0.26473528633399623,
      "q3_ms": 0.2650355170958508,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 64.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.0022947462275624275,
      "rel_err_parts": [
        0.0022947462275624275
      ],
      "rel_vs_probjax_nosplit": 4.091013761353679e-05,
      "compile_s": 0.04687794900382869,
      "inner": 64,
      "median_ms": 0.6179269842050417,
      "q1_ms": 0.6178854296194913,
      "q3_ms": 0.6179917266990742,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_err": 0.002256842562928796,
      "rel_err_parts": [
        0.002256842562928796
      ],
      "rel_vs_probjax_nosplit": 0.003035714151337743,
      "compile_s": 0.0004126100102439523,
      "inner": 174,
      "median_ms": 0.2805317298612215,
      "q1_ms": 0.2800985488250461,
      "q3_ms": 0.2805644597947844,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 5.0781402587890625,
      "xla_arg_mb": 64.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.0022052691783756018,
      "rel_err_parts": [
        0.0022052691783756018
      ],
      "rel_vs_probjax_nosplit": 0.002877879189327359,
      "compile_s": 0.36005078299785964,
      "inner": 195,
      "median_ms": 0.1258829281701205,
      "q1_ms": 0.12460104618908074,
      "q3_ms": 0.12590386665975434,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 64.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.00229478208348155,
      "rel_err_parts": [
        0.00229478208348155
      ],
      "compile_s": 0.13396283499605488,
      "inner": 80,
      "median_ms": 0.4901459125903785,
      "q1_ms": 0.4899126500276907,
      "q3_ms": 0.49017064384315745,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 256.0,
      "xla_arg_mb": 64.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.00229478208348155,
      "rel_err_parts": [
        0.00229478208348155
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.13495519298885483,
      "inner": 47,
      "median_ms": 0.9054434256618248,
      "q1_ms": 0.9052480425143199,
      "q3_ms": 0.9056070212341667,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 48.0,
      "xla_arg_mb": 64.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.002328990027308464,
      "rel_err_parts": [
        0.002328990027308464
      ],
      "rel_vs_probjax_nosplit": 0.0031103475484997034,
      "compile_s": 0.2965616479923483,
      "inner": 120,
      "median_ms": 0.26501993330991047,
      "q1_ms": 0.2649808124867074,
      "q3_ms": 0.26508701248530997,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 64.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.0022947462275624275,
      "rel_err_parts": [
        0.0022947462275624275
      ],
      "rel_vs_probjax_nosplit": 4.091013761353679e-05,
      "compile_s": 0.04500444400764536,
      "inner": 65,
      "median_ms": 0.6186622770306153,
      "q1_ms": 0.6186440692265303,
      "q3_ms": 0.6187231538485951,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_err": 0.002256842562928796,
      "rel_err_parts": [
        0.002256842562928796
      ],
      "rel_vs_probjax_nosplit": 0.003035714151337743,
      "compile_s": 0.0004169489984633401,
      "inner": 171,
      "median_ms": 0.27897112860891704,
      "q1_ms": 0.2784345848148836,
      "q3_ms": 0.279198871358374,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 5.0781402587890625,
      "xla_arg_mb": 64.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.0022052691783756018,
      "rel_err_parts": [
        0.0022052691783756018
      ],
      "rel_vs_probjax_nosplit": 0.002877879189327359,
      "compile_s": 0.35790798800007906,
      "inner": 200,
      "median_ms": 0.12587910998263396,
      "q1_ms": 0.1251261575089302,
      "q3_ms": 0.12592865001352038,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 64.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.00229478208348155,
      "rel_err_parts": [
        0.00229478208348155
      ],
      "compile_s": 0.1350835760094924,
      "inner": 85,
      "median_ms": 0.4904721529729774,
      "q1_ms": 0.4904282411991838,
      "q3_ms": 0.49083329418756283,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 256.0,
      "xla_arg_mb": 64.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.00229478208348155,
      "rel_err_parts": [
        0.00229478208348155
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.1347552239894867,
      "inner": 47,
      "median_ms": 0.9062436596302196,
      "q1_ms": 0.9061057234190384,
      "q3_ms": 0.9063925531239526,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 48.0,
      "xla_arg_mb": 64.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.002328990027308464,
      "rel_err_parts": [
        0.002328990027308464
      ],
      "rel_vs_probjax_nosplit": 0.0031103475484997034,
      "compile_s": 0.29392850599833764,
      "inner": 126,
      "median_ms": 0.2649701507632724,
      "q1_ms": 0.26494521830861073,
      "q3_ms": 0.2650501745746523,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 64.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.0022947462275624275,
      "rel_err_parts": [
        0.0022947462275624275
      ],
      "rel_vs_probjax_nosplit": 4.091013761353679e-05,
      "compile_s": 0.045387920006760396,
      "inner": 65,
      "median_ms": 0.6190531537868083,
      "q1_ms": 0.61901461539002,
      "q3_ms": 0.6190868384482411,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_err": 0.002256842562928796,
      "rel_err_parts": [
        0.002256842562928796
      ],
      "rel_vs_probjax_nosplit": 0.003035714151337743,
      "compile_s": 0.00041745000635273755,
      "inner": 174,
      "median_ms": 0.27876867818707685,
      "q1_ms": 0.27843491667205866,
      "q3_ms": 0.27894737645876644,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 1.7617340087890625,
      "xla_arg_mb": 512.0625,
      "xla_out_mb": 0.0625,
      "rel_vs_probjax_nosplit": 0.0024134826380759478,
      "compile_s": 0.17285013299260754,
      "inner": 67,
      "median_ms": 0.6222600745447258,
      "q1_ms": 0.6221973433015076,
      "q3_ms": 0.622368865748937,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 512.0625,
      "xla_out_mb": 0.0625,
      "compile_s": 0.15322278199892025,
      "inner": 22,
      "median_ms": 2.054178090475034,
      "q1_ms": 2.0424659543840047,
      "q3_ms": 2.0544411361863073,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 2050.0,
      "xla_arg_mb": 512.0625,
      "xla_out_mb": 0.0625,
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.15988953699707054,
      "inner": 8,
      "median_ms": 6.096436625739443,
      "q1_ms": 6.072964625673194,
      "q3_ms": 6.1033860001771245,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 280.0,
      "xla_arg_mb": 512.0625,
      "xla_out_mb": 0.0625,
      "rel_vs_probjax_nosplit": 0.00309058022685349,
      "compile_s": 0.2748125020007137,
      "inner": 23,
      "median_ms": 1.991355739174532,
      "q1_ms": 1.9908985869614813,
      "q3_ms": 1.9916567170698927,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 512.0625,
      "xla_out_mb": 0.0625,
      "rel_vs_probjax_nosplit": 0.0008694572607055306,
      "compile_s": 0.05090148400631733,
      "inner": 66,
      "median_ms": 0.6171369241201319,
      "q1_ms": 0.6168928333808612,
      "q3_ms": 0.6172527651925162,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_vs_probjax_nosplit": 0.0030018561519682407,
      "compile_s": 0.0007986309938132763,
      "inner": 75,
      "median_ms": 0.6294107999807845,
      "q1_ms": 0.6293888934305869,
      "q3_ms": 0.6294805333406355,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 1.7617340087890625,
      "xla_arg_mb": 512.0625,
      "xla_out_mb": 0.0625,
      "rel_vs_probjax_nosplit": 0.0024134826380759478,
      "compile_s": 0.17469249000714626,
      "inner": 66,
      "median_ms": 0.6215551362732764,
      "q1_ms": 0.6214179849132162,
      "q3_ms": 0.6216862273118529,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 512.0625,
      "xla_out_mb": 0.0625,
      "compile_s": 0.15552018900052644,
      "inner": 22,
      "median_ms": 2.054857590833721,
      "q1_ms": 2.053477500091751,
      "q3_ms": 2.0553690228553023,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 2050.0,
      "xla_arg_mb": 512.0625,
      "xla_out_mb": 0.0625,
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.1628512809984386,
      "inner": 8,
      "median_ms": 6.074274624552345,
      "q1_ms": 6.072054874493915,
      "q3_ms": 6.079054936890316,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 280.0,
      "xla_arg_mb": 512.0625,
      "xla_out_mb": 0.0625,
      "rel_vs_probjax_nosplit": 0.00309058022685349,
      "compile_s": 0.28129811500548385,
      "inner": 23,
      "median_ms": 1.9916924350621665,
      "q1_ms": 1.9911461957077417,
      "q3_ms": 1.991984304091286,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 512.0625,
      "xla_out_mb": 0.0625,
      "rel_vs_probjax_nosplit": 0.0008694572607055306,
      "compile_s": 0.04691713399370201,
      "inner": 65,
      "median_ms": 0.6170556923070063,
      "q1_ms": 0.6167892539031947,
      "q3_ms": 0.6173583461741845,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_vs_probjax_nosplit": 0.0030018561519682407,
      "compile_s": 0.0007877710013417527,
      "inner": 76,
      "median_ms": 0.6295608157080622,
      "q1_ms": 0.6294384868516845,
      "q3_ms": 0.6296685986805977,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 1.7617340087890625,
      "xla_arg_mb": 512.0625,
      "xla_out_mb": 0.0625,
      "rel_vs_probjax_nosplit": 0.0024134826380759478,
      "compile_s": 0.17729502500151284,
      "inner": 68,
      "median_ms": 0.6208438233657302,
      "q1_ms": 0.6205671912432402,
      "q3_ms": 0.6211125735931359,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 512.0625,
      "xla_out_mb": 0.0625,
      "compile_s": 0.15653909699176438,
      "inner": 22,
      "median_ms": 2.0541699089914225,
      "q1_ms": 2.0538614087193063,
      "q3_ms": 2.055659136451802,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 2050.0,
      "xla_arg_mb": 512.0625,
      "xla_out_mb": 0.0625,
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.16100221300439443,
      "inner": 8,
      "median_ms": 6.073911374187446,
      "q1_ms": 6.073029812796449,
      "q3_ms": 6.075180187508522,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 280.0,
      "xla_arg_mb": 512.0625,
      "xla_out_mb": 0.0625,
      "rel_vs_probjax_nosplit": 0.00309058022685349,
      "compile_s": 0.27944469699286856,
      "inner": 21,
      "median_ms": 1.992622619062396,
      "q1_ms": 1.9916054765677211,
      "q3_ms": 1.9961001900290807,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 512.0625,
      "xla_out_mb": 0.0625,
      "rel_vs_probjax_nosplit": 0.0008694572607055306,
      "compile_s": 0.04787351400591433,
      "inner": 65,
      "median_ms": 0.6173998768924951,
      "q1_ms": 0.617241653918217,
      "q3_ms": 0.6175216461997479,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_vs_probjax_nosplit": 0.0030018561519682407,
      "compile_s": 0.0007899750053184107,
      "inner": 75,
      "median_ms": 0.6294197598860288,
      "q1_ms": 0.6293753265830067,
      "q3_ms": 0.6294975599545676,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 1.5234527587890625,
      "xla_arg_mb": 512.25,
      "xla_out_mb": 0.25,
      "rel_vs_probjax_nosplit": 0.0024811241310089827,
      "compile_s": 0.17242656199960038,
      "inner": 60,
      "median_ms": 0.6232692666041354,
      "q1_ms": 0.6232222000107868,
      "q3_ms": 0.6233650417319344,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 512.25,
      "xla_out_mb": 0.25,
      "compile_s": 0.15896678100398276,
      "inner": 22,
      "median_ms": 2.070775909074159,
      "q1_ms": 2.0624239544056104,
      "q3_ms": 2.07151886372065,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 2050.0,
      "xla_arg_mb": 512.25,
      "xla_out_mb": 0.25,
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.16550672700395808,
      "inner": 8,
      "median_ms": 6.088418875151547,
      "q1_ms": 6.082098812839831,
      "q3_ms": 6.095722750615096,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 352.0,
      "xla_arg_mb": 512.25,
      "xla_out_mb": 0.25,
      "rel_vs_probjax_nosplit": 0.0031053079292178154,
      "compile_s": 0.37484104100440163,
      "inner": 19,
      "median_ms": 2.423994999844581,
      "q1_ms": 2.422990131507456,
      "q3_ms": 2.4297543158785344,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 512.25,
      "xla_out_mb": 0.25,
      "rel_vs_probjax_nosplit": 2.0777717509190552e-05,
      "compile_s": 0.05284169499645941,
      "inner": 19,
      "median_ms": 2.435089631226698,
      "q1_ms": 2.434262289362347,
      "q3_ms": 2.4357640263348768,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_vs_probjax_nosplit": 0.0030537541024386883,
      "compile_s": 0.004933825999614783,
      "inner": 10,
      "median_ms": 4.887439699086826,
      "q1_ms": 4.886651700508082,
      "q3_ms": 4.946251550427405,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 1.5234527587890625,
      "xla_arg_mb": 512.25,
      "xla_out_mb": 0.25,
      "rel_vs_probjax_nosplit": 0.0024811241310089827,
      "compile_s": 0.17359103399212472,
      "inner": 60,
      "median_ms": 0.6213998999252606,
      "q1_ms": 0.6213420499989297,
      "q3_ms": 0.6215860082496268,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 512.25,
      "xla_out_mb": 0.25,
      "compile_s": 0.15742583500104956,
      "inner": 22,
      "median_ms": 2.0697549089990472,
      "q1_ms": 2.069438159113898,
      "q3_ms": 2.0702275677732276,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 2050.0,
      "xla_arg_mb": 512.25,
      "xla_out_mb": 0.25,
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.16308385200682096,
      "inner": 8,
      "median_ms": 6.086660625442164,
      "q1_ms": 6.076667812521919,
      "q3_ms": 6.092931812418101,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 352.0,
      "xla_arg_mb": 512.25,
      "xla_out_mb": 0.25,
      "rel_vs_probjax_nosplit": 0.0031053079292178154,
      "compile_s": 0.3736206719913753,
      "inner": 19,
      "median_ms": 2.4266319998883104,
      "q1_ms": 2.4242836577621727,
      "q3_ms": 2.4301972106944674,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 512.25,
      "xla_out_mb": 0.25,
      "rel_vs_probjax_nosplit": 2.0777717509190552e-05,
      "compile_s": 0.051340054007596336,
      "inner": 19,
      "median_ms": 2.4467094205503694,
      "q1_ms": 2.4360543424049146,
      "q3_ms": 2.451742210388674,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_vs_probjax_nosplit": 0.0030537541024386883,
      "compile_s": 0.004940248996717855,
      "inner": 10,
      "median_ms": 4.919139799312688,
      "q1_ms": 4.902652550663333,
      "q3_ms": 4.9223187495954335,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 1.5234527587890625,
      "xla_arg_mb": 512.25,
      "xla_out_mb": 0.25,
      "rel_vs_probjax_nosplit": 0.0024811241310089827,
      "compile_s": 0.17065886600175872,
      "inner": 66,
      "median_ms": 0.6221131667244305,
      "q1_ms": 0.6220818182610497,
      "q3_ms": 0.6221976362929016,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 512.25,
      "xla_out_mb": 0.25,
      "compile_s": 0.1577137090062024,
      "inner": 22,
      "median_ms": 2.0729586818064987,
      "q1_ms": 2.0711429771670904,
      "q3_ms": 2.082722613770536,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 2050.0,
      "xla_arg_mb": 512.25,
      "xla_out_mb": 0.25,
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.1610446429986041,
      "inner": 7,
      "median_ms": 6.104025143680961,
      "q1_ms": 6.090118857433222,
      "q3_ms": 6.1140277851206655,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 352.0,
      "xla_arg_mb": 512.25,
      "xla_out_mb": 0.25,
      "rel_vs_probjax_nosplit": 0.0031053079292178154,
      "compile_s": 0.371558869999717,
      "inner": 19,
      "median_ms": 2.424689473915159,
      "q1_ms": 2.4241526575918355,
      "q3_ms": 2.4250129212496656,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 512.25,
      "xla_out_mb": 0.25,
      "rel_vs_probjax_nosplit": 2.0777717509190552e-05,
      "compile_s": 0.048868696991121396,
      "inner": 19,
      "median_ms": 2.4380187893065774,
      "q1_ms": 2.434441552491319,
      "q3_ms": 2.4457920000330837,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_vs_probjax_nosplit": 0.0030537541024386883,
      "compile_s": 0.004963473009411246,
      "inner": 10,
      "median_ms": 4.917543700139504,
      "q1_ms": 4.917089849914191,
      "q3_ms": 4.917929450311931,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 6.0937652587890625,
      "xla_arg_mb": 513.0,
      "xla_out_mb": 1.0,
      "rel_vs_probjax_nosplit": 0.002458289498463273,
      "compile_s": 0.37029828400409315,
      "inner": 40,
      "median_ms": 1.0260774251946714,
      "q1_ms": 1.0233669374429155,
      "q3_ms": 1.085718562535476,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 513.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.13873258999956306,
      "inner": 22,
      "median_ms": 2.074353136280975,
      "q1_ms": 2.067508136282082,
      "q3_ms": 2.0750360226057554,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 2048.0,
      "xla_arg_mb": 513.0,
      "xla_out_mb": 1.0,
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.14125347099616192,
      "inner": 8,
      "median_ms": 6.096986375268898,
      "q1_ms": 6.087731375373551,
      "q3_ms": 6.129531750048045,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 512.0,
      "xla_arg_mb": 513.0,
      "xla_out_mb": 1.0,
      "rel_vs_probjax_nosplit": 0.0031200863886624575,
      "compile_s": 0.3790893199911807,
      "inner": 14,
      "median_ms": 3.242419428715948,
      "q1_ms": 3.24201396454425,
      "q3_ms": 3.244967713856438,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 513.0,
      "xla_out_mb": 1.0,
      "rel_vs_probjax_nosplit": 2.6284174964530393e-05,
      "compile_s": 0.04842730199743528,
      "inner": 19,
      "median_ms": 2.4642120529039713,
      "q1_ms": 2.4610329472986785,
      "q3_ms": 2.465530894895844,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_vs_probjax_nosplit": 0.003088453784584999,
      "compile_s": 0.004962630991940387,
      "inner": 10,
      "median_ms": 4.919365201203618,
      "q1_ms": 4.919100699771661,
      "q3_ms": 4.919395199249266,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 6.0937652587890625,
      "xla_arg_mb": 513.0,
      "xla_out_mb": 1.0,
      "rel_vs_probjax_nosplit": 0.002458289498463273,
      "compile_s": 0.36433367800782435,
      "inner": 45,
      "median_ms": 1.0319724223033213,
      "q1_ms": 1.0226321890109427,
      "q3_ms": 1.032105544517334,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 513.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.13660128699848428,
      "inner": 23,
      "median_ms": 2.0817289566722414,
      "q1_ms": 2.0782791306492467,
      "q3_ms": 2.089222891220782,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 2048.0,
      "xla_arg_mb": 513.0,
      "xla_out_mb": 1.0,
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.1427576479909476,
      "inner": 8,
      "median_ms": 6.1003077498753555,
      "q1_ms": 6.096732812693517,
      "q3_ms": 6.119341251178412,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 512.0,
      "xla_arg_mb": 513.0,
      "xla_out_mb": 1.0,
      "rel_vs_probjax_nosplit": 0.0031200863886624575,
      "compile_s": 0.38278198800981045,
      "inner": 14,
      "median_ms": 3.2441992142204463,
      "q1_ms": 3.242397571739275,
      "q3_ms": 3.2455056072129604,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 513.0,
      "xla_out_mb": 1.0,
      "rel_vs_probjax_nosplit": 2.6284174964530393e-05,
      "compile_s": 0.049431211009505205,
      "inner": 20,
      "median_ms": 2.45647054980509,
      "q1_ms": 2.4559522498748265,
      "q3_ms": 2.45678957508062,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_vs_probjax_nosplit": 0.003088453784584999,
      "compile_s": 0.004966848995536566,
      "inner": 10,
      "median_ms": 4.94724390009651,
      "q1_ms": 4.925622100563487,
      "q3_ms": 4.960511499666609,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 6.0937652587890625,
      "xla_arg_mb": 513.0,
      "xla_out_mb": 1.0,
      "rel_vs_probjax_nosplit": 0.002458289498463273,
      "compile_s": 0.36327965300006326,
      "inner": 44,
      "median_ms": 1.0250388182943095,
      "q1_ms": 1.0245903636545832,
      "q3_ms": 1.0251985682275104,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 513.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.13846049600397237,
      "inner": 21,
      "median_ms": 2.0771068098145493,
      "q1_ms": 2.0757354047548557,
      "q3_ms": 2.0789424289389338,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 2048.0,
      "xla_arg_mb": 513.0,
      "xla_out_mb": 1.0,
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.1397091190010542,
      "inner": 8,
      "median_ms": 6.109567750172573,
      "q1_ms": 6.096558687204379,
      "q3_ms": 6.145393562292156,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 512.0,
      "xla_arg_mb": 513.0,
      "xla_out_mb": 1.0,
      "rel_vs_probjax_nosplit": 0.0031200863886624575,
      "compile_s": 0.3738374410022516,
      "inner": 14,
      "median_ms": 3.247478927992883,
      "q1_ms": 3.244957392910562,
      "q3_ms": 3.249106249963266,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 513.0,
      "xla_out_mb": 1.0,
      "rel_vs_probjax_nosplit": 2.6284174964530393e-05,
      "compile_s": 0.048590651000267826,
      "inner": 20,
      "median_ms": 2.4566022504586726,
      "q1_ms": 2.4545829499402316,
      "q3_ms": 2.4653162250615424,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_vs_probjax_nosplit": 0.003088453784584999,
      "compile_s": 0.004964384002960287,
      "inner": 10,
      "median_ms": 4.958165599964559,
      "q1_ms": 4.946966900752159,
      "q3_ms": 4.9591975497605745,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 1.3540191650390625,
      "xla_arg_mb": 256.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.0022563713137060404,
      "rel_err_parts": [
        0.0022563713137060404
      ],
      "rel_vs_probjax_nosplit": 0.00289671728387475,
      "compile_s": 0.17482091298734304,
      "inner": 117,
      "median_ms": 0.3180591452719532,
      "q1_ms": 0.31739254274383644,
      "q3_ms": 0.3185141923052307,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 256.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.002308015711605549,
      "rel_err_parts": [
        0.002308015711605549
      ],
      "compile_s": 0.1561910780001199,
      "inner": 23,
      "median_ms": 2.0169850866533006,
      "q1_ms": 2.0158401087468043,
      "q3_ms": 2.0191320214982387,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 1024.25,
      "xla_arg_mb": 256.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.002308015711605549,
      "rel_err_parts": [
        0.002308015711605549
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.1557019030005904,
      "inner": 12,
      "median_ms": 3.680470666343657,
      "q1_ms": 3.674381541713956,
      "q3_ms": 3.6868300837037773,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 12.0001220703125,
      "xla_arg_mb": 256.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.002363304840400815,
      "rel_err_parts": [
        0.002363304840400815
      ],
      "rel_vs_probjax_nosplit": 0.003134555649012327,
      "compile_s": 0.09080362199165393,
      "inner": 110,
      "median_ms": 0.37158862727334385,
      "q1_ms": 0.371574831726982,
      "q3_ms": 0.3716842682148457,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.2542724609375,
      "xla_arg_mb": 256.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.0022921692579984665,
      "rel_err_parts": [
        0.0022921692579984665
      ],
      "rel_vs_probjax_nosplit": 0.0029888967983424664,
      "compile_s": 0.04693000800034497,
      "inner": 106,
      "median_ms": 0.31833409428934756,
      "q1_ms": 0.3181861320626922,
      "q3_ms": 0.31837010849486896,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_err": 0.0022799859289079905,
      "rel_err_parts": [
        0.0022799859289079905
      ],
      "rel_vs_probjax_nosplit": 0.0031248414888978004,
      "compile_s": 0.0004930730065098032,
      "inner": 137,
      "median_ms": 0.3477316788426621,
      "q1_ms": 0.3476149270210728,
      "q3_ms": 0.3477754853653083,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 1.3540191650390625,
      "xla_arg_mb": 256.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.0022563713137060404,
      "rel_err_parts": [
        0.0022563713137060404
      ],
      "rel_vs_probjax_nosplit": 0.00289671728387475,
      "compile_s": 0.172221573011484,
      "inner": 109,
      "median_ms": 0.31771278903294153,
      "q1_ms": 0.3176213761710417,
      "q3_ms": 0.31775759177660534,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 256.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.002308015711605549,
      "rel_err_parts": [
        0.002308015711605549
      ],
      "compile_s": 0.15548370000033174,
      "inner": 22,
      "median_ms": 2.0282230907469057,
      "q1_ms": 2.0277098635481456,
      "q3_ms": 2.0286748863740782,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 1024.25,
      "xla_arg_mb": 256.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.002308015711605549,
      "rel_err_parts": [
        0.002308015711605549
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.15731578700069804,
      "inner": 13,
      "median_ms": 3.6807953068861163,
      "q1_ms": 3.67489226930327,
      "q3_ms": 3.691380769081521,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 12.0001220703125,
      "xla_arg_mb": 256.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.002363304840400815,
      "rel_err_parts": [
        0.002363304840400815
      ],
      "rel_vs_probjax_nosplit": 0.003134555649012327,
      "compile_s": 0.09305801799928304,
      "inner": 105,
      "median_ms": 0.3716281142606888,
      "q1_ms": 0.3716033428818697,
      "q3_ms": 0.371659361865438,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.2542724609375,
      "xla_arg_mb": 256.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.0022921692579984665,
      "rel_err_parts": [
        0.0022921692579984665
      ],
      "rel_vs_probjax_nosplit": 0.0029888967983424664,
      "compile_s": 0.04733380199468229,
      "inner": 105,
      "median_ms": 0.31810920959937256,
      "q1_ms": 0.3180704286759941,
      "q3_ms": 0.3182335428954546,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_err": 0.0022799859289079905,
      "rel_err_parts": [
        0.0022799859289079905
      ],
      "rel_vs_probjax_nosplit": 0.0031248414888978004,
      "compile_s": 0.0004917709884466603,
      "inner": 138,
      "median_ms": 0.3463120144484835,
      "q1_ms": 0.3462811195725184,
      "q3_ms": 0.3463315724334477,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 1.3540191650390625,
      "xla_arg_mb": 256.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.0022563713137060404,
      "rel_err_parts": [
        0.0022563713137060404
      ],
      "rel_vs_probjax_nosplit": 0.00289671728387475,
      "compile_s": 0.17714010199415497,
      "inner": 109,
      "median_ms": 0.31739888072479516,
      "q1_ms": 0.3173555046028389,
      "q3_ms": 0.3174341834708835,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 256.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.002308015711605549,
      "rel_err_parts": [
        0.002308015711605549
      ],
      "compile_s": 0.1534406430000672,
      "inner": 23,
      "median_ms": 2.0308330871478617,
      "q1_ms": 2.029940478361981,
      "q3_ms": 2.031726260927668,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 1024.25,
      "xla_arg_mb": 256.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.002308015711605549,
      "rel_err_parts": [
        0.002308015711605549
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.15783711400581524,
      "inner": 13,
      "median_ms": 3.6679880006256727,
      "q1_ms": 3.6677487693780746,
      "q3_ms": 3.668413076621409,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 12.0001220703125,
      "xla_arg_mb": 256.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.002363304840400815,
      "rel_err_parts": [
        0.002363304840400815
      ],
      "rel_vs_probjax_nosplit": 0.003134555649012327,
      "compile_s": 0.0927131150092464,
      "inner": 101,
      "median_ms": 0.3715839604662713,
      "q1_ms": 0.3714559950299865,
      "q3_ms": 0.37169198516818536,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.2542724609375,
      "xla_arg_mb": 256.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.0022921692579984665,
      "rel_err_parts": [
        0.0022921692579984665
      ],
      "rel_vs_probjax_nosplit": 0.0029888967983424664,
      "compile_s": 0.04764868799247779,
      "inner": 118,
      "median_ms": 0.31814245757216714,
      "q1_ms": 0.3180331016557244,
      "q3_ms": 0.3181811736825601,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_err": 0.0022799859289079905,
      "rel_err_parts": [
        0.0022799859289079905
      ],
      "rel_vs_probjax_nosplit": 0.0031248414888978004,
      "compile_s": 0.0004917899932479486,
      "inner": 138,
      "median_ms": 0.3476594203570183,
      "q1_ms": 0.34760032244578365,
      "q3_ms": 0.3477085289584162,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 1.3330230712890625,
      "xla_arg_mb": 256.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.002255770144984126,
      "rel_err_parts": [
        0.002255770144984126
      ],
      "rel_vs_probjax_nosplit": 0.0029406538233160973,
      "compile_s": 0.17520939800306223,
      "inner": 116,
      "median_ms": 0.32016016371801076,
      "q1_ms": 0.3200713361930584,
      "q3_ms": 0.3202160904916598,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 256.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.0023078436497598886,
      "rel_err_parts": [
        0.0023078436497598886
      ],
      "compile_s": 0.1542523899988737,
      "inner": 22,
      "median_ms": 2.033349181477785,
      "q1_ms": 2.03211570491972,
      "q3_ms": 2.0341456818807107,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 1024.25,
      "xla_arg_mb": 256.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.0023078436497598886,
      "rel_err_parts": [
        0.0023078436497598886
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.15649645699886605,
      "inner": 12,
      "median_ms": 3.673184416887428,
      "q1_ms": 3.671330792106649,
      "q3_ms": 3.6905395827488974,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 48.00048828125,
      "xla_arg_mb": 256.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.0023529664613306522,
      "rel_err_parts": [
        0.0023529664613306522
      ],
      "rel_vs_probjax_nosplit": 0.003129406599327922,
      "compile_s": 0.15123296600359026,
      "inner": 82,
      "median_ms": 0.519723499941969,
      "q1_ms": 0.5197153781099165,
      "q3_ms": 0.5197715123003723,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 256.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.002308053895831108,
      "rel_err_parts": [
        0.002308053895831108
      ],
      "rel_vs_probjax_nosplit": 3.643652962637134e-05,
      "compile_s": 0.0499894769891398,
      "inner": 19,
      "median_ms": 2.4596229471315287,
      "q1_ms": 2.4582242105297105,
      "q3_ms": 2.467011078567195,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_err": 0.002281938912346959,
      "rel_err_parts": [
        0.002281938912346959
      ],
      "rel_vs_probjax_nosplit": 0.0030344196129590273,
      "compile_s": 0.0011558380065253004,
      "inner": 48,
      "median_ms": 1.0719384790718323,
      "q1_ms": 1.0681131144944325,
      "q3_ms": 1.076203541591288,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 1.3330230712890625,
      "xla_arg_mb": 256.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.002255770144984126,
      "rel_err_parts": [
        0.002255770144984126
      ],
      "rel_vs_probjax_nosplit": 0.0029406538233160973,
      "compile_s": 0.17164208700705785,
      "inner": 107,
      "median_ms": 0.31888390652983334,
      "q1_ms": 0.3187621355592056,
      "q3_ms": 0.3189508551318268,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 256.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.0023078436497598886,
      "rel_err_parts": [
        0.0023078436497598886
      ],
      "compile_s": 0.1526168639975367,
      "inner": 22,
      "median_ms": 2.038584090769291,
      "q1_ms": 2.038028931896075,
      "q3_ms": 2.0411507727906364,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 1024.25,
      "xla_arg_mb": 256.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.0023078436497598886,
      "rel_err_parts": [
        0.0023078436497598886
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.1585789569944609,
      "inner": 13,
      "median_ms": 3.676322999616297,
      "q1_ms": 3.6752938081358133,
      "q3_ms": 3.6825602303276983,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 48.00048828125,
      "xla_arg_mb": 256.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.0023529664613306522,
      "rel_err_parts": [
        0.0023529664613306522
      ],
      "rel_vs_probjax_nosplit": 0.003129406599327922,
      "compile_s": 0.15373330700094812,
      "inner": 76,
      "median_ms": 0.5200773289528259,
      "q1_ms": 0.5199774079397015,
      "q3_ms": 0.520221427665858,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 256.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.002308053895831108,
      "rel_err_parts": [
        0.002308053895831108
      ],
      "rel_vs_probjax_nosplit": 3.643652962637134e-05,
      "compile_s": 0.05310151701269206,
      "inner": 19,
      "median_ms": 2.4536700527134694,
      "q1_ms": 2.4534003417724497,
      "q3_ms": 2.456905921127672,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_err": 0.002281938912346959,
      "rel_err_parts": [
        0.002281938912346959
      ],
      "rel_vs_probjax_nosplit": 0.0030344196129590273,
      "compile_s": 0.0011590739886742085,
      "inner": 48,
      "median_ms": 1.0685898541851202,
      "q1_ms": 1.0676762396239308,
      "q3_ms": 1.0715255105727315,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 1.3330230712890625,
      "xla_arg_mb": 256.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.002255770144984126,
      "rel_err_parts": [
        0.002255770144984126
      ],
      "rel_vs_probjax_nosplit": 0.0029406538233160973,
      "compile_s": 0.1786721789976582,
      "inner": 110,
      "median_ms": 0.31838959098835895,
      "q1_ms": 0.3183410545303063,
      "q3_ms": 0.31857872275974264,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 256.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.0023078436497598886,
      "rel_err_parts": [
        0.0023078436497598886
      ],
      "compile_s": 0.1537423440022394,
      "inner": 23,
      "median_ms": 2.0304950436263383,
      "q1_ms": 2.029605065573655,
      "q3_ms": 2.0332805000615064,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 1024.25,
      "xla_arg_mb": 256.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.0023078436497598886,
      "rel_err_parts": [
        0.0023078436497598886
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.15798027500568423,
      "inner": 13,
      "median_ms": 3.667955692015732,
      "q1_ms": 3.6674246926290484,
      "q3_ms": 3.6687513845838393,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 48.00048828125,
      "xla_arg_mb": 256.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.0023529664613306522,
      "rel_err_parts": [
        0.0023529664613306522
      ],
      "rel_vs_probjax_nosplit": 0.003129406599327922,
      "compile_s": 0.15664918499533087,
      "inner": 79,
      "median_ms": 0.5196938228563559,
      "q1_ms": 0.519637063295209,
      "q3_ms": 0.5197864620065127,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 256.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.002308053895831108,
      "rel_err_parts": [
        0.002308053895831108
      ],
      "rel_vs_probjax_nosplit": 3.643652962637134e-05,
      "compile_s": 0.05012853200605605,
      "inner": 18,
      "median_ms": 2.4574882223128545,
      "q1_ms": 2.4569488887209445,
      "q3_ms": 2.4614751944480098,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_err": 0.002281938912346959,
      "rel_err_parts": [
        0.002281938912346959
      ],
      "rel_vs_probjax_nosplit": 0.0030344196129590273,
      "compile_s": 0.0011522300046635792,
      "inner": 48,
      "median_ms": 1.0702688747793825,
      "q1_ms": 1.06905166649085,
      "q3_ms": 1.0726705001313046,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 5.3320465087890625,
      "xla_arg_mb": 256.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.0022598435170948505,
      "rel_err_parts": [
        0.0022598435170948505
      ],
      "rel_vs_probjax_nosplit": 0.002951048780232668,
      "compile_s": 0.3629398699995363,
      "inner": 72,
      "median_ms": 0.5747900555029951,
      "q1_ms": 0.5746939861839767,
      "q3_ms": 0.5768246041447533,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 256.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.002327960915863514,
      "rel_err_parts": [
        0.002327960915863514
      ],
      "compile_s": 0.13772577499912586,
      "inner": 23,
      "median_ms": 1.9601017822582358,
      "q1_ms": 1.9600558478724333,
      "q3_ms": 1.9615905213235554,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 1024.0,
      "xla_arg_mb": 256.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.002327960915863514,
      "rel_err_parts": [
        0.002327960915863514
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.13777880600537173,
      "inner": 13,
      "median_ms": 3.6039916158188134,
      "q1_ms": 3.6030551536196773,
      "q3_ms": 3.620195268572738,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 192.001953125,
      "xla_arg_mb": 256.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.00235678069293499,
      "rel_err_parts": [
        0.00235678069293499
      ],
      "rel_vs_probjax_nosplit": 0.0031457876320928335,
      "compile_s": 0.1638464749994455,
      "inner": 41,
      "median_ms": 1.1172334879364183,
      "q1_ms": 1.1170023170984122,
      "q3_ms": 1.1177324390368217,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 256.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.002328024012967944,
      "rel_err_parts": [
        0.002328024012967944
      ],
      "rel_vs_probjax_nosplit": 4.697022450272925e-05,
      "compile_s": 0.04913373899762519,
      "inner": 18,
      "median_ms": 2.466955055650841,
      "q1_ms": 2.4650024170114193,
      "q3_ms": 2.4675383331972016,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_err": 0.0022993902675807476,
      "rel_err_parts": [
        0.0022993902675807476
      ],
      "rel_vs_probjax_nosplit": 0.003092801431193948,
      "compile_s": 0.001156759011792019,
      "inner": 48,
      "median_ms": 1.0780032916954951,
      "q1_ms": 1.0771680625415079,
      "q3_ms": 1.078191041718431,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 5.3320465087890625,
      "xla_arg_mb": 256.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.0022598435170948505,
      "rel_err_parts": [
        0.0022598435170948505
      ],
      "rel_vs_probjax_nosplit": 0.002951048780232668,
      "compile_s": 0.3661894789984217,
      "inner": 68,
      "median_ms": 0.5794851176153101,
      "q1_ms": 0.5756191250625039,
      "q3_ms": 0.579694639728747,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 256.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.002327960915863514,
      "rel_err_parts": [
        0.002327960915863514
      ],
      "compile_s": 0.1366740349913016,
      "inner": 23,
      "median_ms": 1.957065217024849,
      "q1_ms": 1.9569649998827472,
      "q3_ms": 1.9575961955326973,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 1024.0,
      "xla_arg_mb": 256.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.002327960915863514,
      "rel_err_parts": [
        0.002327960915863514
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.1405088320025243,
      "inner": 12,
      "median_ms": 3.62403050045638,
      "q1_ms": 3.6148289170038574,
      "q3_ms": 3.626128916948801,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 192.001953125,
      "xla_arg_mb": 256.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.00235678069293499,
      "rel_err_parts": [
        0.00235678069293499
      ],
      "rel_vs_probjax_nosplit": 0.0031457876320928335,
      "compile_s": 0.17155858999467455,
      "inner": 40,
      "median_ms": 1.1181779998878483,
      "q1_ms": 1.1177471624250757,
      "q3_ms": 1.1182030249983654,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 256.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.002328024012967944,
      "rel_err_parts": [
        0.002328024012967944
      ],
      "rel_vs_probjax_nosplit": 4.697022450272925e-05,
      "compile_s": 0.04957495200505946,
      "inner": 18,
      "median_ms": 2.4690104999333724,
      "q1_ms": 2.4633852218458845,
      "q3_ms": 2.473057972464529,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_err": 0.0022993902675807476,
      "rel_err_parts": [
        0.0022993902675807476
      ],
      "rel_vs_probjax_nosplit": 0.003092801431193948,
      "compile_s": 0.0011539239931153134,
      "inner": 48,
      "median_ms": 1.0778726248948562,
      "q1_ms": 1.0769174897783764,
      "q3_ms": 1.079462614749597,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 5.3320465087890625,
      "xla_arg_mb": 256.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.0022598435170948505,
      "rel_err_parts": [
        0.0022598435170948505
      ],
      "rel_vs_probjax_nosplit": 0.002951048780232668,
      "compile_s": 0.3632225949986605,
      "inner": 68,
      "median_ms": 0.5788397941026179,
      "q1_ms": 0.5757180659951272,
      "q3_ms": 0.5791183970466269,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 256.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.002327960915863514,
      "rel_err_parts": [
        0.002327960915863514
      ],
      "compile_s": 0.1378895950037986,
      "inner": 23,
      "median_ms": 1.9455255216756917,
      "q1_ms": 1.9450221956029292,
      "q3_ms": 1.9456490433936857,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 1024.0,
      "xla_arg_mb": 256.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.002327960915863514,
      "rel_err_parts": [
        0.002327960915863514
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.1406316450011218,
      "inner": 13,
      "median_ms": 3.6067174615499633,
      "q1_ms": 3.6044531924954546,
      "q3_ms": 3.613788576890679,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 192.001953125,
      "xla_arg_mb": 256.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.00235678069293499,
      "rel_err_parts": [
        0.00235678069293499
      ],
      "rel_vs_probjax_nosplit": 0.0031457876320928335,
      "compile_s": 0.1674743410112569,
      "inner": 39,
      "median_ms": 1.1171776669261118,
      "q1_ms": 1.1169166796324321,
      "q3_ms": 1.1175512179481584,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 256.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.002328024012967944,
      "rel_err_parts": [
        0.002328024012967944
      ],
      "rel_vs_probjax_nosplit": 4.697022450272925e-05,
      "compile_s": 0.04850086999067571,
      "inner": 18,
      "median_ms": 2.4564918890569563,
      "q1_ms": 2.456361389098068,
      "q3_ms": 2.4626326663160905,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_err": 0.0022993902675807476,
      "rel_err_parts": [
        0.0022993902675807476
      ],
      "rel_vs_probjax_nosplit": 0.003092801431193948,
      "compile_s": 0.001152060998720117,
      "inner": 48,
      "median_ms": 1.077740312515137,
      "q1_ms": 1.0762296251414227,
      "q3_ms": 1.079223416733536,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 1.7617340087890625,
      "xla_arg_mb": 2048.0625,
      "xla_out_mb": 0.0625,
      "rel_vs_probjax_nosplit": 0.0024784288834780455,
      "compile_s": 0.1781085740076378,
      "inner": 19,
      "median_ms": 2.433705947049403,
      "q1_ms": 2.4327607369693736,
      "q3_ms": 2.433912394687803,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 2048.0625,
      "xla_out_mb": 0.0625,
      "compile_s": 0.16256995999719948,
      "inner": 6,
      "median_ms": 8.228168500257501,
      "q1_ms": 8.22163425133719,
      "q3_ms": 8.242360750349082,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 8194.0,
      "xla_arg_mb": 2048.0625,
      "xla_out_mb": 0.0625,
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.17873523000162095,
      "inner": 2,
      "median_ms": 24.385144497500733,
      "q1_ms": 24.367857247852953,
      "q3_ms": 24.42282399715623,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 1120.0,
      "xla_arg_mb": 2048.0625,
      "xla_out_mb": 0.0625,
      "rel_vs_probjax_nosplit": 0.0031434306874871254,
      "compile_s": 0.14661787200020626,
      "inner": 6,
      "median_ms": 8.171448331268039,
      "q1_ms": 8.168112750960669,
      "q3_ms": 8.183757998873869,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 2048.0625,
      "xla_out_mb": 0.0625,
      "rel_vs_probjax_nosplit": 0.000729168183170259,
      "compile_s": 0.05252627799927723,
      "inner": 18,
      "median_ms": 2.437899999980194,
      "q1_ms": 2.4371427498408593,
      "q3_ms": 2.438033888918451,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_vs_probjax_nosplit": 0.0030760576482862234,
      "compile_s": 0.00261675201181788,
      "inner": 20,
      "median_ms": 2.4489181996614207,
      "q1_ms": 2.4481943750288337,
      "q3_ms": 2.451389899943024,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 1.7617340087890625,
      "xla_arg_mb": 2048.0625,
      "xla_out_mb": 0.0625,
      "rel_vs_probjax_nosplit": 0.0024784288834780455,
      "compile_s": 0.17627239901048597,
      "inner": 19,
      "median_ms": 2.4383215786656365,
      "q1_ms": 2.436242842170606,
      "q3_ms": 2.4435888684054494,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 2048.0625,
      "xla_out_mb": 0.0625,
      "compile_s": 0.1593698740034597,
      "inner": 6,
      "median_ms": 8.232690667985784,
      "q1_ms": 8.225515750382328,
      "q3_ms": 8.237171081418637,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 8194.0,
      "xla_arg_mb": 2048.0625,
      "xla_out_mb": 0.0625,
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.18049805500777438,
      "inner": 2,
      "median_ms": 24.387779994867742,
      "q1_ms": 24.379562248213915,
      "q3_ms": 24.484457502694568,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 1120.0,
      "xla_arg_mb": 2048.0625,
      "xla_out_mb": 0.0625,
      "rel_vs_probjax_nosplit": 0.0031434306874871254,
      "compile_s": 0.14357372200174723,
      "inner": 6,
      "median_ms": 8.187770998726288,
      "q1_ms": 8.180791084062852,
      "q3_ms": 8.207360249798512,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 2048.0625,
      "xla_out_mb": 0.0625,
      "rel_vs_probjax_nosplit": 0.000729168183170259,
      "compile_s": 0.05030383299163077,
      "inner": 18,
      "median_ms": 2.440457667237044,
      "q1_ms": 2.4378830554875703,
      "q3_ms": 2.4450343331005167,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_vs_probjax_nosplit": 0.0030760576482862234,
      "compile_s": 0.002604468012577854,
      "inner": 20,
      "median_ms": 2.451652900344925,
      "q1_ms": 2.448950799953309,
      "q3_ms": 2.4522642750525847,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 1.7617340087890625,
      "xla_arg_mb": 2048.0625,
      "xla_out_mb": 0.0625,
      "rel_vs_probjax_nosplit": 0.0024784288834780455,
      "compile_s": 0.17808279499877244,
      "inner": 19,
      "median_ms": 2.433275631598295,
      "q1_ms": 2.4324546313054225,
      "q3_ms": 2.434218236885499,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 2048.0625,
      "xla_out_mb": 0.0625,
      "compile_s": 0.1622357869928237,
      "inner": 6,
      "median_ms": 8.247852664984142,
      "q1_ms": 8.215715251329431,
      "q3_ms": 8.28340858303515,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 8194.0,
      "xla_arg_mb": 2048.0625,
      "xla_out_mb": 0.0625,
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.17744512800709344,
      "inner": 2,
      "median_ms": 24.39519399922574,
      "q1_ms": 24.386981000134256,
      "q3_ms": 24.409805504546966,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 1120.0,
      "xla_arg_mb": 2048.0625,
      "xla_out_mb": 0.0625,
      "rel_vs_probjax_nosplit": 0.0031434306874871254,
      "compile_s": 0.15200889400148299,
      "inner": 6,
      "median_ms": 8.198238666712617,
      "q1_ms": 8.18834275075157,
      "q3_ms": 8.216417833561234,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 2048.0625,
      "xla_out_mb": 0.0625,
      "rel_vs_probjax_nosplit": 0.000729168183170259,
      "compile_s": 0.05120620100933593,
      "inner": 17,
      "median_ms": 2.4373016469112105,
      "q1_ms": 2.4369395297767102,
      "q3_ms": 2.464682588384569,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_vs_probjax_nosplit": 0.0030760576482862234,
      "compile_s": 0.002604559005703777,
      "inner": 20,
      "median_ms": 2.4522394996893127,
      "q1_ms": 2.44893899980525,
      "q3_ms": 2.4528183250367874,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 1.5234527587890625,
      "xla_arg_mb": 2048.25,
      "xla_out_mb": 0.25,
      "rel_vs_probjax_nosplit": 0.002506550867110491,
      "compile_s": 0.17695033201016486,
      "inner": 19,
      "median_ms": 2.4379534734180197,
      "q1_ms": 2.4378421838168594,
      "q3_ms": 2.4393882891959757,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 2048.25,
      "xla_out_mb": 0.25,
      "compile_s": 0.16471912599809002,
      "inner": 5,
      "median_ms": 8.3200623979792,
      "q1_ms": 8.293018200492952,
      "q3_ms": 8.342971799720544,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 8194.0,
      "xla_arg_mb": 2048.25,
      "xla_out_mb": 0.25,
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.1795894860115368,
      "inner": 2,
      "median_ms": 24.42435999546433,
      "q1_ms": 24.387341498368187,
      "q3_ms": 24.46694350146572,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 1408.0,
      "xla_arg_mb": 2048.25,
      "xla_out_mb": 0.25,
      "rel_vs_probjax_nosplit": 0.003125615883618593,
      "compile_s": 0.17602233600337058,
      "inner": 5,
      "median_ms": 9.351289601181634,
      "q1_ms": 9.3391270012944,
      "q3_ms": 9.377839100488927,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 2048.25,
      "xla_out_mb": 0.25,
      "rel_vs_probjax_nosplit": 3.318980452604592e-05,
      "compile_s": 0.05826782299845945,
      "inner": 5,
      "median_ms": 9.889272000873461,
      "q1_ms": 9.880074599641375,
      "q3_ms": 9.895640300237574,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_vs_probjax_nosplit": 0.003107517259195447,
      "compile_s": 0.019369531000847928,
      "inner": 2,
      "median_ms": 19.81178300047759,
      "q1_ms": 19.804864750767592,
      "q3_ms": 19.825952247629175,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 1.5234527587890625,
      "xla_arg_mb": 2048.25,
      "xla_out_mb": 0.25,
      "rel_vs_probjax_nosplit": 0.002506550867110491,
      "compile_s": 0.176602567997179,
      "inner": 19,
      "median_ms": 2.4381002624163797,
      "q1_ms": 2.437443447170632,
      "q3_ms": 2.43874539477187,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 2048.25,
      "xla_out_mb": 0.25,
      "compile_s": 0.16308588199899532,
      "inner": 5,
      "median_ms": 8.325272999354638,
      "q1_ms": 8.296478498959914,
      "q3_ms": 8.339189800608438,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 8194.0,
      "xla_arg_mb": 2048.25,
      "xla_out_mb": 0.25,
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.17863711799145676,
      "inner": 2,
      "median_ms": 24.377813497267198,
      "q1_ms": 24.366714751522522,
      "q3_ms": 24.43042724917177,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 1408.0,
      "xla_arg_mb": 2048.25,
      "xla_out_mb": 0.25,
      "rel_vs_probjax_nosplit": 0.003125615883618593,
      "compile_s": 0.1789671919977991,
      "inner": 5,
      "median_ms": 9.370074199978262,
      "q1_ms": 9.337235100974794,
      "q3_ms": 9.371910500340164,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 2048.25,
      "xla_out_mb": 0.25,
      "rel_vs_probjax_nosplit": 3.318980452604592e-05,
      "compile_s": 0.05639635300030932,
      "inner": 5,
      "median_ms": 9.877053601667285,
      "q1_ms": 9.810043599281926,
      "q3_ms": 9.911190999264363,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_vs_probjax_nosplit": 0.003107517259195447,
      "compile_s": 0.01938578199769836,
      "inner": 2,
      "median_ms": 19.816436499240808,
      "q1_ms": 19.810615747701377,
      "q3_ms": 19.817917000182206,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 1.5234527587890625,
      "xla_arg_mb": 2048.25,
      "xla_out_mb": 0.25,
      "rel_vs_probjax_nosplit": 0.002506550867110491,
      "compile_s": 0.17259084800025448,
      "inner": 19,
      "median_ms": 2.438959736824281,
      "q1_ms": 2.438104710324153,
      "q3_ms": 2.4392822107122476,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 2048.25,
      "xla_out_mb": 0.25,
      "compile_s": 0.16256468500068877,
      "inner": 5,
      "median_ms": 8.350410999264568,
      "q1_ms": 8.340563101228327,
      "q3_ms": 8.366400699014775,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 8194.0,
      "xla_arg_mb": 2048.25,
      "xla_out_mb": 0.25,
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.18392686800507363,
      "inner": 2,
      "median_ms": 24.416446998657193,
      "q1_ms": 24.3995864984754,
      "q3_ms": 24.444563496217597,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 1408.0,
      "xla_arg_mb": 2048.25,
      "xla_out_mb": 0.25,
      "rel_vs_probjax_nosplit": 0.003125615883618593,
      "compile_s": 0.18154432899609674,
      "inner": 5,
      "median_ms": 9.342800598824397,
      "q1_ms": 9.339734999230132,
      "q3_ms": 9.355251799570397,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 2048.25,
      "xla_out_mb": 0.25,
      "rel_vs_probjax_nosplit": 3.318980452604592e-05,
      "compile_s": 0.05711307800083887,
      "inner": 5,
      "median_ms": 9.937660198193043,
      "q1_ms": 9.908544199424796,
      "q3_ms": 9.95252410066314,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_vs_probjax_nosplit": 0.003107517259195447,
      "compile_s": 0.019387605003430508,
      "inner": 2,
      "median_ms": 19.813566002994776,
      "q1_ms": 19.810946498182602,
      "q3_ms": 19.820669753244147,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 6.0937652587890625,
      "xla_arg_mb": 2049.0,
      "xla_out_mb": 1.0,
      "rel_vs_probjax_nosplit": 0.0025066551752388477,
      "compile_s": 0.37099373700038996,
      "inner": 12,
      "median_ms": 3.997713665739866,
      "q1_ms": 3.9836199163498045,
      "q3_ms": 4.010761291283416,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 2049.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.14194181399943773,
      "inner": 6,
      "median_ms": 8.27922616736032,
      "q1_ms": 8.26626058308951,
      "q3_ms": 8.306990250275703,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 8192.0,
      "xla_arg_mb": 2049.0,
      "xla_out_mb": 1.0,
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.16125780300353654,
      "inner": 2,
      "median_ms": 24.388268495386,
      "q1_ms": 24.345767000340857,
      "q3_ms": 24.469263502396643,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 2048.0,
      "xla_arg_mb": 2049.0,
      "xla_out_mb": 1.0,
      "rel_vs_probjax_nosplit": 0.0031410623341798782,
      "compile_s": 0.2077805430017179,
      "inner": 3,
      "median_ms": 14.19058200068927,
      "q1_ms": 14.149943167770592,
      "q3_ms": 14.2017169976801,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 2049.0,
      "xla_out_mb": 1.0,
      "rel_vs_probjax_nosplit": 4.824485949939117e-05,
      "compile_s": 0.058997765998356044,
      "inner": 5,
      "median_ms": 10.020461000385694,
      "q1_ms": 9.990583200124092,
      "q3_ms": 10.070420999545604,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_vs_probjax_nosplit": 0.0031295339576900005,
      "compile_s": 0.019368358989595436,
      "inner": 2,
      "median_ms": 19.815389503492042,
      "q1_ms": 19.814194998616586,
      "q3_ms": 19.817553751636297,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 6.0937652587890625,
      "xla_arg_mb": 2049.0,
      "xla_out_mb": 1.0,
      "rel_vs_probjax_nosplit": 0.0025066551752388477,
      "compile_s": 0.3726693260105094,
      "inner": 12,
      "median_ms": 4.015614167049837,
      "q1_ms": 4.007483041277737,
      "q3_ms": 4.01935583310357,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 2049.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.14327846199739724,
      "inner": 6,
      "median_ms": 8.342181832025139,
      "q1_ms": 8.323163250073168,
      "q3_ms": 8.3438734181982,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 8192.0,
      "xla_arg_mb": 2049.0,
      "xla_out_mb": 1.0,
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.16044775999034755,
      "inner": 2,
      "median_ms": 24.406816497503314,
      "q1_ms": 24.358290753298206,
      "q3_ms": 24.481937751261285,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 2048.0,
      "xla_arg_mb": 2049.0,
      "xla_out_mb": 1.0,
      "rel_vs_probjax_nosplit": 0.0031410623341798782,
      "compile_s": 0.20015594600408804,
      "inner": 3,
      "median_ms": 14.170270000856059,
      "q1_ms": 14.155671999712165,
      "q3_ms": 14.189350833476055,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 2049.0,
      "xla_out_mb": 1.0,
      "rel_vs_probjax_nosplit": 4.824485949939117e-05,
      "compile_s": 0.05668944799981546,
      "inner": 5,
      "median_ms": 9.986352201667614,
      "q1_ms": 9.982426799251698,
      "q3_ms": 9.989835901069455,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_vs_probjax_nosplit": 0.0031295339576900005,
      "compile_s": 0.01938920900283847,
      "inner": 2,
      "median_ms": 19.81599599821493,
      "q1_ms": 19.813282997347414,
      "q3_ms": 19.817648997559445,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 6.0937652587890625,
      "xla_arg_mb": 2049.0,
      "xla_out_mb": 1.0,
      "rel_vs_probjax_nosplit": 0.0025066551752388477,
      "compile_s": 0.36771923099877313,
      "inner": 12,
      "median_ms": 4.021729083130292,
      "q1_ms": 4.019059416653666,
      "q3_ms": 4.028219333728581,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 2049.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.14174983999691904,
      "inner": 6,
      "median_ms": 8.324837331504872,
      "q1_ms": 8.323911250045057,
      "q3_ms": 8.346519251063,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 8192.0,
      "xla_arg_mb": 2049.0,
      "xla_out_mb": 1.0,
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.1604693609988317,
      "inner": 2,
      "median_ms": 24.46290400257567,
      "q1_ms": 24.34666599947377,
      "q3_ms": 24.51627499976894,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 2048.0,
      "xla_arg_mb": 2049.0,
      "xla_out_mb": 1.0,
      "rel_vs_probjax_nosplit": 0.0031410623341798782,
      "compile_s": 0.204780446991208,
      "inner": 3,
      "median_ms": 14.165213666274212,
      "q1_ms": 14.15374333494886,
      "q3_ms": 14.212641333870124,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 2049.0,
      "xla_out_mb": 1.0,
      "rel_vs_probjax_nosplit": 4.824485949939117e-05,
      "compile_s": 0.05924193799728528,
      "inner": 5,
      "median_ms": 9.985791001236066,
      "q1_ms": 9.972965699853376,
      "q3_ms": 9.998948201246094,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_vs_probjax_nosplit": 0.0031295339576900005,
      "compile_s": 0.019369219997315668,
      "inner": 2,
      "median_ms": 19.816978005110286,
      "q1_ms": 19.814330247754697,
      "q3_ms": 19.818686003418406,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h8"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 1.6367340087890625,
      "xla_arg_mb": 64.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.002198614878579974,
      "rel_err_parts": [
        0.002198614878579974
      ],
      "rel_vs_probjax_nosplit": 0.0026371402200311422,
      "compile_s": 0.1700978699955158,
      "inner": 200,
      "median_ms": 0.09157652995781973,
      "q1_ms": 0.09157071999652544,
      "q3_ms": 0.09160085250186967,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 64.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.0022587531711906195,
      "rel_err_parts": [
        0.0022587531711906195
      ],
      "compile_s": 0.15422423400741536,
      "inner": 185,
      "median_ms": 0.13883248108401391,
      "q1_ms": 0.13881285133696086,
      "q3_ms": 0.13885937298402645,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 64.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.0022587531711906195,
      "rel_err_parts": [
        0.0022587531711906195
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.1467545199993765,
      "inner": 193,
      "median_ms": 0.13878172018608667,
      "q1_ms": 0.13876030828690378,
      "q3_ms": 0.13882390155132446,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 1.0000152587890625,
      "xla_arg_mb": 64.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.0016579708317294717,
      "rel_err_parts": [
        0.0016579708317294717
      ],
      "rel_vs_probjax_nosplit": 0.0024238598998636007,
      "compile_s": 0.07455571499303915,
      "inner": 200,
      "median_ms": 0.1125752150255721,
      "q1_ms": 0.1125184100237675,
      "q3_ms": 0.1126359074987704,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.0638427734375,
      "xla_arg_mb": 64.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.002212806837633252,
      "rel_err_parts": [
        0.002212806837633252
      ],
      "rel_vs_probjax_nosplit": 0.0026995744556188583,
      "compile_s": 0.04913260899775196,
      "inner": 140,
      "median_ms": 0.08958937864268332,
      "q1_ms": 0.08952805001172237,
      "q3_ms": 0.08964047504248031,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_err": 0.0021837472449988127,
      "rel_err_parts": [
        0.0021837472449988127
      ],
      "rel_vs_probjax_nosplit": 0.002820101100951433,
      "compile_s": 0.0002504940057406202,
      "inner": 200,
      "median_ms": 0.09550882496114355,
      "q1_ms": 0.09494172496488318,
      "q3_ms": 0.09552953502861783,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 1.6367340087890625,
      "xla_arg_mb": 64.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.002198614878579974,
      "rel_err_parts": [
        0.002198614878579974
      ],
      "rel_vs_probjax_nosplit": 0.0026371402200311422,
      "compile_s": 0.16954459300905,
      "inner": 200,
      "median_ms": 0.0900062650180189,
      "q1_ms": 0.08995205997052835,
      "q3_ms": 0.09003320996271214,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 64.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.0022587531711906195,
      "rel_err_parts": [
        0.0022587531711906195
      ],
      "compile_s": 0.1508478049945552,
      "inner": 200,
      "median_ms": 0.13916217001678888,
      "q1_ms": 0.13901434253057232,
      "q3_ms": 0.1392666950050625,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 64.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.0022587531711906195,
      "rel_err_parts": [
        0.0022587531711906195
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.14743888500379398,
      "inner": 174,
      "median_ms": 0.13930785058911932,
      "q1_ms": 0.1392567212641207,
      "q3_ms": 0.1393204885093367,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 1.0000152587890625,
      "xla_arg_mb": 64.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.0016579708317294717,
      "rel_err_parts": [
        0.0016579708317294717
      ],
      "rel_vs_probjax_nosplit": 0.0024238598998636007,
      "compile_s": 0.07120126800145954,
      "inner": 200,
      "median_ms": 0.11258047496085055,
      "q1_ms": 0.1124734475160949,
      "q3_ms": 0.11260619747190503,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.0638427734375,
      "xla_arg_mb": 64.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.002212806837633252,
      "rel_err_parts": [
        0.002212806837633252
      ],
      "rel_vs_probjax_nosplit": 0.0026995744556188583,
      "compile_s": 0.04631244200572837,
      "inner": 200,
      "median_ms": 0.09000891499454156,
      "q1_ms": 0.08998930250527337,
      "q3_ms": 0.0900739899952896,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_err": 0.0021837472449988127,
      "rel_err_parts": [
        0.0021837472449988127
      ],
      "rel_vs_probjax_nosplit": 0.002820101100951433,
      "compile_s": 0.00025132599694188684,
      "inner": 200,
      "median_ms": 0.09567804001562763,
      "q1_ms": 0.09466327497648308,
      "q3_ms": 0.09569693003868451,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 1.6367340087890625,
      "xla_arg_mb": 64.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.002198614878579974,
      "rel_err_parts": [
        0.002198614878579974
      ],
      "rel_vs_probjax_nosplit": 0.0026371402200311422,
      "compile_s": 0.1734129730029963,
      "inner": 200,
      "median_ms": 0.0903438499517506,
      "q1_ms": 0.0903281474893447,
      "q3_ms": 0.09035191749717342,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 64.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.0022587531711906195,
      "rel_err_parts": [
        0.0022587531711906195
      ],
      "compile_s": 0.14671197900315747,
      "inner": 200,
      "median_ms": 0.13884061503631528,
      "q1_ms": 0.13869413749489468,
      "q3_ms": 0.13887465498555684,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 64.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.0022587531711906195,
      "rel_err_parts": [
        0.0022587531711906195
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.14821740800107364,
      "inner": 179,
      "median_ms": 0.13892382115721388,
      "q1_ms": 0.13890221786350998,
      "q3_ms": 0.13893775973766886,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 1.0000152587890625,
      "xla_arg_mb": 64.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.0016579708317294717,
      "rel_err_parts": [
        0.0016579708317294717
      ],
      "rel_vs_probjax_nosplit": 0.0024238598998636007,
      "compile_s": 0.07710533899080474,
      "inner": 200,
      "median_ms": 0.11230360498302616,
      "q1_ms": 0.11223211498872843,
      "q3_ms": 0.11234553003305336,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.0638427734375,
      "xla_arg_mb": 64.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.002212806837633252,
      "rel_err_parts": [
        0.002212806837633252
      ],
      "rel_vs_probjax_nosplit": 0.0026995744556188583,
      "compile_s": 0.04611558098986279,
      "inner": 200,
      "median_ms": 0.08951989497290924,
      "q1_ms": 0.0894635100121377,
      "q3_ms": 0.08956808247603476,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_err": 0.0021837472449988127,
      "rel_err_parts": [
        0.0021837472449988127
      ],
      "rel_vs_probjax_nosplit": 0.002820101100951433,
      "compile_s": 0.0002451439941069111,
      "inner": 200,
      "median_ms": 0.09569847999955527,
      "q1_ms": 0.09449630499148043,
      "q3_ms": 0.09572047496476443,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 1.6250152587890625,
      "xla_arg_mb": 64.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.0021783614065498114,
      "rel_err_parts": [
        0.0021783614065498114
      ],
      "rel_vs_probjax_nosplit": 0.0026152331847697496,
      "compile_s": 0.17204788100207224,
      "inner": 200,
      "median_ms": 0.09012398499180563,
      "q1_ms": 0.09009557750687236,
      "q3_ms": 0.09016891748615308,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 64.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.0022518106270581484,
      "rel_err_parts": [
        0.0022518106270581484
      ],
      "compile_s": 0.15682986500905827,
      "inner": 197,
      "median_ms": 0.1387631776653996,
      "q1_ms": 0.13873009643455975,
      "q3_ms": 0.13879013453435324,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 64.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.0022518106270581484,
      "rel_err_parts": [
        0.0022518106270581484
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.14931349198741373,
      "inner": 191,
      "median_ms": 0.13879094239633957,
      "q1_ms": 0.13877499738283922,
      "q3_ms": 0.13880458375023363,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 3.0,
      "xla_arg_mb": 64.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.0023412490263581276,
      "rel_err_parts": [
        0.0023412490263581276
      ],
      "rel_vs_probjax_nosplit": 0.003107257653027773,
      "compile_s": 0.1452703120012302,
      "inner": 200,
      "median_ms": 0.09464653994655237,
      "q1_ms": 0.0944466375221964,
      "q3_ms": 0.09466985749895684,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 64.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.0022518739569932222,
      "rel_err_parts": [
        0.0022518739569932222
      ],
      "rel_vs_probjax_nosplit": 7.24051205907017e-05,
      "compile_s": 0.04429306200472638,
      "inner": 138,
      "median_ms": 0.1603306957184796,
      "q1_ms": 0.16030093114671257,
      "q3_ms": 0.1603453985834663,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_err": 0.002191473264247179,
      "rel_err_parts": [
        0.002191473264247179
      ],
      "rel_vs_probjax_nosplit": 0.0028353917878121138,
      "compile_s": 0.00025378000282216817,
      "inner": 200,
      "median_ms": 0.09637356000894215,
      "q1_ms": 0.09555415752402041,
      "q3_ms": 0.09641826749430038,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 1.6250152587890625,
      "xla_arg_mb": 64.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.0021783614065498114,
      "rel_err_parts": [
        0.0021783614065498114
      ],
      "rel_vs_probjax_nosplit": 0.0026152331847697496,
      "compile_s": 0.17171502999553923,
      "inner": 200,
      "median_ms": 0.08929281000746414,
      "q1_ms": 0.08917739000025904,
      "q3_ms": 0.08931430249504047,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 64.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.0022518106270581484,
      "rel_err_parts": [
        0.0022518106270581484
      ],
      "compile_s": 0.15039218300080393,
      "inner": 189,
      "median_ms": 0.13906509523395763,
      "q1_ms": 0.13903422219088674,
      "q3_ms": 0.13910586508077188,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 64.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.0022518106270581484,
      "rel_err_parts": [
        0.0022518106270581484
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.149555859999964,
      "inner": 186,
      "median_ms": 0.13906346242536882,
      "q1_ms": 0.13905527149652563,
      "q3_ms": 0.13912516129415942,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 3.0,
      "xla_arg_mb": 64.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.0023412490263581276,
      "rel_err_parts": [
        0.0023412490263581276
      ],
      "rel_vs_probjax_nosplit": 0.003107257653027773,
      "compile_s": 0.13526142299815547,
      "inner": 200,
      "median_ms": 0.09457776002818719,
      "q1_ms": 0.09445600502658635,
      "q3_ms": 0.09461147499678191,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 64.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.0022518739569932222,
      "rel_err_parts": [
        0.0022518739569932222
      ],
      "rel_vs_probjax_nosplit": 7.24051205907017e-05,
      "compile_s": 0.04539324399956968,
      "inner": 141,
      "median_ms": 0.16059963115580142,
      "q1_ms": 0.1605691134674534,
      "q3_ms": 0.16064382976016808,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_err": 0.002191473264247179,
      "rel_err_parts": [
        0.002191473264247179
      ],
      "rel_vs_probjax_nosplit": 0.0028353917878121138,
      "compile_s": 0.0002529889898141846,
      "inner": 200,
      "median_ms": 0.09615889997803606,
      "q1_ms": 0.09539565748127643,
      "q3_ms": 0.09618259751732694,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 1.6250152587890625,
      "xla_arg_mb": 64.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.0021783614065498114,
      "rel_err_parts": [
        0.0021783614065498114
      ],
      "rel_vs_probjax_nosplit": 0.0026152331847697496,
      "compile_s": 0.17510231900087092,
      "inner": 200,
      "median_ms": 0.08927892995416187,
      "q1_ms": 0.08920945503632538,
      "q3_ms": 0.0892835924969404,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 64.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.0022518106270581484,
      "rel_err_parts": [
        0.0022518106270581484
      ],
      "compile_s": 0.14792399300495163,
      "inner": 134,
      "median_ms": 0.13938875368938766,
      "q1_ms": 0.13935035449728966,
      "q3_ms": 0.13940251116778132,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 64.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.0022518106270581484,
      "rel_err_parts": [
        0.0022518106270581484
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.1490355760033708,
      "inner": 185,
      "median_ms": 0.13891019995286558,
      "q1_ms": 0.13879051084340727,
      "q3_ms": 0.13892617296093973,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 3.0,
      "xla_arg_mb": 64.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.0023412490263581276,
      "rel_err_parts": [
        0.0023412490263581276
      ],
      "rel_vs_probjax_nosplit": 0.003107257653027773,
      "compile_s": 0.13642639700265136,
      "inner": 200,
      "median_ms": 0.09434621999389492,
      "q1_ms": 0.09423208255611826,
      "q3_ms": 0.0943498024935252,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 64.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.0022518739569932222,
      "rel_err_parts": [
        0.0022518739569932222
      ],
      "rel_vs_probjax_nosplit": 7.24051205907017e-05,
      "compile_s": 0.045305617997655645,
      "inner": 167,
      "median_ms": 0.15975930536316407,
      "q1_ms": 0.15974259579973854,
      "q3_ms": 0.15979596107123242,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_err": 0.002191473264247179,
      "rel_err_parts": [
        0.002191473264247179
      ],
      "rel_vs_probjax_nosplit": 0.0028353917878121138,
      "compile_s": 0.00024947199563030154,
      "inner": 200,
      "median_ms": 0.09638608004024718,
      "q1_ms": 0.09575979751389241,
      "q3_ms": 0.09643484998377971,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 1.5234527587890625,
      "xla_arg_mb": 64.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.0021824080031365156,
      "rel_err_parts": [
        0.0021824080031365156
      ],
      "rel_vs_probjax_nosplit": 0.0026159496046602726,
      "compile_s": 0.17499444499844685,
      "inner": 200,
      "median_ms": 0.09090827006730251,
      "q1_ms": 0.09089111001230776,
      "q3_ms": 0.0909498974942835,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 64.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.002256848616525531,
      "rel_err_parts": [
        0.002256848616525531
      ],
      "compile_s": 0.13721699299640022,
      "inner": 200,
      "median_ms": 0.1292345899855718,
      "q1_ms": 0.12922316749609308,
      "q3_ms": 0.12958893246832304,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 64.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.002256848616525531,
      "rel_err_parts": [
        0.002256848616525531
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.13005447998875752,
      "inner": 181,
      "median_ms": 0.13024539782251462,
      "q1_ms": 0.13020181218329407,
      "q3_ms": 0.1302497486441746,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 12.0,
      "xla_arg_mb": 64.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.002338483463972807,
      "rel_err_parts": [
        0.002338483463972807
      ],
      "rel_vs_probjax_nosplit": 0.0030901406425982714,
      "compile_s": 0.14165550099278335,
      "inner": 200,
      "median_ms": 0.11817846498161089,
      "q1_ms": 0.11793190249591134,
      "q3_ms": 0.1182112050082651,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 64.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.0022568346466869116,
      "rel_err_parts": [
        0.0022568346466869116
      ],
      "rel_vs_probjax_nosplit": 3.77187789126765e-05,
      "compile_s": 0.0450677680055378,
      "inner": 169,
      "median_ms": 0.16054211834655005,
      "q1_ms": 0.16053746741667346,
      "q3_ms": 0.1605595828835803,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_err": 0.002206334611400962,
      "rel_err_parts": [
        0.002206334611400962
      ],
      "rel_vs_probjax_nosplit": 0.0028899607714265585,
      "compile_s": 0.00024887100153137,
      "inner": 200,
      "median_ms": 0.09860563499387354,
      "q1_ms": 0.09736393003549892,
      "q3_ms": 0.0986147249568603,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 1.5234527587890625,
      "xla_arg_mb": 64.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.0021824080031365156,
      "rel_err_parts": [
        0.0021824080031365156
      ],
      "rel_vs_probjax_nosplit": 0.0026159496046602726,
      "compile_s": 0.1709804910060484,
      "inner": 200,
      "median_ms": 0.08983013001852669,
      "q1_ms": 0.08964758246293059,
      "q3_ms": 0.08987806999357417,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 64.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.002256848616525531,
      "rel_err_parts": [
        0.002256848616525531
      ],
      "compile_s": 0.13034757399873342,
      "inner": 186,
      "median_ms": 0.1304983817008374,
      "q1_ms": 0.13049431721791274,
      "q3_ms": 0.1305083467897969,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 64.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.002256848616525531,
      "rel_err_parts": [
        0.002256848616525531
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.13034540999797173,
      "inner": 180,
      "median_ms": 0.13057630001437953,
      "q1_ms": 0.13054245557416128,
      "q3_ms": 0.13058161667787418,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 12.0,
      "xla_arg_mb": 64.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.002338483463972807,
      "rel_err_parts": [
        0.002338483463972807
      ],
      "rel_vs_probjax_nosplit": 0.0030901406425982714,
      "compile_s": 0.13912244800303597,
      "inner": 200,
      "median_ms": 0.11850178001623135,
      "q1_ms": 0.11837809750431916,
      "q3_ms": 0.1185047350008972,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 64.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.0022568346466869116,
      "rel_err_parts": [
        0.0022568346466869116
      ],
      "rel_vs_probjax_nosplit": 3.77187789126765e-05,
      "compile_s": 0.04335799199179746,
      "inner": 159,
      "median_ms": 0.16104019495567787,
      "q1_ms": 0.16102572952824756,
      "q3_ms": 0.16104570756894798,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_err": 0.002206334611400962,
      "rel_err_parts": [
        0.002206334611400962
      ],
      "rel_vs_probjax_nosplit": 0.0028899607714265585,
      "compile_s": 0.000248591008130461,
      "inner": 200,
      "median_ms": 0.09859255500487052,
      "q1_ms": 0.09745760751684429,
      "q3_ms": 0.09863511248113355,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 1.5234527587890625,
      "xla_arg_mb": 64.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.0021824080031365156,
      "rel_err_parts": [
        0.0021824080031365156
      ],
      "rel_vs_probjax_nosplit": 0.0026159496046602726,
      "compile_s": 0.17011290800292045,
      "inner": 200,
      "median_ms": 0.09005399995658081,
      "q1_ms": 0.08982511750218691,
      "q3_ms": 0.09008791250380455,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 64.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.002256848616525531,
      "rel_err_parts": [
        0.002256848616525531
      ],
      "compile_s": 0.1321662650007056,
      "inner": 179,
      "median_ms": 0.13035062005269918,
      "q1_ms": 0.13034037428494416,
      "q3_ms": 0.13035923742862188,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 64.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.002256848616525531,
      "rel_err_parts": [
        0.002256848616525531
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.13154087199654896,
      "inner": 200,
      "median_ms": 0.13029544497840106,
      "q1_ms": 0.1302761599799851,
      "q3_ms": 0.1303194150023046,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 12.0,
      "xla_arg_mb": 64.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.002338483463972807,
      "rel_err_parts": [
        0.002338483463972807
      ],
      "rel_vs_probjax_nosplit": 0.0030901406425982714,
      "compile_s": 0.14292045499314554,
      "inner": 200,
      "median_ms": 0.11809165494923946,
      "q1_ms": 0.11804271249275189,
      "q3_ms": 0.11813478497060714,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 64.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.0022568346466869116,
      "rel_err_parts": [
        0.0022568346466869116
      ],
      "rel_vs_probjax_nosplit": 3.77187789126765e-05,
      "compile_s": 0.04281004500808194,
      "inner": 164,
      "median_ms": 0.1603935609348743,
      "q1_ms": 0.1603324329530218,
      "q3_ms": 0.1604366890405356,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 64.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_err": 0.002206334611400962,
      "rel_err_parts": [
        0.002206334611400962
      ],
      "rel_vs_probjax_nosplit": 0.0028899607714265585,
      "compile_s": 0.00025039400497917086,
      "inner": 200,
      "median_ms": 0.09885545499855652,
      "q1_ms": 0.09751414250786183,
      "q3_ms": 0.09886960750009166,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 5.0312652587890625,
      "xla_arg_mb": 512.0625,
      "xla_out_mb": 0.0625,
      "rel_err": 0.0022383439354598522,
      "rel_err_parts": [
        0.0022383439354598522
      ],
      "rel_vs_probjax_nosplit": 0.0021074514370411634,
      "compile_s": 0.17567337999935262,
      "inner": 67,
      "median_ms": 0.6277560596979821,
      "q1_ms": 0.627534134303039,
      "q3_ms": 0.6279817088628744,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 512.0625,
      "xla_out_mb": 0.0625,
      "rel_err": 0.002278390806168318,
      "rel_err_parts": [
        0.002278390806168318
      ],
      "compile_s": 0.15153838200785685,
      "inner": 60,
      "median_ms": 0.7214823500059235,
      "q1_ms": 0.7205645331850974,
      "q3_ms": 0.7215615917327038,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 512.0625,
      "xla_out_mb": 0.0625,
      "rel_err": 0.002278390806168318,
      "rel_err_parts": [
        0.002278390806168318
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.15213276700524148,
      "inner": 59,
      "median_ms": 0.7217056100686886,
      "q1_ms": 0.7216803899964208,
      "q3_ms": 0.7217502796389084,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 4.0,
      "xla_arg_mb": 512.0625,
      "xla_out_mb": 0.0625,
      "rel_err": 0.0016591971507295966,
      "rel_err_parts": [
        0.0016591971507295966
      ],
      "rel_vs_probjax_nosplit": 0.002475152024999261,
      "compile_s": 0.079815917997621,
      "inner": 65,
      "median_ms": 0.6680169536803777,
      "q1_ms": 0.6677553769595062,
      "q3_ms": 0.6682654924448939,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 512.0625,
      "xla_out_mb": 0.0625,
      "rel_err": 0.0022775507532060146,
      "rel_err_parts": [
        0.0022775507532060146
      ],
      "rel_vs_probjax_nosplit": 0.001088160090148449,
      "compile_s": 0.04804781699203886,
      "inner": 62,
      "median_ms": 0.6373054354973587,
      "q1_ms": 0.6371826128733735,
      "q3_ms": 0.6373456613331536,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_err": 0.002279111184179783,
      "rel_err_parts": [
        0.002279111184179783
      ],
      "rel_vs_probjax_nosplit": 0.003011618973687291,
      "compile_s": 0.0013511860015569255,
      "inner": 40,
      "median_ms": 1.2652215500565944,
      "q1_ms": 1.2643312498767045,
      "q3_ms": 1.265303962645703,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 5.0312652587890625,
      "xla_arg_mb": 512.0625,
      "xla_out_mb": 0.0625,
      "rel_err": 0.0022383439354598522,
      "rel_err_parts": [
        0.0022383439354598522
      ],
      "rel_vs_probjax_nosplit": 0.0021074514370411634,
      "compile_s": 0.1727818889921764,
      "inner": 66,
      "median_ms": 0.6266624393791779,
      "q1_ms": 0.6264419469957692,
      "q3_ms": 0.626774234908887,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 512.0625,
      "xla_out_mb": 0.0625,
      "rel_err": 0.002278390806168318,
      "rel_err_parts": [
        0.002278390806168318
      ],
      "compile_s": 0.14946005900856107,
      "inner": 57,
      "median_ms": 0.7218341052066535,
      "q1_ms": 0.7216555175929036,
      "q3_ms": 0.7221378332918935,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 512.0625,
      "xla_out_mb": 0.0625,
      "rel_err": 0.002278390806168318,
      "rel_err_parts": [
        0.002278390806168318
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.15185633400687948,
      "inner": 58,
      "median_ms": 0.7217137412978979,
      "q1_ms": 0.7215807241225085,
      "q3_ms": 0.72176157770847,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 4.0,
      "xla_arg_mb": 512.0625,
      "xla_out_mb": 0.0625,
      "rel_err": 0.0016591971507295966,
      "rel_err_parts": [
        0.0016591971507295966
      ],
      "rel_vs_probjax_nosplit": 0.002475152024999261,
      "compile_s": 0.08220209300634451,
      "inner": 66,
      "median_ms": 0.6681884395110308,
      "q1_ms": 0.6680468863068557,
      "q3_ms": 0.6682887121287147,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 512.0625,
      "xla_out_mb": 0.0625,
      "rel_err": 0.0022775507532060146,
      "rel_err_parts": [
        0.0022775507532060146
      ],
      "rel_vs_probjax_nosplit": 0.001088160090148449,
      "compile_s": 0.046652596996864304,
      "inner": 63,
      "median_ms": 0.6369961269055715,
      "q1_ms": 0.6367453254550504,
      "q3_ms": 0.6370865317183503,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_err": 0.002279111184179783,
      "rel_err_parts": [
        0.002279111184179783
      ],
      "rel_vs_probjax_nosplit": 0.003011618973687291,
      "compile_s": 0.001351496990537271,
      "inner": 40,
      "median_ms": 1.2653403002332198,
      "q1_ms": 1.2652251874897047,
      "q3_ms": 1.2697536250925623,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 5.0312652587890625,
      "xla_arg_mb": 512.0625,
      "xla_out_mb": 0.0625,
      "rel_err": 0.0022383439354598522,
      "rel_err_parts": [
        0.0022383439354598522
      ],
      "rel_vs_probjax_nosplit": 0.0021074514370411634,
      "compile_s": 0.17554898500384297,
      "inner": 67,
      "median_ms": 0.6263864477714107,
      "q1_ms": 0.6262312312879766,
      "q3_ms": 0.6264104476804957,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 512.0625,
      "xla_out_mb": 0.0625,
      "rel_err": 0.002278390806168318,
      "rel_err_parts": [
        0.002278390806168318
      ],
      "compile_s": 0.15231303799373563,
      "inner": 59,
      "median_ms": 0.7219856271472901,
      "q1_ms": 0.7215760422819215,
      "q3_ms": 0.7226023135136544,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 512.0625,
      "xla_out_mb": 0.0625,
      "rel_err": 0.002278390806168318,
      "rel_err_parts": [
        0.002278390806168318
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.15075939799135085,
      "inner": 57,
      "median_ms": 0.7218991405240734,
      "q1_ms": 0.721763885960906,
      "q3_ms": 0.723763543774431,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 4.0,
      "xla_arg_mb": 512.0625,
      "xla_out_mb": 0.0625,
      "rel_err": 0.0016591971507295966,
      "rel_err_parts": [
        0.0016591971507295966
      ],
      "rel_vs_probjax_nosplit": 0.002475152024999261,
      "compile_s": 0.07843927398789674,
      "inner": 65,
      "median_ms": 0.6678561845686859,
      "q1_ms": 0.6676601999340794,
      "q3_ms": 0.6679990691526865,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 512.0625,
      "xla_out_mb": 0.0625,
      "rel_err": 0.0022775507532060146,
      "rel_err_parts": [
        0.0022775507532060146
      ],
      "rel_vs_probjax_nosplit": 0.001088160090148449,
      "compile_s": 0.04619186399213504,
      "inner": 62,
      "median_ms": 0.6374775323145573,
      "q1_ms": 0.6371799434184082,
      "q3_ms": 0.6375781209029694,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_err": 0.002279111184179783,
      "rel_err_parts": [
        0.002279111184179783
      ],
      "rel_vs_probjax_nosplit": 0.003011618973687291,
      "compile_s": 0.0013504759990610182,
      "inner": 40,
      "median_ms": 1.2654244499572087,
      "q1_ms": 1.2652650126256049,
      "q3_ms": 1.272228187553992,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 5.0000152587890625,
      "xla_arg_mb": 512.25,
      "xla_out_mb": 0.25,
      "rel_err": 0.0022445658687502146,
      "rel_err_parts": [
        0.0022445658687502146
      ],
      "rel_vs_probjax_nosplit": 0.00215164921246469,
      "compile_s": 0.17585018499812577,
      "inner": 68,
      "median_ms": 0.6289739852517788,
      "q1_ms": 0.6284593162169638,
      "q3_ms": 0.6292906765499056,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 512.25,
      "xla_out_mb": 0.25,
      "rel_err": 0.002270719036459923,
      "rel_err_parts": [
        0.002270719036459923
      ],
      "compile_s": 0.1530280810111435,
      "inner": 60,
      "median_ms": 0.7217560334538575,
      "q1_ms": 0.721182299942787,
      "q3_ms": 0.7218202583317179,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 512.25,
      "xla_out_mb": 0.25,
      "rel_err": 0.002270719036459923,
      "rel_err_parts": [
        0.002270719036459923
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.1511368009960279,
      "inner": 59,
      "median_ms": 0.7218353389130937,
      "q1_ms": 0.7215634915016798,
      "q3_ms": 0.721905915208666,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 280.0,
      "xla_arg_mb": 512.25,
      "xla_out_mb": 0.25,
      "rel_err": 0.002343224361538887,
      "rel_err_parts": [
        0.002343224361538887
      ],
      "rel_vs_probjax_nosplit": 0.003097591223195195,
      "compile_s": 0.18467717799649108,
      "inner": 23,
      "median_ms": 1.9807915219723287,
      "q1_ms": 1.980353978148702,
      "q3_ms": 1.9810958041014837,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 512.25,
      "xla_out_mb": 0.25,
      "rel_err": 0.002270706230774522,
      "rel_err_parts": [
        0.002270706230774522
      ],
      "rel_vs_probjax_nosplit": 3.7419431464513764e-05,
      "compile_s": 0.046669628005474806,
      "inner": 58,
      "median_ms": 0.7134019826229758,
      "q1_ms": 0.7132037584408957,
      "q3_ms": 0.7135563103255318,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_err": 0.0022763360757380724,
      "rel_err_parts": [
        0.0022763360757380724
      ],
      "rel_vs_probjax_nosplit": 0.003007022198289633,
      "compile_s": 0.0014007810095790774,
      "inner": 40,
      "median_ms": 1.2695126748440089,
      "q1_ms": 1.2655706124860444,
      "q3_ms": 1.272983875060163,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 5.0000152587890625,
      "xla_arg_mb": 512.25,
      "xla_out_mb": 0.25,
      "rel_err": 0.0022445658687502146,
      "rel_err_parts": [
        0.0022445658687502146
      ],
      "rel_vs_probjax_nosplit": 0.00215164921246469,
      "compile_s": 0.17581586999585852,
      "inner": 28,
      "median_ms": 0.6293687500666627,
      "q1_ms": 0.6288588571935958,
      "q3_ms": 0.6294317141250108,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 512.25,
      "xla_out_mb": 0.25,
      "rel_err": 0.002270719036459923,
      "rel_err_parts": [
        0.002270719036459923
      ],
      "compile_s": 0.15271890599979088,
      "inner": 60,
      "median_ms": 0.7224346666286389,
      "q1_ms": 0.7220668748535293,
      "q3_ms": 0.7225209917426886,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 512.25,
      "xla_out_mb": 0.25,
      "rel_err": 0.002270719036459923,
      "rel_err_parts": [
        0.002270719036459923
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.160741436004173,
      "inner": 59,
      "median_ms": 0.7224986440663905,
      "q1_ms": 0.7223730763298127,
      "q3_ms": 0.7226240507453301,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 280.0,
      "xla_arg_mb": 512.25,
      "xla_out_mb": 0.25,
      "rel_err": 0.002343224361538887,
      "rel_err_parts": [
        0.002343224361538887
      ],
      "rel_vs_probjax_nosplit": 0.003097591223195195,
      "compile_s": 0.18212839400803205,
      "inner": 23,
      "median_ms": 1.9820186520835545,
      "q1_ms": 1.9816864998732773,
      "q3_ms": 1.9827400000118043,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 512.25,
      "xla_out_mb": 0.25,
      "rel_err": 0.002270706230774522,
      "rel_err_parts": [
        0.002270706230774522
      ],
      "rel_vs_probjax_nosplit": 3.7419431464513764e-05,
      "compile_s": 0.04879536099906545,
      "inner": 59,
      "median_ms": 0.7135651355029194,
      "q1_ms": 0.7131698814376605,
      "q3_ms": 0.7137150678317994,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_err": 0.0022763360757380724,
      "rel_err_parts": [
        0.0022763360757380724
      ],
      "rel_vs_probjax_nosplit": 0.003007022198289633,
      "compile_s": 0.0014016119966981933,
      "inner": 40,
      "median_ms": 1.2688309001532616,
      "q1_ms": 1.2655990374696557,
      "q3_ms": 1.274654662302055,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 5.0000152587890625,
      "xla_arg_mb": 512.25,
      "xla_out_mb": 0.25,
      "rel_err": 0.0022445658687502146,
      "rel_err_parts": [
        0.0022445658687502146
      ],
      "rel_vs_probjax_nosplit": 0.00215164921246469,
      "compile_s": 0.1729589939932339,
      "inner": 66,
      "median_ms": 0.6276698029868192,
      "q1_ms": 0.6275789393949901,
      "q3_ms": 0.6278203105797399,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 512.25,
      "xla_out_mb": 0.25,
      "rel_err": 0.002270719036459923,
      "rel_err_parts": [
        0.002270719036459923
      ],
      "compile_s": 0.15309212201100308,
      "inner": 56,
      "median_ms": 0.7223536426213936,
      "q1_ms": 0.7221265892180131,
      "q3_ms": 0.7225038215048178,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 512.25,
      "xla_out_mb": 0.25,
      "rel_err": 0.002270719036459923,
      "rel_err_parts": [
        0.002270719036459923
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.15098820099956356,
      "inner": 58,
      "median_ms": 0.7220288103341993,
      "q1_ms": 0.7219266464807301,
      "q3_ms": 0.7220688016316452,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 280.0,
      "xla_arg_mb": 512.25,
      "xla_out_mb": 0.25,
      "rel_err": 0.002343224361538887,
      "rel_err_parts": [
        0.002343224361538887
      ],
      "rel_vs_probjax_nosplit": 0.003097591223195195,
      "compile_s": 0.1793354600085877,
      "inner": 23,
      "median_ms": 1.9808703908220986,
      "q1_ms": 1.980641478439793,
      "q3_ms": 1.9812162389153736,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 512.25,
      "xla_out_mb": 0.25,
      "rel_err": 0.002270706230774522,
      "rel_err_parts": [
        0.002270706230774522
      ],
      "rel_vs_probjax_nosplit": 3.7419431464513764e-05,
      "compile_s": 0.04547433699190151,
      "inner": 59,
      "median_ms": 0.7132623560225465,
      "q1_ms": 0.7131870338294687,
      "q3_ms": 0.7136710000695993,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_err": 0.0022763360757380724,
      "rel_err_parts": [
        0.0022763360757380724
      ],
      "rel_vs_probjax_nosplit": 0.003007022198289633,
      "compile_s": 0.0013996589987073094,
      "inner": 40,
      "median_ms": 1.2700499497441342,
      "q1_ms": 1.2658048000957933,
      "q3_ms": 1.2746979873554665,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 4.0625152587890625,
      "xla_arg_mb": 513.0,
      "xla_out_mb": 1.0,
      "rel_err": 0.0022466913796961308,
      "rel_err_parts": [
        0.0022466913796961308
      ],
      "rel_vs_probjax_nosplit": 0.0021634562872350216,
      "compile_s": 0.17313212199951522,
      "inner": 66,
      "median_ms": 0.6338458787826258,
      "q1_ms": 0.6332747272443322,
      "q3_ms": 0.6339109317967526,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 513.0,
      "xla_out_mb": 1.0,
      "rel_err": 0.002270277589559555,
      "rel_err_parts": [
        0.002270277589559555
      ],
      "compile_s": 0.13285245400038548,
      "inner": 59,
      "median_ms": 0.715054389606309,
      "q1_ms": 0.7148763305217549,
      "q3_ms": 0.7180287371985457,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 513.0,
      "xla_out_mb": 1.0,
      "rel_err": 0.002270277589559555,
      "rel_err_parts": [
        0.002270277589559555
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.1316387679980835,
      "inner": 61,
      "median_ms": 0.7152059344483204,
      "q1_ms": 0.7151221720736519,
      "q3_ms": 0.7160752950871715,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 352.0,
      "xla_arg_mb": 513.0,
      "xla_out_mb": 1.0,
      "rel_err": 0.002343859290704131,
      "rel_err_parts": [
        0.002343859290704131
      ],
      "rel_vs_probjax_nosplit": 0.003091159276664257,
      "compile_s": 0.2561418529949151,
      "inner": 21,
      "median_ms": 2.1642710004622736,
      "q1_ms": 2.163908357033506,
      "q3_ms": 2.1654806190608293,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 513.0,
      "xla_out_mb": 1.0,
      "rel_err": 0.0022702727001160383,
      "rel_err_parts": [
        0.0022702727001160383
      ],
      "rel_vs_probjax_nosplit": 3.1741088605485857e-05,
      "compile_s": 0.04648345700115897,
      "inner": 59,
      "median_ms": 0.719353355865894,
      "q1_ms": 0.7191138473730061,
      "q3_ms": 0.7194071101474623,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_err": 0.0022721898276358843,
      "rel_err_parts": [
        0.0022721898276358843
      ],
      "rel_vs_probjax_nosplit": 0.0030189254321157932,
      "compile_s": 0.0013531999866245314,
      "inner": 40,
      "median_ms": 1.2684847501077456,
      "q1_ms": 1.2672998625930632,
      "q3_ms": 1.2760080999214551,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 4.0625152587890625,
      "xla_arg_mb": 513.0,
      "xla_out_mb": 1.0,
      "rel_err": 0.0022466913796961308,
      "rel_err_parts": [
        0.0022466913796961308
      ],
      "rel_vs_probjax_nosplit": 0.0021634562872350216,
      "compile_s": 0.17189818700717296,
      "inner": 64,
      "median_ms": 0.632521390571128,
      "q1_ms": 0.632463398346772,
      "q3_ms": 0.6325534139932643,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 513.0,
      "xla_out_mb": 1.0,
      "rel_err": 0.002270277589559555,
      "rel_err_parts": [
        0.002270277589559555
      ],
      "compile_s": 0.1314419749978697,
      "inner": 58,
      "median_ms": 0.7160352241922299,
      "q1_ms": 0.7154595689920323,
      "q3_ms": 0.7160386810247819,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 513.0,
      "xla_out_mb": 1.0,
      "rel_err": 0.002270277589559555,
      "rel_err_parts": [
        0.002270277589559555
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.13201803600532003,
      "inner": 57,
      "median_ms": 0.7161499999806678,
      "q1_ms": 0.7158498860039879,
      "q3_ms": 0.716285622809409,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 352.0,
      "xla_arg_mb": 513.0,
      "xla_out_mb": 1.0,
      "rel_err": 0.002343859290704131,
      "rel_err_parts": [
        0.002343859290704131
      ],
      "rel_vs_probjax_nosplit": 0.003091159276664257,
      "compile_s": 0.2510848249949049,
      "inner": 22,
      "median_ms": 2.1634937723071994,
      "q1_ms": 2.1630410909975497,
      "q3_ms": 2.1638674320870037,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 513.0,
      "xla_out_mb": 1.0,
      "rel_err": 0.0022702727001160383,
      "rel_err_parts": [
        0.0022702727001160383
      ],
      "rel_vs_probjax_nosplit": 3.1741088605485857e-05,
      "compile_s": 0.04536389799613971,
      "inner": 60,
      "median_ms": 0.7189926499753103,
      "q1_ms": 0.7188597333879443,
      "q3_ms": 0.7193849749455694,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_err": 0.0022721898276358843,
      "rel_err_parts": [
        0.0022721898276358843
      ],
      "rel_vs_probjax_nosplit": 0.0030189254321157932,
      "compile_s": 0.0013771560043096542,
      "inner": 39,
      "median_ms": 1.2765722050785255,
      "q1_ms": 1.2762289869640835,
      "q3_ms": 1.278096499915521,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 4.0625152587890625,
      "xla_arg_mb": 513.0,
      "xla_out_mb": 1.0,
      "rel_err": 0.0022466913796961308,
      "rel_err_parts": [
        0.0022466913796961308
      ],
      "rel_vs_probjax_nosplit": 0.0021634562872350216,
      "compile_s": 0.17201996699441224,
      "inner": 66,
      "median_ms": 0.6326027726048998,
      "q1_ms": 0.6324264620409601,
      "q3_ms": 0.6326479318280083,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 513.0,
      "xla_out_mb": 1.0,
      "rel_err": 0.002270277589559555,
      "rel_err_parts": [
        0.002270277589559555
      ],
      "compile_s": 0.13138413599517662,
      "inner": 54,
      "median_ms": 0.7155767964798508,
      "q1_ms": 0.7154211296572747,
      "q3_ms": 0.7160992684442102,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 513.0,
      "xla_out_mb": 1.0,
      "rel_err": 0.002270277589559555,
      "rel_err_parts": [
        0.002270277589559555
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.13092111000150908,
      "inner": 60,
      "median_ms": 0.7151829333452042,
      "q1_ms": 0.7148705167007089,
      "q3_ms": 0.7158434333784195,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 352.0,
      "xla_arg_mb": 513.0,
      "xla_out_mb": 1.0,
      "rel_err": 0.002343859290704131,
      "rel_err_parts": [
        0.002343859290704131
      ],
      "rel_vs_probjax_nosplit": 0.003091159276664257,
      "compile_s": 0.24574053200194612,
      "inner": 21,
      "median_ms": 2.1640400477086326,
      "q1_ms": 2.1636877143118043,
      "q3_ms": 2.164305309374218,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 513.0,
      "xla_out_mb": 1.0,
      "rel_err": 0.0022702727001160383,
      "rel_err_parts": [
        0.0022702727001160383
      ],
      "rel_vs_probjax_nosplit": 3.1741088605485857e-05,
      "compile_s": 0.04676205400028266,
      "inner": 59,
      "median_ms": 0.7196928305902629,
      "q1_ms": 0.7192336441653039,
      "q3_ms": 0.7198030253382937,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 512.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_err": 0.0022721898276358843,
      "rel_err_parts": [
        0.0022721898276358843
      ],
      "rel_vs_probjax_nosplit": 0.0030189254321157932,
      "compile_s": 0.0013712549989577383,
      "inner": 40,
      "median_ms": 1.277159150049556,
      "q1_ms": 1.276611612411216,
      "q3_ms": 1.2774668626661878,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 1.6367340087890625,
      "xla_arg_mb": 256.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.002216068794950843,
      "rel_err_parts": [
        0.002216068794950843
      ],
      "rel_vs_probjax_nosplit": 0.0026591068599373102,
      "compile_s": 0.17199972899106797,
      "inner": 106,
      "median_ms": 0.31934870749641703,
      "q1_ms": 0.31828310842821284,
      "q3_ms": 0.3195596698376148,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 256.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.002309530507773161,
      "rel_err_parts": [
        0.002309530507773161
      ],
      "compile_s": 0.14904117600235622,
      "inner": 80,
      "median_ms": 0.5225795875958283,
      "q1_ms": 0.5225589125075203,
      "q3_ms": 0.5226075750215387,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 256.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.002309530507773161,
      "rel_err_parts": [
        0.002309530507773161
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.14810032599780243,
      "inner": 76,
      "median_ms": 0.52282021045501,
      "q1_ms": 0.5226316249096079,
      "q3_ms": 0.5233750066269606,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 4.0,
      "xla_arg_mb": 256.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.0016528420383110642,
      "rel_err_parts": [
        0.0016528420383110642
      ],
      "rel_vs_probjax_nosplit": 0.0024188272655010223,
      "compile_s": 0.1054277609946439,
      "inner": 116,
      "median_ms": 0.3416346293745627,
      "q1_ms": 0.3416063922611948,
      "q3_ms": 0.34165043532617534,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.0638427734375,
      "xla_arg_mb": 256.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.002230970887467265,
      "rel_err_parts": [
        0.002230970887467265
      ],
      "rel_vs_probjax_nosplit": 0.0026582214049994946,
      "compile_s": 0.046401742001762614,
      "inner": 105,
      "median_ms": 0.3183051523852295,
      "q1_ms": 0.318292176165241,
      "q3_ms": 0.3183306809487043,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_err": 0.0022304949816316366,
      "rel_err_parts": [
        0.0022304949816316366
      ],
      "rel_vs_probjax_nosplit": 0.0028869658708572388,
      "compile_s": 0.00048236300062853843,
      "inner": 142,
      "median_ms": 0.3229294295675239,
      "q1_ms": 0.322852595045205,
      "q3_ms": 0.323197123269521,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 1.6367340087890625,
      "xla_arg_mb": 256.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.002216068794950843,
      "rel_err_parts": [
        0.002216068794950843
      ],
      "rel_vs_probjax_nosplit": 0.0026591068599373102,
      "compile_s": 0.17206730700854678,
      "inner": 108,
      "median_ms": 0.3184510000951209,
      "q1_ms": 0.3183603611331718,
      "q3_ms": 0.3186713241087057,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 256.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.002309530507773161,
      "rel_err_parts": [
        0.002309530507773161
      ],
      "compile_s": 0.1457437080098316,
      "inner": 76,
      "median_ms": 0.5241073683320888,
      "q1_ms": 0.5240944539860051,
      "q3_ms": 0.5242353156637891,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 256.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.002309530507773161,
      "rel_err_parts": [
        0.002309530507773161
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.15188327500072774,
      "inner": 75,
      "median_ms": 0.5241742400297275,
      "q1_ms": 0.5241464599384926,
      "q3_ms": 0.5243149733481308,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 4.0,
      "xla_arg_mb": 256.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.0016528420383110642,
      "rel_err_parts": [
        0.0016528420383110642
      ],
      "rel_vs_probjax_nosplit": 0.0024188272655010223,
      "compile_s": 0.10956484900088981,
      "inner": 103,
      "median_ms": 0.34178769907086765,
      "q1_ms": 0.3416911165260493,
      "q3_ms": 0.34184830581986353,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.0638427734375,
      "xla_arg_mb": 256.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.002230970887467265,
      "rel_err_parts": [
        0.002230970887467265
      ],
      "rel_vs_probjax_nosplit": 0.0026582214049994946,
      "compile_s": 0.04880951800441835,
      "inner": 102,
      "median_ms": 0.31833127457717,
      "q1_ms": 0.31825828923234334,
      "q3_ms": 0.31839094606088925,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_err": 0.0022304949816316366,
      "rel_err_parts": [
        0.0022304949816316366
      ],
      "rel_vs_probjax_nosplit": 0.0028869658708572388,
      "compile_s": 0.00047411699779331684,
      "inner": 141,
      "median_ms": 0.32309895746264583,
      "q1_ms": 0.3230740532530009,
      "q3_ms": 0.3232272162889824,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 1.6367340087890625,
      "xla_arg_mb": 256.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.002216068794950843,
      "rel_err_parts": [
        0.002216068794950843
      ],
      "rel_vs_probjax_nosplit": 0.0026591068599373102,
      "compile_s": 0.17021127500629518,
      "inner": 113,
      "median_ms": 0.31858185845513753,
      "q1_ms": 0.3183326194324506,
      "q3_ms": 0.3186163406992183,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 256.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.002309530507773161,
      "rel_err_parts": [
        0.002309530507773161
      ],
      "compile_s": 0.1487997190124588,
      "inner": 75,
      "median_ms": 0.5225335333185891,
      "q1_ms": 0.5223575400304981,
      "q3_ms": 0.5226360666953649,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 256.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.002309530507773161,
      "rel_err_parts": [
        0.002309530507773161
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.14929979600128718,
      "inner": 74,
      "median_ms": 0.5220953377854155,
      "q1_ms": 0.5220099121435998,
      "q3_ms": 0.5221556553757133,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 4.0,
      "xla_arg_mb": 256.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.0016528420383110642,
      "rel_err_parts": [
        0.0016528420383110642
      ],
      "rel_vs_probjax_nosplit": 0.0024188272655010223,
      "compile_s": 0.10866215200803708,
      "inner": 110,
      "median_ms": 0.34164791815617884,
      "q1_ms": 0.3416048863569317,
      "q3_ms": 0.3417102681520522,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.0638427734375,
      "xla_arg_mb": 256.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.002230970887467265,
      "rel_err_parts": [
        0.002230970887467265
      ],
      "rel_vs_probjax_nosplit": 0.0026582214049994946,
      "compile_s": 0.04771459600306116,
      "inner": 108,
      "median_ms": 0.3179703704113814,
      "q1_ms": 0.31792987499135994,
      "q3_ms": 0.3180083147998192,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_err": 0.0022304949816316366,
      "rel_err_parts": [
        0.0022304949816316366
      ],
      "rel_vs_probjax_nosplit": 0.0028869658708572388,
      "compile_s": 0.00047912599984556437,
      "inner": 143,
      "median_ms": 0.32298086708496243,
      "q1_ms": 0.3229056538320374,
      "q3_ms": 0.32300237059256165,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 1.6250152587890625,
      "xla_arg_mb": 256.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.0022407067008316517,
      "rel_err_parts": [
        0.0022407067008316517
      ],
      "rel_vs_probjax_nosplit": 0.0027165981009602547,
      "compile_s": 0.17709462900529616,
      "inner": 116,
      "median_ms": 0.3179505948882534,
      "q1_ms": 0.31735839654394454,
      "q3_ms": 0.31811906040369153,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 256.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.002314816229045391,
      "rel_err_parts": [
        0.002314816229045391
      ],
      "compile_s": 0.14874441501160618,
      "inner": 79,
      "median_ms": 0.522471924019164,
      "q1_ms": 0.5221412278773687,
      "q3_ms": 0.5227520696625786,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 256.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.002314816229045391,
      "rel_err_parts": [
        0.002314816229045391
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.15102703499724157,
      "inner": 75,
      "median_ms": 0.5225107066022853,
      "q1_ms": 0.5223394399702859,
      "q3_ms": 0.5227715199968468,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 12.0,
      "xla_arg_mb": 256.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.0023422695230692625,
      "rel_err_parts": [
        0.0023422695230692625
      ],
      "rel_vs_probjax_nosplit": 0.0030810327734798193,
      "compile_s": 0.24412003600446042,
      "inner": 102,
      "median_ms": 0.3523773627296346,
      "q1_ms": 0.3523644411995751,
      "q3_ms": 0.3524311863407539,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 256.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.002314816229045391,
      "rel_err_parts": [
        0.002314816229045391
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.045047218998661265,
      "inner": 64,
      "median_ms": 0.6182863280628226,
      "q1_ms": 0.6180808670706028,
      "q3_ms": 0.6184175937278269,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_err": 0.0022560283541679382,
      "rel_err_parts": [
        0.0022560283541679382
      ],
      "rel_vs_probjax_nosplit": 0.003015496302396059,
      "compile_s": 0.0004789460072061047,
      "inner": 142,
      "median_ms": 0.32373489436945996,
      "q1_ms": 0.3237000810130808,
      "q3_ms": 0.3238223521485695,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 1.6250152587890625,
      "xla_arg_mb": 256.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.0022407067008316517,
      "rel_err_parts": [
        0.0022407067008316517
      ],
      "rel_vs_probjax_nosplit": 0.0027165981009602547,
      "compile_s": 0.1725497599982191,
      "inner": 121,
      "median_ms": 0.31685807430935825,
      "q1_ms": 0.31683521902319506,
      "q3_ms": 0.31691404957418556,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 256.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.002314816229045391,
      "rel_err_parts": [
        0.002314816229045391
      ],
      "compile_s": 0.1502165510028135,
      "inner": 78,
      "median_ms": 0.5221018974015561,
      "q1_ms": 0.5219661282190384,
      "q3_ms": 0.5222825000087958,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 256.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.002314816229045391,
      "rel_err_parts": [
        0.002314816229045391
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.1475675089895958,
      "inner": 73,
      "median_ms": 0.5222615069862812,
      "q1_ms": 0.5221981026969008,
      "q3_ms": 0.5223922466752375,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 12.0,
      "xla_arg_mb": 256.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.0023422695230692625,
      "rel_err_parts": [
        0.0023422695230692625
      ],
      "rel_vs_probjax_nosplit": 0.0030810327734798193,
      "compile_s": 0.24940059801156167,
      "inner": 102,
      "median_ms": 0.3523950393119936,
      "q1_ms": 0.35238316181051893,
      "q3_ms": 0.3524315784578029,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 256.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.002314816229045391,
      "rel_err_parts": [
        0.002314816229045391
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.04661646799650043,
      "inner": 64,
      "median_ms": 0.618938656316459,
      "q1_ms": 0.6188448124930801,
      "q3_ms": 0.6191281561314099,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_err": 0.0022560283541679382,
      "rel_err_parts": [
        0.0022560283541679382
      ],
      "rel_vs_probjax_nosplit": 0.003015496302396059,
      "compile_s": 0.0004778750007972121,
      "inner": 144,
      "median_ms": 0.3237836389315716,
      "q1_ms": 0.32370738194408155,
      "q3_ms": 0.32379553476857836,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 1.6250152587890625,
      "xla_arg_mb": 256.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.0022407067008316517,
      "rel_err_parts": [
        0.0022407067008316517
      ],
      "rel_vs_probjax_nosplit": 0.0027165981009602547,
      "compile_s": 0.17323612800100818,
      "inner": 110,
      "median_ms": 0.31725849998606875,
      "q1_ms": 0.3170961045808243,
      "q3_ms": 0.31730613634201954,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 256.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.002314816229045391,
      "rel_err_parts": [
        0.002314816229045391
      ],
      "compile_s": 0.1543077620008262,
      "inner": 74,
      "median_ms": 0.5226940405476411,
      "q1_ms": 0.522588432480965,
      "q3_ms": 0.523273249937857,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 256.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.002314816229045391,
      "rel_err_parts": [
        0.002314816229045391
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.14728441199986264,
      "inner": 77,
      "median_ms": 0.522398090900176,
      "q1_ms": 0.522319376661424,
      "q3_ms": 0.5227464220793495,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 12.0,
      "xla_arg_mb": 256.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.0023422695230692625,
      "rel_err_parts": [
        0.0023422695230692625
      ],
      "rel_vs_probjax_nosplit": 0.0030810327734798193,
      "compile_s": 0.24840856999799144,
      "inner": 105,
      "median_ms": 0.35242553331911386,
      "q1_ms": 0.3523952333905202,
      "q3_ms": 0.3524496666310976,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 256.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.002314816229045391,
      "rel_err_parts": [
        0.002314816229045391
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.04654046399809886,
      "inner": 64,
      "median_ms": 0.6182417187119427,
      "q1_ms": 0.6182345078968865,
      "q3_ms": 0.6183313438441473,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_err": 0.0022560283541679382,
      "rel_err_parts": [
        0.0022560283541679382
      ],
      "rel_vs_probjax_nosplit": 0.003015496302396059,
      "compile_s": 0.0004760510055348277,
      "inner": 141,
      "median_ms": 0.32393237588686064,
      "q1_ms": 0.32375033328504504,
      "q3_ms": 0.3240619858224044,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 1.5234527587890625,
      "xla_arg_mb": 256.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.0022548669949173927,
      "rel_err_parts": [
        0.0022548669949173927
      ],
      "rel_vs_probjax_nosplit": 0.002735789632424712,
      "compile_s": 0.17088172299554572,
      "inner": 117,
      "median_ms": 0.3183186495762207,
      "q1_ms": 0.3173383760378441,
      "q3_ms": 0.31839314953654874,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 256.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.002306178444996476,
      "rel_err_parts": [
        0.002306178444996476
      ],
      "compile_s": 0.13348364901321474,
      "inner": 79,
      "median_ms": 0.4994675063578223,
      "q1_ms": 0.4993807532731738,
      "q3_ms": 0.49953198101286056,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 256.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.002306178444996476,
      "rel_err_parts": [
        0.002306178444996476
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.12959200599289034,
      "inner": 78,
      "median_ms": 0.499568102569379,
      "q1_ms": 0.4994711923175671,
      "q3_ms": 0.4996219871920915,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 48.0,
      "xla_arg_mb": 256.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.0023478225339204073,
      "rel_err_parts": [
        0.0023478225339204073
      ],
      "rel_vs_probjax_nosplit": 0.003126756055280566,
      "compile_s": 0.25047541099775117,
      "inner": 84,
      "median_ms": 0.4500580237488196,
      "q1_ms": 0.44994591068340617,
      "q3_ms": 0.45007126798390384,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 256.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.002306221751496196,
      "rel_err_parts": [
        0.002306221751496196
      ],
      "rel_vs_probjax_nosplit": 1.8592370906844735e-05,
      "compile_s": 0.04476420300488826,
      "inner": 64,
      "median_ms": 0.618577203113091,
      "q1_ms": 0.6185742108755221,
      "q3_ms": 0.6186921015114422,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_err": 0.002259491477161646,
      "rel_err_parts": [
        0.002259491477161646
      ],
      "rel_vs_probjax_nosplit": 0.0030235869344323874,
      "compile_s": 0.00048301399510819465,
      "inner": 141,
      "median_ms": 0.326351709197381,
      "q1_ms": 0.3263326312236963,
      "q3_ms": 0.32642212410908233,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 1.5234527587890625,
      "xla_arg_mb": 256.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.0022548669949173927,
      "rel_err_parts": [
        0.0022548669949173927
      ],
      "rel_vs_probjax_nosplit": 0.002735789632424712,
      "compile_s": 0.17592819199489895,
      "inner": 116,
      "median_ms": 0.31739091380210277,
      "q1_ms": 0.31725612932169844,
      "q3_ms": 0.3174239526084295,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 256.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.002306178444996476,
      "rel_err_parts": [
        0.002306178444996476
      ],
      "compile_s": 0.13032912000198849,
      "inner": 76,
      "median_ms": 0.5025452632444763,
      "q1_ms": 0.5024406578549563,
      "q3_ms": 0.5027192697812815,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 256.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.002306178444996476,
      "rel_err_parts": [
        0.002306178444996476
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.12872453300224151,
      "inner": 78,
      "median_ms": 0.5027176666771993,
      "q1_ms": 0.5026545898148563,
      "q3_ms": 0.5027842564278473,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 48.0,
      "xla_arg_mb": 256.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.0023478225339204073,
      "rel_err_parts": [
        0.0023478225339204073
      ],
      "rel_vs_probjax_nosplit": 0.003126756055280566,
      "compile_s": 0.2509042430028785,
      "inner": 82,
      "median_ms": 0.4495169146471975,
      "q1_ms": 0.44945631710040135,
      "q3_ms": 0.449556878110587,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 256.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.002306221751496196,
      "rel_err_parts": [
        0.002306221751496196
      ],
      "rel_vs_probjax_nosplit": 1.8592370906844735e-05,
      "compile_s": 0.04453252500388771,
      "inner": 64,
      "median_ms": 0.6181459218623786,
      "q1_ms": 0.618050257799041,
      "q3_ms": 0.6181574141237434,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_err": 0.002259491477161646,
      "rel_err_parts": [
        0.002259491477161646
      ],
      "rel_vs_probjax_nosplit": 0.0030235869344323874,
      "compile_s": 0.00048228200466837734,
      "inner": 141,
      "median_ms": 0.32628790070038577,
      "q1_ms": 0.32622192551609763,
      "q3_ms": 0.32644830852604956,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 1.5234527587890625,
      "xla_arg_mb": 256.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.0022548669949173927,
      "rel_err_parts": [
        0.0022548669949173927
      ],
      "rel_vs_probjax_nosplit": 0.002735789632424712,
      "compile_s": 0.17029741799342446,
      "inner": 112,
      "median_ms": 0.3174325268316482,
      "q1_ms": 0.31713392412971836,
      "q3_ms": 0.3174740758627844,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 256.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.002306178444996476,
      "rel_err_parts": [
        0.002306178444996476
      ],
      "compile_s": 0.13269622799998615,
      "inner": 83,
      "median_ms": 0.5019657229317397,
      "q1_ms": 0.501872662644895,
      "q3_ms": 0.5020146747856087,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 256.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.002306178444996476,
      "rel_err_parts": [
        0.002306178444996476
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.13266910701349843,
      "inner": 77,
      "median_ms": 0.5019292467449151,
      "q1_ms": 0.5018165518942459,
      "q3_ms": 0.501955000044137,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 48.0,
      "xla_arg_mb": 256.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.0023478225339204073,
      "rel_err_parts": [
        0.0023478225339204073
      ],
      "rel_vs_probjax_nosplit": 0.003126756055280566,
      "compile_s": 0.24727345300198067,
      "inner": 82,
      "median_ms": 0.44956517075018093,
      "q1_ms": 0.4495111097639892,
      "q3_ms": 0.4497380670239558,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 256.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.002306221751496196,
      "rel_err_parts": [
        0.002306221751496196
      ],
      "rel_vs_probjax_nosplit": 1.8592370906844735e-05,
      "compile_s": 0.04419124900596216,
      "inner": 63,
      "median_ms": 0.6193314285834306,
      "q1_ms": 0.619156730184812,
      "q3_ms": 0.6196939445097028,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 256.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_err": 0.002259491477161646,
      "rel_err_parts": [
        0.002259491477161646
      ],
      "rel_vs_probjax_nosplit": 0.0030235869344323874,
      "compile_s": 0.0004920710052829236,
      "inner": 140,
      "median_ms": 0.32611922849485253,
      "q1_ms": 0.32609468216833193,
      "q3_ms": 0.32624818933462457,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 5.0312652587890625,
      "xla_arg_mb": 2048.0625,
      "xla_out_mb": 0.0625,
      "rel_vs_probjax_nosplit": 0.00219778623431921,
      "compile_s": 0.17550512199522927,
      "inner": 19,
      "median_ms": 2.4484313685924897,
      "q1_ms": 2.447743210436066,
      "q3_ms": 2.449971105593967,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 2048.0625,
      "xla_out_mb": 0.0625,
      "compile_s": 0.15501906799909193,
      "inner": 16,
      "median_ms": 2.855898374946264,
      "q1_ms": 2.8504135621005844,
      "q3_ms": 2.857596530702722,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 2048.0625,
      "xla_out_mb": 0.0625,
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.1570467749988893,
      "inner": 16,
      "median_ms": 2.8510761248981,
      "q1_ms": 2.8467635625020193,
      "q3_ms": 2.861793531337753,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 32.0,
      "xla_arg_mb": 2048.0625,
      "xla_out_mb": 0.0625,
      "rel_vs_probjax_nosplit": 0.0024736658670008183,
      "compile_s": 0.10504139000840951,
      "inner": 17,
      "median_ms": 2.644986058893504,
      "q1_ms": 2.643759058949346,
      "q3_ms": 2.6560112939264133,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 2048.0625,
      "xla_out_mb": 0.0625,
      "rel_vs_probjax_nosplit": 0.0008966828463599086,
      "compile_s": 0.05006803901051171,
      "inner": 17,
      "median_ms": 2.5190395296222587,
      "q1_ms": 2.51735658821417,
      "q3_ms": 2.5217045292854037,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_vs_probjax_nosplit": 0.0030787340365350246,
      "compile_s": 0.004930798997520469,
      "inner": 10,
      "median_ms": 5.041015600727405,
      "q1_ms": 5.04053729964653,
      "q3_ms": 5.064720500377007,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 5.0312652587890625,
      "xla_arg_mb": 2048.0625,
      "xla_out_mb": 0.0625,
      "rel_vs_probjax_nosplit": 0.00219778623431921,
      "compile_s": 0.1740890220098663,
      "inner": 19,
      "median_ms": 2.4479398945052373,
      "q1_ms": 2.447189526145322,
      "q3_ms": 2.4483878418993155,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 2048.0625,
      "xla_out_mb": 0.0625,
      "compile_s": 0.15258082399668638,
      "inner": 16,
      "median_ms": 2.852385437108751,
      "q1_ms": 2.8479924062594364,
      "q3_ms": 2.863447656181961,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 2048.0625,
      "xla_out_mb": 0.0625,
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.1551328539935639,
      "inner": 16,
      "median_ms": 2.855649187040399,
      "q1_ms": 2.853504406175489,
      "q3_ms": 2.8592925623343035,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 32.0,
      "xla_arg_mb": 2048.0625,
      "xla_out_mb": 0.0625,
      "rel_vs_probjax_nosplit": 0.0024736658670008183,
      "compile_s": 0.10488523400272243,
      "inner": 17,
      "median_ms": 2.64384741175865,
      "q1_ms": 2.6434620004318488,
      "q3_ms": 2.650185264321729,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 2048.0625,
      "xla_out_mb": 0.0625,
      "rel_vs_probjax_nosplit": 0.0008966828463599086,
      "compile_s": 0.05194560199743137,
      "inner": 18,
      "median_ms": 2.5180579994210146,
      "q1_ms": 2.5170285831942847,
      "q3_ms": 2.52032513890299,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_vs_probjax_nosplit": 0.0030787340365350246,
      "compile_s": 0.004935379009111784,
      "inner": 10,
      "median_ms": 5.052898099529557,
      "q1_ms": 5.046753999340581,
      "q3_ms": 5.065834649576573,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 5.0312652587890625,
      "xla_arg_mb": 2048.0625,
      "xla_out_mb": 0.0625,
      "rel_vs_probjax_nosplit": 0.00219778623431921,
      "compile_s": 0.17322050800430588,
      "inner": 19,
      "median_ms": 2.447633052820732,
      "q1_ms": 2.4457903158403687,
      "q3_ms": 2.449809737108329,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 2048.0625,
      "xla_out_mb": 0.0625,
      "compile_s": 0.15762828500010073,
      "inner": 16,
      "median_ms": 2.8524105000542477,
      "q1_ms": 2.8506208436738234,
      "q3_ms": 2.8542921559164824,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 2048.0625,
      "xla_out_mb": 0.0625,
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.15478624698880594,
      "inner": 16,
      "median_ms": 2.852417436770338,
      "q1_ms": 2.8511693749351252,
      "q3_ms": 2.860915343717352,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 32.0,
      "xla_arg_mb": 2048.0625,
      "xla_out_mb": 0.0625,
      "rel_vs_probjax_nosplit": 0.0024736658670008183,
      "compile_s": 0.09901842400722671,
      "inner": 18,
      "median_ms": 2.644142166698455,
      "q1_ms": 2.6435451945872046,
      "q3_ms": 2.645966999908625,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 2048.0625,
      "xla_out_mb": 0.0625,
      "rel_vs_probjax_nosplit": 0.0008966828463599086,
      "compile_s": 0.050389917989377864,
      "inner": 18,
      "median_ms": 2.520405777937008,
      "q1_ms": 2.5190237223190097,
      "q3_ms": 2.5244826942475305,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_vs_probjax_nosplit": 0.0030787340365350246,
      "compile_s": 0.00497354099934455,
      "inner": 10,
      "median_ms": 5.080150699359365,
      "q1_ms": 5.078716949356021,
      "q3_ms": 5.0813655005185865,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 5.0000152587890625,
      "xla_arg_mb": 2048.25,
      "xla_out_mb": 0.25,
      "rel_vs_probjax_nosplit": 0.0021878161933273077,
      "compile_s": 0.17910540499724448,
      "inner": 19,
      "median_ms": 2.4449568422584744,
      "q1_ms": 2.444209710550862,
      "q3_ms": 2.4483229736736236,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 2048.25,
      "xla_out_mb": 0.25,
      "compile_s": 0.15252881599008106,
      "inner": 16,
      "median_ms": 2.844293812813703,
      "q1_ms": 2.843311061951681,
      "q3_ms": 2.844982968781551,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 2048.25,
      "xla_out_mb": 0.25,
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.15386193899030332,
      "inner": 16,
      "median_ms": 2.844653250576812,
      "q1_ms": 2.842456656253489,
      "q3_ms": 2.847219406248769,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 1120.0,
      "xla_arg_mb": 2048.25,
      "xla_out_mb": 0.25,
      "rel_vs_probjax_nosplit": 0.003116343403235078,
      "compile_s": 0.30000447999918833,
      "inner": 6,
      "median_ms": 8.046617333699638,
      "q1_ms": 8.04324374985299,
      "q3_ms": 8.057803749882927,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 2048.25,
      "xla_out_mb": 0.25,
      "rel_vs_probjax_nosplit": 2.6179977794527076e-05,
      "compile_s": 0.04794401000253856,
      "inner": 16,
      "median_ms": 2.751318500486377,
      "q1_ms": 2.7259756247985933,
      "q3_ms": 2.763794031125144,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_vs_probjax_nosplit": 0.0030772541649639606,
      "compile_s": 0.00494828300725203,
      "inner": 10,
      "median_ms": 5.078893399331719,
      "q1_ms": 5.078139900433598,
      "q3_ms": 5.089019949809881,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 5.0000152587890625,
      "xla_arg_mb": 2048.25,
      "xla_out_mb": 0.25,
      "rel_vs_probjax_nosplit": 0.0021878161933273077,
      "compile_s": 0.17393075399741065,
      "inner": 19,
      "median_ms": 2.4560753158048576,
      "q1_ms": 2.4469113685834025,
      "q3_ms": 2.45785026345402,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 2048.25,
      "xla_out_mb": 0.25,
      "compile_s": 0.15562998500536196,
      "inner": 16,
      "median_ms": 2.8459413742893958,
      "q1_ms": 2.8451345001485606,
      "q3_ms": 2.8488405941970996,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 2048.25,
      "xla_out_mb": 0.25,
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.15222610399359837,
      "inner": 16,
      "median_ms": 2.8484523127190187,
      "q1_ms": 2.847095093329699,
      "q3_ms": 2.853355062597984,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 1120.0,
      "xla_arg_mb": 2048.25,
      "xla_out_mb": 0.25,
      "rel_vs_probjax_nosplit": 0.003116343403235078,
      "compile_s": 0.29482115199789405,
      "inner": 6,
      "median_ms": 8.065269665773181,
      "q1_ms": 8.062131166904388,
      "q3_ms": 8.069099834150013,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 2048.25,
      "xla_out_mb": 0.25,
      "rel_vs_probjax_nosplit": 2.6179977794527076e-05,
      "compile_s": 0.05143035700893961,
      "inner": 16,
      "median_ms": 2.7480242497404106,
      "q1_ms": 2.729768406425137,
      "q3_ms": 2.760001250408095,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_vs_probjax_nosplit": 0.0030772541649639606,
      "compile_s": 0.004966477004927583,
      "inner": 10,
      "median_ms": 5.078809200495016,
      "q1_ms": 5.0779450495610945,
      "q3_ms": 5.080088550312212,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 5.0000152587890625,
      "xla_arg_mb": 2048.25,
      "xla_out_mb": 0.25,
      "rel_vs_probjax_nosplit": 0.0021878161933273077,
      "compile_s": 0.1773518659902038,
      "inner": 19,
      "median_ms": 2.4453417894332423,
      "q1_ms": 2.445024131846271,
      "q3_ms": 2.4524054740475933,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 2048.25,
      "xla_out_mb": 0.25,
      "compile_s": 0.15496098800213076,
      "inner": 16,
      "median_ms": 2.8488230618677335,
      "q1_ms": 2.848441436981375,
      "q3_ms": 2.850702249816095,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 2048.25,
      "xla_out_mb": 0.25,
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.15638725699682254,
      "inner": 16,
      "median_ms": 2.849516250535089,
      "q1_ms": 2.847801155894558,
      "q3_ms": 2.865935781301232,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 1120.0,
      "xla_arg_mb": 2048.25,
      "xla_out_mb": 0.25,
      "rel_vs_probjax_nosplit": 0.003116343403235078,
      "compile_s": 0.2986065350123681,
      "inner": 6,
      "median_ms": 8.069507166510448,
      "q1_ms": 8.067016832986457,
      "q3_ms": 8.076851667283336,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 2048.25,
      "xla_out_mb": 0.25,
      "rel_vs_probjax_nosplit": 2.6179977794527076e-05,
      "compile_s": 0.048304281997843646,
      "inner": 16,
      "median_ms": 2.7558358751775813,
      "q1_ms": 2.732483219006099,
      "q3_ms": 2.774311124539963,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_vs_probjax_nosplit": 0.0030772541649639606,
      "compile_s": 0.0050090480071958154,
      "inner": 10,
      "median_ms": 5.081961100222543,
      "q1_ms": 5.079214400029741,
      "q3_ms": 5.083938350435346,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 4.0625152587890625,
      "xla_arg_mb": 2049.0,
      "xla_out_mb": 1.0,
      "rel_vs_probjax_nosplit": 0.0022026244550943375,
      "compile_s": 0.17653765600698534,
      "inner": 19,
      "median_ms": 2.454402684203447,
      "q1_ms": 2.453697105106824,
      "q3_ms": 2.4558654208943294,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 2049.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.13871046899294015,
      "inner": 16,
      "median_ms": 2.843973249582632,
      "q1_ms": 2.8414763123691955,
      "q3_ms": 2.850016906450037,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 2049.0,
      "xla_out_mb": 1.0,
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.1335536209953716,
      "inner": 17,
      "median_ms": 2.8442601769176474,
      "q1_ms": 2.8392185588874033,
      "q3_ms": 2.8489222356310004,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 1408.0,
      "xla_arg_mb": 2049.0,
      "xla_out_mb": 1.0,
      "rel_vs_probjax_nosplit": 0.003111846512183547,
      "compile_s": 0.313228433005861,
      "inner": 5,
      "median_ms": 8.77612619951833,
      "q1_ms": 8.761170999787282,
      "q3_ms": 8.782519200758543,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 2049.0,
      "xla_out_mb": 1.0,
      "rel_vs_probjax_nosplit": 2.6732266633189283e-05,
      "compile_s": 0.04757613499532454,
      "inner": 16,
      "median_ms": 2.716193999731331,
      "q1_ms": 2.708951500153489,
      "q3_ms": 2.7304926252327277,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_vs_probjax_nosplit": 0.003072373103350401,
      "compile_s": 0.004975924995960668,
      "inner": 10,
      "median_ms": 5.080296999949496,
      "q1_ms": 5.074318149854662,
      "q3_ms": 5.080549500416964,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 4.0625152587890625,
      "xla_arg_mb": 2049.0,
      "xla_out_mb": 1.0,
      "rel_vs_probjax_nosplit": 0.0022026244550943375,
      "compile_s": 0.17700545099796727,
      "inner": 19,
      "median_ms": 2.455509999938505,
      "q1_ms": 2.4542906053523583,
      "q3_ms": 2.4558058415239334,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 2049.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.1342741050029872,
      "inner": 17,
      "median_ms": 2.84659223518033,
      "q1_ms": 2.845916529408629,
      "q3_ms": 2.8510742644464377,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 2049.0,
      "xla_out_mb": 1.0,
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.1372841599950334,
      "inner": 16,
      "median_ms": 2.8497629373305244,
      "q1_ms": 2.844704937615461,
      "q3_ms": 2.852436749435583,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 1408.0,
      "xla_arg_mb": 2049.0,
      "xla_out_mb": 1.0,
      "rel_vs_probjax_nosplit": 0.003111846512183547,
      "compile_s": 0.31153267400804907,
      "inner": 5,
      "median_ms": 8.784307600581087,
      "q1_ms": 8.76849280030001,
      "q3_ms": 8.790057599253487,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 2049.0,
      "xla_out_mb": 1.0,
      "rel_vs_probjax_nosplit": 2.6732266633189283e-05,
      "compile_s": 0.047533042990835384,
      "inner": 17,
      "median_ms": 2.7171885885763913,
      "q1_ms": 2.716903293998364,
      "q3_ms": 2.7219390296308763,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_vs_probjax_nosplit": 0.003072373103350401,
      "compile_s": 0.004980192999937572,
      "inner": 10,
      "median_ms": 5.0848065002355725,
      "q1_ms": 5.0827441504225135,
      "q3_ms": 5.085401149699464,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 4.0625152587890625,
      "xla_arg_mb": 2049.0,
      "xla_out_mb": 1.0,
      "rel_vs_probjax_nosplit": 0.0022026244550943375,
      "compile_s": 0.17415320398868062,
      "inner": 19,
      "median_ms": 2.4546557894944656,
      "q1_ms": 2.454400052604462,
      "q3_ms": 2.4560552368424586,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 2049.0,
      "xla_out_mb": 1.0,
      "compile_s": 0.13735493501008023,
      "inner": 16,
      "median_ms": 2.8495425003711716,
      "q1_ms": 2.8484764688982978,
      "q3_ms": 2.8509301564554335,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 2049.0,
      "xla_out_mb": 1.0,
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.1334199379925849,
      "inner": 17,
      "median_ms": 2.848299588282209,
      "q1_ms": 2.8439268824192423,
      "q3_ms": 2.8599870000856322,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 1408.0,
      "xla_arg_mb": 2049.0,
      "xla_out_mb": 1.0,
      "rel_vs_probjax_nosplit": 0.003111846512183547,
      "compile_s": 0.31093388100271113,
      "inner": 5,
      "median_ms": 8.765135399880819,
      "q1_ms": 8.760471601272002,
      "q3_ms": 8.77015590085648,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 2049.0,
      "xla_out_mb": 1.0,
      "rel_vs_probjax_nosplit": 2.6732266633189283e-05,
      "compile_s": 0.049199585992028005,
      "inner": 16,
      "median_ms": 2.7294471256027464,
      "q1_ms": 2.7133182184115867,
      "q3_ms": 2.732373968228785,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 2048.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_vs_probjax_nosplit": 0.003072373103350401,
      "compile_s": 0.0049807640025392175,
      "inner": 10,
      "median_ms": 5.0826564998715185,
      "q1_ms": 5.080304999137297,
      "q3_ms": 5.08519775030436,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 1024.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 1.6367340087890625,
      "xla_arg_mb": 1024.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.00224181916564703,
      "rel_err_parts": [
        0.00224181916564703
      ],
      "rel_vs_probjax_nosplit": 0.002772044623270631,
      "compile_s": 0.17405770299956203,
      "inner": 37,
      "median_ms": 1.227495351387817,
      "q1_ms": 1.2270247701766297,
      "q3_ms": 1.227865945799528,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 1024.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 1024.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.002282705856487155,
      "rel_err_parts": [
        0.002282705856487155
      ],
      "compile_s": 0.15223992998653557,
      "inner": 22,
      "median_ms": 2.057944954272402,
      "q1_ms": 2.0566065226872032,
      "q3_ms": 2.062499500342793,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 1024.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 1024.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.002282705856487155,
      "rel_err_parts": [
        0.002282705856487155
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.1496883509971667,
      "inner": 23,
      "median_ms": 2.056791869108565,
      "q1_ms": 2.0564039783177734,
      "q3_ms": 2.057021434745326,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 1024.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 12.0,
      "xla_arg_mb": 1024.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.0016560620861127973,
      "rel_err_parts": [
        0.0016560620861127973
      ],
      "rel_vs_probjax_nosplit": 0.0024661370553076267,
      "compile_s": 0.0896057030040538,
      "inner": 33,
      "median_ms": 1.3852821513975153,
      "q1_ms": 1.384920409085074,
      "q3_ms": 1.3867499393668505,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 1024.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.0638427734375,
      "xla_arg_mb": 1024.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.0022924216464161873,
      "rel_err_parts": [
        0.0022924216464161873
      ],
      "rel_vs_probjax_nosplit": 0.002760650124400854,
      "compile_s": 0.04848621699784417,
      "inner": 32,
      "median_ms": 1.2359333122731186,
      "q1_ms": 1.235400453197144,
      "q3_ms": 1.2367484530386719,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 1024.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_err": 0.002305536298081279,
      "rel_err_parts": [
        0.002305536298081279
      ],
      "rel_vs_probjax_nosplit": 0.0030674005392938852,
      "compile_s": 0.001386733987601474,
      "inner": 39,
      "median_ms": 1.2342172054699263,
      "q1_ms": 1.2334594743711373,
      "q3_ms": 1.23442206416094,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 1024.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 1.6367340087890625,
      "xla_arg_mb": 1024.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.00224181916564703,
      "rel_err_parts": [
        0.00224181916564703
      ],
      "rel_vs_probjax_nosplit": 0.002772044623270631,
      "compile_s": 0.16882524199900217,
      "inner": 36,
      "median_ms": 1.2282179446224797,
      "q1_ms": 1.2279520138286495,
      "q3_ms": 1.2287081667131536,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 1024.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 1024.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.002282705856487155,
      "rel_err_parts": [
        0.002282705856487155
      ],
      "compile_s": 0.14834389700263273,
      "inner": 19,
      "median_ms": 2.0653580523733246,
      "q1_ms": 2.0612452897598574,
      "q3_ms": 2.0689820261108443,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 1024.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 1024.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.002282705856487155,
      "rel_err_parts": [
        0.002282705856487155
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.14835730199411046,
      "inner": 22,
      "median_ms": 2.060180545165415,
      "q1_ms": 2.060099022690503,
      "q3_ms": 2.0605334998849267,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 1024.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 12.0,
      "xla_arg_mb": 1024.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.0016560620861127973,
      "rel_err_parts": [
        0.0016560620861127973
      ],
      "rel_vs_probjax_nosplit": 0.0024661370553076267,
      "compile_s": 0.0854266859969357,
      "inner": 33,
      "median_ms": 1.3855383938679358,
      "q1_ms": 1.384859817996833,
      "q3_ms": 1.3858310758774732,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 1024.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.0638427734375,
      "xla_arg_mb": 1024.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.0022924216464161873,
      "rel_err_parts": [
        0.0022924216464161873
      ],
      "rel_vs_probjax_nosplit": 0.002760650124400854,
      "compile_s": 0.048296687004039995,
      "inner": 35,
      "median_ms": 1.2339815139837031,
      "q1_ms": 1.2331257572181389,
      "q3_ms": 1.2348793573827215,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 1024.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_err": 0.002305536298081279,
      "rel_err_parts": [
        0.002305536298081279
      ],
      "rel_vs_probjax_nosplit": 0.0030674005392938852,
      "compile_s": 0.0013930759887443855,
      "inner": 39,
      "median_ms": 1.2342600769046934,
      "q1_ms": 1.2338236155277953,
      "q3_ms": 1.2343014616187875,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 1024.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 1.6367340087890625,
      "xla_arg_mb": 1024.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.00224181916564703,
      "rel_err_parts": [
        0.00224181916564703
      ],
      "rel_vs_probjax_nosplit": 0.002772044623270631,
      "compile_s": 0.17598352700588293,
      "inner": 36,
      "median_ms": 1.2279552223238472,
      "q1_ms": 1.2275686666624905,
      "q3_ms": 1.2281679723350356,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 1024.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 1024.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.002282705856487155,
      "rel_err_parts": [
        0.002282705856487155
      ],
      "compile_s": 0.15170886499981862,
      "inner": 22,
      "median_ms": 2.073721636076119,
      "q1_ms": 2.073102954454953,
      "q3_ms": 2.0739120002624323,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 1024.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 1024.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.002282705856487155,
      "rel_err_parts": [
        0.002282705856487155
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.14839742900221609,
      "inner": 22,
      "median_ms": 2.060417818111918,
      "q1_ms": 2.0600944772013463,
      "q3_ms": 2.061179272542094,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 1024.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 12.0,
      "xla_arg_mb": 1024.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.0016560620861127973,
      "rel_err_parts": [
        0.0016560620861127973
      ],
      "rel_vs_probjax_nosplit": 0.0024661370553076267,
      "compile_s": 0.08620829400024377,
      "inner": 33,
      "median_ms": 1.3866835757984188,
      "q1_ms": 1.3863156363544422,
      "q3_ms": 1.3870899697808041,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 1024.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.0638427734375,
      "xla_arg_mb": 1024.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.0022924216464161873,
      "rel_err_parts": [
        0.0022924216464161873
      ],
      "rel_vs_probjax_nosplit": 0.002760650124400854,
      "compile_s": 0.04758154299634043,
      "inner": 33,
      "median_ms": 1.2357777273846846,
      "q1_ms": 1.2356752726411675,
      "q3_ms": 1.2364602424151432,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 1024.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_err": 0.002305536298081279,
      "rel_err_parts": [
        0.002305536298081279
      ],
      "rel_vs_probjax_nosplit": 0.0030674005392938852,
      "compile_s": 0.0013891779963159934,
      "inner": 39,
      "median_ms": 1.234036076661701,
      "q1_ms": 1.233490807816195,
      "q3_ms": 1.2341613078322737,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 1024.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 1.6250152587890625,
      "xla_arg_mb": 1024.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.002288058865815401,
      "rel_err_parts": [
        0.002288058865815401
      ],
      "rel_vs_probjax_nosplit": 0.002747800201177597,
      "compile_s": 0.17583499700413086,
      "inner": 37,
      "median_ms": 1.2261501349810813,
      "q1_ms": 1.2259133379572902,
      "q3_ms": 1.2262208108802803,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 1024.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 1024.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.002311257878318429,
      "rel_err_parts": [
        0.002311257878318429
      ],
      "compile_s": 0.1498058540018974,
      "inner": 22,
      "median_ms": 2.0585028184789484,
      "q1_ms": 2.0580451590913897,
      "q3_ms": 2.0589780456661675,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 1024.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 1024.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.002311257878318429,
      "rel_err_parts": [
        0.002311257878318429
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.14958418400783557,
      "inner": 22,
      "median_ms": 2.0588721361656845,
      "q1_ms": 2.0584809773770925,
      "q3_ms": 2.0667108861247967,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 1024.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 48.00048828125,
      "xla_arg_mb": 1024.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.0023400033824145794,
      "rel_err_parts": [
        0.0023400033824145794
      ],
      "rel_vs_probjax_nosplit": 0.0031450316309928894,
      "compile_s": 0.09790346400404815,
      "inner": 32,
      "median_ms": 1.4388823128683725,
      "q1_ms": 1.4388172344297345,
      "q3_ms": 1.4390565781923215,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 1024.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 1024.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.0023113626521080732,
      "rel_err_parts": [
        0.0023113626521080732
      ],
      "rel_vs_probjax_nosplit": 7.500975334551185e-05,
      "compile_s": 0.04936242300027516,
      "inner": 18,
      "median_ms": 2.4572122775781384,
      "q1_ms": 2.455194888979248,
      "q3_ms": 2.458143277989519,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 1024.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_err": 0.002299167448654771,
      "rel_err_parts": [
        0.002299167448654771
      ],
      "rel_vs_probjax_nosplit": 0.0030602167826145887,
      "compile_s": 0.001390751000144519,
      "inner": 39,
      "median_ms": 1.2343278974115561,
      "q1_ms": 1.2342499357122834,
      "q3_ms": 1.2347615512016301,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 1024.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 1.6250152587890625,
      "xla_arg_mb": 1024.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.002288058865815401,
      "rel_err_parts": [
        0.002288058865815401
      ],
      "rel_vs_probjax_nosplit": 0.002747800201177597,
      "compile_s": 0.17623154599277768,
      "inner": 37,
      "median_ms": 1.226485351327102,
      "q1_ms": 1.2262256621741816,
      "q3_ms": 1.2266784322344244,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 1024.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 1024.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.002311257878318429,
      "rel_err_parts": [
        0.002311257878318429
      ],
      "compile_s": 0.15169742300349753,
      "inner": 21,
      "median_ms": 2.05597028529848,
      "q1_ms": 2.0558384047555074,
      "q3_ms": 2.0604712138135386,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 1024.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 1024.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.002311257878318429,
      "rel_err_parts": [
        0.002311257878318429
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.14913535499363206,
      "inner": 23,
      "median_ms": 2.0573730004759017,
      "q1_ms": 2.056336456482076,
      "q3_ms": 2.0578763259159727,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 1024.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 48.00048828125,
      "xla_arg_mb": 1024.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.0023400033824145794,
      "rel_err_parts": [
        0.0023400033824145794
      ],
      "rel_vs_probjax_nosplit": 0.0031450316309928894,
      "compile_s": 0.09374474399373867,
      "inner": 32,
      "median_ms": 1.43903668731582,
      "q1_ms": 1.4385570468675724,
      "q3_ms": 1.4394836562132696,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 1024.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 1024.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.0023113626521080732,
      "rel_err_parts": [
        0.0023113626521080732
      ],
      "rel_vs_probjax_nosplit": 7.500975334551185e-05,
      "compile_s": 0.051927738008089364,
      "inner": 18,
      "median_ms": 2.4595456108929485,
      "q1_ms": 2.4581527500736735,
      "q3_ms": 2.464251194598748,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 1024.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_err": 0.002299167448654771,
      "rel_err_parts": [
        0.002299167448654771
      ],
      "rel_vs_probjax_nosplit": 0.0030602167826145887,
      "compile_s": 0.0013852910051355138,
      "inner": 39,
      "median_ms": 1.234055589660644,
      "q1_ms": 1.2338836024401902,
      "q3_ms": 1.2342315769787184,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 1024.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 1.6250152587890625,
      "xla_arg_mb": 1024.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.002288058865815401,
      "rel_err_parts": [
        0.002288058865815401
      ],
      "rel_vs_probjax_nosplit": 0.002747800201177597,
      "compile_s": 0.17529279000882525,
      "inner": 36,
      "median_ms": 1.226452083048773,
      "q1_ms": 1.22606066664351,
      "q3_ms": 1.2266059861253273,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 1024.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 1024.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.002311257878318429,
      "rel_err_parts": [
        0.002311257878318429
      ],
      "compile_s": 0.1514403460023459,
      "inner": 22,
      "median_ms": 2.058477318231863,
      "q1_ms": 2.0573893862919332,
      "q3_ms": 2.064472545581256,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 1024.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 1024.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.002311257878318429,
      "rel_err_parts": [
        0.002311257878318429
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.1515051299938932,
      "inner": 22,
      "median_ms": 2.0575988634267226,
      "q1_ms": 2.057261181603694,
      "q3_ms": 2.057976386409295,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 1024.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 48.00048828125,
      "xla_arg_mb": 1024.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.0023400033824145794,
      "rel_err_parts": [
        0.0023400033824145794
      ],
      "rel_vs_probjax_nosplit": 0.0031450316309928894,
      "compile_s": 0.09508995999931358,
      "inner": 32,
      "median_ms": 1.4389330626727315,
      "q1_ms": 1.438729406118,
      "q3_ms": 1.439372187405752,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 1024.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 1024.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.0023113626521080732,
      "rel_err_parts": [
        0.0023113626521080732
      ],
      "rel_vs_probjax_nosplit": 7.500975334551185e-05,
      "compile_s": 0.04965200199512765,
      "inner": 17,
      "median_ms": 2.4626315888085895,
      "q1_ms": 2.4617449115098946,
      "q3_ms": 2.4720057939275115,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 1024.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_err": 0.002299167448654771,
      "rel_err_parts": [
        0.002299167448654771
      ],
      "rel_vs_probjax_nosplit": 0.0030602167826145887,
      "compile_s": 0.00138261700340081,
      "inner": 39,
      "median_ms": 1.2340430257609114,
      "q1_ms": 1.2336973590068876,
      "q3_ms": 1.2345073590926732,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 1024.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 1.5234527587890625,
      "xla_arg_mb": 1024.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.0022924006916582584,
      "rel_err_parts": [
        0.0022924006916582584
      ],
      "rel_vs_probjax_nosplit": 0.0027632759883999825,
      "compile_s": 0.17464612699404825,
      "inner": 36,
      "median_ms": 1.227502944756351,
      "q1_ms": 1.2268116525875588,
      "q3_ms": 1.2281684028291946,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 1024.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 1024.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.0023226167540997267,
      "rel_err_parts": [
        0.0023226167540997267
      ],
      "compile_s": 0.13245748699409887,
      "inner": 23,
      "median_ms": 1.9867467395353902,
      "q1_ms": 1.9860336520334543,
      "q3_ms": 1.9870553480122117,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 1024.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 1024.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.0023226167540997267,
      "rel_err_parts": [
        0.0023226167540997267
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.13028558799123857,
      "inner": 23,
      "median_ms": 1.9860453910001998,
      "q1_ms": 1.9856694346779715,
      "q3_ms": 1.9861083694748864,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 1024.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 192.001953125,
      "xla_arg_mb": 1024.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.002348241861909628,
      "rel_err_parts": [
        0.002348241861909628
      ],
      "rel_vs_probjax_nosplit": 0.003144636284559965,
      "compile_s": 0.1285002389922738,
      "inner": 23,
      "median_ms": 2.0353722173477644,
      "q1_ms": 2.0349847390721347,
      "q3_ms": 2.035806325970359,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 1024.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 1024.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.0023226651828736067,
      "rel_err_parts": [
        0.0023226651828736067
      ],
      "rel_vs_probjax_nosplit": 4.228383113513701e-05,
      "compile_s": 0.04968214899417944,
      "inner": 19,
      "median_ms": 2.4563822111016824,
      "q1_ms": 2.4559328942957923,
      "q3_ms": 2.4630413685563806,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 1024.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_err": 0.002301033353433013,
      "rel_err_parts": [
        0.002301033353433013
      ],
      "rel_vs_probjax_nosplit": 0.0030783195979893208,
      "compile_s": 0.0013902200007578358,
      "inner": 39,
      "median_ms": 1.2373477179491414,
      "q1_ms": 1.2367975768239166,
      "q3_ms": 1.2374728334506448,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 1024.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 1.5234527587890625,
      "xla_arg_mb": 1024.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.0022924006916582584,
      "rel_err_parts": [
        0.0022924006916582584
      ],
      "rel_vs_probjax_nosplit": 0.0027632759883999825,
      "compile_s": 0.1701818100118544,
      "inner": 36,
      "median_ms": 1.2275105000298936,
      "q1_ms": 1.2272889442910027,
      "q3_ms": 1.2277488472641886,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 1024.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 1024.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.0023226167540997267,
      "rel_err_parts": [
        0.0023226167540997267
      ],
      "compile_s": 0.1319077859952813,
      "inner": 23,
      "median_ms": 1.9810363478229747,
      "q1_ms": 1.980692456645952,
      "q3_ms": 1.9819846736778426,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 1024.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 1024.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.0023226167540997267,
      "rel_err_parts": [
        0.0023226167540997267
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.13183935699635185,
      "inner": 23,
      "median_ms": 1.981616565583112,
      "q1_ms": 1.9812855218491596,
      "q3_ms": 1.9862690653257151,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 1024.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 192.001953125,
      "xla_arg_mb": 1024.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.002348241861909628,
      "rel_err_parts": [
        0.002348241861909628
      ],
      "rel_vs_probjax_nosplit": 0.003144636284559965,
      "compile_s": 0.12868503900244832,
      "inner": 22,
      "median_ms": 2.0363286365119913,
      "q1_ms": 2.0361944774446203,
      "q3_ms": 2.036450408592249,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 1024.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 1024.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.0023226651828736067,
      "rel_err_parts": [
        0.0023226651828736067
      ],
      "rel_vs_probjax_nosplit": 4.228383113513701e-05,
      "compile_s": 0.0487273320031818,
      "inner": 18,
      "median_ms": 2.4713401111916435,
      "q1_ms": 2.4710491112879636,
      "q3_ms": 2.4770423608262924,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 1024.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_err": 0.002301033353433013,
      "rel_err_parts": [
        0.002301033353433013
      ],
      "rel_vs_probjax_nosplit": 0.0030783195979893208,
      "compile_s": 0.0013896599994041026,
      "inner": 39,
      "median_ms": 1.2370576923086236,
      "q1_ms": 1.2366724872546724,
      "q3_ms": 1.2373317947957472,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 1024.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 1.5234527587890625,
      "xla_arg_mb": 1024.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.0022924006916582584,
      "rel_err_parts": [
        0.0022924006916582584
      ],
      "rel_vs_probjax_nosplit": 0.0027632759883999825,
      "compile_s": 0.1729097310017096,
      "inner": 37,
      "median_ms": 1.228275486249225,
      "q1_ms": 1.2278923378609172,
      "q3_ms": 1.2285565810727985,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 1024.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 1024.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.0023226167540997267,
      "rel_err_parts": [
        0.0023226167540997267
      ],
      "compile_s": 0.13158719899365678,
      "inner": 23,
      "median_ms": 2.0013896523185237,
      "q1_ms": 2.001202999963425,
      "q3_ms": 2.001671717469277,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 1024.0,
      "impl": "flex_repeat",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 1024.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.0023226167540997267,
      "rel_err_parts": [
        0.0023226167540997267
      ],
      "rel_vs_probjax_nosplit": 0.0,
      "compile_s": 0.13324588800605852,
      "inner": 23,
      "median_ms": 2.0017930434298012,
      "q1_ms": 2.001271391434231,
      "q3_ms": 2.0020472176205972,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 1024.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 192.001953125,
      "xla_arg_mb": 1024.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.002348241861909628,
      "rel_err_parts": [
        0.002348241861909628
      ],
      "rel_vs_probjax_nosplit": 0.003144636284559965,
      "compile_s": 0.12756954799988307,
      "inner": 23,
      "median_ms": 2.036425912671763,
      "q1_ms": 2.0351738045143932,
      "q3_ms": 2.036773565518872,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 1024.0,
      "impl": "jax_cudnn",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.000244140625,
      "xla_arg_mb": 1024.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.0023226651828736067,
      "rel_err_parts": [
        0.0023226651828736067
      ],
      "rel_vs_probjax_nosplit": 4.228383113513701e-05,
      "compile_s": 0.04718467399652582,
      "inner": 18,
      "median_ms": 2.47430800017254,
      "q1_ms": 2.4729826945986133,
      "q3_ms": 2.4761974168035925,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 1024.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_err": 0.002301033353433013,
      "rel_err_parts": [
        0.002301033353433013
      ],
      "rel_vs_probjax_nosplit": 0.0030783195979893208,
      "compile_s": 0.0013890789996366948,
      "inner": 39,
      "median_ms": 1.2368581027168637,
      "q1_ms": 1.2365003589925787,
      "q3_ms": 1.2371738078236842,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 8192.0,
      "error": "skipped: K+V > 4.5 GB",
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 8192.0,
      "error": "skipped: K+V > 4.5 GB",
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "kv_mb": 8192.0,
      "error": "skipped: K+V > 4.5 GB",
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_h32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 16.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.5312652587890625,
      "xla_arg_mb": 16.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.00041898307972587645,
      "rel_err_parts": [
        0.00041898307972587645
      ],
      "rel_vs_probjax_nosplit": 0.0002272034325869754,
      "compile_s": 0.17417054699035361,
      "inner": 200,
      "median_ms": 0.04568560005282052,
      "q1_ms": 0.045501727508963086,
      "q3_ms": 0.04633700500562554,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 16.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 16.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.00042372700409032404,
      "rel_err_parts": [
        0.00042372700409032404
      ],
      "compile_s": 0.14053422999859322,
      "inner": 134,
      "median_ms": 0.2286648507005841,
      "q1_ms": 0.22394434699211824,
      "q3_ms": 0.22870675375495703,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 16.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 1.0,
      "xla_arg_mb": 16.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.0004352576797828078,
      "rel_err_parts": [
        0.0004352576797828078
      ],
      "rel_vs_probjax_nosplit": 0.00027709564892575145,
      "compile_s": 0.13482505700085312,
      "inner": 200,
      "median_ms": 0.041323089972138405,
      "q1_ms": 0.04081713250343455,
      "q3_ms": 0.04976924999937182,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 16.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_err": 9.284203770221211e-07,
      "rel_err_parts": [
        9.284203770221211e-07
      ],
      "rel_vs_probjax_nosplit": 0.0004238006367813796,
      "compile_s": 0.000856791011756286,
      "inner": 69,
      "median_ms": 0.6798272317636461,
      "q1_ms": 0.6796927827172508,
      "q3_ms": 0.6818666159338993,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 16.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.5312652587890625,
      "xla_arg_mb": 16.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.00041898307972587645,
      "rel_err_parts": [
        0.00041898307972587645
      ],
      "rel_vs_probjax_nosplit": 0.0002272034325869754,
      "compile_s": 0.17186403200321365,
      "inner": 200,
      "median_ms": 0.03650492501037661,
      "q1_ms": 0.03643080999609083,
      "q3_ms": 0.03672724498755997,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 16.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 16.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.00042372700409032404,
      "rel_err_parts": [
        0.00042372700409032404
      ],
      "compile_s": 0.14255927200429142,
      "inner": 158,
      "median_ms": 0.2064200189470054,
      "q1_ms": 0.20638282908697,
      "q3_ms": 0.20643172154197165,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 16.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 1.0,
      "xla_arg_mb": 16.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.0004352576797828078,
      "rel_err_parts": [
        0.0004352576797828078
      ],
      "rel_vs_probjax_nosplit": 0.00027709564892575145,
      "compile_s": 0.13838181699975394,
      "inner": 200,
      "median_ms": 0.039318750059464946,
      "q1_ms": 0.0389102500048466,
      "q3_ms": 0.04897251499642152,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 16.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_err": 9.284203770221211e-07,
      "rel_err_parts": [
        9.284203770221211e-07
      ],
      "rel_vs_probjax_nosplit": 0.0004238006367813796,
      "compile_s": 0.0008502179989591241,
      "inner": 69,
      "median_ms": 0.6849501303910022,
      "q1_ms": 0.6848993767904914,
      "q3_ms": 0.6849833694961516,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 16.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.5312652587890625,
      "xla_arg_mb": 16.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.00041898307972587645,
      "rel_err_parts": [
        0.00041898307972587645
      ],
      "rel_vs_probjax_nosplit": 0.0002272034325869754,
      "compile_s": 0.1768939199973829,
      "inner": 200,
      "median_ms": 0.03997744497610256,
      "q1_ms": 0.03980241497629322,
      "q3_ms": 0.04060881499754032,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 16.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 16.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.00042372700409032404,
      "rel_err_parts": [
        0.00042372700409032404
      ],
      "compile_s": 0.14435584101011045,
      "inner": 135,
      "median_ms": 0.20641139266512323,
      "q1_ms": 0.2063775444252382,
      "q3_ms": 0.2064479000249843,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 16.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 1.0,
      "xla_arg_mb": 16.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.0004352576797828078,
      "rel_err_parts": [
        0.0004352576797828078
      ],
      "rel_vs_probjax_nosplit": 0.00027709564892575145,
      "compile_s": 0.135293173007085,
      "inner": 200,
      "median_ms": 0.04470464002224617,
      "q1_ms": 0.04327398248278769,
      "q3_ms": 0.045182517533248756,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 16.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_err": 9.284203770221211e-07,
      "rel_err_parts": [
        9.284203770221211e-07
      ],
      "rel_vs_probjax_nosplit": 0.0004238006367813796,
      "compile_s": 0.0008635239937575534,
      "inner": 70,
      "median_ms": 0.6825067285847451,
      "q1_ms": 0.6824650786126897,
      "q3_ms": 0.6825347143084011,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 16.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.5156402587890625,
      "xla_arg_mb": 16.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.00042241925257258117,
      "rel_err_parts": [
        0.00042241925257258117
      ],
      "rel_vs_probjax_nosplit": 0.00023395738389808685,
      "compile_s": 0.1766294690023642,
      "inner": 200,
      "median_ms": 0.0365900850010803,
      "q1_ms": 0.03630189250543481,
      "q3_ms": 0.03691302248626016,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 16.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 16.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.00043457368155941367,
      "rel_err_parts": [
        0.00043457368155941367
      ],
      "compile_s": 0.14581581299717072,
      "inner": 143,
      "median_ms": 0.23421758740661677,
      "q1_ms": 0.2300460664437387,
      "q3_ms": 0.2342530838527287,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 16.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 4.0,
      "xla_arg_mb": 16.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.00043248559813946486,
      "rel_err_parts": [
        0.00043248559813946486
      ],
      "rel_vs_probjax_nosplit": 0.00028231620672158897,
      "compile_s": 0.3003191559982952,
      "inner": 200,
      "median_ms": 0.04491504005272873,
      "q1_ms": 0.0448524700186681,
      "q3_ms": 0.04666084252676228,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 16.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_err": 1.1089206282122177e-06,
      "rel_err_parts": [
        1.1089206282122177e-06
      ],
      "rel_vs_probjax_nosplit": 0.0004346004279796034,
      "compile_s": 0.0008498079987475649,
      "inner": 69,
      "median_ms": 0.6820078986862242,
      "q1_ms": 0.6819320217313151,
      "q3_ms": 0.6820440434292296,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 16.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.5156402587890625,
      "xla_arg_mb": 16.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.00042241925257258117,
      "rel_err_parts": [
        0.00042241925257258117
      ],
      "rel_vs_probjax_nosplit": 0.00023395738389808685,
      "compile_s": 0.17736961001355667,
      "inner": 200,
      "median_ms": 0.037149445051909424,
      "q1_ms": 0.037091760023031384,
      "q3_ms": 0.04548604501906084,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 16.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 16.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.00043457368155941367,
      "rel_err_parts": [
        0.00043457368155941367
      ],
      "compile_s": 0.1436706639942713,
      "inner": 140,
      "median_ms": 0.21161219992791303,
      "q1_ms": 0.2115442178494829,
      "q3_ms": 0.21169600004213862,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 16.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 4.0,
      "xla_arg_mb": 16.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.00043248559813946486,
      "rel_err_parts": [
        0.00043248559813946486
      ],
      "rel_vs_probjax_nosplit": 0.00028231620672158897,
      "compile_s": 0.29885155899683014,
      "inner": 200,
      "median_ms": 0.045584304971271195,
      "q1_ms": 0.045344950012804475,
      "q3_ms": 0.0456662099895766,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 16.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_err": 1.1089206282122177e-06,
      "rel_err_parts": [
        1.1089206282122177e-06
      ],
      "rel_vs_probjax_nosplit": 0.0004346004279796034,
      "compile_s": 0.00085387600120157,
      "inner": 69,
      "median_ms": 0.683483884052333,
      "q1_ms": 0.6834178841571608,
      "q3_ms": 0.6836095362550755,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 16.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.5156402587890625,
      "xla_arg_mb": 16.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.00042241925257258117,
      "rel_err_parts": [
        0.00042241925257258117
      ],
      "rel_vs_probjax_nosplit": 0.00023395738389808685,
      "compile_s": 0.17369123999378644,
      "inner": 200,
      "median_ms": 0.03743984503671527,
      "q1_ms": 0.03725166254298529,
      "q3_ms": 0.03814895750110736,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 16.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 16.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.00043457368155941367,
      "rel_err_parts": [
        0.00043457368155941367
      ],
      "compile_s": 0.14306414700695314,
      "inner": 141,
      "median_ms": 0.21180780144187075,
      "q1_ms": 0.21176201069284997,
      "q3_ms": 0.21183945743056637,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 16.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 4.0,
      "xla_arg_mb": 16.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.00043248559813946486,
      "rel_err_parts": [
        0.00043248559813946486
      ],
      "rel_vs_probjax_nosplit": 0.00028231620672158897,
      "compile_s": 0.29431088699493557,
      "inner": 200,
      "median_ms": 0.04517813504207879,
      "q1_ms": 0.04509450500336243,
      "q3_ms": 0.045522940017690416,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 16.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_err": 1.1089206282122177e-06,
      "rel_err_parts": [
        1.1089206282122177e-06
      ],
      "rel_vs_probjax_nosplit": 0.0004346004279796034,
      "compile_s": 0.0008552979998057708,
      "inner": 69,
      "median_ms": 0.6847670143622014,
      "q1_ms": 0.6846394637672136,
      "q3_ms": 0.6849239131121937,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 16.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 2.0625152587890625,
      "xla_arg_mb": 16.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.0004194006323814392,
      "rel_err_parts": [
        0.0004194006323814392
      ],
      "rel_vs_probjax_nosplit": 0.00023686086933594197,
      "compile_s": 0.37020188799942844,
      "inner": 200,
      "median_ms": 0.044282495000516064,
      "q1_ms": 0.0442282900257851,
      "q3_ms": 0.04432397003256483,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 16.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 16.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.0004301099688746035,
      "rel_err_parts": [
        0.0004301099688746035
      ],
      "compile_s": 0.12549106500227936,
      "inner": 139,
      "median_ms": 0.21935038129537762,
      "q1_ms": 0.2192673813061715,
      "q3_ms": 0.2193948165630828,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 16.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 16.0,
      "xla_arg_mb": 16.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.00043545031803660095,
      "rel_err_parts": [
        0.00043545031803660095
      ],
      "rel_vs_probjax_nosplit": 0.0002821694070007652,
      "compile_s": 0.2217752440046752,
      "inner": 200,
      "median_ms": 0.08855962500092573,
      "q1_ms": 0.08813331754936371,
      "q3_ms": 0.08868400997016579,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 16.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_err": 1.1093359262304148e-06,
      "rel_err_parts": [
        1.1093359262304148e-06
      ],
      "rel_vs_probjax_nosplit": 0.000430161424446851,
      "compile_s": 0.000872481003170833,
      "inner": 69,
      "median_ms": 0.6828493332368848,
      "q1_ms": 0.6828173986371215,
      "q3_ms": 0.6830790434437601,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 16.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 2.0625152587890625,
      "xla_arg_mb": 16.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.0004194006323814392,
      "rel_err_parts": [
        0.0004194006323814392
      ],
      "rel_vs_probjax_nosplit": 0.00023686086933594197,
      "compile_s": 0.37433098199835513,
      "inner": 200,
      "median_ms": 0.0417038649902679,
      "q1_ms": 0.04161564251262462,
      "q3_ms": 0.043659162511175964,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 16.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 16.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.0004301099688746035,
      "rel_err_parts": [
        0.0004301099688746035
      ],
      "compile_s": 0.12692985800094903,
      "inner": 152,
      "median_ms": 0.19673390794962056,
      "q1_ms": 0.19667662500547808,
      "q3_ms": 0.19677879275772475,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 16.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 16.0,
      "xla_arg_mb": 16.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.00043545031803660095,
      "rel_err_parts": [
        0.00043545031803660095
      ],
      "rel_vs_probjax_nosplit": 0.0002821694070007652,
      "compile_s": 0.23004477299400605,
      "inner": 200,
      "median_ms": 0.08888428499631118,
      "q1_ms": 0.0885947649658192,
      "q3_ms": 0.08896423998521641,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 16.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_err": 1.1093359262304148e-06,
      "rel_err_parts": [
        1.1093359262304148e-06
      ],
      "rel_vs_probjax_nosplit": 0.000430161424446851,
      "compile_s": 0.000853074001497589,
      "inner": 69,
      "median_ms": 0.6862906376372321,
      "q1_ms": 0.6862360289391211,
      "q3_ms": 0.6865130145224216,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 16.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 2.0625152587890625,
      "xla_arg_mb": 16.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.0004194006323814392,
      "rel_err_parts": [
        0.0004194006323814392
      ],
      "rel_vs_probjax_nosplit": 0.00023686086933594197,
      "compile_s": 0.3666717590094777,
      "inner": 200,
      "median_ms": 0.04197101494355593,
      "q1_ms": 0.041728232536115684,
      "q3_ms": 0.04211701998428907,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 16.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 16.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.0004301099688746035,
      "rel_err_parts": [
        0.0004301099688746035
      ],
      "compile_s": 0.13006725400919095,
      "inner": 150,
      "median_ms": 0.19848537999981394,
      "q1_ms": 0.19846330333772735,
      "q3_ms": 0.19848851661663502,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 16.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 16.0,
      "xla_arg_mb": 16.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.00043545031803660095,
      "rel_err_parts": [
        0.00043545031803660095
      ],
      "rel_vs_probjax_nosplit": 0.0002821694070007652,
      "compile_s": 0.22351145998982247,
      "inner": 200,
      "median_ms": 0.0889094849844696,
      "q1_ms": 0.08859751997079002,
      "q3_ms": 0.08898084499378456,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 16.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_err": 1.1093359262304148e-06,
      "rel_err_parts": [
        1.1093359262304148e-06
      ],
      "rel_vs_probjax_nosplit": 0.000430161424446851,
      "compile_s": 0.0008610589866293594,
      "inner": 69,
      "median_ms": 0.6857140291738661,
      "q1_ms": 0.6855869057514044,
      "q3_ms": 0.6857918695828445,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 64.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.6977691650390625,
      "xla_arg_mb": 64.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.0004088225250598043,
      "rel_err_parts": [
        0.0004088225250598043
      ],
      "rel_vs_probjax_nosplit": 0.0002505243173800409,
      "compile_s": 0.17597207600192633,
      "inner": 200,
      "median_ms": 0.08930563497415278,
      "q1_ms": 0.08925606496632099,
      "q3_ms": 0.08931347496400122,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 64.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 64.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.000419473071815446,
      "rel_err_parts": [
        0.000419473071815446
      ],
      "compile_s": 0.14512998599093407,
      "inner": 55,
      "median_ms": 0.7935723273972558,
      "q1_ms": 0.792869372790764,
      "q3_ms": 0.7937509455421771,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 64.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 4.0,
      "xla_arg_mb": 64.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.0004198154201731086,
      "rel_err_parts": [
        0.0004198154201731086
      ],
      "rel_vs_probjax_nosplit": 0.0002820993249770254,
      "compile_s": 0.2454588089894969,
      "inner": 200,
      "median_ms": 0.09902407495246734,
      "q1_ms": 0.09895045750454301,
      "q3_ms": 0.09903812748234486,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 64.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_err": 9.867621884041e-07,
      "rel_err_parts": [
        9.867621884041e-07
      ],
      "rel_vs_probjax_nosplit": 0.00041956911445595324,
      "compile_s": 0.0028737479879055172,
      "inner": 18,
      "median_ms": 2.697645667164276,
      "q1_ms": 2.6973400559250473,
      "q3_ms": 2.698034111138744,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 64.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.6977691650390625,
      "xla_arg_mb": 64.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.0004088225250598043,
      "rel_err_parts": [
        0.0004088225250598043
      ],
      "rel_vs_probjax_nosplit": 0.0002505243173800409,
      "compile_s": 0.1723213380028028,
      "inner": 200,
      "median_ms": 0.08837351997499354,
      "q1_ms": 0.08815255252557108,
      "q3_ms": 0.08840272748784628,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 64.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 64.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.000419473071815446,
      "rel_err_parts": [
        0.000419473071815446
      ],
      "compile_s": 0.14320025499910116,
      "inner": 55,
      "median_ms": 0.795136563564566,
      "q1_ms": 0.7950518546170894,
      "q3_ms": 0.7952864817608233,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 64.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 4.0,
      "xla_arg_mb": 64.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.0004198154201731086,
      "rel_err_parts": [
        0.0004198154201731086
      ],
      "rel_vs_probjax_nosplit": 0.0002820993249770254,
      "compile_s": 0.2544972510077059,
      "inner": 200,
      "median_ms": 0.0992489999771351,
      "q1_ms": 0.09918435502186185,
      "q3_ms": 0.09931314500136068,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 64.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_err": 9.867621884041e-07,
      "rel_err_parts": [
        9.867621884041e-07
      ],
      "rel_vs_probjax_nosplit": 0.00041956911445595324,
      "compile_s": 0.0028668750019278377,
      "inner": 18,
      "median_ms": 2.699786888519561,
      "q1_ms": 2.699230027954197,
      "q3_ms": 2.699992277889073,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 64.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.6977691650390625,
      "xla_arg_mb": 64.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.0004088225250598043,
      "rel_err_parts": [
        0.0004088225250598043
      ],
      "rel_vs_probjax_nosplit": 0.0002505243173800409,
      "compile_s": 0.17339325600187294,
      "inner": 200,
      "median_ms": 0.08828274498227984,
      "q1_ms": 0.08821136500046123,
      "q3_ms": 0.08834379001200432,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 64.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 64.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.000419473071815446,
      "rel_err_parts": [
        0.000419473071815446
      ],
      "compile_s": 0.1410772969975369,
      "inner": 53,
      "median_ms": 0.7920458867271729,
      "q1_ms": 0.7919695094899925,
      "q3_ms": 0.7921807547012027,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 64.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 4.0,
      "xla_arg_mb": 64.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.0004198154201731086,
      "rel_err_parts": [
        0.0004198154201731086
      ],
      "rel_vs_probjax_nosplit": 0.0002820993249770254,
      "compile_s": 0.24505056699854322,
      "inner": 200,
      "median_ms": 0.09924088502884842,
      "q1_ms": 0.09913057499943534,
      "q3_ms": 0.0992714174935827,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 64.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_err": 9.867621884041e-07,
      "rel_err_parts": [
        9.867621884041e-07
      ],
      "rel_vs_probjax_nosplit": 0.00041956911445595324,
      "compile_s": 0.002868358002160676,
      "inner": 18,
      "median_ms": 2.6974775553551606,
      "q1_ms": 2.69732611114806,
      "q3_ms": 2.69794644474233,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 64.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.6767730712890625,
      "xla_arg_mb": 64.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.0004139773081988096,
      "rel_err_parts": [
        0.0004139773081988096
      ],
      "rel_vs_probjax_nosplit": 0.0002570152282714844,
      "compile_s": 0.17942535900510848,
      "inner": 200,
      "median_ms": 0.08916061000491027,
      "q1_ms": 0.08914661004382651,
      "q3_ms": 0.08918485749745741,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 64.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 64.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.0004236093955114484,
      "rel_err_parts": [
        0.0004236093955114484
      ],
      "compile_s": 0.14310736900370102,
      "inner": 52,
      "median_ms": 0.8174930192651048,
      "q1_ms": 0.8173232692150543,
      "q3_ms": 0.8178114039765205,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 64.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 16.0,
      "xla_arg_mb": 64.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.00042680473416112363,
      "rel_err_parts": [
        0.00042680473416112363
      ],
      "rel_vs_probjax_nosplit": 0.00028978692716918886,
      "compile_s": 0.29777018299500924,
      "inner": 187,
      "median_ms": 0.17987094124789066,
      "q1_ms": 0.17961947594297412,
      "q3_ms": 0.179994973252256,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 64.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_err": 1.3565161225415068e-06,
      "rel_err_parts": [
        1.3565161225415068e-06
      ],
      "rel_vs_probjax_nosplit": 0.00042365072295069695,
      "compile_s": 0.0028780860011465847,
      "inner": 18,
      "median_ms": 2.712887833088947,
      "q1_ms": 2.712488999855446,
      "q3_ms": 2.713078943795861,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 64.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.6767730712890625,
      "xla_arg_mb": 64.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.0004139773081988096,
      "rel_err_parts": [
        0.0004139773081988096
      ],
      "rel_vs_probjax_nosplit": 0.0002570152282714844,
      "compile_s": 0.179716892002034,
      "inner": 200,
      "median_ms": 0.0883371000236366,
      "q1_ms": 0.08815921250061365,
      "q3_ms": 0.08837256998958765,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 64.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 64.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.0004236093955114484,
      "rel_err_parts": [
        0.0004236093955114484
      ],
      "compile_s": 0.14874224101367872,
      "inner": 51,
      "median_ms": 0.8229147645704192,
      "q1_ms": 0.8226390490045898,
      "q3_ms": 0.8236471373546759,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 64.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 16.0,
      "xla_arg_mb": 64.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.00042680473416112363,
      "rel_err_parts": [
        0.00042680473416112363
      ],
      "rel_vs_probjax_nosplit": 0.00028978692716918886,
      "compile_s": 0.30363549100002274,
      "inner": 155,
      "median_ms": 0.1800296838477915,
      "q1_ms": 0.17993140321660547,
      "q3_ms": 0.18003282255508127,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 64.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_err": 1.3565161225415068e-06,
      "rel_err_parts": [
        1.3565161225415068e-06
      ],
      "rel_vs_probjax_nosplit": 0.00042365072295069695,
      "compile_s": 0.0028894569986732677,
      "inner": 18,
      "median_ms": 2.7219226111709656,
      "q1_ms": 2.721767638932862,
      "q3_ms": 2.721943833522447,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 64.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.6767730712890625,
      "xla_arg_mb": 64.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.0004139773081988096,
      "rel_err_parts": [
        0.0004139773081988096
      ],
      "rel_vs_probjax_nosplit": 0.0002570152282714844,
      "compile_s": 0.1764783740072744,
      "inner": 200,
      "median_ms": 0.08838163994369097,
      "q1_ms": 0.08827724250295432,
      "q3_ms": 0.08839676251227502,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 64.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 64.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.0004236093955114484,
      "rel_err_parts": [
        0.0004236093955114484
      ],
      "compile_s": 0.14702986099291593,
      "inner": 52,
      "median_ms": 0.8223811153304548,
      "q1_ms": 0.822092769270127,
      "q3_ms": 0.8225222500186646,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 64.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 16.0,
      "xla_arg_mb": 64.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.00042680473416112363,
      "rel_err_parts": [
        0.00042680473416112363
      ],
      "rel_vs_probjax_nosplit": 0.00028978692716918886,
      "compile_s": 0.30624737399921287,
      "inner": 159,
      "median_ms": 0.18005447796071475,
      "q1_ms": 0.17969599999480076,
      "q3_ms": 0.18009974530813103,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 64.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_err": 1.3565161225415068e-06,
      "rel_err_parts": [
        1.3565161225415068e-06
      ],
      "rel_vs_probjax_nosplit": 0.00042365072295069695,
      "compile_s": 0.0028615350020118058,
      "inner": 18,
      "median_ms": 2.69262333313236,
      "q1_ms": 2.6922128607919955,
      "q3_ms": 2.693071389027561,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 64.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 2.5781402587890625,
      "xla_arg_mb": 64.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.0004149123269598931,
      "rel_err_parts": [
        0.0004149123269598931
      ],
      "rel_vs_probjax_nosplit": 0.0002553444937802851,
      "compile_s": 0.3649903780024033,
      "inner": 200,
      "median_ms": 0.10701594503188971,
      "q1_ms": 0.10660539501259336,
      "q3_ms": 0.10736743250163272,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 64.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 64.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.0004272123333066702,
      "rel_err_parts": [
        0.0004272123333066702
      ],
      "compile_s": 0.12824687000829726,
      "inner": 56,
      "median_ms": 0.7729778034055406,
      "q1_ms": 0.7729232411293196,
      "q3_ms": 0.7730331785101693,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 64.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 64.0,
      "xla_arg_mb": 64.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.0004275477840565145,
      "rel_err_parts": [
        0.0004275477840565145
      ],
      "rel_vs_probjax_nosplit": 0.00028908264357596636,
      "compile_s": 0.2996704690012848,
      "inner": 111,
      "median_ms": 0.3417393423119761,
      "q1_ms": 0.3416790945643165,
      "q3_ms": 0.34193556748177467,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 64.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_err": 1.3555304576584604e-06,
      "rel_err_parts": [
        1.3555304576584604e-06
      ],
      "rel_vs_probjax_nosplit": 0.0004272261285223067,
      "compile_s": 0.002874188998248428,
      "inner": 18,
      "median_ms": 2.700170444667391,
      "q1_ms": 2.6999165836766927,
      "q3_ms": 2.7002725275653776,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 64.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 2.5781402587890625,
      "xla_arg_mb": 64.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.0004149123269598931,
      "rel_err_parts": [
        0.0004149123269598931
      ],
      "rel_vs_probjax_nosplit": 0.0002553444937802851,
      "compile_s": 0.3686268979945453,
      "inner": 200,
      "median_ms": 0.10769307504233439,
      "q1_ms": 0.10444272753375117,
      "q3_ms": 0.10802846249134745,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 64.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 64.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.0004272123333066702,
      "rel_err_parts": [
        0.0004272123333066702
      ],
      "compile_s": 0.12620953499572352,
      "inner": 55,
      "median_ms": 0.7716698000546206,
      "q1_ms": 0.7715876453651369,
      "q3_ms": 0.7718193544529971,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 64.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 64.0,
      "xla_arg_mb": 64.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.0004275477840565145,
      "rel_err_parts": [
        0.0004275477840565145
      ],
      "rel_vs_probjax_nosplit": 0.00028908264357596636,
      "compile_s": 0.2977458669920452,
      "inner": 100,
      "median_ms": 0.3422361800039653,
      "q1_ms": 0.34194217492768075,
      "q3_ms": 0.342271894987789,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 64.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_err": 1.3555304576584604e-06,
      "rel_err_parts": [
        1.3555304576584604e-06
      ],
      "rel_vs_probjax_nosplit": 0.0004272261285223067,
      "compile_s": 0.0028615649935090914,
      "inner": 18,
      "median_ms": 2.6999293889578744,
      "q1_ms": 2.6998806667203707,
      "q3_ms": 2.700170416372,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 64.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 2.5781402587890625,
      "xla_arg_mb": 64.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.0004149123269598931,
      "rel_err_parts": [
        0.0004149123269598931
      ],
      "rel_vs_probjax_nosplit": 0.0002553444937802851,
      "compile_s": 0.365344638004899,
      "inner": 200,
      "median_ms": 0.10814443499839399,
      "q1_ms": 0.10684424498322187,
      "q3_ms": 0.10822560750966659,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 64.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 64.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.0004272123333066702,
      "rel_err_parts": [
        0.0004272123333066702
      ],
      "compile_s": 0.12714478399720974,
      "inner": 57,
      "median_ms": 0.7732120350722159,
      "q1_ms": 0.7730982105018931,
      "q3_ms": 0.7732894474317793,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 64.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 64.0,
      "xla_arg_mb": 64.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.0004275477840565145,
      "rel_err_parts": [
        0.0004275477840565145
      ],
      "rel_vs_probjax_nosplit": 0.00028908264357596636,
      "compile_s": 0.2943244429916376,
      "inner": 113,
      "median_ms": 0.3418620088316945,
      "q1_ms": 0.34178300442347154,
      "q3_ms": 0.3419181725280195,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 64.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_err": 1.3555304576584604e-06,
      "rel_err_parts": [
        1.3555304576584604e-06
      ],
      "rel_vs_probjax_nosplit": 0.0004272261285223067,
      "compile_s": 0.0028889060049550608,
      "inner": 18,
      "median_ms": 2.7249127775171975,
      "q1_ms": 2.7246823333876415,
      "q3_ms": 2.725151277748713,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 256.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.6977691650390625,
      "xla_arg_mb": 256.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.0003962305490858853,
      "rel_err_parts": [
        0.0003962305490858853
      ],
      "rel_vs_probjax_nosplit": 0.0003244838153477758,
      "compile_s": 0.17948092499864288,
      "inner": 109,
      "median_ms": 0.3192431558833517,
      "q1_ms": 0.31896069733460153,
      "q3_ms": 0.31948228436896,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 256.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 256.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.0004612490301951766,
      "rel_err_parts": [
        0.0004612490301951766
      ],
      "compile_s": 0.14535051300481427,
      "inner": 15,
      "median_ms": 3.186850799829699,
      "q1_ms": 3.1858115330881747,
      "q3_ms": 3.190411300359604,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 256.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 8.25,
      "xla_arg_mb": 256.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.00040911766700446606,
      "rel_err_parts": [
        0.00040911766700446606
      ],
      "rel_vs_probjax_nosplit": 0.00034066138323396444,
      "compile_s": 0.11288453100132756,
      "inner": 106,
      "median_ms": 0.38101520751194395,
      "q1_ms": 0.38085211321540213,
      "q3_ms": 0.38122271699940713,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 256.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_err": 1.3800058695778716e-06,
      "rel_err_parts": [
        1.3800058695778716e-06
      ],
      "rel_vs_probjax_nosplit": 0.00046113660209812224,
      "compile_s": 0.010933779005426913,
      "inner": 4,
      "median_ms": 10.767494251922471,
      "q1_ms": 10.766453751784866,
      "q3_ms": 10.774221998872235,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 256.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.6977691650390625,
      "xla_arg_mb": 256.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.0003962305490858853,
      "rel_err_parts": [
        0.0003962305490858853
      ],
      "rel_vs_probjax_nosplit": 0.0003244838153477758,
      "compile_s": 0.17754921900632326,
      "inner": 120,
      "median_ms": 0.3181293000428317,
      "q1_ms": 0.31793029999486555,
      "q3_ms": 0.31818173338251654,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 256.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 256.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.0004612490301951766,
      "rel_err_parts": [
        0.0004612490301951766
      ],
      "compile_s": 0.15001326600031462,
      "inner": 14,
      "median_ms": 3.1855469288919784,
      "q1_ms": 3.1853423928883524,
      "q3_ms": 3.190730749858111,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 256.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 8.25,
      "xla_arg_mb": 256.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.00040911766700446606,
      "rel_err_parts": [
        0.00040911766700446606
      ],
      "rel_vs_probjax_nosplit": 0.00034066138323396444,
      "compile_s": 0.11160491000919137,
      "inner": 103,
      "median_ms": 0.3811992621230124,
      "q1_ms": 0.3811581165120639,
      "q3_ms": 0.3812821893907577,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 256.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_err": 1.3800058695778716e-06,
      "rel_err_parts": [
        1.3800058695778716e-06
      ],
      "rel_vs_probjax_nosplit": 0.00046113660209812224,
      "compile_s": 0.011387136997655034,
      "inner": 4,
      "median_ms": 11.337973748595687,
      "q1_ms": 10.864376499739592,
      "q3_ms": 11.340496124830679,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 256.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.6977691650390625,
      "xla_arg_mb": 256.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.0003962305490858853,
      "rel_err_parts": [
        0.0003962305490858853
      ],
      "rel_vs_probjax_nosplit": 0.0003244838153477758,
      "compile_s": 0.1764558810100425,
      "inner": 116,
      "median_ms": 0.31800690515504765,
      "q1_ms": 0.3176207887761472,
      "q3_ms": 0.3182236508517486,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 256.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 256.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.0004612490301951766,
      "rel_err_parts": [
        0.0004612490301951766
      ],
      "compile_s": 0.14737681800033897,
      "inner": 13,
      "median_ms": 3.2666273081513983,
      "q1_ms": 3.265384923277172,
      "q3_ms": 3.2756601918103674,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 256.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 8.25,
      "xla_arg_mb": 256.0078125,
      "xla_out_mb": 0.0078125,
      "rel_err": 0.00040911766700446606,
      "rel_err_parts": [
        0.00040911766700446606
      ],
      "rel_vs_probjax_nosplit": 0.00034066138323396444,
      "compile_s": 0.10990287800086662,
      "inner": 105,
      "median_ms": 0.3811469237830135,
      "q1_ms": 0.3811342333349222,
      "q3_ms": 0.38116620001216817,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 256.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_err": 1.3800058695778716e-06,
      "rel_err_parts": [
        1.3800058695778716e-06
      ],
      "rel_vs_probjax_nosplit": 0.00046113660209812224,
      "compile_s": 0.011049938999349251,
      "inner": 4,
      "median_ms": 10.883977502089692,
      "q1_ms": 10.882491875236155,
      "q3_ms": 10.884753624850418,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 256.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.6767730712890625,
      "xla_arg_mb": 256.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.0004115557239856571,
      "rel_err_parts": [
        0.0004115557239856571
      ],
      "rel_vs_probjax_nosplit": 0.0003260743396822363,
      "compile_s": 0.1785671760007972,
      "inner": 122,
      "median_ms": 0.31820850813719365,
      "q1_ms": 0.3177935450665126,
      "q3_ms": 0.31897848766377257,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 256.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 256.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.0004626582667697221,
      "rel_err_parts": [
        0.0004626582667697221
      ],
      "compile_s": 0.14747424199595116,
      "inner": 14,
      "median_ms": 3.242062500378649,
      "q1_ms": 3.2406154995052407,
      "q3_ms": 3.246751393037682,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 256.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 64.00048828125,
      "xla_arg_mb": 256.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.0004197948146611452,
      "rel_err_parts": [
        0.0004197948146611452
      ],
      "rel_vs_probjax_nosplit": 0.0003466318594291806,
      "compile_s": 0.1312441409972962,
      "inner": 77,
      "median_ms": 0.5576810909150919,
      "q1_ms": 0.5574797403162599,
      "q3_ms": 0.5578190129282037,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 256.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_err": 1.8776778460960486e-06,
      "rel_err_parts": [
        1.8776778460960486e-06
      ],
      "rel_vs_probjax_nosplit": 0.00046258093789219856,
      "compile_s": 0.011007849010638893,
      "inner": 4,
      "median_ms": 10.835450499143917,
      "q1_ms": 10.835171122380416,
      "q3_ms": 10.83597387332702,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 256.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.6767730712890625,
      "xla_arg_mb": 256.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.0004115557239856571,
      "rel_err_parts": [
        0.0004115557239856571
      ],
      "rel_vs_probjax_nosplit": 0.0003260743396822363,
      "compile_s": 0.1768518700118875,
      "inner": 121,
      "median_ms": 0.3177969587085507,
      "q1_ms": 0.31766886367079655,
      "q3_ms": 0.31783831813796004,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 256.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 256.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.0004626582667697221,
      "rel_err_parts": [
        0.0004626582667697221
      ],
      "compile_s": 0.14690145799249876,
      "inner": 14,
      "median_ms": 3.2707182150001506,
      "q1_ms": 3.267044071565449,
      "q3_ms": 3.270917785682416,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 256.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 64.00048828125,
      "xla_arg_mb": 256.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.0004197948146611452,
      "rel_err_parts": [
        0.0004197948146611452
      ],
      "rel_vs_probjax_nosplit": 0.0003466318594291806,
      "compile_s": 0.13113384300959297,
      "inner": 76,
      "median_ms": 0.5579851579233562,
      "q1_ms": 0.5579503619006381,
      "q3_ms": 0.558112710534356,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 256.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_err": 1.8776778460960486e-06,
      "rel_err_parts": [
        1.8776778460960486e-06
      ],
      "rel_vs_probjax_nosplit": 0.00046258093789219856,
      "compile_s": 0.011027156011550687,
      "inner": 4,
      "median_ms": 10.840750503120944,
      "q1_ms": 10.840529876077198,
      "q3_ms": 10.842145624337718,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 256.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.6767730712890625,
      "xla_arg_mb": 256.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.0004115557239856571,
      "rel_err_parts": [
        0.0004115557239856571
      ],
      "rel_vs_probjax_nosplit": 0.0003260743396822363,
      "compile_s": 0.17530795800848864,
      "inner": 109,
      "median_ms": 0.3180447431845129,
      "q1_ms": 0.3179729999728723,
      "q3_ms": 0.31811170181987913,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 256.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.25,
      "xla_arg_mb": 256.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.0004626582667697221,
      "rel_err_parts": [
        0.0004626582667697221
      ],
      "compile_s": 0.14448697000625543,
      "inner": 14,
      "median_ms": 3.2643962133858753,
      "q1_ms": 3.2594235718923823,
      "q3_ms": 3.2710656071555735,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 256.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 64.00048828125,
      "xla_arg_mb": 256.03125,
      "xla_out_mb": 0.03125,
      "rel_err": 0.0004197948146611452,
      "rel_err_parts": [
        0.0004197948146611452
      ],
      "rel_vs_probjax_nosplit": 0.0003466318594291806,
      "compile_s": 0.12920289899921045,
      "inner": 74,
      "median_ms": 0.5579202971185205,
      "q1_ms": 0.557815364732665,
      "q3_ms": 0.5581583918191845,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 256.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_err": 1.8776778460960486e-06,
      "rel_err_parts": [
        1.8776778460960486e-06
      ],
      "rel_vs_probjax_nosplit": 0.00046258093789219856,
      "compile_s": 0.010999072997947223,
      "inner": 4,
      "median_ms": 10.836825502337888,
      "q1_ms": 10.835680875970866,
      "q3_ms": 10.838333249921561,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 256.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 2.7070465087890625,
      "xla_arg_mb": 256.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.00041627694736234844,
      "rel_err_parts": [
        0.00041627694736234844
      ],
      "rel_vs_probjax_nosplit": 0.0003233308088965714,
      "compile_s": 0.3686582279915456,
      "inner": 78,
      "median_ms": 0.4860604999404854,
      "q1_ms": 0.4855067564699894,
      "q3_ms": 0.4893170448784263,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 256.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 256.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.00047000282211229205,
      "rel_err_parts": [
        0.00047000282211229205
      ],
      "compile_s": 0.12962588900700212,
      "inner": 15,
      "median_ms": 3.0821593332802877,
      "q1_ms": 3.081186899604897,
      "q3_ms": 3.0903916330620027,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 256.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 0,
      "xla_temp_mb": 256.001953125,
      "xla_arg_mb": 256.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.00042401067912578583,
      "rel_err_parts": [
        0.00042401067912578583
      ],
      "rel_vs_probjax_nosplit": 0.00034850952215492725,
      "compile_s": 0.18279896301100962,
      "inner": 36,
      "median_ms": 1.278624694451638,
      "q1_ms": 1.2783578194001974,
      "q3_ms": 1.2789094028751908,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 256.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 0,
      "rel_err": 1.8688363070396008e-06,
      "rel_err_parts": [
        1.8688363070396008e-06
      ],
      "rel_vs_probjax_nosplit": 0.00046989857219159603,
      "compile_s": 0.011022456994396634,
      "inner": 4,
      "median_ms": 10.853294250409817,
      "q1_ms": 10.852849500224693,
      "q3_ms": 10.853697373022442,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 256.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 2.7070465087890625,
      "xla_arg_mb": 256.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.00041627694736234844,
      "rel_err_parts": [
        0.00041627694736234844
      ],
      "rel_vs_probjax_nosplit": 0.0003233308088965714,
      "compile_s": 0.3690667010087054,
      "inner": 82,
      "median_ms": 0.484952378184705,
      "q1_ms": 0.48452353049225805,
      "q3_ms": 0.48512045120843117,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 256.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 256.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.00047000282211229205,
      "rel_err_parts": [
        0.00047000282211229205
      ],
      "compile_s": 0.1315573530009715,
      "inner": 13,
      "median_ms": 3.110941922596584,
      "q1_ms": 3.10930765352588,
      "q3_ms": 3.118640731442541,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 256.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 1,
      "xla_temp_mb": 256.001953125,
      "xla_arg_mb": 256.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.00042401067912578583,
      "rel_err_parts": [
        0.00042401067912578583
      ],
      "rel_vs_probjax_nosplit": 0.00034850952215492725,
      "compile_s": 0.17453047700109892,
      "inner": 36,
      "median_ms": 1.2790048612158798,
      "q1_ms": 1.278611610865078,
      "q3_ms": 1.2793495417200678,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 256.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 1,
      "rel_err": 1.8688363070396008e-06,
      "rel_err_parts": [
        1.8688363070396008e-06
      ],
      "rel_vs_probjax_nosplit": 0.00046989857219159603,
      "compile_s": 0.011031172994989902,
      "inner": 4,
      "median_ms": 10.874892501306022,
      "q1_ms": 10.87340737649356,
      "q3_ms": 10.876570626351167,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 256.0,
      "impl": "probjax_splitkv",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 2.7070465087890625,
      "xla_arg_mb": 256.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.00041627694736234844,
      "rel_err_parts": [
        0.00041627694736234844
      ],
      "rel_vs_probjax_nosplit": 0.0003233308088965714,
      "compile_s": 0.37712713300425094,
      "inner": 79,
      "median_ms": 0.4851386202292965,
      "q1_ms": 0.48463462017140574,
      "q3_ms": 0.4854701266733081,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 256.0,
      "impl": "probjax_nosplit",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 256.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.00047000282211229205,
      "rel_err_parts": [
        0.00047000282211229205
      ],
      "compile_s": 0.12916846299776807,
      "inner": 15,
      "median_ms": 3.1088171992450953,
      "q1_ms": 3.1077484997998304,
      "q3_ms": 3.1170585003565066,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 256.0,
      "impl": "jax_xla",
      "mode": "fwd",
      "round": 2,
      "xla_temp_mb": 256.001953125,
      "xla_arg_mb": 256.125,
      "xla_out_mb": 0.125,
      "rel_err": 0.00042401067912578583,
      "rel_err_parts": [
        0.00042401067912578583
      ],
      "rel_vs_probjax_nosplit": 0.00034850952215492725,
      "compile_s": 0.18760387500515208,
      "inner": 35,
      "median_ms": 1.2791275142392675,
      "q1_ms": 1.2789644999429584,
      "q3_ms": 1.279348943041571,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "decode",
      "dtype": "float32",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 64,
      "kv_mb": 256.0,
      "impl": "torch_sdpa",
      "mode": "fwd",
      "round": 2,
      "rel_err": 1.8688363070396008e-06,
      "rel_err_parts": [
        1.8688363070396008e-06
      ],
      "rel_vs_probjax_nosplit": 0.00046989857219159603,
      "compile_s": 0.011020182006177492,
      "inner": 4,
      "median_ms": 10.8639644968207,
      "q1_ms": 10.859427002287703,
      "q3_ms": 10.868715875403723,
      "tag": null,
      "device_arg": "gpu",
      "run": "decode2_f32"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 0,
      "compile_s": 0.15547055000206456,
      "inner": 172,
      "median_ms": 0.1523691918369285,
      "q1_ms": 0.15223387789203202,
      "q3_ms": 0.1523893750049336,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 0,
      "compile_s": 0.19270268399850465,
      "inner": 200,
      "median_ms": 0.06165120998048223,
      "q1_ms": 0.061462080011551734,
      "q3_ms": 0.06178987499879441,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 0,
      "compile_s": 0.23888300299586263,
      "inner": 200,
      "median_ms": 0.05760069994721562,
      "q1_ms": 0.05732668250857387,
      "q3_ms": 0.057829712495731656,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 0,
      "compile_s": 0.2391689439973561,
      "inner": 200,
      "median_ms": 0.04977429998689331,
      "q1_ms": 0.0497057674874668,
      "q3_ms": 0.050032562467094976,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 0,
      "compile_s": 0.2420159689936554,
      "inner": 200,
      "median_ms": 0.04355769000540022,
      "q1_ms": 0.04032337252283469,
      "q3_ms": 0.04579098749672994,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 32,
      "round": 0,
      "compile_s": 0.2547070040018298,
      "inner": 200,
      "median_ms": 0.040951160044642165,
      "q1_ms": 0.03802347000601003,
      "q3_ms": 0.04848461001529358,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 64,
      "round": 0,
      "compile_s": 0.25285638400237076,
      "inner": 200,
      "median_ms": 0.038830295015941374,
      "q1_ms": 0.03880028998537455,
      "q3_ms": 0.0410429574912996,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 1,
      "compile_s": 0.16080200498981867,
      "inner": 181,
      "median_ms": 0.15225671821863015,
      "q1_ms": 0.15220839500643354,
      "q3_ms": 0.15228276244861214,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 1,
      "compile_s": 0.19707999400270637,
      "inner": 200,
      "median_ms": 0.06082409498048946,
      "q1_ms": 0.060817085031885654,
      "q3_ms": 0.06084932250814745,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 1,
      "compile_s": 0.19800276900059544,
      "inner": 200,
      "median_ms": 0.040549545010435395,
      "q1_ms": 0.040402844970230944,
      "q3_ms": 0.04110482503165258,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 1,
      "compile_s": 0.194964793001418,
      "inner": 200,
      "median_ms": 0.04476366993912961,
      "q1_ms": 0.04474074998142896,
      "q3_ms": 0.04536545751761878,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 1,
      "compile_s": 0.19558829099696595,
      "inner": 200,
      "median_ms": 0.0384638000105042,
      "q1_ms": 0.03788243255257839,
      "q3_ms": 0.042882939997070935,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 32,
      "round": 1,
      "compile_s": 0.20012801900156774,
      "inner": 200,
      "median_ms": 0.04680293001001701,
      "q1_ms": 0.03999495249445317,
      "q3_ms": 0.05076233999716351,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 64,
      "round": 1,
      "compile_s": 0.18558453999867197,
      "inner": 200,
      "median_ms": 0.038204614975256845,
      "q1_ms": 0.03798910496698227,
      "q3_ms": 0.04635004250303609,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 0,
      "compile_s": 0.15870765400177334,
      "inner": 183,
      "median_ms": 0.15250332244796566,
      "q1_ms": 0.15235881147304386,
      "q3_ms": 0.15252524590224048,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 0,
      "compile_s": 0.23459973100398201,
      "inner": 200,
      "median_ms": 0.06177589501021429,
      "q1_ms": 0.061773942470608745,
      "q3_ms": 0.06178987248858902,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 0,
      "compile_s": 0.1983943490049569,
      "inner": 200,
      "median_ms": 0.040389394998783246,
      "q1_ms": 0.0402849999954924,
      "q3_ms": 0.04115955252927961,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 0,
      "compile_s": 0.20794875200954266,
      "inner": 200,
      "median_ms": 0.056251150017487817,
      "q1_ms": 0.04572108246065909,
      "q3_ms": 0.06157306252134731,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 0,
      "compile_s": 0.2028946220088983,
      "inner": 200,
      "median_ms": 0.0380035350099206,
      "q1_ms": 0.0377368274712353,
      "q3_ms": 0.04035488502267981,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 32,
      "round": 0,
      "compile_s": 0.20280616500531323,
      "inner": 200,
      "median_ms": 0.041035215035662986,
      "q1_ms": 0.03990738499851432,
      "q3_ms": 0.04474543500691652,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 64,
      "round": 0,
      "compile_s": 0.18422318498778623,
      "inner": 200,
      "median_ms": 0.03922163501556497,
      "q1_ms": 0.039020582516968716,
      "q3_ms": 0.03950647504098015,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 1,
      "compile_s": 0.16242073600005824,
      "inner": 200,
      "median_ms": 0.15227988500555512,
      "q1_ms": 0.1474031100224238,
      "q3_ms": 0.15235464998113457,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 1,
      "compile_s": 0.19803396800125483,
      "inner": 200,
      "median_ms": 0.06158778996905312,
      "q1_ms": 0.06158228003187105,
      "q3_ms": 0.061611337478097994,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 1,
      "compile_s": 0.1988667530094972,
      "inner": 200,
      "median_ms": 0.0469398849963909,
      "q1_ms": 0.04132694499276113,
      "q3_ms": 0.05139275748661021,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 1,
      "compile_s": 0.2047491700068349,
      "inner": 200,
      "median_ms": 0.057612119999248534,
      "q1_ms": 0.05450126245705178,
      "q3_ms": 0.06625772501138272,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 1,
      "compile_s": 0.21288702300807927,
      "inner": 200,
      "median_ms": 0.04178563496679999,
      "q1_ms": 0.041441132489126176,
      "q3_ms": 0.04925183500745334,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 32,
      "round": 1,
      "compile_s": 0.1968290190125117,
      "inner": 200,
      "median_ms": 0.03729418996954337,
      "q1_ms": 0.03682800997921731,
      "q3_ms": 0.039607345024705864,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 64,
      "round": 1,
      "compile_s": 0.18719401399721391,
      "inner": 200,
      "median_ms": 0.03858292999211699,
      "q1_ms": 0.03847915500955423,
      "q3_ms": 0.038657497498206794,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 0,
      "compile_s": 0.14167142899532337,
      "inner": 200,
      "median_ms": 0.14204828505171463,
      "q1_ms": 0.1420363350189291,
      "q3_ms": 0.14205496750946622,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 0,
      "compile_s": 0.410326265002368,
      "inner": 200,
      "median_ms": 0.11830512004962657,
      "q1_ms": 0.1182966799751739,
      "q3_ms": 0.11832315500214463,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 0,
      "compile_s": 0.3925274490029551,
      "inner": 200,
      "median_ms": 0.06734688497090247,
      "q1_ms": 0.06732384248607559,
      "q3_ms": 0.06737120500474703,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 0,
      "compile_s": 0.38725049800996203,
      "inner": 200,
      "median_ms": 0.05664954995154403,
      "q1_ms": 0.0447800775145879,
      "q3_ms": 0.06051007752830628,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 0,
      "compile_s": 0.42412757099373266,
      "inner": 200,
      "median_ms": 0.04696548501669895,
      "q1_ms": 0.04679882251366507,
      "q3_ms": 0.047175482504826505,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 32,
      "round": 0,
      "compile_s": 0.46118003499577753,
      "inner": 200,
      "median_ms": 0.06328534997010138,
      "q1_ms": 0.06326756749331253,
      "q3_ms": 0.06330426753265783,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 64,
      "round": 0,
      "compile_s": 0.4592598619929049,
      "inner": 200,
      "median_ms": 0.07646187499631196,
      "q1_ms": 0.07610955999552971,
      "q3_ms": 0.07657488753466168,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 1,
      "compile_s": 0.1391789740009699,
      "inner": 162,
      "median_ms": 0.13001567905446812,
      "q1_ms": 0.1299861141887959,
      "q3_ms": 0.13009010801374923,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 1,
      "compile_s": 0.3693909379944671,
      "inner": 184,
      "median_ms": 0.11822791843176517,
      "q1_ms": 0.11821085053784572,
      "q3_ms": 0.11829252715593296,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 1,
      "compile_s": 0.3726558919879608,
      "inner": 200,
      "median_ms": 0.06728677006321959,
      "q1_ms": 0.06728496748110047,
      "q3_ms": 0.06738703748851549,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 1,
      "compile_s": 0.37921255400578957,
      "inner": 200,
      "median_ms": 0.05584788501437288,
      "q1_ms": 0.05202933749387739,
      "q3_ms": 0.06161451750813285,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 1,
      "compile_s": 0.3909266109985765,
      "inner": 200,
      "median_ms": 0.048047980017145164,
      "q1_ms": 0.04797742250957526,
      "q3_ms": 0.048701444975449704,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 32,
      "round": 1,
      "compile_s": 0.38201526600460056,
      "inner": 200,
      "median_ms": 0.05811527502373792,
      "q1_ms": 0.05778835500677815,
      "q3_ms": 0.05828231503983261,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 64,
      "round": 1,
      "compile_s": 0.3818156170018483,
      "inner": 200,
      "median_ms": 0.0765390200103866,
      "q1_ms": 0.0759163200200419,
      "q3_ms": 0.07655247252841946,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 0,
      "compile_s": 0.15672281899605878,
      "inner": 77,
      "median_ms": 0.5198147791579922,
      "q1_ms": 0.5193590454896306,
      "q3_ms": 0.5204657533216628,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 0,
      "compile_s": 0.20073303300887346,
      "inner": 192,
      "median_ms": 0.16442340105034722,
      "q1_ms": 0.16439444012424548,
      "q3_ms": 0.1644506692552265,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 0,
      "compile_s": 0.20476584100106265,
      "inner": 174,
      "median_ms": 0.16784765516761047,
      "q1_ms": 0.16781820108760792,
      "q3_ms": 0.1678933707201698,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 0,
      "compile_s": 0.2653524300112622,
      "inner": 157,
      "median_ms": 0.18038035670345778,
      "q1_ms": 0.18032879296506332,
      "q3_ms": 0.18040961787070497,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 0,
      "compile_s": 0.2702799909893656,
      "inner": 158,
      "median_ms": 0.18135004432258917,
      "q1_ms": 0.18130876901235338,
      "q3_ms": 0.18143010449329028,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 32,
      "round": 0,
      "compile_s": 0.29270229100075085,
      "inner": 162,
      "median_ms": 0.20718527778143775,
      "q1_ms": 0.20713008640917524,
      "q3_ms": 0.20719669447754752,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 64,
      "round": 0,
      "compile_s": 0.26811018599255476,
      "inner": 142,
      "median_ms": 0.2454159577200058,
      "q1_ms": 0.245321348596866,
      "q3_ms": 0.2454847534876612,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 1,
      "compile_s": 0.16464455200184602,
      "inner": 75,
      "median_ms": 0.5225033867948999,
      "q1_ms": 0.5222253867153389,
      "q3_ms": 0.5227058199428333,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 1,
      "compile_s": 0.20178653601033147,
      "inner": 162,
      "median_ms": 0.16449181485869405,
      "q1_ms": 0.16447715744208025,
      "q3_ms": 0.1645727994108722,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 1,
      "compile_s": 0.1980566109996289,
      "inner": 168,
      "median_ms": 0.1678795655434009,
      "q1_ms": 0.16782649107058423,
      "q3_ms": 0.16792616965436274,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 1,
      "compile_s": 0.19725381200260017,
      "inner": 158,
      "median_ms": 0.18028099367268313,
      "q1_ms": 0.18020119617326197,
      "q3_ms": 0.18032770251149394,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 1,
      "compile_s": 0.2092758619983215,
      "inner": 177,
      "median_ms": 0.18126647457372566,
      "q1_ms": 0.18124810738474562,
      "q3_ms": 0.18130332771011254,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 32,
      "round": 1,
      "compile_s": 0.20088543200108688,
      "inner": 143,
      "median_ms": 0.20732326572272483,
      "q1_ms": 0.2073106223980284,
      "q3_ms": 0.20739301747904595,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 64,
      "round": 1,
      "compile_s": 0.18700979600544088,
      "inner": 142,
      "median_ms": 0.2454044647813237,
      "q1_ms": 0.2452207324021286,
      "q3_ms": 0.24559711266807477,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 0,
      "compile_s": 0.15659778099507093,
      "inner": 75,
      "median_ms": 0.5250672800078368,
      "q1_ms": 0.5218369066521215,
      "q3_ms": 0.5254585600535695,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 0,
      "compile_s": 0.24936020099266898,
      "inner": 181,
      "median_ms": 0.16776439222476308,
      "q1_ms": 0.16774186465402974,
      "q3_ms": 0.16783012985862863,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 0,
      "compile_s": 0.207859753005323,
      "inner": 146,
      "median_ms": 0.1718900548190529,
      "q1_ms": 0.1718573287079168,
      "q3_ms": 0.1719372054164885,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 0,
      "compile_s": 0.2078608650044771,
      "inner": 146,
      "median_ms": 0.18402362333639358,
      "q1_ms": 0.18393764040498256,
      "q3_ms": 0.1840473013457185,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 0,
      "compile_s": 0.24176219900255091,
      "inner": 168,
      "median_ms": 0.18826264282064845,
      "q1_ms": 0.18817333633099528,
      "q3_ms": 0.18827522619955872,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 32,
      "round": 0,
      "compile_s": 0.2395264600054361,
      "inner": 154,
      "median_ms": 0.20945279869589614,
      "q1_ms": 0.2093627305142733,
      "q3_ms": 0.20948743830168887,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 64,
      "round": 0,
      "compile_s": 0.19202302901248913,
      "inner": 141,
      "median_ms": 0.24623524821578052,
      "q1_ms": 0.24621549997675227,
      "q3_ms": 0.24635043622288483,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 1,
      "compile_s": 0.16532385800383054,
      "inner": 76,
      "median_ms": 0.5242401447934466,
      "q1_ms": 0.5239600131493429,
      "q3_ms": 0.5250691447145073,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 1,
      "compile_s": 0.20275002899870742,
      "inner": 165,
      "median_ms": 0.16686975149875224,
      "q1_ms": 0.16684588785706594,
      "q3_ms": 0.16703636969340907,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 1,
      "compile_s": 0.2011957490030909,
      "inner": 171,
      "median_ms": 0.17072456141345596,
      "q1_ms": 0.17052415205398574,
      "q3_ms": 0.1707719590862783,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 1,
      "compile_s": 0.20179329899838194,
      "inner": 167,
      "median_ms": 0.18308444306237656,
      "q1_ms": 0.18298731734983142,
      "q3_ms": 0.1832339192028185,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 1,
      "compile_s": 0.20339040900580585,
      "inner": 161,
      "median_ms": 0.18798662734619973,
      "q1_ms": 0.18786963975448912,
      "q3_ms": 0.1880173695325036,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 32,
      "round": 1,
      "compile_s": 0.19887928701064084,
      "inner": 135,
      "median_ms": 0.20930871843895013,
      "q1_ms": 0.2091630332870409,
      "q3_ms": 0.20933821851266893,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 64,
      "round": 1,
      "compile_s": 0.19612634999793954,
      "inner": 146,
      "median_ms": 0.24645706165856238,
      "q1_ms": 0.2463225582064359,
      "q3_ms": 0.2465064349613784,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 0,
      "compile_s": 0.13978022200171836,
      "inner": 74,
      "median_ms": 0.5205600271536339,
      "q1_ms": 0.5177032838440248,
      "q3_ms": 0.5210214528503727,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 0,
      "compile_s": 0.4218301649962086,
      "inner": 133,
      "median_ms": 0.2640379924078549,
      "q1_ms": 0.2629226503615386,
      "q3_ms": 0.26416620676172897,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 0,
      "compile_s": 0.3906353000056697,
      "inner": 112,
      "median_ms": 0.2804604911424186,
      "q1_ms": 0.2801592901440537,
      "q3_ms": 0.280709973204856,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 0,
      "compile_s": 0.3800110039883293,
      "inner": 121,
      "median_ms": 0.27967552887275815,
      "q1_ms": 0.27934453311358953,
      "q3_ms": 0.27995361982121264,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 0,
      "compile_s": 0.41957028000615537,
      "inner": 120,
      "median_ms": 0.33170308330833603,
      "q1_ms": 0.33164463335803396,
      "q3_ms": 0.3319593541770397,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 32,
      "round": 0,
      "compile_s": 0.4666588279942516,
      "inner": 110,
      "median_ms": 0.36495694538726553,
      "q1_ms": 0.3648358999105933,
      "q3_ms": 0.36531630910096946,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 64,
      "round": 0,
      "compile_s": 0.4667722119920654,
      "inner": 78,
      "median_ms": 0.5007185128096563,
      "q1_ms": 0.5004386922593813,
      "q3_ms": 0.5008079743465015,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 1,
      "compile_s": 0.14703507399826776,
      "inner": 77,
      "median_ms": 0.5210529738438449,
      "q1_ms": 0.5203082466984845,
      "q3_ms": 0.5220109544877363,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 1,
      "compile_s": 0.38713639200432226,
      "inner": 136,
      "median_ms": 0.2642241912219427,
      "q1_ms": 0.2639043162108625,
      "q3_ms": 0.26477209929212486,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 1,
      "compile_s": 0.3789227859961102,
      "inner": 128,
      "median_ms": 0.2817713202603045,
      "q1_ms": 0.28157000389228415,
      "q3_ms": 0.28210022657049194,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 1,
      "compile_s": 0.3784204560070066,
      "inner": 136,
      "median_ms": 0.27967399268270715,
      "q1_ms": 0.27966788601428133,
      "q3_ms": 0.27976626835665874,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 1,
      "compile_s": 0.3780040679994272,
      "inner": 109,
      "median_ms": 0.33356138524266143,
      "q1_ms": 0.3318759036536645,
      "q3_ms": 0.33360780738753376,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 32,
      "round": 1,
      "compile_s": 0.3790413910028292,
      "inner": 111,
      "median_ms": 0.3665925946447428,
      "q1_ms": 0.36596306751750673,
      "q3_ms": 0.3666662027548604,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 4096,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 64,
      "round": 1,
      "compile_s": 0.38021724500868004,
      "inner": 81,
      "median_ms": 0.5007481481305058,
      "q1_ms": 0.5006713332410004,
      "q3_ms": 0.5011895926055545,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 0,
      "compile_s": 0.15606473399384413,
      "inner": 71,
      "median_ms": 0.5135286478883206,
      "q1_ms": 0.5126955210227966,
      "q3_ms": 0.5366650914599542,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 0,
      "compile_s": 0.1959035589970881,
      "inner": 136,
      "median_ms": 0.2187186323948588,
      "q1_ms": 0.21862426103111562,
      "q3_ms": 0.21873726467781715,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 0,
      "compile_s": 0.19905659201322123,
      "inner": 173,
      "median_ms": 0.1175111156149439,
      "q1_ms": 0.1175049480001679,
      "q3_ms": 0.11756529479261993,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 0,
      "compile_s": 0.19567486600135453,
      "inner": 200,
      "median_ms": 0.08790557498286944,
      "q1_ms": 0.0878790225033299,
      "q3_ms": 0.08796050002274569,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 0,
      "compile_s": 0.20695644400257152,
      "inner": 200,
      "median_ms": 0.08892119498341344,
      "q1_ms": 0.08889196498785168,
      "q3_ms": 0.08899034750356805,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 32,
      "round": 0,
      "compile_s": 0.19991961399500724,
      "inner": 200,
      "median_ms": 0.0927703249908518,
      "q1_ms": 0.09271218750654953,
      "q3_ms": 0.09288601751904935,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 64,
      "round": 0,
      "compile_s": 0.21425733499927446,
      "inner": 200,
      "median_ms": 0.10192359004577156,
      "q1_ms": 0.10178612497838913,
      "q3_ms": 0.10196496499702334,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 1,
      "compile_s": 0.1580381569947349,
      "inner": 81,
      "median_ms": 0.5128637777589676,
      "q1_ms": 0.5124673455914792,
      "q3_ms": 0.5134232283855881,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 1,
      "compile_s": 0.19880989599914756,
      "inner": 147,
      "median_ms": 0.21942224485014064,
      "q1_ms": 0.21939494892191397,
      "q3_ms": 0.2195439013362653,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 1,
      "compile_s": 0.204060517004109,
      "inner": 200,
      "median_ms": 0.11779581000155304,
      "q1_ms": 0.11778741998568876,
      "q3_ms": 0.11781446999520995,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 1,
      "compile_s": 0.20040017399878707,
      "inner": 200,
      "median_ms": 0.08798476999800187,
      "q1_ms": 0.08795779252977809,
      "q3_ms": 0.08802514748822432,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 1,
      "compile_s": 0.20241980400169268,
      "inner": 200,
      "median_ms": 0.08898325999325607,
      "q1_ms": 0.0889539800118655,
      "q3_ms": 0.08900079250452109,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 32,
      "round": 1,
      "compile_s": 0.20780960899719503,
      "inner": 200,
      "median_ms": 0.09283910498197656,
      "q1_ms": 0.09275016000174219,
      "q3_ms": 0.09288346249377355,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 64,
      "round": 1,
      "compile_s": 0.20049602500512265,
      "inner": 200,
      "median_ms": 0.10178301999985706,
      "q1_ms": 0.10174517501582159,
      "q3_ms": 0.10193255249760114,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 0,
      "compile_s": 0.1619046099949628,
      "inner": 77,
      "median_ms": 0.5132403895984609,
      "q1_ms": 0.5127289090547469,
      "q3_ms": 0.5133652467319936,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 0,
      "compile_s": 0.2012535370013211,
      "inner": 135,
      "median_ms": 0.2199933776864782,
      "q1_ms": 0.2199728518866727,
      "q3_ms": 0.2201027629241623,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 0,
      "compile_s": 0.2026417839952046,
      "inner": 200,
      "median_ms": 0.11801116503193043,
      "q1_ms": 0.11800230000517331,
      "q3_ms": 0.1180393700269633,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 0,
      "compile_s": 0.20583246900059748,
      "inner": 200,
      "median_ms": 0.08860053996613715,
      "q1_ms": 0.0885626174567733,
      "q3_ms": 0.08865864750987384,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 0,
      "compile_s": 0.19631475600181147,
      "inner": 200,
      "median_ms": 0.08857073000399396,
      "q1_ms": 0.08855001502524829,
      "q3_ms": 0.08867349752108566,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 32,
      "round": 0,
      "compile_s": 0.20934371001203544,
      "inner": 200,
      "median_ms": 0.09269428002880886,
      "q1_ms": 0.09265512748243054,
      "q3_ms": 0.09275973003241234,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 64,
      "round": 0,
      "compile_s": 0.203799592985888,
      "inner": 200,
      "median_ms": 0.10202452998782974,
      "q1_ms": 0.10193903999606846,
      "q3_ms": 0.10206385250057792,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 1,
      "compile_s": 0.16052965899871197,
      "inner": 82,
      "median_ms": 0.513300414640093,
      "q1_ms": 0.5127661646749802,
      "q3_ms": 0.5137794938580156,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 1,
      "compile_s": 0.19945784099400043,
      "inner": 149,
      "median_ms": 0.21963597315554226,
      "q1_ms": 0.21959589928077633,
      "q3_ms": 0.2198525569764991,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 1,
      "compile_s": 0.19846949199563824,
      "inner": 200,
      "median_ms": 0.11857598503411282,
      "q1_ms": 0.1185322775199893,
      "q3_ms": 0.11862352497701067,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 1,
      "compile_s": 0.2197187040001154,
      "inner": 200,
      "median_ms": 0.08913824996852782,
      "q1_ms": 0.08911829248972936,
      "q3_ms": 0.08920938998926431,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 1,
      "compile_s": 0.20031444200139958,
      "inner": 200,
      "median_ms": 0.08849453493894543,
      "q1_ms": 0.08831757502775872,
      "q3_ms": 0.08853260504110949,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 32,
      "round": 1,
      "compile_s": 0.230894943990279,
      "inner": 200,
      "median_ms": 0.09269768495869357,
      "q1_ms": 0.09258860503905453,
      "q3_ms": 0.09277202745579416,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 64,
      "round": 1,
      "compile_s": 0.2045950379979331,
      "inner": 200,
      "median_ms": 0.10204752506979275,
      "q1_ms": 0.10192904752329923,
      "q3_ms": 0.10205593749560649,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 0,
      "compile_s": 0.14417383199906908,
      "inner": 80,
      "median_ms": 0.4913596500045969,
      "q1_ms": 0.49114830617327243,
      "q3_ms": 0.49168789373652544,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 0,
      "compile_s": 0.3714025910012424,
      "inner": 91,
      "median_ms": 0.43685190110329214,
      "q1_ms": 0.4366970440189886,
      "q3_ms": 0.437145697795351,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 0,
      "compile_s": 0.38131467999482993,
      "inner": 132,
      "median_ms": 0.22685993176172345,
      "q1_ms": 0.2268478673971858,
      "q3_ms": 0.22688092427746265,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 0,
      "compile_s": 0.3947096759948181,
      "inner": 200,
      "median_ms": 0.12736305994621944,
      "q1_ms": 0.12652035999053624,
      "q3_ms": 0.12778405249264324,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 0,
      "compile_s": 0.37420001199643593,
      "inner": 195,
      "median_ms": 0.13198760518720612,
      "q1_ms": 0.13062183846695685,
      "q3_ms": 0.13200812562643433,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 32,
      "round": 0,
      "compile_s": 0.37211253501300234,
      "inner": 175,
      "median_ms": 0.1446432457305491,
      "q1_ms": 0.14420816566728586,
      "q3_ms": 0.14498437995955882,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 64,
      "round": 0,
      "compile_s": 0.37580200300726574,
      "inner": 167,
      "median_ms": 0.1578292514728717,
      "q1_ms": 0.15777216770778227,
      "q3_ms": 0.15795571554569293,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 1,
      "compile_s": 0.147130704994197,
      "inner": 79,
      "median_ms": 0.49101341758188616,
      "q1_ms": 0.49096378486727277,
      "q3_ms": 0.49128786716359685,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 1,
      "compile_s": 0.38603103200148325,
      "inner": 88,
      "median_ms": 0.43679380690842995,
      "q1_ms": 0.4365705397669279,
      "q3_ms": 0.43690565342496324,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 1,
      "compile_s": 0.3736731960088946,
      "inner": 148,
      "median_ms": 0.2266081486399147,
      "q1_ms": 0.22652684797444758,
      "q3_ms": 0.22665692567652268,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 1,
      "compile_s": 0.3775991519942181,
      "inner": 184,
      "median_ms": 0.12792963586790668,
      "q1_ms": 0.12648089404651924,
      "q3_ms": 0.1280589863792715,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 1,
      "compile_s": 0.38909433600201737,
      "inner": 200,
      "median_ms": 0.1320864399895072,
      "q1_ms": 0.1307931675182772,
      "q3_ms": 0.13238943749456666,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 32,
      "round": 1,
      "compile_s": 0.37307640799554065,
      "inner": 181,
      "median_ms": 0.14487872932824045,
      "q1_ms": 0.1448247596915446,
      "q3_ms": 0.14504030937308354,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 64,
      "round": 1,
      "compile_s": 0.3943708360020537,
      "inner": 165,
      "median_ms": 0.15790667878746084,
      "q1_ms": 0.1572498575445603,
      "q3_ms": 0.15794077577690285,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 0,
      "compile_s": 0.1640271750075044,
      "inner": 22,
      "median_ms": 2.061597227028952,
      "q1_ms": 2.0599957497324795,
      "q3_ms": 2.069313613570888,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 0,
      "compile_s": 0.21170072900713421,
      "inner": 71,
      "median_ms": 0.6192638590016966,
      "q1_ms": 0.6185324788200144,
      "q3_ms": 0.6196513521143983,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 0,
      "compile_s": 0.20033361799141858,
      "inner": 63,
      "median_ms": 0.6248762381599389,
      "q1_ms": 0.6245204682012886,
      "q3_ms": 0.6249920080003831,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 0,
      "compile_s": 0.2032542020024266,
      "inner": 60,
      "median_ms": 0.6482917167886626,
      "q1_ms": 0.6480427666247124,
      "q3_ms": 0.6485146583145251,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 0,
      "compile_s": 0.20545348100131378,
      "inner": 64,
      "median_ms": 0.6500447032067314,
      "q1_ms": 0.6496797266208887,
      "q3_ms": 0.6515209374811093,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 32,
      "round": 0,
      "compile_s": 0.21122061100322753,
      "inner": 63,
      "median_ms": 0.6696940159484271,
      "q1_ms": 0.6694552221377972,
      "q3_ms": 0.6705871428638357,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 64,
      "round": 0,
      "compile_s": 0.20216061300016008,
      "inner": 60,
      "median_ms": 0.7134020500719392,
      "q1_ms": 0.7128991084755398,
      "q3_ms": 0.7139621999764738,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 1,
      "compile_s": 0.17144474500673823,
      "inner": 23,
      "median_ms": 2.0685474351605238,
      "q1_ms": 2.066195347731042,
      "q3_ms": 2.072032695413446,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 1,
      "compile_s": 0.20650857700093184,
      "inner": 66,
      "median_ms": 0.6185476970241518,
      "q1_ms": 0.6182606363533221,
      "q3_ms": 0.6192362727068195,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 1,
      "compile_s": 0.20755695999832824,
      "inner": 69,
      "median_ms": 0.6229947536232192,
      "q1_ms": 0.6228404927769324,
      "q3_ms": 0.6244085144725682,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 1,
      "compile_s": 0.20961440200335346,
      "inner": 64,
      "median_ms": 0.6488211406576738,
      "q1_ms": 0.6484968515678702,
      "q3_ms": 0.6492385000456125,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 1,
      "compile_s": 0.2176612430048408,
      "inner": 65,
      "median_ms": 0.6494145076990557,
      "q1_ms": 0.6492801076651981,
      "q3_ms": 0.6504470846266486,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 32,
      "round": 1,
      "compile_s": 0.21353355499741156,
      "inner": 59,
      "median_ms": 0.6695539153859776,
      "q1_ms": 0.6692033303976532,
      "q3_ms": 0.6701162287677338,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 64,
      "round": 1,
      "compile_s": 0.21063336900260765,
      "inner": 61,
      "median_ms": 0.7128207541173935,
      "q1_ms": 0.7124584098775932,
      "q3_ms": 0.7130447704751487,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 0,
      "compile_s": 0.15972435699950438,
      "inner": 22,
      "median_ms": 2.0701242273472333,
      "q1_ms": 2.0577020910753743,
      "q3_ms": 2.0714961819149638,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 0,
      "compile_s": 0.20228811400011182,
      "inner": 54,
      "median_ms": 0.6214875740826957,
      "q1_ms": 0.6213483146792795,
      "q3_ms": 0.621625805554989,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 0,
      "compile_s": 0.21134327200707048,
      "inner": 65,
      "median_ms": 0.6283005692799074,
      "q1_ms": 0.6277664769396114,
      "q3_ms": 0.6290808923390264,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 0,
      "compile_s": 0.20429490000242367,
      "inner": 61,
      "median_ms": 0.658192491964208,
      "q1_ms": 0.6575029180909996,
      "q3_ms": 0.6587286721905464,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 0,
      "compile_s": 0.20158574600645807,
      "inner": 64,
      "median_ms": 0.6563982967691118,
      "q1_ms": 0.6560736094343156,
      "q3_ms": 0.658536382843522,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 32,
      "round": 0,
      "compile_s": 0.20675532500899862,
      "inner": 63,
      "median_ms": 0.6715675713383549,
      "q1_ms": 0.6714794603000869,
      "q3_ms": 0.6718185080292945,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 64,
      "round": 0,
      "compile_s": 0.20287901200936176,
      "inner": 58,
      "median_ms": 0.7138486895389083,
      "q1_ms": 0.7135094999504308,
      "q3_ms": 0.7154030948907443,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 1,
      "compile_s": 0.16879874998994637,
      "inner": 22,
      "median_ms": 2.078187227280895,
      "q1_ms": 2.076939613537186,
      "q3_ms": 2.0833834092958234,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 1,
      "compile_s": 0.20473335900169332,
      "inner": 66,
      "median_ms": 0.6210890150572541,
      "q1_ms": 0.6209441286927992,
      "q3_ms": 0.6212887954351379,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 1,
      "compile_s": 0.19731874599528965,
      "inner": 68,
      "median_ms": 0.6281541763575446,
      "q1_ms": 0.6278650956065737,
      "q3_ms": 0.6283789411741857,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 1,
      "compile_s": 0.2080301160021918,
      "inner": 66,
      "median_ms": 0.6579168484401607,
      "q1_ms": 0.6574880226804508,
      "q3_ms": 0.6585945908833917,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 1,
      "compile_s": 0.20500374199764337,
      "inner": 66,
      "median_ms": 0.6564464999951959,
      "q1_ms": 0.6563534470015756,
      "q3_ms": 0.657355568171514,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 32,
      "round": 1,
      "compile_s": 0.2110618899896508,
      "inner": 64,
      "median_ms": 0.6725075156737148,
      "q1_ms": 0.6721660702169174,
      "q3_ms": 0.6736315078796906,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 64,
      "round": 1,
      "compile_s": 0.22486493799078744,
      "inner": 61,
      "median_ms": 0.7134039836676149,
      "q1_ms": 0.7132637212157906,
      "q3_ms": 0.7139529754239761,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 0,
      "compile_s": 0.1507243230007589,
      "inner": 22,
      "median_ms": 2.0676409092795276,
      "q1_ms": 2.065383204766973,
      "q3_ms": 2.0717493637427897,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 0,
      "compile_s": 0.3890685759979533,
      "inner": 47,
      "median_ms": 0.9957922553207647,
      "q1_ms": 0.9930227552491062,
      "q3_ms": 0.9973313296560514,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 0,
      "compile_s": 0.39786549599375576,
      "inner": 42,
      "median_ms": 1.0285099283681205,
      "q1_ms": 1.0206611668330705,
      "q3_ms": 1.0310071666026488,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 0,
      "compile_s": 0.3784972310095327,
      "inner": 48,
      "median_ms": 0.9458129373645837,
      "q1_ms": 0.9454161562037673,
      "q3_ms": 0.9487281458859798,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 0,
      "compile_s": 0.3923288140067598,
      "inner": 49,
      "median_ms": 0.9659726326520156,
      "q1_ms": 0.9643207448213932,
      "q3_ms": 0.9660750612074851,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 32,
      "round": 0,
      "compile_s": 0.3847218150040135,
      "inner": 47,
      "median_ms": 0.9871333189794119,
      "q1_ms": 0.9868366913461285,
      "q3_ms": 0.9914714257242753,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 64,
      "round": 0,
      "compile_s": 0.3861007429950405,
      "inner": 42,
      "median_ms": 1.1356068093196623,
      "q1_ms": 1.1354755953929963,
      "q3_ms": 1.1378083334074334,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 1,
      "compile_s": 0.15527630200085696,
      "inner": 22,
      "median_ms": 2.076464863917367,
      "q1_ms": 2.0737738635613248,
      "q3_ms": 2.0819069772683587,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 1,
      "compile_s": 0.3799806669994723,
      "inner": 45,
      "median_ms": 0.9925053112157104,
      "q1_ms": 0.9922223222222076,
      "q3_ms": 0.9933200666435167,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 1,
      "compile_s": 0.39158544699603226,
      "inner": 45,
      "median_ms": 1.0275257777215705,
      "q1_ms": 1.025894577711976,
      "q3_ms": 1.0284903667828378,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 1,
      "compile_s": 0.39225783899019007,
      "inner": 48,
      "median_ms": 0.9482605000812327,
      "q1_ms": 0.9468594168235238,
      "q3_ms": 0.948747875099798,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 1,
      "compile_s": 0.39264170500973705,
      "inner": 49,
      "median_ms": 0.9644542448341429,
      "q1_ms": 0.9635013162291476,
      "q3_ms": 0.9653266122545667,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 32,
      "round": 1,
      "compile_s": 0.386445726006059,
      "inner": 48,
      "median_ms": 0.9869371876144821,
      "q1_ms": 0.9864903022389626,
      "q3_ms": 0.9880497188229734,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 64,
      "round": 1,
      "compile_s": 0.39792177300842013,
      "inner": 41,
      "median_ms": 1.1343903173749312,
      "q1_ms": 1.1341623293247256,
      "q3_ms": 1.1362980731064454,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 0,
      "compile_s": 0.1612708710017614,
      "inner": 22,
      "median_ms": 2.038187454507517,
      "q1_ms": 2.0358111363417595,
      "q3_ms": 2.039426204811951,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 0,
      "compile_s": 0.2087496970052598,
      "inner": 51,
      "median_ms": 0.8586570783175856,
      "q1_ms": 0.8577554706112007,
      "q3_ms": 0.8588743529638604,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 0,
      "compile_s": 0.20814609400986228,
      "inner": 91,
      "median_ms": 0.4397375823528209,
      "q1_ms": 0.43940722525304,
      "q3_ms": 0.4401869505322743,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 0,
      "compile_s": 0.20333141800074372,
      "inner": 111,
      "median_ms": 0.31567465758836316,
      "q1_ms": 0.31558065772211796,
      "q3_ms": 0.31578432429563,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 0,
      "compile_s": 0.20782061899080873,
      "inner": 115,
      "median_ms": 0.31564573911727284,
      "q1_ms": 0.3154555521652346,
      "q3_ms": 0.31569203483345715,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 32,
      "round": 0,
      "compile_s": 0.21540435300266836,
      "inner": 114,
      "median_ms": 0.32027228939624736,
      "q1_ms": 0.3201734122596559,
      "q3_ms": 0.3205629648602968,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 64,
      "round": 0,
      "compile_s": 0.21214482799405232,
      "inner": 110,
      "median_ms": 0.33938858184476633,
      "q1_ms": 0.3391855636436958,
      "q3_ms": 0.3396213000418026,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 1,
      "compile_s": 0.16993436800839845,
      "inner": 23,
      "median_ms": 2.037184173747411,
      "q1_ms": 2.03199891320369,
      "q3_ms": 2.038215674237226,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 1,
      "compile_s": 0.2074023879977176,
      "inner": 51,
      "median_ms": 0.8534085097969236,
      "q1_ms": 0.8532915295183878,
      "q3_ms": 0.8537423922212375,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 1,
      "compile_s": 0.19238337001297623,
      "inner": 89,
      "median_ms": 0.43794387651298694,
      "q1_ms": 0.43787052251563374,
      "q3_ms": 0.4381333820173894,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 1,
      "compile_s": 0.20023937900259625,
      "inner": 118,
      "median_ms": 0.3150190762934764,
      "q1_ms": 0.3149637542896849,
      "q3_ms": 0.3151805634596543,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 1,
      "compile_s": 0.20357927700388245,
      "inner": 113,
      "median_ms": 0.315666106139079,
      "q1_ms": 0.3156044380198464,
      "q3_ms": 0.3157330885404861,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 32,
      "round": 1,
      "compile_s": 0.1999551920016529,
      "inner": 111,
      "median_ms": 0.32068341435562514,
      "q1_ms": 0.32054684685987866,
      "q3_ms": 0.32076437380489986,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 64,
      "round": 1,
      "compile_s": 0.19987951799703296,
      "inner": 114,
      "median_ms": 0.3394198157414059,
      "q1_ms": 0.3393628640128953,
      "q3_ms": 0.33954192986129783,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 0,
      "compile_s": 0.1575172219891101,
      "inner": 23,
      "median_ms": 2.0430256522756634,
      "q1_ms": 2.041622565150686,
      "q3_ms": 2.0464094350640867,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 0,
      "compile_s": 0.2030461389949778,
      "inner": 50,
      "median_ms": 0.8526640001218766,
      "q1_ms": 0.852259130042512,
      "q3_ms": 0.8529211800487246,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 0,
      "compile_s": 0.2107963479938917,
      "inner": 86,
      "median_ms": 0.4376084301796699,
      "q1_ms": 0.4374989128040778,
      "q3_ms": 0.43839660463353813,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 0,
      "compile_s": 0.21262323399423622,
      "inner": 110,
      "median_ms": 0.315619072617582,
      "q1_ms": 0.31548049545645795,
      "q3_ms": 0.31583125455654226,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 0,
      "compile_s": 0.21019055201031733,
      "inner": 118,
      "median_ms": 0.3164490677220984,
      "q1_ms": 0.3164100974464742,
      "q3_ms": 0.31675799569495855,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 32,
      "round": 0,
      "compile_s": 0.2093351040093694,
      "inner": 114,
      "median_ms": 0.3213443070355999,
      "q1_ms": 0.3212998376591458,
      "q3_ms": 0.3214174298990289,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 64,
      "round": 0,
      "compile_s": 0.20852812800148968,
      "inner": 102,
      "median_ms": 0.3414718725538228,
      "q1_ms": 0.3411137402187759,
      "q3_ms": 0.34165044602488315,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 1,
      "compile_s": 0.16104166800505482,
      "inner": 22,
      "median_ms": 2.0372620909421877,
      "q1_ms": 2.035825045326915,
      "q3_ms": 2.038154681940796,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 1,
      "compile_s": 0.2062815179961035,
      "inner": 45,
      "median_ms": 0.8536775997425947,
      "q1_ms": 0.8535539333630975,
      "q3_ms": 0.8537878222543643,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 1,
      "compile_s": 0.2016540650074603,
      "inner": 89,
      "median_ms": 0.4373869663784594,
      "q1_ms": 0.4372151909929731,
      "q3_ms": 0.4376962079026925,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 1,
      "compile_s": 0.19702539099671412,
      "inner": 107,
      "median_ms": 0.3150745606970338,
      "q1_ms": 0.315029378496241,
      "q3_ms": 0.3151927290486485,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 1,
      "compile_s": 0.2020286340120947,
      "inner": 109,
      "median_ms": 0.31633320178541807,
      "q1_ms": 0.3162640780520224,
      "q3_ms": 0.31711652297142107,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 32,
      "round": 1,
      "compile_s": 0.20962545300426427,
      "inner": 110,
      "median_ms": 0.3216350910158574,
      "q1_ms": 0.32151222734352236,
      "q3_ms": 0.32187472725953825,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 64,
      "round": 1,
      "compile_s": 0.20252790799713694,
      "inner": 109,
      "median_ms": 0.34080141284310794,
      "q1_ms": 0.3406741972633535,
      "q3_ms": 0.34091883031128045,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 0,
      "compile_s": 0.13645175599958748,
      "inner": 23,
      "median_ms": 1.9722256522230888,
      "q1_ms": 1.972007413037911,
      "q3_ms": 1.9728209128674203,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 0,
      "compile_s": 0.39144264700007625,
      "inner": 26,
      "median_ms": 1.7179452692149566,
      "q1_ms": 1.7172564232448795,
      "q3_ms": 1.721991730636756,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 0,
      "compile_s": 0.3771271900041029,
      "inner": 50,
      "median_ms": 0.8680100200581364,
      "q1_ms": 0.8677815900591668,
      "q3_ms": 0.8685583600890823,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 0,
      "compile_s": 0.3739212159998715,
      "inner": 85,
      "median_ms": 0.46948437640607793,
      "q1_ms": 0.4688849529723966,
      "q3_ms": 0.4701870529542622,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 0,
      "compile_s": 0.37944557500304654,
      "inner": 81,
      "median_ms": 0.4762605926378366,
      "q1_ms": 0.47615403713684323,
      "q3_ms": 0.47638094443238604,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 32,
      "round": 0,
      "compile_s": 0.38671623800473753,
      "inner": 85,
      "median_ms": 0.48551422352080836,
      "q1_ms": 0.48540908824446993,
      "q3_ms": 0.4862791412640089,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 64,
      "round": 0,
      "compile_s": 0.3838911829952849,
      "inner": 89,
      "median_ms": 0.4624023706864126,
      "q1_ms": 0.4621882697474473,
      "q3_ms": 0.46329063487963834,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 1,
      "compile_s": 0.1528290939895669,
      "inner": 23,
      "median_ms": 1.9748187826379484,
      "q1_ms": 1.9714565869830751,
      "q3_ms": 1.974906804009705,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 1,
      "compile_s": 0.38305063400184736,
      "inner": 26,
      "median_ms": 1.7245419613363293,
      "q1_ms": 1.7241899230821915,
      "q3_ms": 1.724705903772981,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 1,
      "compile_s": 0.3838218430028064,
      "inner": 48,
      "median_ms": 0.868680166604463,
      "q1_ms": 0.8683076979044321,
      "q3_ms": 0.8701169582915705,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 1,
      "compile_s": 0.37334145899512805,
      "inner": 90,
      "median_ms": 0.46982511105145225,
      "q1_ms": 0.46930645556939354,
      "q3_ms": 0.4699057056212849,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 1,
      "compile_s": 0.37910857699171174,
      "inner": 83,
      "median_ms": 0.47292790364429177,
      "q1_ms": 0.471501114391806,
      "q3_ms": 0.47311223503370514,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 32,
      "round": 1,
      "compile_s": 0.37804408399097156,
      "inner": 87,
      "median_ms": 0.48554183907897863,
      "q1_ms": 0.48538027577110215,
      "q3_ms": 0.4861079137665273,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 64,
      "round": 1,
      "compile_s": 0.40887448900321033,
      "inner": 93,
      "median_ms": 0.46257365583216836,
      "q1_ms": 0.46227718277844365,
      "q3_ms": 0.46319995164340466,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 0,
      "compile_s": 0.16667146599502303,
      "inner": 6,
      "median_ms": 8.227431998723963,
      "q1_ms": 8.194897249874582,
      "q3_ms": 8.235671083336154,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 0,
      "compile_s": 0.21413410099921748,
      "inner": 19,
      "median_ms": 2.43566705275474,
      "q1_ms": 2.4353540789697465,
      "q3_ms": 2.4388116312605375,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 0,
      "compile_s": 0.20980546399368905,
      "inner": 19,
      "median_ms": 2.4460355787979147,
      "q1_ms": 2.4421521581978975,
      "q3_ms": 2.446306605155482,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 0,
      "compile_s": 0.20032399900082964,
      "inner": 19,
      "median_ms": 2.517692210140491,
      "q1_ms": 2.511888552889056,
      "q3_ms": 2.5184164474730526,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 0,
      "compile_s": 0.2004892410041066,
      "inner": 18,
      "median_ms": 2.5182448330775107,
      "q1_ms": 2.5176304168578802,
      "q3_ms": 2.521920999748141,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 32,
      "round": 0,
      "compile_s": 0.1996465380070731,
      "inner": 19,
      "median_ms": 2.5082569474378893,
      "q1_ms": 2.506606710468125,
      "q3_ms": 2.508407262978021,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 64,
      "round": 0,
      "compile_s": 0.20274112100014463,
      "inner": 18,
      "median_ms": 2.545550444564368,
      "q1_ms": 2.5449128334003035,
      "q3_ms": 2.5507085277543715,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 1,
      "compile_s": 0.17323656400549226,
      "inner": 6,
      "median_ms": 8.26666966653041,
      "q1_ms": 8.245483666541986,
      "q3_ms": 8.279629083214484,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 1,
      "compile_s": 0.20678696400136687,
      "inner": 19,
      "median_ms": 2.436537631805741,
      "q1_ms": 2.4355897893608356,
      "q3_ms": 2.4383800524734207,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 1,
      "compile_s": 0.2003379249945283,
      "inner": 19,
      "median_ms": 2.4945346843818887,
      "q1_ms": 2.4847211320057365,
      "q3_ms": 2.512323078821952,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 1,
      "compile_s": 0.2046905689931009,
      "inner": 19,
      "median_ms": 2.5148963682831416,
      "q1_ms": 2.511642578419445,
      "q3_ms": 2.5269852895259026,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 1,
      "compile_s": 0.2213952150050318,
      "inner": 18,
      "median_ms": 2.524373166832245,
      "q1_ms": 2.517503749711776,
      "q3_ms": 2.527543861005041,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 32,
      "round": 1,
      "compile_s": 0.20311646099435166,
      "inner": 19,
      "median_ms": 2.549191736953186,
      "q1_ms": 2.5306358423365864,
      "q3_ms": 2.6019238947284116,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 64,
      "round": 1,
      "compile_s": 0.21844799099199008,
      "inner": 18,
      "median_ms": 2.757613777551645,
      "q1_ms": 2.628106583061809,
      "q3_ms": 2.898525305984852,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 0,
      "compile_s": 0.17189091799082235,
      "inner": 6,
      "median_ms": 8.327264166534102,
      "q1_ms": 8.307158749327453,
      "q3_ms": 8.351190333390454,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 0,
      "compile_s": 0.2042922050022753,
      "inner": 17,
      "median_ms": 2.448648647473687,
      "q1_ms": 2.4431125583599647,
      "q3_ms": 2.5004501469940057,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 0,
      "compile_s": 0.20646288999705575,
      "inner": 19,
      "median_ms": 2.6117854209396207,
      "q1_ms": 2.5618438682469886,
      "q3_ms": 2.7774305524776297,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 0,
      "compile_s": 0.20117101199866738,
      "inner": 18,
      "median_ms": 2.5768307225209557,
      "q1_ms": 2.5388575001367526,
      "q3_ms": 2.5798169165985505,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 0,
      "compile_s": 0.20079435899970122,
      "inner": 19,
      "median_ms": 2.5590888950606123,
      "q1_ms": 2.525102789437178,
      "q3_ms": 2.5826236841997963,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 32,
      "round": 0,
      "compile_s": 0.20263330699526705,
      "inner": 17,
      "median_ms": 2.5476405297762112,
      "q1_ms": 2.5366123822202686,
      "q3_ms": 2.559931176673511,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 64,
      "round": 0,
      "compile_s": 0.20327020299737342,
      "inner": 18,
      "median_ms": 2.5949677777437805,
      "q1_ms": 2.587162444089901,
      "q3_ms": 2.6068277224516754,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 1,
      "compile_s": 0.17953507699712645,
      "inner": 5,
      "median_ms": 8.408214000519365,
      "q1_ms": 8.326315099839121,
      "q3_ms": 8.584374199563172,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 1,
      "compile_s": 0.20248039900616277,
      "inner": 19,
      "median_ms": 2.4911240524440807,
      "q1_ms": 2.480860657478364,
      "q3_ms": 2.4945162631889906,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 1,
      "compile_s": 0.20459228300023824,
      "inner": 19,
      "median_ms": 2.518316000197573,
      "q1_ms": 2.4944418948603255,
      "q3_ms": 2.7343988951638734,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 1,
      "compile_s": 0.19782891999057028,
      "inner": 18,
      "median_ms": 2.5705126107722105,
      "q1_ms": 2.5631122500474,
      "q3_ms": 2.593255666498509,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 1,
      "compile_s": 0.1980231080087833,
      "inner": 18,
      "median_ms": 2.564762888747888,
      "q1_ms": 2.5509740274477126,
      "q3_ms": 2.594160971865575,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 32,
      "round": 1,
      "compile_s": 0.20440996799152344,
      "inner": 19,
      "median_ms": 2.529653737407276,
      "q1_ms": 2.5242485001064394,
      "q3_ms": 2.5555353682935484,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 4,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 64,
      "round": 1,
      "compile_s": 0.20127673099341337,
      "inner": 18,
      "median_ms": 2.5682706109364517,
      "q1_ms": 2.55581288891133,
      "q3_ms": 2.5795046668564385,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 0,
      "compile_s": 0.14474085498659406,
      "inner": 6,
      "median_ms": 8.747303666799175,
      "q1_ms": 8.50047766653006,
      "q3_ms": 9.157552999871161,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 0,
      "compile_s": 0.3702558530058013,
      "inner": 12,
      "median_ms": 3.979295167179468,
      "q1_ms": 3.948156666713961,
      "q3_ms": 4.120802333394143,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 0,
      "compile_s": 0.39549414099019486,
      "inner": 11,
      "median_ms": 4.224274728172036,
      "q1_ms": 4.087693181645591,
      "q3_ms": 4.4194115451484155,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 0,
      "compile_s": 0.3779730090027442,
      "inner": 14,
      "median_ms": 3.641651714653043,
      "q1_ms": 3.623837285496328,
      "q3_ms": 3.6553293570510244,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 0,
      "compile_s": 0.4011678820097586,
      "inner": 14,
      "median_ms": 3.533677500464754,
      "q1_ms": 3.488096249514326,
      "q3_ms": 3.8031185358704533,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 32,
      "round": 0,
      "compile_s": 0.39913473700289614,
      "inner": 15,
      "median_ms": 3.509648466327538,
      "q1_ms": 3.463438333225592,
      "q3_ms": 3.7708862330570505,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 64,
      "round": 0,
      "compile_s": 0.38553467199380975,
      "inner": 14,
      "median_ms": 3.6378960004575283,
      "q1_ms": 3.6130482497225915,
      "q3_ms": 3.664256249716605,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 1,
      "compile_s": 0.16725417799898423,
      "inner": 5,
      "median_ms": 8.350428799167275,
      "q1_ms": 8.339223600341938,
      "q3_ms": 8.405778399901465,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 1,
      "compile_s": 0.39539523400890175,
      "inner": 12,
      "median_ms": 4.346176582960955,
      "q1_ms": 4.069562499959527,
      "q3_ms": 4.358066124647545,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 1,
      "compile_s": 0.3957887479919009,
      "inner": 12,
      "median_ms": 4.055454999615904,
      "q1_ms": 4.0206549995976575,
      "q3_ms": 4.269551667069512,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 1,
      "compile_s": 0.37618283300253097,
      "inner": 14,
      "median_ms": 3.651200499526957,
      "q1_ms": 3.612655035145248,
      "q3_ms": 3.6907563210531538,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 1,
      "compile_s": 0.426782031005132,
      "inner": 15,
      "median_ms": 3.5166730667697266,
      "q1_ms": 3.5039146333777658,
      "q3_ms": 3.7382451662172875,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 32,
      "round": 1,
      "compile_s": 0.38461823000398,
      "inner": 15,
      "median_ms": 3.4707902668742463,
      "q1_ms": 3.4387179335074807,
      "q3_ms": 3.5547588001160575,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 64,
      "round": 1,
      "compile_s": 0.38873797199630644,
      "inner": 14,
      "median_ms": 3.6303467142196104,
      "q1_ms": 3.5908721074520145,
      "q3_ms": 3.658343249948562,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep_h8",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 1024,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 0,
      "round": 0,
      "compile_s": 0.14911859600397293,
      "inner": 200,
      "median_ms": 0.047411124978680164,
      "q1_ms": 0.047391817533934955,
      "q3_ms": 0.04750663498271024,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 1024,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 2,
      "round": 0,
      "compile_s": 0.1933701870002551,
      "inner": 200,
      "median_ms": 0.04521978000411764,
      "q1_ms": 0.04503635747823864,
      "q3_ms": 0.04558081749564735,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 1024,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 4,
      "round": 0,
      "compile_s": 0.24108325400447939,
      "inner": 200,
      "median_ms": 0.049683524994179606,
      "q1_ms": 0.04923242246150039,
      "q3_ms": 0.051069272521999665,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 1024,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 8,
      "round": 0,
      "compile_s": 0.2337823449925054,
      "inner": 200,
      "median_ms": 0.04226063501846511,
      "q1_ms": 0.042006902513094246,
      "q3_ms": 0.04364968997833785,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 1024,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 16,
      "round": 0,
      "compile_s": 0.22799678199226037,
      "inner": 200,
      "median_ms": 0.04832475504372269,
      "q1_ms": 0.041673625019029714,
      "q3_ms": 0.04918926497339271,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 1024,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 0,
      "round": 1,
      "compile_s": 0.15140814600454178,
      "inner": 200,
      "median_ms": 0.04745706501125824,
      "q1_ms": 0.04743937999592163,
      "q3_ms": 0.04815370499272831,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 1024,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 2,
      "round": 1,
      "compile_s": 0.19421380199491978,
      "inner": 200,
      "median_ms": 0.04302799003198743,
      "q1_ms": 0.04130836001422722,
      "q3_ms": 0.04373109248263063,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 1024,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 4,
      "round": 1,
      "compile_s": 0.19274748899624683,
      "inner": 200,
      "median_ms": 0.05075344997749198,
      "q1_ms": 0.04498378748394316,
      "q3_ms": 0.05132718248205492,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 1024,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 8,
      "round": 1,
      "compile_s": 0.19242689201200847,
      "inner": 200,
      "median_ms": 0.057574905004003085,
      "q1_ms": 0.05743669251387473,
      "q3_ms": 0.057943499996326864,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 1024,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 16,
      "round": 1,
      "compile_s": 0.17899781100277323,
      "inner": 200,
      "median_ms": 0.05288572996505536,
      "q1_ms": 0.05222475501795998,
      "q3_ms": 0.05399626003054436,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 32,
      "lkv": 1024,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 0,
      "round": 0,
      "compile_s": 0.1603905859956285,
      "inner": 200,
      "median_ms": 0.05548500004806556,
      "q1_ms": 0.055466440026066266,
      "q3_ms": 0.055505060008727014,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 32,
      "lkv": 1024,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 2,
      "round": 0,
      "compile_s": 0.356983728997875,
      "inner": 200,
      "median_ms": 0.03948934499931056,
      "q1_ms": 0.03902709504473023,
      "q3_ms": 0.03967985750932712,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 32,
      "lkv": 1024,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 4,
      "round": 0,
      "compile_s": 0.33956866900552996,
      "inner": 200,
      "median_ms": 0.038426934988819994,
      "q1_ms": 0.03834923751128372,
      "q3_ms": 0.038599107501795515,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 32,
      "lkv": 1024,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 8,
      "round": 0,
      "compile_s": 0.3191082579869544,
      "inner": 200,
      "median_ms": 0.03999970998847857,
      "q1_ms": 0.039845764986239374,
      "q3_ms": 0.040109742476488464,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 32,
      "lkv": 1024,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 16,
      "round": 0,
      "compile_s": 0.33091755499481224,
      "inner": 200,
      "median_ms": 0.04904141002043616,
      "q1_ms": 0.048918480024440214,
      "q3_ms": 0.04912759752187412,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 32,
      "lkv": 1024,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 0,
      "round": 1,
      "compile_s": 0.15921263799828012,
      "inner": 200,
      "median_ms": 0.05539758501981851,
      "q1_ms": 0.055391320020135026,
      "q3_ms": 0.05540492253203411,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 32,
      "lkv": 1024,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 2,
      "round": 1,
      "compile_s": 0.24361471300653648,
      "inner": 200,
      "median_ms": 0.03745424502994865,
      "q1_ms": 0.037033797525509726,
      "q3_ms": 0.048615154992148746,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 32,
      "lkv": 1024,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 4,
      "round": 1,
      "compile_s": 0.24488064799515996,
      "inner": 200,
      "median_ms": 0.038126864965306595,
      "q1_ms": 0.03809522750088945,
      "q3_ms": 0.03822946000582306,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 32,
      "lkv": 1024,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 8,
      "round": 1,
      "compile_s": 0.2413879609957803,
      "inner": 200,
      "median_ms": 0.039958885026862845,
      "q1_ms": 0.039862074991106056,
      "q3_ms": 0.04004507001809543,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 32,
      "lkv": 1024,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 16,
      "round": 1,
      "compile_s": 0.23617798699706327,
      "inner": 200,
      "median_ms": 0.04898890998447314,
      "q1_ms": 0.04894623245490948,
      "q3_ms": 0.04912967997370288,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 64,
      "lkv": 1024,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 0,
      "round": 0,
      "compile_s": 0.21696933399653062,
      "inner": 200,
      "median_ms": 0.0818219649954699,
      "q1_ms": 0.08181472250726074,
      "q3_ms": 0.08184362999600125,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 64,
      "lkv": 1024,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 2,
      "round": 0,
      "compile_s": 0.4747910389996832,
      "inner": 200,
      "median_ms": 0.04253395498380996,
      "q1_ms": 0.04236220251186751,
      "q3_ms": 0.044951324998692144,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 64,
      "lkv": 1024,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 4,
      "round": 0,
      "compile_s": 0.441623302991502,
      "inner": 200,
      "median_ms": 0.04857999003434088,
      "q1_ms": 0.048456727527081966,
      "q3_ms": 0.048759880010038614,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 64,
      "lkv": 1024,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 8,
      "round": 0,
      "compile_s": 0.4573202540050261,
      "inner": 200,
      "median_ms": 0.057433085021330044,
      "q1_ms": 0.05727162748371484,
      "q3_ms": 0.05791782503365539,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 64,
      "lkv": 1024,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 16,
      "round": 0,
      "compile_s": 0.46138428000267595,
      "inner": 200,
      "median_ms": 0.07159586995840073,
      "q1_ms": 0.07132723749236902,
      "q3_ms": 0.07171354500314919,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 64,
      "lkv": 1024,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 0,
      "round": 1,
      "compile_s": 0.2209783770085778,
      "inner": 200,
      "median_ms": 0.07359025003097486,
      "q1_ms": 0.07357634000072721,
      "q3_ms": 0.073603870041552,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 64,
      "lkv": 1024,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 2,
      "round": 1,
      "compile_s": 0.36017442400043365,
      "inner": 200,
      "median_ms": 0.04255073501553852,
      "q1_ms": 0.04243644249072531,
      "q3_ms": 0.04430345001310343,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 64,
      "lkv": 1024,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 4,
      "round": 1,
      "compile_s": 0.3612601680069929,
      "inner": 200,
      "median_ms": 0.04865157505264506,
      "q1_ms": 0.048550657484156545,
      "q3_ms": 0.04985430005035596,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 64,
      "lkv": 1024,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 8,
      "round": 1,
      "compile_s": 0.36233090300811455,
      "inner": 200,
      "median_ms": 0.057651095048640855,
      "q1_ms": 0.057520247537468094,
      "q3_ms": 0.058040104995598085,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 64,
      "lkv": 1024,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 16,
      "round": 1,
      "compile_s": 0.3679205150110647,
      "inner": 200,
      "median_ms": 0.07313082503969781,
      "q1_ms": 0.0731047775116167,
      "q3_ms": 0.07317212748603197,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 1024,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 0,
      "round": 0,
      "compile_s": 0.1505427999945823,
      "inner": 163,
      "median_ms": 0.19475432518838212,
      "q1_ms": 0.1947191349524245,
      "q3_ms": 0.19476787735132411,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 1024,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 2,
      "round": 0,
      "compile_s": 0.18810693100385834,
      "inner": 169,
      "median_ms": 0.17893300593221703,
      "q1_ms": 0.17886693192785053,
      "q3_ms": 0.17899442306539243,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 1024,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 4,
      "round": 0,
      "compile_s": 0.19182624599488918,
      "inner": 181,
      "median_ms": 0.18090870720527305,
      "q1_ms": 0.18090046129417117,
      "q3_ms": 0.1809958618414686,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 1024,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 8,
      "round": 0,
      "compile_s": 0.23645180500170682,
      "inner": 149,
      "median_ms": 0.19016755032259347,
      "q1_ms": 0.19015413086866062,
      "q3_ms": 0.19019629865659762,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 1024,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 16,
      "round": 0,
      "compile_s": 0.21966565398906823,
      "inner": 153,
      "median_ms": 0.20849541180806686,
      "q1_ms": 0.20846545098393585,
      "q3_ms": 0.20852602613683535,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 1024,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 0,
      "round": 1,
      "compile_s": 0.15144054700795095,
      "inner": 164,
      "median_ms": 0.19512444513810162,
      "q1_ms": 0.19510483226196546,
      "q3_ms": 0.19515330486905352,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 1024,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 2,
      "round": 1,
      "compile_s": 0.19124652999744285,
      "inner": 159,
      "median_ms": 0.17909964772802336,
      "q1_ms": 0.17898612576624695,
      "q3_ms": 0.17925752200546483,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 1024,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 4,
      "round": 1,
      "compile_s": 0.19144282100023702,
      "inner": 154,
      "median_ms": 0.18106839612395464,
      "q1_ms": 0.18103440256650233,
      "q3_ms": 0.18109262983470417,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 1024,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 8,
      "round": 1,
      "compile_s": 0.189555881996057,
      "inner": 162,
      "median_ms": 0.19079762344596402,
      "q1_ms": 0.19078256171965619,
      "q3_ms": 0.19083756786939274,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 1024,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 16,
      "round": 1,
      "compile_s": 0.1781898429908324,
      "inner": 145,
      "median_ms": 0.2085825310947759,
      "q1_ms": 0.20854939658451696,
      "q3_ms": 0.2086255758607375,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 32,
      "lkv": 1024,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 0,
      "round": 0,
      "compile_s": 0.15629231400089338,
      "inner": 145,
      "median_ms": 0.2063746344306153,
      "q1_ms": 0.2061497551541583,
      "q3_ms": 0.20638747930582935,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 32,
      "lkv": 1024,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 2,
      "round": 0,
      "compile_s": 0.44123369599401485,
      "inner": 154,
      "median_ms": 0.2069748701284301,
      "q1_ms": 0.20695332795504007,
      "q3_ms": 0.2070857564784139,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 32,
      "lkv": 1024,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 4,
      "round": 0,
      "compile_s": 0.37473645800491795,
      "inner": 146,
      "median_ms": 0.2169240069135462,
      "q1_ms": 0.21655594517216914,
      "q3_ms": 0.2169570479839434,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 32,
      "lkv": 1024,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 8,
      "round": 0,
      "compile_s": 0.3320630410016747,
      "inner": 136,
      "median_ms": 0.2556194558565039,
      "q1_ms": 0.2555514228173278,
      "q3_ms": 0.2556300256829298,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 32,
      "lkv": 1024,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 16,
      "round": 0,
      "compile_s": 0.3325512050068937,
      "inner": 114,
      "median_ms": 0.33820041234752063,
      "q1_ms": 0.3381275000400931,
      "q3_ms": 0.33824447809979974,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 32,
      "lkv": 1024,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 0,
      "round": 1,
      "compile_s": 0.16518813000584487,
      "inner": 142,
      "median_ms": 0.20668093663555057,
      "q1_ms": 0.2061963556039239,
      "q3_ms": 0.2066954683113194,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 32,
      "lkv": 1024,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 2,
      "round": 1,
      "compile_s": 0.24264079100976232,
      "inner": 157,
      "median_ms": 0.20714926748064935,
      "q1_ms": 0.20714343313586003,
      "q3_ms": 0.2071666305571065,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 32,
      "lkv": 1024,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 4,
      "round": 1,
      "compile_s": 0.23795086999598425,
      "inner": 139,
      "median_ms": 0.21724594968853506,
      "q1_ms": 0.21708438851537262,
      "q3_ms": 0.21726861869348482,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 32,
      "lkv": 1024,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 8,
      "round": 1,
      "compile_s": 0.23657206300413236,
      "inner": 122,
      "median_ms": 0.2558468770642276,
      "q1_ms": 0.2557705532827178,
      "q3_ms": 0.2559949507933018,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 32,
      "lkv": 1024,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 16,
      "round": 1,
      "compile_s": 0.2350316190131707,
      "inner": 105,
      "median_ms": 0.33831706664170186,
      "q1_ms": 0.33824755235054044,
      "q3_ms": 0.3383711190378692,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 64,
      "lkv": 1024,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 0,
      "round": 0,
      "compile_s": 0.220284223993076,
      "inner": 114,
      "median_ms": 0.302191149099248,
      "q1_ms": 0.3013181842362157,
      "q3_ms": 0.3023351973592582,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 64,
      "lkv": 1024,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 2,
      "round": 0,
      "compile_s": 0.5719595420086989,
      "inner": 124,
      "median_ms": 0.279521782255371,
      "q1_ms": 0.2791286088693361,
      "q3_ms": 0.27969043953963857,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 64,
      "lkv": 1024,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 4,
      "round": 0,
      "compile_s": 0.45525186200393364,
      "inner": 120,
      "median_ms": 0.3079906416436037,
      "q1_ms": 0.30797707083062653,
      "q3_ms": 0.30832172499989,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 64,
      "lkv": 1024,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 8,
      "round": 0,
      "compile_s": 0.45376980799483135,
      "inner": 101,
      "median_ms": 0.3723813861142844,
      "q1_ms": 0.3721603168039767,
      "q3_ms": 0.37241600000583064,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 64,
      "lkv": 1024,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 16,
      "round": 0,
      "compile_s": 0.46426718300790526,
      "inner": 76,
      "median_ms": 0.5101627368850657,
      "q1_ms": 0.5101324079173777,
      "q3_ms": 0.5103113026531279,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 64,
      "lkv": 1024,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 0,
      "round": 1,
      "compile_s": 0.22766891100036446,
      "inner": 111,
      "median_ms": 0.3041577477454887,
      "q1_ms": 0.30317123427930465,
      "q3_ms": 0.30424728371341386,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 64,
      "lkv": 1024,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 2,
      "round": 1,
      "compile_s": 0.357081002992345,
      "inner": 121,
      "median_ms": 0.2793273471213092,
      "q1_ms": 0.2777903099146995,
      "q3_ms": 0.2793491694620187,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 64,
      "lkv": 1024,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 4,
      "round": 1,
      "compile_s": 0.35866599100700114,
      "inner": 117,
      "median_ms": 0.309657162323029,
      "q1_ms": 0.30939555986889866,
      "q3_ms": 0.3096858546576606,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 64,
      "lkv": 1024,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 8,
      "round": 1,
      "compile_s": 0.3614535230008187,
      "inner": 101,
      "median_ms": 0.37198806931333855,
      "q1_ms": 0.37122667816958604,
      "q3_ms": 0.3719976386383849,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 64,
      "lkv": 1024,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 16,
      "round": 1,
      "compile_s": 0.3644811299891444,
      "inner": 75,
      "median_ms": 0.5102632800117135,
      "q1_ms": 0.5101894133258611,
      "q3_ms": 0.5103250666676711,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 2048,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 0,
      "round": 0,
      "compile_s": 0.14660541199555155,
      "inner": 200,
      "median_ms": 0.07496874997741543,
      "q1_ms": 0.07495049252611352,
      "q3_ms": 0.07497786747990176,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 2048,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 2,
      "round": 0,
      "compile_s": 0.19011439800669905,
      "inner": 200,
      "median_ms": 0.04910477997327689,
      "q1_ms": 0.048890775033214595,
      "q3_ms": 0.0492995249805972,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 2048,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 4,
      "round": 0,
      "compile_s": 0.18962887999077793,
      "inner": 200,
      "median_ms": 0.05242496503342409,
      "q1_ms": 0.05227480247413041,
      "q3_ms": 0.05254423751466675,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 2048,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 8,
      "round": 0,
      "compile_s": 0.19025760899239685,
      "inner": 200,
      "median_ms": 0.053384074999485165,
      "q1_ms": 0.053349434965639375,
      "q3_ms": 0.05341277999832528,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 2048,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 16,
      "round": 0,
      "compile_s": 0.1906475770083489,
      "inner": 200,
      "median_ms": 0.05527159504708834,
      "q1_ms": 0.05525180749827996,
      "q3_ms": 0.055277857500186656,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 2048,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 0,
      "round": 1,
      "compile_s": 0.15066862700041384,
      "inner": 200,
      "median_ms": 0.07496414502384141,
      "q1_ms": 0.07493844499549596,
      "q3_ms": 0.07498620998376282,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 2048,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 2,
      "round": 1,
      "compile_s": 0.18822402200021315,
      "inner": 200,
      "median_ms": 0.048678019957151264,
      "q1_ms": 0.04863959999056533,
      "q3_ms": 0.04882710494712228,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 2048,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 4,
      "round": 1,
      "compile_s": 0.1952155769977253,
      "inner": 200,
      "median_ms": 0.05241639504674822,
      "q1_ms": 0.052259420008340385,
      "q3_ms": 0.05244234504061751,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 2048,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 8,
      "round": 1,
      "compile_s": 0.20184958500612993,
      "inner": 200,
      "median_ms": 0.053617015000781976,
      "q1_ms": 0.0535519199911505,
      "q3_ms": 0.05371397499402519,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 2048,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 16,
      "round": 1,
      "compile_s": 0.18878651600971352,
      "inner": 200,
      "median_ms": 0.05526779001229443,
      "q1_ms": 0.055192795007315,
      "q3_ms": 0.055469267535954714,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 32,
      "lkv": 2048,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 0,
      "round": 0,
      "compile_s": 0.16029649699339643,
      "inner": 200,
      "median_ms": 0.10351284501666669,
      "q1_ms": 0.10349227999540744,
      "q3_ms": 0.10358597999584163,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 32,
      "lkv": 2048,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 2,
      "round": 0,
      "compile_s": 0.24249609799881,
      "inner": 200,
      "median_ms": 0.052859580027870834,
      "q1_ms": 0.052701759959745687,
      "q3_ms": 0.05291000499710208,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 32,
      "lkv": 2048,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 4,
      "round": 0,
      "compile_s": 0.24351959399064071,
      "inner": 200,
      "median_ms": 0.055679365032119676,
      "q1_ms": 0.05553757000598125,
      "q3_ms": 0.055775350010662805,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 32,
      "lkv": 2048,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 8,
      "round": 0,
      "compile_s": 0.2435668620018987,
      "inner": 200,
      "median_ms": 0.059884064976358786,
      "q1_ms": 0.05975567750283517,
      "q3_ms": 0.059987794993503485,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 32,
      "lkv": 2048,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 16,
      "round": 0,
      "compile_s": 0.23869527799251955,
      "inner": 200,
      "median_ms": 0.07042696503049228,
      "q1_ms": 0.07029443499050103,
      "q3_ms": 0.07046650996926473,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 32,
      "lkv": 2048,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 0,
      "round": 1,
      "compile_s": 0.15921756699390244,
      "inner": 200,
      "median_ms": 0.09248558504623361,
      "q1_ms": 0.09244721499271691,
      "q3_ms": 0.09249066999473143,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 32,
      "lkv": 2048,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 2,
      "round": 1,
      "compile_s": 0.23852127899590414,
      "inner": 200,
      "median_ms": 0.05260315003397409,
      "q1_ms": 0.05243282750598155,
      "q3_ms": 0.05268239499855554,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 32,
      "lkv": 2048,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 4,
      "round": 1,
      "compile_s": 0.24610251899866853,
      "inner": 200,
      "median_ms": 0.055770889957784675,
      "q1_ms": 0.055532362530357204,
      "q3_ms": 0.055854700003692415,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 32,
      "lkv": 2048,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 8,
      "round": 1,
      "compile_s": 0.24314101699565072,
      "inner": 200,
      "median_ms": 0.06012853002175689,
      "q1_ms": 0.06009061002259841,
      "q3_ms": 0.06030040753103094,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 32,
      "lkv": 2048,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 16,
      "round": 1,
      "compile_s": 0.24131462299556006,
      "inner": 200,
      "median_ms": 0.07073554501403123,
      "q1_ms": 0.07054044745018473,
      "q3_ms": 0.07206856502307346,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 64,
      "lkv": 2048,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 0,
      "round": 0,
      "compile_s": 0.2138336219941266,
      "inner": 172,
      "median_ms": 0.1562359243824937,
      "q1_ms": 0.15618787209711818,
      "q3_ms": 0.15624911628587734,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 64,
      "lkv": 2048,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 2,
      "round": 0,
      "compile_s": 0.35740520599938463,
      "inner": 200,
      "median_ms": 0.07059583003865555,
      "q1_ms": 0.06983398998272605,
      "q3_ms": 0.07075711251673056,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 64,
      "lkv": 2048,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 4,
      "round": 0,
      "compile_s": 0.35801501899550203,
      "inner": 200,
      "median_ms": 0.07655649998923764,
      "q1_ms": 0.0758938024591771,
      "q3_ms": 0.0767116499991971,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 64,
      "lkv": 2048,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 8,
      "round": 0,
      "compile_s": 0.3635771810077131,
      "inner": 200,
      "median_ms": 0.0876780399994459,
      "q1_ms": 0.08704284253326477,
      "q3_ms": 0.08772395252890419,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 64,
      "lkv": 2048,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 16,
      "round": 0,
      "compile_s": 0.3596166890056338,
      "inner": 200,
      "median_ms": 0.10433790499519091,
      "q1_ms": 0.10383908746007364,
      "q3_ms": 0.104343865023111,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 64,
      "lkv": 2048,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 0,
      "round": 1,
      "compile_s": 0.22075652799685486,
      "inner": 174,
      "median_ms": 0.14014306328334342,
      "q1_ms": 0.1400875631986111,
      "q3_ms": 0.14017076431753533,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 64,
      "lkv": 2048,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 2,
      "round": 1,
      "compile_s": 0.35666504698747303,
      "inner": 200,
      "median_ms": 0.07076149500790052,
      "q1_ms": 0.06942508996871766,
      "q3_ms": 0.07077031001244904,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 64,
      "lkv": 2048,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 4,
      "round": 1,
      "compile_s": 0.3553154739929596,
      "inner": 200,
      "median_ms": 0.07656742498511449,
      "q1_ms": 0.07603566999023315,
      "q3_ms": 0.07677847250306513,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 64,
      "lkv": 2048,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 8,
      "round": 1,
      "compile_s": 0.35949343600077555,
      "inner": 200,
      "median_ms": 0.08765224498347379,
      "q1_ms": 0.0868930324941175,
      "q3_ms": 0.08766038252360886,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 64,
      "lkv": 2048,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 16,
      "round": 1,
      "compile_s": 0.36490753600082826,
      "inner": 200,
      "median_ms": 0.10431145499751437,
      "q1_ms": 0.10345706497901119,
      "q3_ms": 0.10431503502331907,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 2048,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 0,
      "round": 0,
      "compile_s": 0.1493561540119117,
      "inner": 102,
      "median_ms": 0.3717636764300622,
      "q1_ms": 0.371679838138459,
      "q3_ms": 0.37184657847711944,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 2048,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 2,
      "round": 0,
      "compile_s": 0.1963783359969966,
      "inner": 114,
      "median_ms": 0.3363468070012076,
      "q1_ms": 0.3361542061483896,
      "q3_ms": 0.3364079780830683,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 2048,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 4,
      "round": 0,
      "compile_s": 0.19427064900810365,
      "inner": 105,
      "median_ms": 0.33732557134325836,
      "q1_ms": 0.33722681435479207,
      "q3_ms": 0.3373824427918797,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 2048,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 8,
      "round": 0,
      "compile_s": 0.18858988500142004,
      "inner": 101,
      "median_ms": 0.3472373860895099,
      "q1_ms": 0.3471274305933782,
      "q3_ms": 0.3472596583720804,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 2048,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 16,
      "round": 0,
      "compile_s": 0.19285676500294358,
      "inner": 100,
      "median_ms": 0.364376609941246,
      "q1_ms": 0.3643212000315543,
      "q3_ms": 0.36448381491936743,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 2048,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 0,
      "round": 1,
      "compile_s": 0.15578741900390014,
      "inner": 104,
      "median_ms": 0.3719612595887156,
      "q1_ms": 0.3719239663181585,
      "q3_ms": 0.3719889038200983,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 2048,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 2,
      "round": 1,
      "compile_s": 0.19713822400080971,
      "inner": 106,
      "median_ms": 0.3363866509182185,
      "q1_ms": 0.33629671706189845,
      "q3_ms": 0.33646098587646606,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 2048,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 4,
      "round": 1,
      "compile_s": 0.19435264500498306,
      "inner": 109,
      "median_ms": 0.33735664221757183,
      "q1_ms": 0.33714964222472693,
      "q3_ms": 0.3373773623425556,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 2048,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 8,
      "round": 1,
      "compile_s": 0.19210778900014702,
      "inner": 105,
      "median_ms": 0.3466642761070814,
      "q1_ms": 0.3465542143593276,
      "q3_ms": 0.3467311714692707,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 2048,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 16,
      "round": 1,
      "compile_s": 0.19337741000344977,
      "inner": 103,
      "median_ms": 0.36435372812687017,
      "q1_ms": 0.3643058738004398,
      "q3_ms": 0.3645115534665325,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 32,
      "lkv": 2048,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 0,
      "round": 0,
      "compile_s": 0.15931014200032223,
      "inner": 95,
      "median_ms": 0.3982887263287251,
      "q1_ms": 0.3972481789165421,
      "q3_ms": 0.3985364052871438,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 32,
      "lkv": 2048,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 2,
      "round": 0,
      "compile_s": 0.23938908999843989,
      "inner": 100,
      "median_ms": 0.3757555400079582,
      "q1_ms": 0.37555916504061315,
      "q3_ms": 0.3757688649784541,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 32,
      "lkv": 2048,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 4,
      "round": 0,
      "compile_s": 0.24367559900565539,
      "inner": 103,
      "median_ms": 0.376485417556307,
      "q1_ms": 0.37635614561737574,
      "q3_ms": 0.3765165436791382,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 32,
      "lkv": 2048,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 8,
      "round": 0,
      "compile_s": 0.2448760989937,
      "inner": 95,
      "median_ms": 0.41326568416001175,
      "q1_ms": 0.41323156842018327,
      "q3_ms": 0.41328255266355546,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 32,
      "lkv": 2048,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 16,
      "round": 0,
      "compile_s": 0.24198379799781833,
      "inner": 85,
      "median_ms": 0.490981388172847,
      "q1_ms": 0.49095451172701465,
      "q3_ms": 0.49108228824265737,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 32,
      "lkv": 2048,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 0,
      "round": 1,
      "compile_s": 0.1613307840016205,
      "inner": 102,
      "median_ms": 0.398629882279327,
      "q1_ms": 0.39855297547867796,
      "q3_ms": 0.3987549804579056,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 32,
      "lkv": 2048,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 2,
      "round": 1,
      "compile_s": 0.239818442001706,
      "inner": 98,
      "median_ms": 0.3760814490343197,
      "q1_ms": 0.37590693875467784,
      "q3_ms": 0.37630319898848286,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 32,
      "lkv": 2048,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 4,
      "round": 1,
      "compile_s": 0.2396237939974526,
      "inner": 101,
      "median_ms": 0.37657595046514664,
      "q1_ms": 0.376474668270687,
      "q3_ms": 0.376640079170696,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 32,
      "lkv": 2048,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 8,
      "round": 1,
      "compile_s": 0.24314008498913608,
      "inner": 90,
      "median_ms": 0.4135630222865277,
      "q1_ms": 0.4134921610077274,
      "q3_ms": 0.41363872215798536,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 32,
      "lkv": 2048,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 16,
      "round": 1,
      "compile_s": 0.24430985799699556,
      "inner": 82,
      "median_ms": 0.49116342681993885,
      "q1_ms": 0.491124567083827,
      "q3_ms": 0.4912381951815006,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 64,
      "lkv": 2048,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 0,
      "round": 0,
      "compile_s": 0.22617074800655246,
      "inner": 73,
      "median_ms": 0.5933395479263559,
      "q1_ms": 0.5932551438279118,
      "q3_ms": 0.5934087945106886,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 64,
      "lkv": 2048,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 2,
      "round": 0,
      "compile_s": 0.3566306610009633,
      "inner": 87,
      "median_ms": 0.49574528743871543,
      "q1_ms": 0.4956792930814159,
      "q3_ms": 0.49583268967047656,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 64,
      "lkv": 2048,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 4,
      "round": 0,
      "compile_s": 0.363698690009187,
      "inner": 81,
      "median_ms": 0.5198558889484654,
      "q1_ms": 0.5188716171642789,
      "q3_ms": 0.521728629623626,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 64,
      "lkv": 2048,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 8,
      "round": 0,
      "compile_s": 0.3650810840044869,
      "inner": 72,
      "median_ms": 0.5877763610947618,
      "q1_ms": 0.587679222311029,
      "q3_ms": 0.5878357013797438,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 64,
      "lkv": 2048,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 16,
      "round": 0,
      "compile_s": 0.3633510119980201,
      "inner": 62,
      "median_ms": 0.7143131935604728,
      "q1_ms": 0.7140523226670696,
      "q3_ms": 0.7143384193055211,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 64,
      "lkv": 2048,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 0,
      "round": 1,
      "compile_s": 0.22820275199774187,
      "inner": 71,
      "median_ms": 0.5935143239133741,
      "q1_ms": 0.5932655986841164,
      "q3_ms": 0.5941479788637314,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 64,
      "lkv": 2048,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 2,
      "round": 1,
      "compile_s": 0.3574296820006566,
      "inner": 83,
      "median_ms": 0.49608689154279473,
      "q1_ms": 0.4959030481685143,
      "q3_ms": 0.49612413252372833,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 64,
      "lkv": 2048,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 4,
      "round": 1,
      "compile_s": 0.36087429799954407,
      "inner": 79,
      "median_ms": 0.5190458227556105,
      "q1_ms": 0.5189517151514257,
      "q3_ms": 0.5211595631362222,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 64,
      "lkv": 2048,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 8,
      "round": 1,
      "compile_s": 0.36335068198968656,
      "inner": 71,
      "median_ms": 0.5876952817115907,
      "q1_ms": 0.5872238168745421,
      "q3_ms": 0.5902719155567187,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 64,
      "lkv": 2048,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 16,
      "round": 1,
      "compile_s": 0.3723179130029166,
      "inner": 63,
      "median_ms": 0.7143457778655584,
      "q1_ms": 0.713682777824856,
      "q3_ms": 0.7146119921041343,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 0,
      "round": 0,
      "compile_s": 0.16054559900658205,
      "inner": 76,
      "median_ms": 0.5209925525138571,
      "q1_ms": 0.5208974276971083,
      "q3_ms": 0.5210199145183556,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 2,
      "round": 0,
      "compile_s": 0.2102560259954771,
      "inner": 52,
      "median_ms": 0.3134162884984667,
      "q1_ms": 0.31335145188047764,
      "q3_ms": 0.3137572212626065,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 4,
      "round": 0,
      "compile_s": 0.24303233200043906,
      "inner": 120,
      "median_ms": 0.31709719999829156,
      "q1_ms": 0.31691936255811015,
      "q3_ms": 0.3175097332738611,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 8,
      "round": 0,
      "compile_s": 0.20026405600947328,
      "inner": 110,
      "median_ms": 0.3207536091477695,
      "q1_ms": 0.3205561909205492,
      "q3_ms": 0.32079191362506454,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 16,
      "round": 0,
      "compile_s": 0.21404680600971915,
      "inner": 115,
      "median_ms": 0.3352985565510133,
      "q1_ms": 0.33512645225172216,
      "q3_ms": 0.3383727652481352,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 0,
      "round": 1,
      "compile_s": 0.1679307469894411,
      "inner": 78,
      "median_ms": 0.5214613332869844,
      "q1_ms": 0.521307519263176,
      "q3_ms": 0.521684448661593,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 2,
      "round": 1,
      "compile_s": 0.20757054700516164,
      "inner": 110,
      "median_ms": 0.3131306453724392,
      "q1_ms": 0.3128211500139018,
      "q3_ms": 0.3135122772867114,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 4,
      "round": 1,
      "compile_s": 0.20834347800700925,
      "inner": 121,
      "median_ms": 0.3173332066084682,
      "q1_ms": 0.31726783465484376,
      "q3_ms": 0.3179152603073561,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 8,
      "round": 1,
      "compile_s": 0.21315924800001085,
      "inner": 118,
      "median_ms": 0.321185652495133,
      "q1_ms": 0.3205688516865879,
      "q3_ms": 0.3212700889350144,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 16,
      "round": 1,
      "compile_s": 0.20440525900630746,
      "inner": 114,
      "median_ms": 0.33508134207263457,
      "q1_ms": 0.33499732895155865,
      "q3_ms": 0.3351720437773079,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 32,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 0,
      "round": 0,
      "compile_s": 0.17531026700453367,
      "inner": 62,
      "median_ms": 0.6947626773786221,
      "q1_ms": 0.6946825242560026,
      "q3_ms": 0.6953591370303559,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 32,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 2,
      "round": 0,
      "compile_s": 0.26500379099161364,
      "inner": 115,
      "median_ms": 0.3307092695689315,
      "q1_ms": 0.33061496088402753,
      "q3_ms": 0.3319694434732968,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 32,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 4,
      "round": 0,
      "compile_s": 0.264696678990731,
      "inner": 103,
      "median_ms": 0.3260946699793369,
      "q1_ms": 0.3258297038326736,
      "q3_ms": 0.3266221213830382,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 32,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 8,
      "round": 0,
      "compile_s": 0.24638324999250472,
      "inner": 117,
      "median_ms": 0.3344319743635611,
      "q1_ms": 0.3343405640827348,
      "q3_ms": 0.3345211153540513,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 32,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 16,
      "round": 0,
      "compile_s": 0.2546181660000002,
      "inner": 104,
      "median_ms": 0.36882907681082056,
      "q1_ms": 0.368777629773831,
      "q3_ms": 0.36940097113704096,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 32,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 0,
      "round": 1,
      "compile_s": 0.1733857259969227,
      "inner": 58,
      "median_ms": 0.6995245173035575,
      "q1_ms": 0.6993723274780645,
      "q3_ms": 0.7011888706544815,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 32,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 2,
      "round": 1,
      "compile_s": 0.26080580000416376,
      "inner": 114,
      "median_ms": 0.3317319474023718,
      "q1_ms": 0.32940345613098,
      "q3_ms": 0.33179574121165567,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 32,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 4,
      "round": 1,
      "compile_s": 0.2531916970037855,
      "inner": 114,
      "median_ms": 0.326256763129781,
      "q1_ms": 0.32601155703958823,
      "q3_ms": 0.32628527626736303,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 32,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 8,
      "round": 1,
      "compile_s": 0.2518657499895198,
      "inner": 108,
      "median_ms": 0.3343668981364721,
      "q1_ms": 0.33409693983062255,
      "q3_ms": 0.3346067500118537,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 32,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 16,
      "round": 1,
      "compile_s": 0.24821996499667875,
      "inner": 107,
      "median_ms": 0.36865983185005835,
      "q1_ms": 0.3685658738054074,
      "q3_ms": 0.3734369532811177,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 64,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 0,
      "round": 0,
      "compile_s": 0.23179672900005244,
      "inner": 42,
      "median_ms": 1.0643365236492606,
      "q1_ms": 1.0637365713981646,
      "q3_ms": 1.072471654922646,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 64,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 2,
      "round": 0,
      "compile_s": 0.36701247800374404,
      "inner": 84,
      "median_ms": 0.47633944060570294,
      "q1_ms": 0.4690566785852536,
      "q3_ms": 0.4852335238129377,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 64,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 4,
      "round": 0,
      "compile_s": 0.37649120698915794,
      "inner": 81,
      "median_ms": 0.4729837778147977,
      "q1_ms": 0.47274950617222966,
      "q3_ms": 0.47319553090750965,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 64,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 8,
      "round": 0,
      "compile_s": 0.3662805740023032,
      "inner": 79,
      "median_ms": 0.4852606455907461,
      "q1_ms": 0.48520649367754787,
      "q3_ms": 0.4856041520229788,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 64,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 16,
      "round": 0,
      "compile_s": 0.40349862100265455,
      "inner": 92,
      "median_ms": 0.4606428152803109,
      "q1_ms": 0.46054001635955344,
      "q3_ms": 0.4608801684852275,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 64,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 0,
      "round": 1,
      "compile_s": 0.24076505299308337,
      "inner": 41,
      "median_ms": 1.063613536577422,
      "q1_ms": 1.0628138413999184,
      "q3_ms": 1.064185451158243,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 64,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 2,
      "round": 1,
      "compile_s": 0.3652829670027131,
      "inner": 86,
      "median_ms": 0.4678889535446002,
      "q1_ms": 0.46681820940859836,
      "q3_ms": 0.46953201156209745,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 64,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 4,
      "round": 1,
      "compile_s": 0.41278063799836673,
      "inner": 89,
      "median_ms": 0.473866089910342,
      "q1_ms": 0.4724066910037007,
      "q3_ms": 0.4770220562278063,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 64,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 8,
      "round": 1,
      "compile_s": 0.3958399340044707,
      "inner": 87,
      "median_ms": 0.48559424136203294,
      "q1_ms": 0.4836475114831862,
      "q3_ms": 0.49906748274392065,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 64,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 16,
      "round": 1,
      "compile_s": 0.4187444690032862,
      "inner": 93,
      "median_ms": 0.47215188181238066,
      "q1_ms": 0.46149721506529895,
      "q3_ms": 0.4982343333164951,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 0,
      "round": 0,
      "compile_s": 0.1603480260091601,
      "inner": 16,
      "median_ms": 2.881170938053401,
      "q1_ms": 2.8544366246023856,
      "q3_ms": 2.883857281631208,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 2,
      "round": 0,
      "compile_s": 0.20879562399932183,
      "inner": 19,
      "median_ms": 2.5715445257159635,
      "q1_ms": 2.556522447297252,
      "q3_ms": 2.5825682896538638,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 4,
      "round": 0,
      "compile_s": 0.2464045700035058,
      "inner": 18,
      "median_ms": 2.523238722157354,
      "q1_ms": 2.521763193928766,
      "q3_ms": 2.5424782219892625,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 8,
      "round": 0,
      "compile_s": 0.20476285499171354,
      "inner": 19,
      "median_ms": 2.5420951577344617,
      "q1_ms": 2.504919341988984,
      "q3_ms": 2.6158976577694135,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 16,
      "round": 0,
      "compile_s": 0.21197766200930346,
      "inner": 19,
      "median_ms": 2.518429894974504,
      "q1_ms": 2.516734316151614,
      "q3_ms": 2.5329093945114627,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 0,
      "round": 1,
      "compile_s": 0.1617619499884313,
      "inner": 16,
      "median_ms": 2.9142071880414733,
      "q1_ms": 2.8538751876112656,
      "q3_ms": 3.0246410624386044,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 2,
      "round": 1,
      "compile_s": 0.19738843700906727,
      "inner": 19,
      "median_ms": 2.535076631925461,
      "q1_ms": 2.532449815541775,
      "q3_ms": 2.5354441842002986,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 4,
      "round": 1,
      "compile_s": 0.21515666499908548,
      "inner": 19,
      "median_ms": 2.5188248423785953,
      "q1_ms": 2.514234578963001,
      "q3_ms": 2.5246524471042093,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 8,
      "round": 1,
      "compile_s": 0.20218804600881413,
      "inner": 19,
      "median_ms": 2.4846095789383216,
      "q1_ms": 2.4840965262755113,
      "q3_ms": 2.4886369738899368,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 16,
      "round": 1,
      "compile_s": 0.22496041900012642,
      "inner": 19,
      "median_ms": 2.515664631286391,
      "q1_ms": 2.511000315854816,
      "q3_ms": 2.5390298952085066,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 32,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 0,
      "round": 0,
      "compile_s": 0.17025201900105458,
      "inner": 15,
      "median_ms": 3.086097333774281,
      "q1_ms": 3.0840865001664497,
      "q3_ms": 3.0908686336867204,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 32,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 2,
      "round": 0,
      "compile_s": 0.2717438590043457,
      "inner": 17,
      "median_ms": 2.724791528581751,
      "q1_ms": 2.724088147207273,
      "q3_ms": 2.753887470616439,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 32,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 4,
      "round": 0,
      "compile_s": 0.2650449380016653,
      "inner": 18,
      "median_ms": 2.562827555165212,
      "q1_ms": 2.562010750100146,
      "q3_ms": 2.623253500334815,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 32,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 8,
      "round": 0,
      "compile_s": 0.2583513769932324,
      "inner": 16,
      "median_ms": 2.566513874626253,
      "q1_ms": 2.565372687058698,
      "q3_ms": 2.567569625625765,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 32,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 16,
      "round": 0,
      "compile_s": 0.28893230100220535,
      "inner": 18,
      "median_ms": 2.6926631116010764,
      "q1_ms": 2.6646687777732345,
      "q3_ms": 2.78508600038347,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 32,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 0,
      "round": 1,
      "compile_s": 0.17385798900795635,
      "inner": 14,
      "median_ms": 3.083707785533209,
      "q1_ms": 3.0813175710915988,
      "q3_ms": 3.096075856904333,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 32,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 2,
      "round": 1,
      "compile_s": 0.25773131400637794,
      "inner": 17,
      "median_ms": 2.726965059135931,
      "q1_ms": 2.723660587970003,
      "q3_ms": 2.769323441230774,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 32,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 4,
      "round": 1,
      "compile_s": 0.2628657470049802,
      "inner": 17,
      "median_ms": 2.5670248826773947,
      "q1_ms": 2.5615489704411147,
      "q3_ms": 2.5714096174102883,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 32,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 8,
      "round": 1,
      "compile_s": 0.25918031499895733,
      "inner": 18,
      "median_ms": 2.5673994993364126,
      "q1_ms": 2.566616361339887,
      "q3_ms": 2.571723833372946,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 32,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 16,
      "round": 1,
      "compile_s": 0.26004923800064716,
      "inner": 17,
      "median_ms": 2.6583517054164343,
      "q1_ms": 2.6569738530430613,
      "q3_ms": 2.713602441123358,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 64,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 0,
      "round": 0,
      "compile_s": 0.23439000500366092,
      "inner": 10,
      "median_ms": 4.689376900205389,
      "q1_ms": 4.681988350057509,
      "q3_ms": 4.690876250242582,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 64,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 2,
      "round": 0,
      "compile_s": 0.3749347220000345,
      "inner": 14,
      "median_ms": 3.5729539293762564,
      "q1_ms": 3.529617643445298,
      "q3_ms": 3.5896315358903457,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 64,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 4,
      "round": 0,
      "compile_s": 0.38615548600500915,
      "inner": 15,
      "median_ms": 3.5965942001591125,
      "q1_ms": 3.480763466116817,
      "q3_ms": 3.6531382002673736,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 64,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 8,
      "round": 0,
      "compile_s": 0.39637666000635363,
      "inner": 14,
      "median_ms": 3.4584730705578943,
      "q1_ms": 3.4573573934072295,
      "q3_ms": 3.4620484288357796,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 64,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 16,
      "round": 0,
      "compile_s": 0.4086764450039482,
      "inner": 14,
      "median_ms": 3.6213382856138714,
      "q1_ms": 3.611533642210166,
      "q3_ms": 3.639860464837901,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 64,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 0,
      "round": 1,
      "compile_s": 0.23918670800048858,
      "inner": 10,
      "median_ms": 4.705236900190357,
      "q1_ms": 4.7022432496305555,
      "q3_ms": 4.833158349356381,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 64,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 2,
      "round": 1,
      "compile_s": 0.3712287520029349,
      "inner": 14,
      "median_ms": 3.548172785875587,
      "q1_ms": 3.527662535426706,
      "q3_ms": 3.5598183567344677,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 64,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 4,
      "round": 1,
      "compile_s": 0.3749926620075712,
      "inner": 15,
      "median_ms": 3.5150506669500223,
      "q1_ms": 3.4537634002238824,
      "q3_ms": 3.5834997994243167,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 64,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 8,
      "round": 1,
      "compile_s": 0.3820958379947115,
      "inner": 15,
      "median_ms": 3.467373800231144,
      "q1_ms": 3.4664016668102704,
      "q3_ms": 3.4808279330415344,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 64,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "num_splits": 16,
      "round": 1,
      "compile_s": 0.3877009100106079,
      "inner": 14,
      "median_ms": 3.607752570782655,
      "q1_ms": 3.6043275717280006,
      "q3_ms": 3.6102995005161835,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 0,
      "compile_s": 0.15461903899267782,
      "inner": 200,
      "median_ms": 0.05798036996566225,
      "q1_ms": 0.053132347493374255,
      "q3_ms": 0.058565302497299854,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 0,
      "compile_s": 0.19989336500293575,
      "inner": 200,
      "median_ms": 0.049157385001308285,
      "q1_ms": 0.04750457752379589,
      "q3_ms": 0.0513134074935806,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 0,
      "compile_s": 0.24525129899848253,
      "inner": 200,
      "median_ms": 0.049296644938294776,
      "q1_ms": 0.04855669496464543,
      "q3_ms": 0.0500538000051165,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 0,
      "compile_s": 0.24512445900472812,
      "inner": 200,
      "median_ms": 0.05237386496446561,
      "q1_ms": 0.05178778250410687,
      "q3_ms": 0.0607441949978238,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 0,
      "compile_s": 0.24272234499221668,
      "inner": 200,
      "median_ms": 0.050858450049418025,
      "q1_ms": 0.050030429993057624,
      "q3_ms": 0.051347799999348354,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 1,
      "compile_s": 0.15985816899046768,
      "inner": 200,
      "median_ms": 0.05842325997946318,
      "q1_ms": 0.05578236250585178,
      "q3_ms": 0.05894461999560008,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 1,
      "compile_s": 0.23250889699556865,
      "inner": 200,
      "median_ms": 0.04629987997759599,
      "q1_ms": 0.04602598248311551,
      "q3_ms": 0.047602737504348624,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 1,
      "compile_s": 0.20000052799878176,
      "inner": 200,
      "median_ms": 0.04601007996825501,
      "q1_ms": 0.04494842251006048,
      "q3_ms": 0.04829359746508999,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 1,
      "compile_s": 0.20188226700702216,
      "inner": 200,
      "median_ms": 0.04587171504681464,
      "q1_ms": 0.045777437480865046,
      "q3_ms": 0.046283369993034285,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 1,
      "compile_s": 0.18166898400522768,
      "inner": 200,
      "median_ms": 0.04097479999472853,
      "q1_ms": 0.04074887747265166,
      "q3_ms": 0.042504845005169045,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 32,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 0,
      "compile_s": 0.15838321899354924,
      "inner": 200,
      "median_ms": 0.05121652502566576,
      "q1_ms": 0.0511772000027122,
      "q3_ms": 0.05122757000208367,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 32,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 0,
      "compile_s": 0.3916663010022603,
      "inner": 200,
      "median_ms": 0.04208726000797469,
      "q1_ms": 0.042079364975506905,
      "q3_ms": 0.04217574751237407,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 32,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 0,
      "compile_s": 0.3580699729936896,
      "inner": 200,
      "median_ms": 0.04628213995601982,
      "q1_ms": 0.04593949252011953,
      "q3_ms": 0.04984325001714751,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 32,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 0,
      "compile_s": 0.36241738699027337,
      "inner": 200,
      "median_ms": 0.04184259501926135,
      "q1_ms": 0.041650355015008245,
      "q3_ms": 0.04224741001962684,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 32,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 0,
      "compile_s": 0.39490787200338673,
      "inner": 200,
      "median_ms": 0.03876522503560409,
      "q1_ms": 0.038719462500012014,
      "q3_ms": 0.0388503075009794,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 32,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 1,
      "compile_s": 0.16273207499762066,
      "inner": 200,
      "median_ms": 0.05107886499899905,
      "q1_ms": 0.05105729745991994,
      "q3_ms": 0.05110992246045498,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 32,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 1,
      "compile_s": 0.3620229410007596,
      "inner": 200,
      "median_ms": 0.04187845996057149,
      "q1_ms": 0.04187435497442493,
      "q3_ms": 0.04192562751995865,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 32,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 1,
      "compile_s": 0.36188674200093374,
      "inner": 179,
      "median_ms": 0.04620695526987398,
      "q1_ms": 0.046058041944335,
      "q3_ms": 0.0464336369163877,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 32,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 1,
      "compile_s": 0.35902162399725057,
      "inner": 200,
      "median_ms": 0.040419555007247254,
      "q1_ms": 0.04027627997857053,
      "q3_ms": 0.040905997520894744,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 32,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 1,
      "compile_s": 0.37155638200056273,
      "inner": 200,
      "median_ms": 0.03872219494951423,
      "q1_ms": 0.03866348251904128,
      "q3_ms": 0.04012038749351632,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 64,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 0,
      "compile_s": 0.22098038099647965,
      "inner": 200,
      "median_ms": 0.07613426001626067,
      "q1_ms": 0.07612248253280995,
      "q3_ms": 0.07615479498781497,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 64,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 0,
      "compile_s": 0.39166943600866944,
      "inner": 200,
      "median_ms": 0.05818606005050242,
      "q1_ms": 0.05773590499302372,
      "q3_ms": 0.05842736500198953,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 64,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 0,
      "compile_s": 0.3618986850051442,
      "inner": 200,
      "median_ms": 0.0531686149770394,
      "q1_ms": 0.049940887474804185,
      "q3_ms": 0.0592615650020889,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 64,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 0,
      "compile_s": 0.3629388029949041,
      "inner": 200,
      "median_ms": 0.05594000998826232,
      "q1_ms": 0.055920797494763974,
      "q3_ms": 0.05598629748419626,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 64,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 0,
      "compile_s": 0.3975232880038675,
      "inner": 200,
      "median_ms": 0.066925290011568,
      "q1_ms": 0.06683652001811424,
      "q3_ms": 0.06695469000987941,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 64,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 1,
      "compile_s": 0.2227378949901322,
      "inner": 200,
      "median_ms": 0.06740499498846475,
      "q1_ms": 0.06737784497090615,
      "q3_ms": 0.06742715999280335,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 64,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 1,
      "compile_s": 0.3635909369913861,
      "inner": 200,
      "median_ms": 0.04106356995180249,
      "q1_ms": 0.04103379003936425,
      "q3_ms": 0.0411880075262161,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 64,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 1,
      "compile_s": 0.36366509599611163,
      "inner": 200,
      "median_ms": 0.044768380030291155,
      "q1_ms": 0.044756229981430806,
      "q3_ms": 0.04477411501284223,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 64,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 1,
      "compile_s": 0.3641065720003098,
      "inner": 200,
      "median_ms": 0.055737530055921525,
      "q1_ms": 0.055694695001875516,
      "q3_ms": 0.05579310749453725,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 64,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 1,
      "compile_s": 0.3591756949899718,
      "inner": 200,
      "median_ms": 0.06661099498160183,
      "q1_ms": 0.06655954501184169,
      "q3_ms": 0.06662960498942994,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 0,
      "compile_s": 0.1515796199964825,
      "inner": 200,
      "median_ms": 0.13757201500993688,
      "q1_ms": 0.1375139049559948,
      "q3_ms": 0.13764116996753728,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 0,
      "compile_s": 0.19143124899710529,
      "inner": 200,
      "median_ms": 0.050462750004953705,
      "q1_ms": 0.05041806747613009,
      "q3_ms": 0.050553772489365656,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 0,
      "compile_s": 0.19447620899882168,
      "inner": 200,
      "median_ms": 0.052897604982717894,
      "q1_ms": 0.05281251749693183,
      "q3_ms": 0.05290572251396952,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 0,
      "compile_s": 0.24628637700516265,
      "inner": 200,
      "median_ms": 0.05812269002490211,
      "q1_ms": 0.05804388998512877,
      "q3_ms": 0.058467740018386394,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 0,
      "compile_s": 0.24103886100056116,
      "inner": 200,
      "median_ms": 0.060638495051534846,
      "q1_ms": 0.06035746500856476,
      "q3_ms": 0.06075709497963544,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 1,
      "compile_s": 0.15372466800909024,
      "inner": 190,
      "median_ms": 0.1389587316416988,
      "q1_ms": 0.1389384842217336,
      "q3_ms": 0.13896157631135891,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 1,
      "compile_s": 0.1945639640034642,
      "inner": 200,
      "median_ms": 0.05081531504401937,
      "q1_ms": 0.050785835010174196,
      "q3_ms": 0.051224517483206,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 1,
      "compile_s": 0.1915328809991479,
      "inner": 200,
      "median_ms": 0.05287676496664062,
      "q1_ms": 0.052776627489947714,
      "q3_ms": 0.052912482533429284,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 1,
      "compile_s": 0.19915728199703153,
      "inner": 200,
      "median_ms": 0.0573122550122207,
      "q1_ms": 0.05725672748667421,
      "q3_ms": 0.05888478000997566,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 1,
      "compile_s": 0.18413922699983232,
      "inner": 200,
      "median_ms": 0.06053855497157201,
      "q1_ms": 0.060324825026327744,
      "q3_ms": 0.060561352474906016,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 32,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 0,
      "compile_s": 0.15803894799319096,
      "inner": 158,
      "median_ms": 0.18279176583746448,
      "q1_ms": 0.1827568860479212,
      "q3_ms": 0.18290028798863986,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 32,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 0,
      "compile_s": 0.40706152000348084,
      "inner": 200,
      "median_ms": 0.1513461200374877,
      "q1_ms": 0.15107556002476485,
      "q3_ms": 0.15138850252697011,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 32,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 0,
      "compile_s": 0.397275591007201,
      "inner": 184,
      "median_ms": 0.15622174451772994,
      "q1_ms": 0.15615095379990124,
      "q3_ms": 0.15635154890649416,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 32,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 0,
      "compile_s": 0.3987124569976004,
      "inner": 159,
      "median_ms": 0.18659854718153043,
      "q1_ms": 0.18555553142734035,
      "q3_ms": 0.1867066131818325,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 32,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 0,
      "compile_s": 0.406472495989874,
      "inner": 133,
      "median_ms": 0.23398506010907064,
      "q1_ms": 0.2339473194981001,
      "q3_ms": 0.23414276692800967,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 32,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 1,
      "compile_s": 0.1696818299969891,
      "inner": 162,
      "median_ms": 0.1827816543223934,
      "q1_ms": 0.182695197486563,
      "q3_ms": 0.18359715124214193,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 32,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 1,
      "compile_s": 0.37388873400050215,
      "inner": 200,
      "median_ms": 0.15208421500574332,
      "q1_ms": 0.15198600500298198,
      "q3_ms": 0.1522613250199356,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 32,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 1,
      "compile_s": 0.3733470699953614,
      "inner": 178,
      "median_ms": 0.15647546625290631,
      "q1_ms": 0.15646240447832305,
      "q3_ms": 0.15650847469740337,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 32,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 1,
      "compile_s": 0.37293489000876434,
      "inner": 160,
      "median_ms": 0.1865452250058297,
      "q1_ms": 0.1861619656210678,
      "q3_ms": 0.18673955937629216,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 32,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 1,
      "compile_s": 0.3902894060011022,
      "inner": 132,
      "median_ms": 0.23434757575600862,
      "q1_ms": 0.23389377273622702,
      "q3_ms": 0.23451126512416962,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 64,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 0,
      "compile_s": 0.23143576800066512,
      "inner": 117,
      "median_ms": 0.29091812813511264,
      "q1_ms": 0.2908749272922476,
      "q3_ms": 0.2910962051455266,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 64,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 0,
      "compile_s": 0.41959228200721554,
      "inner": 120,
      "median_ms": 0.2628542750850708,
      "q1_ms": 0.26278517919611966,
      "q3_ms": 0.2632480958103164,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 64,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 0,
      "compile_s": 0.4161731450003572,
      "inner": 118,
      "median_ms": 0.28654743208928757,
      "q1_ms": 0.28600113981125197,
      "q3_ms": 0.2867198262858931,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 64,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 0,
      "compile_s": 0.4192300470022019,
      "inner": 108,
      "median_ms": 0.3585651665397799,
      "q1_ms": 0.35852355554029436,
      "q3_ms": 0.3586936065268539,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 64,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 0,
      "compile_s": 0.44299056999443565,
      "inner": 75,
      "median_ms": 0.5693748265427228,
      "q1_ms": 0.5689470200256135,
      "q3_ms": 0.5696548999791655,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 64,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 1,
      "compile_s": 0.23337960400385782,
      "inner": 114,
      "median_ms": 0.29077288593071626,
      "q1_ms": 0.29074322363860056,
      "q3_ms": 0.2910263947820855,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 64,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 1,
      "compile_s": 0.37833727899123915,
      "inner": 121,
      "median_ms": 0.2622058759886033,
      "q1_ms": 0.26182424377367164,
      "q3_ms": 0.2630171982989038,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 64,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 1,
      "compile_s": 0.37687378100235946,
      "inner": 115,
      "median_ms": 0.28617285209728166,
      "q1_ms": 0.28606970432599116,
      "q3_ms": 0.28656820871654654,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 64,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 1,
      "compile_s": 0.38640894599666353,
      "inner": 107,
      "median_ms": 0.3586958410549108,
      "q1_ms": 0.35853049063415343,
      "q3_ms": 0.35889941126502023,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 64,
      "lkv": 1024,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 1,
      "compile_s": 0.3887994670076296,
      "inner": 70,
      "median_ms": 0.5694203427992761,
      "q1_ms": 0.5693784786022401,
      "q3_ms": 0.5699156999721059,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 2048,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 0,
      "compile_s": 0.15763775899540633,
      "inner": 200,
      "median_ms": 0.07300679499167018,
      "q1_ms": 0.07298965752852382,
      "q3_ms": 0.07302254500245908,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 2048,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 0,
      "compile_s": 0.20005992999358568,
      "inner": 200,
      "median_ms": 0.06049437499314081,
      "q1_ms": 0.05868259748240234,
      "q3_ms": 0.06345675246848259,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 2048,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 0,
      "compile_s": 0.20380942200426944,
      "inner": 200,
      "median_ms": 0.05736240003898274,
      "q1_ms": 0.05681872247805586,
      "q3_ms": 0.05860129251232138,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 2048,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 0,
      "compile_s": 0.2070751690043835,
      "inner": 200,
      "median_ms": 0.04974945004505571,
      "q1_ms": 0.04894252502708696,
      "q3_ms": 0.05011213252146263,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 2048,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 0,
      "compile_s": 0.20572888299648184,
      "inner": 200,
      "median_ms": 0.0470122250408167,
      "q1_ms": 0.04640111495973542,
      "q3_ms": 0.0471904625010211,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 2048,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 1,
      "compile_s": 0.16252568500931375,
      "inner": 200,
      "median_ms": 0.08139445999404415,
      "q1_ms": 0.0813561849645339,
      "q3_ms": 0.0814439224996022,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 2048,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 1,
      "compile_s": 0.2003548680077074,
      "inner": 200,
      "median_ms": 0.04096804499567952,
      "q1_ms": 0.03984471502917586,
      "q3_ms": 0.042035612495965324,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 2048,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 1,
      "compile_s": 0.20434202899923548,
      "inner": 200,
      "median_ms": 0.05930086495936848,
      "q1_ms": 0.058007394982269034,
      "q3_ms": 0.06426235002436442,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 2048,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 1,
      "compile_s": 0.21024885200313292,
      "inner": 200,
      "median_ms": 0.045635615024366416,
      "q1_ms": 0.04513840249273926,
      "q3_ms": 0.04807285247807158,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 2048,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 1,
      "compile_s": 0.20899818699399475,
      "inner": 200,
      "median_ms": 0.044608475000131875,
      "q1_ms": 0.043090955005027354,
      "q3_ms": 0.04769338494952535,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 32,
      "lkv": 2048,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 0,
      "compile_s": 0.17450864899728913,
      "inner": 200,
      "median_ms": 0.10162491998926271,
      "q1_ms": 0.10157583001273451,
      "q3_ms": 0.10164641251321882,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 32,
      "lkv": 2048,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 0,
      "compile_s": 0.38808835299278144,
      "inner": 200,
      "median_ms": 0.06669820999377407,
      "q1_ms": 0.06666209250397515,
      "q3_ms": 0.06673848496575374,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 32,
      "lkv": 2048,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 0,
      "compile_s": 0.3793866440100828,
      "inner": 200,
      "median_ms": 0.04707854997832328,
      "q1_ms": 0.0439535874829744,
      "q3_ms": 0.0511413575077313,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 32,
      "lkv": 2048,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 0,
      "compile_s": 0.38598281001031864,
      "inner": 200,
      "median_ms": 0.045782399974996224,
      "q1_ms": 0.045518772494688164,
      "q3_ms": 0.05050071998994099,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 32,
      "lkv": 2048,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 0,
      "compile_s": 0.37322838499676436,
      "inner": 200,
      "median_ms": 0.05462662498757709,
      "q1_ms": 0.054414100013673306,
      "q3_ms": 0.05464974252390675,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 32,
      "lkv": 2048,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 1,
      "compile_s": 0.1738576589996228,
      "inner": 200,
      "median_ms": 0.10134584495972376,
      "q1_ms": 0.10127633751835674,
      "q3_ms": 0.1016940749832429,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 32,
      "lkv": 2048,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 1,
      "compile_s": 0.36715130000084173,
      "inner": 200,
      "median_ms": 0.06714019495120738,
      "q1_ms": 0.06712501497531775,
      "q3_ms": 0.06717451000440633,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 32,
      "lkv": 2048,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 1,
      "compile_s": 0.3845956860022852,
      "inner": 200,
      "median_ms": 0.0492523599677952,
      "q1_ms": 0.04844458249863237,
      "q3_ms": 0.05242603750957642,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 32,
      "lkv": 2048,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 1,
      "compile_s": 0.37860504499985836,
      "inner": 200,
      "median_ms": 0.046060524982749484,
      "q1_ms": 0.046040284978516866,
      "q3_ms": 0.046094287499727216,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 32,
      "lkv": 2048,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 1,
      "compile_s": 0.37288536700361874,
      "inner": 200,
      "median_ms": 0.05450919998111203,
      "q1_ms": 0.054262810008367524,
      "q3_ms": 0.05526240001927363,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 64,
      "lkv": 2048,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 0,
      "compile_s": 0.2473632540059043,
      "inner": 200,
      "median_ms": 0.15140848503506277,
      "q1_ms": 0.13863340001989854,
      "q3_ms": 0.151468450021639,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 64,
      "lkv": 2048,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 0,
      "compile_s": 0.38618193700676784,
      "inner": 200,
      "median_ms": 0.06801054005336482,
      "q1_ms": 0.06798834750952665,
      "q3_ms": 0.0680202350122272,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 64,
      "lkv": 2048,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 0,
      "compile_s": 0.3920729599922197,
      "inner": 200,
      "median_ms": 0.0717709550372092,
      "q1_ms": 0.07176636754593346,
      "q3_ms": 0.0717980824992992,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 64,
      "lkv": 2048,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 0,
      "compile_s": 0.38111777800077107,
      "inner": 200,
      "median_ms": 0.08227101498050615,
      "q1_ms": 0.08201818498491775,
      "q3_ms": 0.08229553001001477,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 64,
      "lkv": 2048,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 0,
      "compile_s": 0.38100122699688654,
      "inner": 200,
      "median_ms": 0.0972488949628314,
      "q1_ms": 0.09682449253887171,
      "q3_ms": 0.09736233500007074,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 64,
      "lkv": 2048,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 1,
      "compile_s": 0.23620270400715526,
      "inner": 178,
      "median_ms": 0.1353427415891561,
      "q1_ms": 0.13532106745157527,
      "q3_ms": 0.13544022752622176,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 64,
      "lkv": 2048,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 1,
      "compile_s": 0.3748602319974452,
      "inner": 200,
      "median_ms": 0.06862850998004433,
      "q1_ms": 0.06861368252430111,
      "q3_ms": 0.06870420253108023,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 64,
      "lkv": 2048,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 1,
      "compile_s": 0.38625001499895006,
      "inner": 200,
      "median_ms": 0.07212958495074417,
      "q1_ms": 0.07204033750895178,
      "q3_ms": 0.07249494752613828,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 64,
      "lkv": 2048,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 1,
      "compile_s": 0.38401264300046023,
      "inner": 200,
      "median_ms": 0.0823816200136207,
      "q1_ms": 0.0820114224916324,
      "q3_ms": 0.08240872499300167,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 64,
      "lkv": 2048,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 1,
      "compile_s": 0.4025104609900154,
      "inner": 200,
      "median_ms": 0.09718226996483281,
      "q1_ms": 0.09685314249509247,
      "q3_ms": 0.09726134499942418,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 2048,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 0,
      "compile_s": 0.15870858400012366,
      "inner": 129,
      "median_ms": 0.266543782993884,
      "q1_ms": 0.2663867364122258,
      "q3_ms": 0.2670613449841175,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 2048,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 0,
      "compile_s": 0.20667660499748308,
      "inner": 200,
      "median_ms": 0.08929605006414931,
      "q1_ms": 0.08927428996685194,
      "q3_ms": 0.08935381250921637,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 2048,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 0,
      "compile_s": 0.20545591600239277,
      "inner": 200,
      "median_ms": 0.09154901505098678,
      "q1_ms": 0.09153729250101605,
      "q3_ms": 0.09159009503491689,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 2048,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 0,
      "compile_s": 0.2148468990053516,
      "inner": 200,
      "median_ms": 0.09883243503281847,
      "q1_ms": 0.09880947251076577,
      "q3_ms": 0.09895376999338623,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 2048,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 0,
      "compile_s": 0.21893617400201038,
      "inner": 200,
      "median_ms": 0.10161931000766344,
      "q1_ms": 0.10155291503906483,
      "q3_ms": 0.10174667750106892,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 2048,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 1,
      "compile_s": 0.17191856099816505,
      "inner": 127,
      "median_ms": 0.26760837803393106,
      "q1_ms": 0.2673432361802057,
      "q3_ms": 0.2678047401457475,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 2048,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 1,
      "compile_s": 0.20645077699737158,
      "inner": 200,
      "median_ms": 0.08936853999330197,
      "q1_ms": 0.08927518752898322,
      "q3_ms": 0.08945711000706069,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 2048,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 1,
      "compile_s": 0.2002475439949194,
      "inner": 200,
      "median_ms": 0.0919302849797532,
      "q1_ms": 0.09192099001666065,
      "q3_ms": 0.09195327751513105,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 2048,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 1,
      "compile_s": 0.20391212700633332,
      "inner": 192,
      "median_ms": 0.09911455723946953,
      "q1_ms": 0.09907813542516428,
      "q3_ms": 0.09913198698541237,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 2048,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 1,
      "compile_s": 0.20506350400683004,
      "inner": 200,
      "median_ms": 0.10189257998717949,
      "q1_ms": 0.10184839746216312,
      "q3_ms": 0.1019124150116113,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 32,
      "lkv": 2048,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 0,
      "compile_s": 0.16454932300257497,
      "inner": 106,
      "median_ms": 0.3587368394979627,
      "q1_ms": 0.35778145752803264,
      "q3_ms": 0.35927403772596866,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 32,
      "lkv": 2048,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 0,
      "compile_s": 0.3790725500002736,
      "inner": 130,
      "median_ms": 0.268694007760272,
      "q1_ms": 0.2684872961253859,
      "q3_ms": 0.2687406268915556,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 32,
      "lkv": 2048,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 0,
      "compile_s": 0.4008503610093612,
      "inner": 130,
      "median_ms": 0.26123358458370116,
      "q1_ms": 0.2610613308206666,
      "q3_ms": 0.26137661922033517,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 32,
      "lkv": 2048,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 0,
      "compile_s": 0.38433646599878557,
      "inner": 130,
      "median_ms": 0.2866868769230608,
      "q1_ms": 0.2845516038248136,
      "q3_ms": 0.2867801307449834,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 32,
      "lkv": 2048,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 0,
      "compile_s": 0.3920598549884744,
      "inner": 103,
      "median_ms": 0.35529941740305643,
      "q1_ms": 0.3548743009303913,
      "q3_ms": 0.3554650291476118,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 32,
      "lkv": 2048,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 1,
      "compile_s": 0.181272023008205,
      "inner": 102,
      "median_ms": 0.3590649508873421,
      "q1_ms": 0.3588166324405785,
      "q3_ms": 0.35917451960438695,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 32,
      "lkv": 2048,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 1,
      "compile_s": 0.3739785849902546,
      "inner": 115,
      "median_ms": 0.26885358257340675,
      "q1_ms": 0.2687443695005799,
      "q3_ms": 0.2688794043775567,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 32,
      "lkv": 2048,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 1,
      "compile_s": 0.371978801995283,
      "inner": 126,
      "median_ms": 0.26128985706554164,
      "q1_ms": 0.2604428213782653,
      "q3_ms": 0.2613124405408079,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 32,
      "lkv": 2048,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 1,
      "compile_s": 0.3636498970008688,
      "inner": 120,
      "median_ms": 0.2868319583891813,
      "q1_ms": 0.2860244250162699,
      "q3_ms": 0.2869149875550647,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 32,
      "lkv": 2048,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 1,
      "compile_s": 0.35586287900514435,
      "inner": 106,
      "median_ms": 0.355246660375379,
      "q1_ms": 0.3551540330419194,
      "q3_ms": 0.35526683962605193,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 64,
      "lkv": 2048,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 0,
      "compile_s": 0.22439224399568047,
      "inner": 69,
      "median_ms": 0.5700910724176472,
      "q1_ms": 0.5697652319424273,
      "q3_ms": 0.5703021231675993,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 64,
      "lkv": 2048,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 0,
      "compile_s": 0.3658890629885718,
      "inner": 81,
      "median_ms": 0.4671740248284797,
      "q1_ms": 0.467087993815823,
      "q3_ms": 0.467275827152965,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 64,
      "lkv": 2048,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 0,
      "compile_s": 0.36646124499384314,
      "inner": 85,
      "median_ms": 0.47820534108562723,
      "q1_ms": 0.4781468176481533,
      "q3_ms": 0.47834537065629024,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 64,
      "lkv": 2048,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 0,
      "compile_s": 0.35798790799162816,
      "inner": 74,
      "median_ms": 0.5461935135236992,
      "q1_ms": 0.5459789256722591,
      "q3_ms": 0.5463155134379662,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 64,
      "lkv": 2048,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 0,
      "compile_s": 0.36224622299778275,
      "inner": 52,
      "median_ms": 0.8329531922824501,
      "q1_ms": 0.8325558269112103,
      "q3_ms": 0.8333176538144471,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 64,
      "lkv": 2048,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 1,
      "compile_s": 0.23053899100341368,
      "inner": 68,
      "median_ms": 0.5700927352282794,
      "q1_ms": 0.5699558530656813,
      "q3_ms": 0.5737639410868082,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 64,
      "lkv": 2048,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 1,
      "compile_s": 0.3590971969970269,
      "inner": 80,
      "median_ms": 0.4668407749704784,
      "q1_ms": 0.4668161000154214,
      "q3_ms": 0.4672510500313365,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 64,
      "lkv": 2048,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 1,
      "compile_s": 0.3582846600038465,
      "inner": 90,
      "median_ms": 0.4779215442896303,
      "q1_ms": 0.47782992219759357,
      "q3_ms": 0.4779503667022153,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 64,
      "lkv": 2048,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 1,
      "compile_s": 0.3609368059987901,
      "inner": 74,
      "median_ms": 0.5462905946596702,
      "q1_ms": 0.5460867634266489,
      "q3_ms": 0.5463811080640094,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 64,
      "lkv": 2048,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 1,
      "compile_s": 0.3653684279997833,
      "inner": 52,
      "median_ms": 0.832627788510693,
      "q1_ms": 0.8323827117307631,
      "q3_ms": 0.8328270961590953,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 0,
      "compile_s": 0.14744614200026263,
      "inner": 75,
      "median_ms": 0.5148920401309927,
      "q1_ms": 0.5138445999667358,
      "q3_ms": 0.5150791333289817,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 0,
      "compile_s": 0.18964361801045015,
      "inner": 136,
      "median_ms": 0.21984318385982668,
      "q1_ms": 0.21980281251350842,
      "q3_ms": 0.21988881615149747,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 0,
      "compile_s": 0.19092547200853005,
      "inner": 200,
      "median_ms": 0.11841503001051024,
      "q1_ms": 0.11837362500955351,
      "q3_ms": 0.1185321524826577,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 0,
      "compile_s": 0.1966181099996902,
      "inner": 200,
      "median_ms": 0.08841904505970888,
      "q1_ms": 0.0883990299917059,
      "q3_ms": 0.08842898747388972,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 0,
      "compile_s": 0.18812084798992146,
      "inner": 200,
      "median_ms": 0.08770975502557121,
      "q1_ms": 0.08766020746406866,
      "q3_ms": 0.08772312496148515,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 1,
      "compile_s": 0.15636820800136775,
      "inner": 80,
      "median_ms": 0.5142595875440747,
      "q1_ms": 0.5141438750797533,
      "q3_ms": 0.5144272249708592,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 1,
      "compile_s": 0.18950089700229,
      "inner": 139,
      "median_ms": 0.22068158997701842,
      "q1_ms": 0.22065848556620224,
      "q3_ms": 0.2207122948845654,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 1,
      "compile_s": 0.1963910000049509,
      "inner": 200,
      "median_ms": 0.11898360506165773,
      "q1_ms": 0.11896213996806182,
      "q3_ms": 0.11899462500878144,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 1,
      "compile_s": 0.18745744301122613,
      "inner": 200,
      "median_ms": 0.08893021004041657,
      "q1_ms": 0.08892244499293156,
      "q3_ms": 0.08894441001757514,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 1,
      "compile_s": 0.19168836499738973,
      "inner": 200,
      "median_ms": 0.08795672001724597,
      "q1_ms": 0.08791471500444459,
      "q3_ms": 0.08799323750281474,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 32,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 0,
      "compile_s": 0.15655442999559455,
      "inner": 58,
      "median_ms": 0.6859459139503025,
      "q1_ms": 0.6858482156530954,
      "q3_ms": 0.6860008362654582,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 32,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 0,
      "compile_s": 0.35837619299127255,
      "inner": 90,
      "median_ms": 0.44236840007619727,
      "q1_ms": 0.4423305500596244,
      "q3_ms": 0.44239383334772553,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 32,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 0,
      "compile_s": 0.3604242669971427,
      "inner": 142,
      "median_ms": 0.23455534504995432,
      "q1_ms": 0.23451597883995673,
      "q3_ms": 0.23462198949115626,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 32,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 0,
      "compile_s": 0.35606138400908094,
      "inner": 96,
      "median_ms": 0.23784686451714757,
      "q1_ms": 0.23707493232905108,
      "q3_ms": 0.23867582293253992,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 32,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 0,
      "compile_s": 0.358533400009037,
      "inner": 137,
      "median_ms": 0.24864490513169768,
      "q1_ms": 0.24852884670355813,
      "q3_ms": 0.24938831755034896,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 32,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 1,
      "compile_s": 0.16849890300363768,
      "inner": 61,
      "median_ms": 0.686691164088695,
      "q1_ms": 0.686151696799407,
      "q3_ms": 0.6867162950077383,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 32,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 1,
      "compile_s": 0.3571385620016372,
      "inner": 82,
      "median_ms": 0.4424314754986681,
      "q1_ms": 0.44239861585438345,
      "q3_ms": 0.4424446220005715,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 32,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 1,
      "compile_s": 0.36331115700886585,
      "inner": 131,
      "median_ms": 0.23467288557386723,
      "q1_ms": 0.2344413015072568,
      "q3_ms": 0.2347839732840124,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 32,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 1,
      "compile_s": 0.3585150650033029,
      "inner": 130,
      "median_ms": 0.23805404612525868,
      "q1_ms": 0.23781566924299113,
      "q3_ms": 0.23930255763564046,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 32,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 1,
      "compile_s": 0.36077043099794537,
      "inner": 133,
      "median_ms": 0.25028551880995,
      "q1_ms": 0.249343996272769,
      "q3_ms": 0.2504811090783958,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 64,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 0,
      "compile_s": 0.22220044799905736,
      "inner": 42,
      "median_ms": 1.0425750951981172,
      "q1_ms": 1.042419702571351,
      "q3_ms": 1.042894523769584,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 64,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 0,
      "compile_s": 0.3567675400117878,
      "inner": 85,
      "median_ms": 0.44600467064061805,
      "q1_ms": 0.445921511752257,
      "q3_ms": 0.4460778706261943,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 64,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 0,
      "compile_s": 0.35648106900043786,
      "inner": 88,
      "median_ms": 0.4490261476762085,
      "q1_ms": 0.44895994315910764,
      "q3_ms": 0.44925282393457816,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 64,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 0,
      "compile_s": 0.35553302500920836,
      "inner": 82,
      "median_ms": 0.46477424390415273,
      "q1_ms": 0.4645820487739246,
      "q3_ms": 0.4648286037369938,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 64,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 0,
      "compile_s": 0.35809043199697044,
      "inner": 88,
      "median_ms": 0.43581695453295566,
      "q1_ms": 0.43550164214552334,
      "q3_ms": 0.43588822725062276,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 64,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 1,
      "compile_s": 0.22089070100628305,
      "inner": 41,
      "median_ms": 1.0424742681718226,
      "q1_ms": 1.0423581831074864,
      "q3_ms": 1.042728524558564,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 64,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 1,
      "compile_s": 0.3598498000064865,
      "inner": 89,
      "median_ms": 0.4459242360977707,
      "q1_ms": 0.44590032582950856,
      "q3_ms": 0.44614961240022033,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 64,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 1,
      "compile_s": 0.35730054799932986,
      "inner": 83,
      "median_ms": 0.44955598808970887,
      "q1_ms": 0.44896868076834673,
      "q3_ms": 0.45214758425704143,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 64,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 1,
      "compile_s": 0.3615236149926204,
      "inner": 87,
      "median_ms": 0.4642155287712384,
      "q1_ms": 0.46414885632814196,
      "q3_ms": 0.46436857469259474,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 64,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 1,
      "compile_s": 0.36283817199000623,
      "inner": 92,
      "median_ms": 0.43560055438201345,
      "q1_ms": 0.43558640762690315,
      "q3_ms": 0.4357497064515685,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 0,
      "compile_s": 0.15268802600621711,
      "inner": 22,
      "median_ms": 2.0551045004553585,
      "q1_ms": 2.0524125681152907,
      "q3_ms": 2.0613082273154184,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 0,
      "compile_s": 0.19290680000267457,
      "inner": 66,
      "median_ms": 0.6180686212963227,
      "q1_ms": 0.6180319544386513,
      "q3_ms": 0.6180973788573421,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 0,
      "compile_s": 0.19048843599739484,
      "inner": 65,
      "median_ms": 0.6226646767642635,
      "q1_ms": 0.6225342152622313,
      "q3_ms": 0.6227381231908042,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 0,
      "compile_s": 0.19252906600013375,
      "inner": 63,
      "median_ms": 0.6480869524816554,
      "q1_ms": 0.6478194681842561,
      "q3_ms": 0.6483960318118556,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 0,
      "compile_s": 0.1915158889896702,
      "inner": 63,
      "median_ms": 0.6493868729901604,
      "q1_ms": 0.6492985396309652,
      "q3_ms": 0.6495506030559126,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 1,
      "compile_s": 0.15703998898970895,
      "inner": 23,
      "median_ms": 2.0561718265525997,
      "q1_ms": 2.053456695727579,
      "q3_ms": 2.06522480426508,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 1,
      "compile_s": 0.18907267799659166,
      "inner": 64,
      "median_ms": 0.6181560936511232,
      "q1_ms": 0.6180519375220683,
      "q3_ms": 0.6181732421737252,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 1,
      "compile_s": 0.19613818200014066,
      "inner": 65,
      "median_ms": 0.622469076859908,
      "q1_ms": 0.6224197538148684,
      "q3_ms": 0.6225954770343378,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 1,
      "compile_s": 0.1929947160097072,
      "inner": 66,
      "median_ms": 0.6475231211604269,
      "q1_ms": 0.6473314621843835,
      "q3_ms": 0.6477357120743501,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 1,
      "compile_s": 0.19546716200420633,
      "inner": 61,
      "median_ms": 0.6491637869344903,
      "q1_ms": 0.6491553197475531,
      "q3_ms": 0.6494083525049958,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 32,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 0,
      "compile_s": 0.15951839700574055,
      "inner": 17,
      "median_ms": 2.8307832356350606,
      "q1_ms": 2.8289509115411953,
      "q3_ms": 2.8394828820155986,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 32,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 0,
      "compile_s": 0.3602820980013348,
      "inner": 24,
      "median_ms": 1.9114646668943653,
      "q1_ms": 1.9066681040082283,
      "q3_ms": 1.9116147501942276,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 32,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 0,
      "compile_s": 0.36023597000166774,
      "inner": 27,
      "median_ms": 1.720905407435364,
      "q1_ms": 1.7208291667177238,
      "q3_ms": 1.721113426255114,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 32,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 0,
      "compile_s": 0.36211869200633373,
      "inner": 28,
      "median_ms": 1.6657879282969847,
      "q1_ms": 1.6655333750220183,
      "q3_ms": 1.6659804644275158,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 32,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 0,
      "compile_s": 0.37032318000274245,
      "inner": 28,
      "median_ms": 1.709302785457112,
      "q1_ms": 1.7086821249644606,
      "q3_ms": 1.7094201071553732,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 32,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 1,
      "compile_s": 0.16368963700369932,
      "inner": 17,
      "median_ms": 2.8370580000816568,
      "q1_ms": 2.831877911590067,
      "q3_ms": 2.8389297650017613,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 32,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 1,
      "compile_s": 0.36010789799911436,
      "inner": 24,
      "median_ms": 1.9126018335858437,
      "q1_ms": 1.8999833749451986,
      "q3_ms": 1.914949604118495,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 32,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 1,
      "compile_s": 0.3626290170068387,
      "inner": 26,
      "median_ms": 1.7225840000229744,
      "q1_ms": 1.7217036922654816,
      "q3_ms": 1.7232400577295968,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 32,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 1,
      "compile_s": 0.3675711950054392,
      "inner": 28,
      "median_ms": 1.666903285825226,
      "q1_ms": 1.6660906608844277,
      "q3_ms": 1.6697209464057647,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 32,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 1,
      "compile_s": 0.36091422400204465,
      "inner": 28,
      "median_ms": 1.7093918572105136,
      "q1_ms": 1.70881721429136,
      "q3_ms": 1.7097489640686296,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 64,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 0,
      "compile_s": 0.22930965501291212,
      "inner": 10,
      "median_ms": 4.520004500227515,
      "q1_ms": 4.518432049371768,
      "q3_ms": 4.525808000471443,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 64,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 0,
      "compile_s": 0.36533854200388305,
      "inner": 13,
      "median_ms": 3.3253776144612437,
      "q1_ms": 3.32507853790258,
      "q3_ms": 3.327232615377467,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 64,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 0,
      "compile_s": 0.36358366200875025,
      "inner": 15,
      "median_ms": 3.1764715997269377,
      "q1_ms": 3.175293666330011,
      "q3_ms": 3.1891755667553907,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 64,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 0,
      "compile_s": 0.3579011939873453,
      "inner": 15,
      "median_ms": 3.167198066754888,
      "q1_ms": 3.16504596703453,
      "q3_ms": 3.1731141997928107,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 64,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 0,
      "compile_s": 0.36623224199865945,
      "inner": 14,
      "median_ms": 3.4795851430057416,
      "q1_ms": 3.47849167831425,
      "q3_ms": 3.4806800712041355,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 64,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 0,
      "round": 1,
      "compile_s": 0.2288413489877712,
      "inner": 10,
      "median_ms": 4.5233458004076965,
      "q1_ms": 4.520072149898624,
      "q3_ms": 4.549739849608159,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 64,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 2,
      "round": 1,
      "compile_s": 0.35986437799874693,
      "inner": 14,
      "median_ms": 3.3438205001792602,
      "q1_ms": 3.3412176427581084,
      "q3_ms": 3.347588035727053,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 64,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 4,
      "round": 1,
      "compile_s": 0.36848984299285803,
      "inner": 15,
      "median_ms": 3.181628000068789,
      "q1_ms": 3.176942132995464,
      "q3_ms": 3.187152733153198,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 64,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 8,
      "round": 1,
      "compile_s": 0.3645817899960093,
      "inner": 15,
      "median_ms": 3.1650619998496645,
      "q1_ms": 3.1608427331472435,
      "q3_ms": 3.1654914671283527,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "splitsweep",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 64,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "num_splits": 16,
      "round": 1,
      "compile_s": 0.36240771900338586,
      "inner": 14,
      "median_ms": 3.4770059282891452,
      "q1_ms": 3.4768431786298084,
      "q3_ms": 3.487368428717933,
      "tag": null,
      "device_arg": "gpu",
      "run": "sweep2",
      "note": "split-KV kernel with block_k=64 (before the 32-key tile change)"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "block_k": 32,
      "round": 0,
      "compile_s": 0.23622803699981887,
      "inner": 24,
      "median_ms": 0.3827868749795016,
      "q1_ms": 0.37666297945785726,
      "q3_ms": 0.4123190001337207,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "block_k": 64,
      "round": 0,
      "compile_s": 0.2159888770111138,
      "inner": 102,
      "median_ms": 0.3681941078884491,
      "q1_ms": 0.3678239901965557,
      "q3_ms": 0.368426902009649,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "block_k": 32,
      "round": 1,
      "compile_s": 0.18957505300932098,
      "inner": 106,
      "median_ms": 0.3180935471494384,
      "q1_ms": 0.31800073112156857,
      "q3_ms": 0.31819444340280306,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "block_k": 64,
      "round": 1,
      "compile_s": 0.2130194979981752,
      "inner": 106,
      "median_ms": 0.3567875283583559,
      "q1_ms": 0.35676030659848246,
      "q3_ms": 0.35706399532226085,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "block_k": 32,
      "round": 2,
      "compile_s": 0.17853228500462137,
      "inner": 108,
      "median_ms": 0.31866663886359203,
      "q1_ms": 0.31842826852139466,
      "q3_ms": 0.31870235647215006,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "block_k": 64,
      "round": 2,
      "compile_s": 0.21902995101117995,
      "inner": 110,
      "median_ms": 0.356997745455539,
      "q1_ms": 0.3566631590233142,
      "q3_ms": 0.3572018635829657,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "block_k": 32,
      "round": 0,
      "compile_s": 0.2693811830104096,
      "inner": 121,
      "median_ms": 0.31715976028637727,
      "q1_ms": 0.31696050824724475,
      "q3_ms": 0.31717359506879156,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "block_k": 64,
      "round": 0,
      "compile_s": 0.20260160400357563,
      "inner": 110,
      "median_ms": 0.3562765636491928,
      "q1_ms": 0.3559186635935806,
      "q3_ms": 0.35650412722712976,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "block_k": 32,
      "round": 1,
      "compile_s": 0.19892297199112363,
      "inner": 117,
      "median_ms": 0.3173639744114226,
      "q1_ms": 0.31725441453417236,
      "q3_ms": 0.3182399102822981,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "block_k": 64,
      "round": 1,
      "compile_s": 0.19731353600218426,
      "inner": 109,
      "median_ms": 0.3558975137629196,
      "q1_ms": 0.35580436234766144,
      "q3_ms": 0.35608511925685654,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "block_k": 32,
      "round": 2,
      "compile_s": 0.19737839899607934,
      "inner": 115,
      "median_ms": 0.31724320876452583,
      "q1_ms": 0.31713660870952043,
      "q3_ms": 0.31727430861159833,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "block_k": 64,
      "round": 2,
      "compile_s": 0.20410156100115273,
      "inner": 104,
      "median_ms": 0.3563575288493527,
      "q1_ms": 0.3558518077485272,
      "q3_ms": 0.3565250961121189,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "block_k": 32,
      "round": 0,
      "compile_s": 0.17684430100780446,
      "inner": 19,
      "median_ms": 2.4468972107187206,
      "q1_ms": 2.4460068946753286,
      "q3_ms": 2.4471068421629942,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "block_k": 64,
      "round": 0,
      "compile_s": 0.2519827799987979,
      "inner": 18,
      "median_ms": 2.5332788330464004,
      "q1_ms": 2.530798361274517,
      "q3_ms": 2.537342944479961,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "block_k": 32,
      "round": 1,
      "compile_s": 0.19790225999895483,
      "inner": 19,
      "median_ms": 2.4477931055068773,
      "q1_ms": 2.4469488949092155,
      "q3_ms": 2.4500768949541163,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "block_k": 64,
      "round": 1,
      "compile_s": 0.2095608829986304,
      "inner": 18,
      "median_ms": 2.528143555132879,
      "q1_ms": 2.5273932773820382,
      "q3_ms": 2.531085000050047,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "block_k": 32,
      "round": 2,
      "compile_s": 0.18750811100471765,
      "inner": 19,
      "median_ms": 2.4449071581904334,
      "q1_ms": 2.4445778418387483,
      "q3_ms": 2.445693158063621,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "block_k": 64,
      "round": 2,
      "compile_s": 0.19891855400055647,
      "inner": 18,
      "median_ms": 2.529137111398288,
      "q1_ms": 2.528326695028227,
      "q3_ms": 2.529591055564298,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "block_k": 32,
      "round": 0,
      "compile_s": 0.3043903439975111,
      "inner": 19,
      "median_ms": 2.4597525794218345,
      "q1_ms": 2.4574015524673376,
      "q3_ms": 2.460196605457704,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "block_k": 64,
      "round": 0,
      "compile_s": 0.2072131500026444,
      "inner": 18,
      "median_ms": 2.5424846114927075,
      "q1_ms": 2.5397246394681536,
      "q3_ms": 2.545528722192911,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "block_k": 32,
      "round": 1,
      "compile_s": 0.18389383201429155,
      "inner": 19,
      "median_ms": 2.4577303679267826,
      "q1_ms": 2.454714631695472,
      "q3_ms": 2.458188079273034,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "block_k": 64,
      "round": 1,
      "compile_s": 0.2161534580081934,
      "inner": 18,
      "median_ms": 2.5488533331857375,
      "q1_ms": 2.5467571386899284,
      "q3_ms": 2.5501919440608214,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "block_k": 32,
      "round": 2,
      "compile_s": 0.18812702300783712,
      "inner": 19,
      "median_ms": 2.4538625263008536,
      "q1_ms": 2.452685026295704,
      "q3_ms": 2.4542585264211274,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "block_k": 64,
      "round": 2,
      "compile_s": 0.2071040230075596,
      "inner": 18,
      "median_ms": 2.5490787221416316,
      "q1_ms": 2.547244722033308,
      "q3_ms": 2.5551004445232035,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "block_k": 32,
      "round": 0,
      "compile_s": 0.1764021139970282,
      "inner": 37,
      "median_ms": 1.2288960267166684,
      "q1_ms": 1.2282942297059818,
      "q3_ms": 1.231090324172574,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "block_k": 64,
      "round": 0,
      "compile_s": 0.199679785000626,
      "inner": 33,
      "median_ms": 1.3922622121313135,
      "q1_ms": 1.3895994548378936,
      "q3_ms": 1.3929016063092137,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "block_k": 32,
      "round": 1,
      "compile_s": 0.18167077600082848,
      "inner": 37,
      "median_ms": 1.2305277840995759,
      "q1_ms": 1.2293344460320785,
      "q3_ms": 1.2310209862629857,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "block_k": 64,
      "round": 1,
      "compile_s": 0.20586178300436586,
      "inner": 32,
      "median_ms": 1.3909590311413922,
      "q1_ms": 1.3889192812257534,
      "q3_ms": 1.3912838905980607,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "block_k": 32,
      "round": 2,
      "compile_s": 0.1825773210002808,
      "inner": 36,
      "median_ms": 1.22877722242265,
      "q1_ms": 1.227860902872635,
      "q3_ms": 1.229169638867865,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "block_k": 64,
      "round": 2,
      "compile_s": 0.19953637199068908,
      "inner": 32,
      "median_ms": 1.3908695314057695,
      "q1_ms": 1.3887885313579318,
      "q3_ms": 1.4056543591323134,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "block_k": 32,
      "round": 0,
      "compile_s": 0.18014070999925025,
      "inner": 36,
      "median_ms": 1.2301879445052084,
      "q1_ms": 1.2287311526759164,
      "q3_ms": 1.2320781944355177,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "block_k": 64,
      "round": 0,
      "compile_s": 0.2082855190092232,
      "inner": 33,
      "median_ms": 1.3830314243078288,
      "q1_ms": 1.3792152727326155,
      "q3_ms": 1.3858485910918028,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "block_k": 32,
      "round": 1,
      "compile_s": 0.1840389370045159,
      "inner": 25,
      "median_ms": 1.2322731199674308,
      "q1_ms": 1.2296787998639047,
      "q3_ms": 1.232674459752161,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "block_k": 64,
      "round": 1,
      "compile_s": 0.21135910600423813,
      "inner": 33,
      "median_ms": 1.3854663332246921,
      "q1_ms": 1.3826663332910165,
      "q3_ms": 1.3858704392491332,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "block_k": 32,
      "round": 2,
      "compile_s": 0.18031665300077293,
      "inner": 37,
      "median_ms": 1.2311757565004047,
      "q1_ms": 1.2299320540576546,
      "q3_ms": 1.231289229778938,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "block_k": 64,
      "round": 2,
      "compile_s": 0.1978892360057216,
      "inner": 32,
      "median_ms": 1.3894965936742665,
      "q1_ms": 1.3823806248183246,
      "q3_ms": 1.3998072031427,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "block_k": 32,
      "round": 0,
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 4.00GiB with allocator GPU_0_bfc on device 0. [tf-allocator-allocation-error=''] [executable_name='jit_splitkv_forward']",
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "block_k": 64,
      "round": 0,
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 4.00GiB with allocator GPU_0_bfc on device 0. [executable_name='jit_splitkv_forward'] [tf-allocator-allocation-error='']",
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "block_k": 32,
      "round": 1,
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 4.00GiB with allocator GPU_0_bfc on device 0. [tf-allocator-allocation-error=''] [executable_name='jit_splitkv_forward']",
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "block_k": 64,
      "round": 1,
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 4.00GiB with allocator GPU_0_bfc on device 0. [executable_name='jit_splitkv_forward'] [tf-allocator-allocation-error='']",
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "block_k": 32,
      "round": 2,
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 4.00GiB with allocator GPU_0_bfc on device 0. [executable_name='jit_splitkv_forward'] [tf-allocator-allocation-error='']",
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "block_k": 64,
      "round": 2,
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 4.00GiB with allocator GPU_0_bfc on device 0. [tf-allocator-allocation-error=''] [executable_name='jit_splitkv_forward']",
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "block_k": 32,
      "round": 0,
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 4.00GiB with allocator GPU_0_bfc on device 0. [executable_name='jit_splitkv_forward'] [tf-allocator-allocation-error='']",
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "block_k": 64,
      "round": 0,
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 4.00GiB with allocator GPU_0_bfc on device 0. [tf-allocator-allocation-error=''] [executable_name='jit_splitkv_forward']",
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "block_k": 32,
      "round": 1,
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 4.00GiB with allocator GPU_0_bfc on device 0. [tf-allocator-allocation-error=''] [executable_name='jit_splitkv_forward']",
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "block_k": 64,
      "round": 1,
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 4.00GiB with allocator GPU_0_bfc on device 0. [tf-allocator-allocation-error=''] [executable_name='jit_splitkv_forward']",
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "block_k": 32,
      "round": 2,
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 4.00GiB with allocator GPU_0_bfc on device 0. [executable_name='jit_splitkv_forward'] [tf-allocator-allocation-error='']",
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 32,
      "head_dim": 128,
      "block_k": 64,
      "round": 2,
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 4.00GiB with allocator GPU_0_bfc on device 0. [executable_name='jit_splitkv_forward'] [tf-allocator-allocation-error='']",
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "block_k": 32,
      "round": 0,
      "compile_s": 0.22726108800270595,
      "inner": 200,
      "median_ms": 0.08999770499940496,
      "q1_ms": 0.08996832999400795,
      "q3_ms": 0.09002846498333383,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "block_k": 64,
      "round": 0,
      "compile_s": 0.24531999199825805,
      "inner": 200,
      "median_ms": 0.08963091495388653,
      "q1_ms": 0.08960356248280732,
      "q3_ms": 0.08965728753537405,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "block_k": 32,
      "round": 1,
      "compile_s": 0.17640343800303526,
      "inner": 200,
      "median_ms": 0.08868311997503042,
      "q1_ms": 0.0886792124947533,
      "q3_ms": 0.08875487998011522,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "block_k": 64,
      "round": 1,
      "compile_s": 0.192264241995872,
      "inner": 200,
      "median_ms": 0.08931201002269518,
      "q1_ms": 0.08927381000830792,
      "q3_ms": 0.08944871751737082,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "block_k": 32,
      "round": 2,
      "compile_s": 0.1692654299986316,
      "inner": 200,
      "median_ms": 0.08850418002111837,
      "q1_ms": 0.08846047749102581,
      "q3_ms": 0.08857416250975803,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "block_k": 64,
      "round": 2,
      "compile_s": 0.19321490998845547,
      "inner": 200,
      "median_ms": 0.0890838749910472,
      "q1_ms": 0.08906742248655064,
      "q3_ms": 0.08909855496312957,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "block_k": 32,
      "round": 0,
      "compile_s": 0.3517274920013733,
      "inner": 180,
      "median_ms": 0.14800096113402914,
      "q1_ms": 0.14763863332013394,
      "q3_ms": 0.14818953891840972,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "block_k": 64,
      "round": 0,
      "compile_s": 0.4000049099995522,
      "inner": 193,
      "median_ms": 0.12448956999008501,
      "q1_ms": 0.12406863729406747,
      "q3_ms": 0.12474141712179843,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "block_k": 32,
      "round": 1,
      "compile_s": 0.30543467999086715,
      "inner": 176,
      "median_ms": 0.14808511365448465,
      "q1_ms": 0.14664927555879165,
      "q3_ms": 0.14813589773025757,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "block_k": 64,
      "round": 1,
      "compile_s": 0.37159589100338053,
      "inner": 200,
      "median_ms": 0.12451547998352909,
      "q1_ms": 0.12263103748409776,
      "q3_ms": 0.12506038496212568,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "block_k": 32,
      "round": 2,
      "compile_s": 0.30964331500581466,
      "inner": 194,
      "median_ms": 0.14790777835908547,
      "q1_ms": 0.1478949999397731,
      "q3_ms": 0.14797835308517154,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "block_k": 64,
      "round": 2,
      "compile_s": 0.3619816159916809,
      "inner": 193,
      "median_ms": 0.12553745596922952,
      "q1_ms": 0.12329930829854172,
      "q3_ms": 0.12559217098199188,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "block_k": 32,
      "round": 0,
      "compile_s": 0.1753059610055061,
      "inner": 66,
      "median_ms": 0.6189896816839323,
      "q1_ms": 0.6189064090104887,
      "q3_ms": 0.6191287273285332,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "block_k": 64,
      "round": 0,
      "compile_s": 0.19566212099744007,
      "inner": 60,
      "median_ms": 0.6987601334306722,
      "q1_ms": 0.698353033294552,
      "q3_ms": 0.6993309582564203,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "block_k": 32,
      "round": 1,
      "compile_s": 0.1709303009993164,
      "inner": 63,
      "median_ms": 0.6191171428543472,
      "q1_ms": 0.6189646429404438,
      "q3_ms": 0.6194135793420072,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "block_k": 64,
      "round": 1,
      "compile_s": 0.19832060199405532,
      "inner": 58,
      "median_ms": 0.6987289656033931,
      "q1_ms": 0.6983540258480734,
      "q3_ms": 0.6987569483091396,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "block_k": 32,
      "round": 2,
      "compile_s": 0.1728903599869227,
      "inner": 66,
      "median_ms": 0.6192313484268996,
      "q1_ms": 0.6190651968636578,
      "q3_ms": 0.6193254621669466,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "block_k": 64,
      "round": 2,
      "compile_s": 0.19723717200395186,
      "inner": 59,
      "median_ms": 0.6987213729821556,
      "q1_ms": 0.6979767456838182,
      "q3_ms": 0.6992303050060015,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "block_k": 32,
      "round": 0,
      "compile_s": 0.35514633198908996,
      "inner": 47,
      "median_ms": 0.9529390424970497,
      "q1_ms": 0.950850202067894,
      "q3_ms": 0.9536444146973775,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "block_k": 64,
      "round": 0,
      "compile_s": 0.3688183729973389,
      "inner": 46,
      "median_ms": 1.024869630308113,
      "q1_ms": 1.0232055000138327,
      "q3_ms": 1.026851750003002,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "block_k": 32,
      "round": 1,
      "compile_s": 0.3093794440064812,
      "inner": 44,
      "median_ms": 0.957279704545561,
      "q1_ms": 0.9562739317268345,
      "q3_ms": 0.9577163180603583,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "block_k": 64,
      "round": 1,
      "compile_s": 0.36314396500529256,
      "inner": 43,
      "median_ms": 1.0242151396977173,
      "q1_ms": 1.0235433021376215,
      "q3_ms": 1.0272137323925057,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "block_k": 32,
      "round": 2,
      "compile_s": 0.31099750599241816,
      "inner": 47,
      "median_ms": 0.9540935744224315,
      "q1_ms": 0.9530985850513853,
      "q3_ms": 0.9558721382958577,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 16384,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "block_k": 64,
      "round": 2,
      "compile_s": 0.363463089990546,
      "inner": 44,
      "median_ms": 1.024512386224656,
      "q1_ms": 1.0241176362548932,
      "q3_ms": 1.026753897762815,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "block_k": 32,
      "round": 0,
      "compile_s": 0.1702257080032723,
      "inner": 111,
      "median_ms": 0.31742807199784157,
      "q1_ms": 0.31737139189473207,
      "q3_ms": 0.3174853873379352,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "block_k": 64,
      "round": 0,
      "compile_s": 0.19262838100257795,
      "inner": 91,
      "median_ms": 0.38204604397133857,
      "q1_ms": 0.381853093411057,
      "q3_ms": 0.38233857688552875,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "block_k": 32,
      "round": 1,
      "compile_s": 0.17083329601155128,
      "inner": 111,
      "median_ms": 0.3176207837978009,
      "q1_ms": 0.3176009684810541,
      "q3_ms": 0.3177151486418118,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "block_k": 64,
      "round": 1,
      "compile_s": 0.19066195999039337,
      "inner": 101,
      "median_ms": 0.38173068317672293,
      "q1_ms": 0.381361569380526,
      "q3_ms": 0.3819314109266385,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "block_k": 32,
      "round": 2,
      "compile_s": 0.17236503500316758,
      "inner": 115,
      "median_ms": 0.31773056531244,
      "q1_ms": 0.3176135608911231,
      "q3_ms": 0.3177719434935843,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "block_k": 64,
      "round": 2,
      "compile_s": 0.19011280999984592,
      "inner": 94,
      "median_ms": 0.38196846806850443,
      "q1_ms": 0.38192620218553125,
      "q3_ms": 0.3821078138029639,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "block_k": 32,
      "round": 0,
      "compile_s": 0.30485330898955,
      "inner": 78,
      "median_ms": 0.5341244742829496,
      "q1_ms": 0.5338687948003494,
      "q3_ms": 0.5344414871965105,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "block_k": 64,
      "round": 0,
      "compile_s": 0.3611195349949412,
      "inner": 71,
      "median_ms": 0.5770890986215664,
      "q1_ms": 0.5761939577445139,
      "q3_ms": 0.5772163168958414,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "block_k": 32,
      "round": 1,
      "compile_s": 0.3052743270091014,
      "inner": 75,
      "median_ms": 0.5347516665157551,
      "q1_ms": 0.5345735334170361,
      "q3_ms": 0.5348449866869487,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "block_k": 64,
      "round": 1,
      "compile_s": 0.3666207760106772,
      "inner": 69,
      "median_ms": 0.5757104492817154,
      "q1_ms": 0.5753150796500858,
      "q3_ms": 0.5759626013889869,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "block_k": 32,
      "round": 2,
      "compile_s": 0.30520015601359773,
      "inner": 73,
      "median_ms": 0.5327591232792793,
      "q1_ms": 0.5317881095585133,
      "q3_ms": 0.5330216028521871,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 1,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "block_k": 64,
      "round": 2,
      "compile_s": 0.36112993500137236,
      "inner": 69,
      "median_ms": 0.5757020433362493,
      "q1_ms": 0.5746741522486538,
      "q3_ms": 0.5758670652183988,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "block_k": 32,
      "round": 0,
      "compile_s": 0.17369429000245873,
      "inner": 19,
      "median_ms": 2.43245836846702,
      "q1_ms": 2.431542894842814,
      "q3_ms": 2.4333850788503986,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "block_k": 64,
      "round": 0,
      "compile_s": 0.19323966799129266,
      "inner": 17,
      "median_ms": 2.75395558855962,
      "q1_ms": 2.7491196468712156,
      "q3_ms": 2.7626611468697186,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "block_k": 32,
      "round": 1,
      "compile_s": 0.1733779319911264,
      "inner": 19,
      "median_ms": 2.4330737367946362,
      "q1_ms": 2.432349420958321,
      "q3_ms": 2.433618973315644,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "block_k": 64,
      "round": 1,
      "compile_s": 0.19349970000621397,
      "inner": 17,
      "median_ms": 2.747704881852876,
      "q1_ms": 2.746412176587482,
      "q3_ms": 2.754872617729501,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "block_k": 32,
      "round": 2,
      "compile_s": 0.1759631739987526,
      "inner": 19,
      "median_ms": 2.432399263182704,
      "q1_ms": 2.4323383948189745,
      "q3_ms": 2.4384855262867764,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 1,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "block_k": 64,
      "round": 2,
      "compile_s": 0.19829339100397192,
      "inner": 17,
      "median_ms": 2.7465244707213166,
      "q1_ms": 2.7442424706190223,
      "q3_ms": 2.7522352645653028,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "block_k": 32,
      "round": 0,
      "compile_s": 0.3114373089920264,
      "inner": 13,
      "median_ms": 3.7399564618961167,
      "q1_ms": 3.7309013848873573,
      "q3_ms": 3.7508198074647225,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "block_k": 64,
      "round": 0,
      "compile_s": 0.37709937400359195,
      "inner": 12,
      "median_ms": 3.9865389177672723,
      "q1_ms": 3.979993208001057,
      "q3_ms": 4.0031716252997285,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "block_k": 32,
      "round": 1,
      "compile_s": 0.3144876709993696,
      "inner": 13,
      "median_ms": 3.7413592307040324,
      "q1_ms": 3.734915500479894,
      "q3_ms": 3.7472087692899203,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "block_k": 64,
      "round": 1,
      "compile_s": 0.36730369699944276,
      "inner": 12,
      "median_ms": 3.9967474998168955,
      "q1_ms": 3.9895922078964454,
      "q3_ms": 4.007388416842636,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "block_k": 32,
      "round": 2,
      "compile_s": 0.3117956369969761,
      "inner": 13,
      "median_ms": 3.748455538208453,
      "q1_ms": 3.7408443850626307,
      "q3_ms": 3.753328461383577,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "suite": "blockk",
      "dtype": "bfloat16",
      "batch": 8,
      "lq": 16,
      "lkv": 65536,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "block_k": 64,
      "round": 2,
      "compile_s": 0.3731309350114316,
      "inner": 12,
      "median_ms": 3.994679333118256,
      "q1_ms": 3.982891583291348,
      "q3_ms": 4.0034725831598434,
      "tag": null,
      "device_arg": "gpu",
      "run": "blockk"
    },
    {
      "lq": 4096,
      "lkv": 4096,
      "batch": 1,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "dtype": "bfloat16",
      "mode": "fwd",
      "impl": "flex_native",
      "suite": "peak",
      "peak_mb": 144.000732421875,
      "inputs_mb": 112.000244140625,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 32.0,
      "tag": null,
      "device_arg": "gpu",
      "run": "peak"
    },
    {
      "lq": 4096,
      "lkv": 4096,
      "batch": 1,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "dtype": "bfloat16",
      "mode": "fwd",
      "impl": "flex_repeat",
      "suite": "peak",
      "peak_mb": 240.000732421875,
      "inputs_mb": 112.000244140625,
      "xla_temp_mb": 64.0,
      "xla_arg_mb": 48.0,
      "xla_out_mb": 32.0,
      "tag": null,
      "device_arg": "gpu",
      "run": "peak"
    },
    {
      "lq": 4096,
      "lkv": 4096,
      "batch": 1,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "dtype": "bfloat16",
      "mode": "fwd",
      "impl": "torch_sdpa",
      "suite": "peak",
      "peak_mb": 144.5009765625,
      "inputs_mb": 112.0,
      "tag": null,
      "device_arg": "gpu",
      "run": "peak"
    },
    {
      "lq": 4096,
      "lkv": 4096,
      "batch": 1,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "dtype": "bfloat16",
      "mode": "fwd_bwd",
      "impl": "flex_native",
      "suite": "peak",
      "peak_mb": 416.0009765625,
      "inputs_mb": 112.000244140625,
      "xla_temp_mb": 128.0002670288086,
      "xla_arg_mb": 112.0,
      "xla_out_mb": 48.000022888183594,
      "tag": null,
      "device_arg": "gpu",
      "run": "peak"
    },
    {
      "lq": 4096,
      "lkv": 4096,
      "batch": 1,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "dtype": "bfloat16",
      "mode": "fwd_bwd",
      "impl": "flex_repeat",
      "suite": "peak",
      "peak_mb": 416.0009765625,
      "inputs_mb": 112.000244140625,
      "xla_temp_mb": 192.0002670288086,
      "xla_arg_mb": 112.0,
      "xla_out_mb": 48.000022888183594,
      "tag": null,
      "device_arg": "gpu",
      "run": "peak"
    },
    {
      "lq": 4096,
      "lkv": 4096,
      "batch": 1,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "dtype": "bfloat16",
      "mode": "fwd_bwd",
      "impl": "torch_sdpa",
      "suite": "peak",
      "peak_mb": 417.001953125,
      "inputs_mb": 112.0,
      "tag": null,
      "device_arg": "gpu",
      "run": "peak"
    },
    {
      "lq": 4096,
      "lkv": 4096,
      "batch": 1,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "dtype": "bfloat16",
      "mode": "fwd",
      "impl": "flex_native",
      "suite": "peak",
      "peak_mb": 131.000244140625,
      "inputs_mb": 98.999755859375,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 34.0,
      "xla_out_mb": 32.0,
      "tag": null,
      "device_arg": "gpu",
      "run": "peak"
    },
    {
      "lq": 4096,
      "lkv": 4096,
      "batch": 1,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "dtype": "bfloat16",
      "mode": "fwd",
      "impl": "flex_repeat",
      "suite": "peak",
      "peak_mb": 227.000244140625,
      "inputs_mb": 98.999755859375,
      "xla_temp_mb": 64.0,
      "xla_arg_mb": 34.0,
      "xla_out_mb": 32.0,
      "tag": null,
      "device_arg": "gpu",
      "run": "peak"
    },
    {
      "lq": 4096,
      "lkv": 4096,
      "batch": 1,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "dtype": "bfloat16",
      "mode": "fwd",
      "impl": "torch_sdpa",
      "suite": "peak",
      "peak_mb": 130.5009765625,
      "inputs_mb": 98.0,
      "tag": null,
      "device_arg": "gpu",
      "run": "peak"
    },
    {
      "lq": 4096,
      "lkv": 4096,
      "batch": 1,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "dtype": "bfloat16",
      "mode": "fwd_bwd",
      "impl": "flex_native",
      "suite": "peak",
      "peak_mb": 389.00048828125,
      "inputs_mb": 98.999755859375,
      "xla_temp_mb": 128.0002670288086,
      "xla_arg_mb": 98.0,
      "xla_out_mb": 34.000022888183594,
      "tag": null,
      "device_arg": "gpu",
      "run": "peak"
    },
    {
      "lq": 4096,
      "lkv": 4096,
      "batch": 1,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "dtype": "bfloat16",
      "mode": "fwd_bwd",
      "impl": "flex_repeat",
      "suite": "peak",
      "peak_mb": 389.00048828125,
      "inputs_mb": 98.999755859375,
      "xla_temp_mb": 192.0002670288086,
      "xla_arg_mb": 98.0,
      "xla_out_mb": 34.000022888183594,
      "tag": null,
      "device_arg": "gpu",
      "run": "peak"
    },
    {
      "lq": 4096,
      "lkv": 4096,
      "batch": 1,
      "hq": 32,
      "hkv": 1,
      "head_dim": 128,
      "dtype": "bfloat16",
      "mode": "fwd_bwd",
      "impl": "torch_sdpa",
      "suite": "peak",
      "peak_mb": 389.001953125,
      "inputs_mb": 98.0,
      "tag": null,
      "device_arg": "gpu",
      "run": "peak"
    },
    {
      "lq": 1,
      "lkv": 65536,
      "batch": 8,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "dtype": "bfloat16",
      "mode": "fwd",
      "impl": "probjax_splitkv",
      "suite": "peak",
      "peak_mb": 2050.0126953125,
      "inputs_mb": 2048.187744140625,
      "xla_temp_mb": 1.7617340087890625,
      "xla_arg_mb": 2048.0625,
      "xla_out_mb": 0.0625,
      "tag": null,
      "device_arg": "gpu",
      "run": "peak"
    },
    {
      "lq": 1,
      "lkv": 65536,
      "batch": 8,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "dtype": "bfloat16",
      "mode": "fwd",
      "impl": "probjax_nosplit",
      "suite": "peak",
      "peak_mb": 2050.250732421875,
      "inputs_mb": 2048.187744140625,
      "xla_temp_mb": 2.0,
      "xla_arg_mb": 2048.0625,
      "xla_out_mb": 0.0625,
      "tag": null,
      "device_arg": "gpu",
      "run": "peak"
    },
    {
      "lq": 1,
      "lkv": 65536,
      "batch": 8,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "dtype": "bfloat16",
      "mode": "fwd",
      "impl": "flex_repeat",
      "suite": "peak",
      "peak_mb": 10242.250732421875,
      "inputs_mb": 2048.187744140625,
      "xla_temp_mb": 8194.0,
      "xla_arg_mb": 2048.0625,
      "xla_out_mb": 0.0625,
      "tag": null,
      "device_arg": "gpu",
      "run": "peak"
    },
    {
      "lq": 1,
      "lkv": 65536,
      "batch": 8,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "dtype": "bfloat16",
      "mode": "fwd",
      "impl": "torch_sdpa",
      "suite": "peak",
      "peak_mb": 2048.8818359375,
      "inputs_mb": 2048.1875,
      "tag": null,
      "device_arg": "gpu",
      "run": "peak"
    },
    {
      "lq": 16,
      "lkv": 65536,
      "batch": 8,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "dtype": "bfloat16",
      "mode": "fwd",
      "impl": "probjax_splitkv",
      "suite": "peak",
      "peak_mb": 2088.999755859375,
      "inputs_mb": 2053.999267578125,
      "xla_temp_mb": 6.0937652587890625,
      "xla_arg_mb": 2049.0,
      "xla_out_mb": 1.0,
      "tag": null,
      "device_arg": "gpu",
      "run": "peak"
    },
    {
      "lq": 16,
      "lkv": 65536,
      "batch": 8,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "dtype": "bfloat16",
      "mode": "fwd",
      "impl": "probjax_nosplit",
      "suite": "peak",
      "peak_mb": 2054.999755859375,
      "inputs_mb": 2053.999267578125,
      "xla_temp_mb": 0.0,
      "xla_arg_mb": 2049.0,
      "xla_out_mb": 1.0,
      "tag": null,
      "device_arg": "gpu",
      "run": "peak"
    },
    {
      "lq": 16,
      "lkv": 65536,
      "batch": 8,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "dtype": "bfloat16",
      "mode": "fwd",
      "impl": "flex_repeat",
      "suite": "peak",
      "peak_mb": 10246.999755859375,
      "inputs_mb": 2053.999267578125,
      "xla_temp_mb": 8192.0,
      "xla_arg_mb": 2049.0,
      "xla_out_mb": 1.0,
      "tag": null,
      "device_arg": "gpu",
      "run": "peak"
    },
    {
      "lq": 16,
      "lkv": 65536,
      "batch": 8,
      "hq": 32,
      "hkv": 8,
      "head_dim": 128,
      "dtype": "bfloat16",
      "mode": "fwd",
      "impl": "torch_sdpa",
      "suite": "peak",
      "peak_mb": 2052.0166015625,
      "inputs_mb": 2051.0,
      "tag": null,
      "device_arg": "gpu",
      "run": "peak"
    }
  ]
}