{ "device": { "jax": "0.11.2", "backend": "cpu", "devices": [ "cpu:0" ], "python": "3.14.8", "device_kind": "cpu", "gpu": "NVIDIA GeForce RTX 3090", "driver": "615.78.08", "gpu_memory_gb": 24.0, "torch": "2.14.1+cu130", "torch_cuda": "13.0" }, "config": { "note": "CPU (JAX_PLATFORMS=cpu). decode rows tagged before/after the flex CPU padding fix; flex_repeat = previous GQA path (repeat K/V).", "runs": { "cpu_gqa": { "device": "cpu", "suite": "gqa", "output": "/tmp/claude-1003/-home-manuel-src/80785a0e-56be-4da7-a2a8-e1d6effd6199/scratchpad/attn-ext/cpu_gqa.json", "append": false, "tag": null, "repeats": 3, "batch": 1, "batches": [ 1, 8 ], "heads": 32, "kv_heads": [ 8, 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", "torch_sdpa" ], "baseline_impl": "flex_repeat", "ref_max_seq": 1024, "max_kv_gb": 4.5, "child": null }, "cpu_decode_before": { "device": "cpu", "suite": "decode", "output": "/tmp/claude-1003/-home-manuel-src/80785a0e-56be-4da7-a2a8-e1d6effd6199/scratchpad/attn-ext/cpu_decode_before.json", "append": false, "tag": "before", "repeats": 3, "batch": 1, "batches": [ 1 ], "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": [ "float32" ], "modes": [ "fwd" ], "impls": [ "flex_native", "flex_repeat", "jax_xla", "torch_sdpa" ], "baseline_impl": "jax_xla", "ref_max_seq": 1024, "max_kv_gb": 4.5, "child": null }, "cpu_decode_after": { "device": "cpu", "suite": "decode", "output": "/tmp/claude-1003/-home-manuel-src/80785a0e-56be-4da7-a2a8-e1d6effd6199/scratchpad/attn-ext/cpu_decode_after.json", "append": false, "tag": "after", "repeats": 3, "batch": 1, "batches": [ 1 ], "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": [ "float32" ], "modes": [ "fwd" ], "impls": [ "flex_native", "flex_repeat", "jax_xla", "torch_sdpa" ], "baseline_impl": "jax_xla", "ref_max_seq": 1024, "max_kv_gb": 4.5, "child": null } } }, "results": [ { "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": 32.59417724609375, "xla_arg_mb": 12.0, "xla_out_mb": 8.0, "rel_err": 6.438401101149793e-07, "rel_err_parts": [ 6.438401101149793e-07 ], "rel_vs_flex_repeat": 0.0, "compile_s": 0.24075318200630136, "inner": 1, "median_ms": 51.074707007501274, "q1_ms": 50.59345100016799, "q3_ms": 51.868579495931044, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 46.09417724609375, "xla_arg_mb": 12.0, "xla_out_mb": 8.0, "rel_err": 6.438401101149793e-07, "rel_err_parts": [ 6.438401101149793e-07 ], "compile_s": 0.19752433999383356, "inner": 1, "median_ms": 59.49386600696016, "q1_ms": 58.71151099563576, "q3_ms": 60.26082750031492, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 386.0, "xla_arg_mb": 12.0, "xla_out_mb": 8.0, "rel_err": 1.4813868176588585e-07, "rel_err_parts": [ 1.4813868176588585e-07 ], "rel_vs_flex_repeat": 6.438485797843896e-07, "compile_s": 0.1803702430042904, "inner": 1, "median_ms": 175.12316499778535, "q1_ms": 169.1476589985541, "q3_ms": 178.3509425004013, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 6.281644004957343e-07, "rel_err_parts": [ 6.281644004957343e-07 ], "rel_vs_flex_repeat": 3.333002496219706e-07, "compile_s": 0.11928493500454351, "inner": 3, "median_ms": 12.216531332038963, "q1_ms": 10.622623832508301, "q3_ms": 12.567012500464141, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 32.59417724609375, "xla_arg_mb": 12.0, "xla_out_mb": 8.0, "rel_err": 6.438401101149793e-07, "rel_err_parts": [ 6.438401101149793e-07 ], "rel_vs_flex_repeat": 0.0, "compile_s": 0.18966681300662458, "inner": 1, "median_ms": 52.91534999560099, "q1_ms": 51.87365349411266, "q3_ms": 53.379758006485645, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 46.09417724609375, "xla_arg_mb": 12.0, "xla_out_mb": 8.0, "rel_err": 6.438401101149793e-07, "rel_err_parts": [ 6.438401101149793e-07 ], "compile_s": 0.1937846740038367, "inner": 1, "median_ms": 59.09279799379874, "q1_ms": 58.628699494875036, "q3_ms": 59.91457749041729, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 386.0, "xla_arg_mb": 12.0, "xla_out_mb": 8.0, "rel_err": 1.4813868176588585e-07, "rel_err_parts": [ 1.4813868176588585e-07 ], "rel_vs_flex_repeat": 6.438485797843896e-07, "compile_s": 0.18694285899982788, "inner": 1, "median_ms": 172.45246199308895, "q1_ms": 162.68088849756168, "q3_ms": 177.8352615001495, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 6.281644004957343e-07, "rel_err_parts": [ 6.281644004957343e-07 ], "rel_vs_flex_repeat": 3.333002496219706e-07, "compile_s": 0.01087153000116814, "inner": 4, "median_ms": 9.50632649983163, "q1_ms": 9.367557875521015, "q3_ms": 9.929959376677289, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 32.59417724609375, "xla_arg_mb": 12.0, "xla_out_mb": 8.0, "rel_err": 6.438401101149793e-07, "rel_err_parts": [ 6.438401101149793e-07 ], "rel_vs_flex_repeat": 0.0, "compile_s": 0.19117982500756625, "inner": 1, "median_ms": 51.58591399958823, "q1_ms": 51.29668649897212, "q3_ms": 52.58217950904509, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 46.09417724609375, "xla_arg_mb": 12.0, "xla_out_mb": 8.0, "rel_err": 6.438401101149793e-07, "rel_err_parts": [ 6.438401101149793e-07 ], "compile_s": 0.2018017029913608, "inner": 1, "median_ms": 60.29908999335021, "q1_ms": 59.34611200063955, "q3_ms": 61.40878999576671, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 386.0, "xla_arg_mb": 12.0, "xla_out_mb": 8.0, "rel_err": 1.4813868176588585e-07, "rel_err_parts": [ 1.4813868176588585e-07 ], "rel_vs_flex_repeat": 6.438485797843896e-07, "compile_s": 0.1745634659891948, "inner": 1, "median_ms": 178.4758790017804, "q1_ms": 164.6505704993615, "q3_ms": 179.87553650164045, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 6.281644004957343e-07, "rel_err_parts": [ 6.281644004957343e-07 ], "rel_vs_flex_repeat": 3.333002496219706e-07, "compile_s": 0.012880711001344025, "inner": 3, "median_ms": 12.196626664566187, "q1_ms": 12.007781166175846, "q3_ms": 12.503457330846384, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 66.25069427490234, "xla_arg_mb": 20.0, "xla_out_mb": 12.000022888183594, "residual_mb": 20.125003814697266, "residual_shapes": [ "float32[1, 8, 4, 1024, 64]", "float32[1, 8, 1024, 64]", "float32[1, 8, 1024, 64]", "uint32[]", "float32[1, 8, 4, 1024, 64]", "float32[1, 32, 1024]" ], "rel_err": 8.295533575619629e-07, "rel_err_parts": [ 6.693323371109727e-07, 8.242420790338656e-07, 8.295533575619629e-07 ], "rel_vs_flex_repeat": 6.003033945489733e-07, "compile_s": 0.4158515199960675, "inner": 1, "median_ms": 205.41209400107618, "q1_ms": 202.86003499495564, "q3_ms": 208.26174550165888, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 96.25069427490234, "xla_arg_mb": 20.0, "xla_out_mb": 12.000022888183594, "residual_mb": 32.125003814697266, "residual_shapes": [ "float32[1, 32, 1, 1024, 64]", "float32[1, 32, 1024, 64]", "float32[1, 32, 1024, 64]", "uint32[]", "float32[1, 32, 1, 1024, 64]", "float32[1, 32, 1024]" ], "rel_err": 6.72211740493367e-07, "rel_err_parts": [ 6.693323371109727e-07, 6.645150278927758e-07, 6.72211740493367e-07 ], "compile_s": 0.4077428879973013, "inner": 1, "median_ms": 197.56983598927036, "q1_ms": 193.79469849809539, "q3_ms": 200.71767699846532, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 522.125, "xla_arg_mb": 20.0, "xla_out_mb": 12.000022888183594, "rel_err": 7.794450880282966e-07, "rel_err_parts": [ 1.5355949756212794e-07, 7.794450880282966e-07, 1.730932552845843e-07 ], "rel_vs_flex_repeat": 6.870205311315658e-07, "compile_s": 0.34366165199026, "inner": 1, "median_ms": 264.17311299883295, "q1_ms": 261.5086324949516, "q3_ms": 266.39383049769094, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 6.720209171362512e-07, "rel_err_parts": [ 6.676874590993975e-07, 6.664754437224474e-07, 6.720209171362512e-07 ], "rel_vs_flex_repeat": 3.569370790046378e-07, "compile_s": 0.058839368997723795, "inner": 1, "median_ms": 51.150740997400135, "q1_ms": 49.75300600199262, "q3_ms": 52.58335699909367, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 66.25069427490234, "xla_arg_mb": 20.0, "xla_out_mb": 12.000022888183594, "residual_mb": 20.125003814697266, "residual_shapes": [ "float32[1, 8, 4, 1024, 64]", "float32[1, 8, 1024, 64]", "float32[1, 8, 1024, 64]", "uint32[]", "float32[1, 8, 4, 1024, 64]", "float32[1, 32, 1024]" ], "rel_err": 8.295533575619629e-07, "rel_err_parts": [ 6.693323371109727e-07, 8.242420790338656e-07, 8.295533575619629e-07 ], "rel_vs_flex_repeat": 6.003033945489733e-07, "compile_s": 0.3907647239975631, "inner": 1, "median_ms": 193.77270199765917, "q1_ms": 192.3389684961876, "q3_ms": 196.27717449475313, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 96.25069427490234, "xla_arg_mb": 20.0, "xla_out_mb": 12.000022888183594, "residual_mb": 32.125003814697266, "residual_shapes": [ "float32[1, 32, 1, 1024, 64]", "float32[1, 32, 1024, 64]", "float32[1, 32, 1024, 64]", "uint32[]", "float32[1, 32, 1, 1024, 64]", "float32[1, 32, 1024]" ], "rel_err": 6.72211740493367e-07, "rel_err_parts": [ 6.693323371109727e-07, 6.645150278927758e-07, 6.72211740493367e-07 ], "compile_s": 0.38894650401198305, "inner": 1, "median_ms": 186.02660500619095, "q1_ms": 182.7790504976292, "q3_ms": 186.97922299907077, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 522.125, "xla_arg_mb": 20.0, "xla_out_mb": 12.000022888183594, "rel_err": 7.794450880282966e-07, "rel_err_parts": [ 1.5355949756212794e-07, 7.794450880282966e-07, 1.730932552845843e-07 ], "rel_vs_flex_repeat": 6.870205311315658e-07, "compile_s": 0.3230882750067394, "inner": 1, "median_ms": 267.6247650088044, "q1_ms": 264.5903279990307, "q3_ms": 269.57994399708696, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 6.720209171362512e-07, "rel_err_parts": [ 6.676874590993975e-07, 6.664754437224474e-07, 6.720209171362512e-07 ], "rel_vs_flex_repeat": 3.569370790046378e-07, "compile_s": 0.05212820999440737, "inner": 1, "median_ms": 56.789589987602085, "q1_ms": 51.94045499956701, "q3_ms": 59.89723949460313, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 66.25069427490234, "xla_arg_mb": 20.0, "xla_out_mb": 12.000022888183594, "residual_mb": 20.125003814697266, "residual_shapes": [ "float32[1, 8, 4, 1024, 64]", "float32[1, 8, 1024, 64]", "float32[1, 8, 1024, 64]", "uint32[]", "float32[1, 8, 4, 1024, 64]", "float32[1, 32, 1024]" ], "rel_err": 8.295533575619629e-07, "rel_err_parts": [ 6.693323371109727e-07, 8.242420790338656e-07, 8.295533575619629e-07 ], "rel_vs_flex_repeat": 6.003033945489733e-07, "compile_s": 0.4136450049991254, "inner": 1, "median_ms": 212.86989499640185, "q1_ms": 209.22447249904508, "q3_ms": 217.89558849559398, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 96.25069427490234, "xla_arg_mb": 20.0, "xla_out_mb": 12.000022888183594, "residual_mb": 32.125003814697266, "residual_shapes": [ "float32[1, 32, 1, 1024, 64]", "float32[1, 32, 1024, 64]", "float32[1, 32, 1024, 64]", "uint32[]", "float32[1, 32, 1, 1024, 64]", "float32[1, 32, 1024]" ], "rel_err": 6.72211740493367e-07, "rel_err_parts": [ 6.693323371109727e-07, 6.645150278927758e-07, 6.72211740493367e-07 ], "compile_s": 0.41531148900685366, "inner": 1, "median_ms": 196.71194200054742, "q1_ms": 193.94635000207927, "q3_ms": 239.54605299513787, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 522.125, "xla_arg_mb": 20.0, "xla_out_mb": 12.000022888183594, "rel_err": 7.794450880282966e-07, "rel_err_parts": [ 1.5355949756212794e-07, 7.794450880282966e-07, 1.730932552845843e-07 ], "rel_vs_flex_repeat": 6.870205311315658e-07, "compile_s": 0.3369707519887015, "inner": 1, "median_ms": 267.248791991733, "q1_ms": 265.0188795014401, "q3_ms": 283.51402349653654, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 6.720209171362512e-07, "rel_err_parts": [ 6.676874590993975e-07, 6.664754437224474e-07, 6.720209171362512e-07 ], "rel_vs_flex_repeat": 3.569370790046378e-07, "compile_s": 0.05594437000399921, "inner": 1, "median_ms": 47.16899699997157, "q1_ms": 47.09914499835577, "q3_ms": 48.50896750576794, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 28.15667724609375, "xla_arg_mb": 8.5, "xla_out_mb": 8.0, "rel_err": 6.482335379587312e-07, "rel_err_parts": [ 6.482335379587312e-07 ], "rel_vs_flex_repeat": 0.0, "compile_s": 0.18472648499300703, "inner": 1, "median_ms": 43.37837500497699, "q1_ms": 42.173750494839624, "q3_ms": 45.49534049874637, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 46.09417724609375, "xla_arg_mb": 8.5, "xla_out_mb": 8.0, "rel_err": 6.482335379587312e-07, "rel_err_parts": [ 6.482335379587312e-07 ], "compile_s": 0.19945989099505823, "inner": 1, "median_ms": 59.34409299516119, "q1_ms": 59.19333799829474, "q3_ms": 69.42406450252747, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 384.0, "xla_arg_mb": 8.5, "xla_out_mb": 8.0, "rel_err": 1.4823216076820245e-07, "rel_err_parts": [ 1.4823216076820245e-07 ], "rel_vs_flex_repeat": 6.483219294750597e-07, "compile_s": 0.19199748199025635, "inner": 1, "median_ms": 158.10697901179083, "q1_ms": 156.76301199709997, "q3_ms": 163.86215850070585, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 6.325635695247911e-07, "rel_err_parts": [ 6.325635695247911e-07 ], "rel_vs_flex_repeat": 3.371636410065548e-07, "compile_s": 0.01675471199268941, "inner": 3, "median_ms": 14.287636333998913, "q1_ms": 14.200099166676711, "q3_ms": 14.518396332277916, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 28.15667724609375, "xla_arg_mb": 8.5, "xla_out_mb": 8.0, "rel_err": 6.482335379587312e-07, "rel_err_parts": [ 6.482335379587312e-07 ], "rel_vs_flex_repeat": 0.0, "compile_s": 0.1836980299995048, "inner": 1, "median_ms": 41.90499799733516, "q1_ms": 41.693952996865846, "q3_ms": 43.279262506985106, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 46.09417724609375, "xla_arg_mb": 8.5, "xla_out_mb": 8.0, "rel_err": 6.482335379587312e-07, "rel_err_parts": [ 6.482335379587312e-07 ], "compile_s": 0.20071452599950135, "inner": 1, "median_ms": 58.261786005459726, "q1_ms": 57.463388999167364, "q3_ms": 58.48598000011407, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 384.0, "xla_arg_mb": 8.5, "xla_out_mb": 8.0, "rel_err": 1.4823216076820245e-07, "rel_err_parts": [ 1.4823216076820245e-07 ], "rel_vs_flex_repeat": 6.483219294750597e-07, "compile_s": 0.1787851240078453, "inner": 1, "median_ms": 169.4387210009154, "q1_ms": 165.05198500817642, "q3_ms": 172.18017699633492, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 6.325635695247911e-07, "rel_err_parts": [ 6.325635695247911e-07 ], "rel_vs_flex_repeat": 3.371636410065548e-07, "compile_s": 0.015360053002950735, "inner": 3, "median_ms": 12.443788000382483, "q1_ms": 12.387703497855302, "q3_ms": 12.52782516530715, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 28.15667724609375, "xla_arg_mb": 8.5, "xla_out_mb": 8.0, "rel_err": 6.482335379587312e-07, "rel_err_parts": [ 6.482335379587312e-07 ], "rel_vs_flex_repeat": 0.0, "compile_s": 0.18095240399998147, "inner": 1, "median_ms": 42.88302799977828, "q1_ms": 41.81563800375443, "q3_ms": 43.307420004566666, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 46.09417724609375, "xla_arg_mb": 8.5, "xla_out_mb": 8.0, "rel_err": 6.482335379587312e-07, "rel_err_parts": [ 6.482335379587312e-07 ], "compile_s": 0.19405508499767166, "inner": 1, "median_ms": 56.959622001159005, "q1_ms": 56.7031360042165, "q3_ms": 57.85070199635811, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 384.0, "xla_arg_mb": 8.5, "xla_out_mb": 8.0, "rel_err": 1.4823216076820245e-07, "rel_err_parts": [ 1.4823216076820245e-07 ], "rel_vs_flex_repeat": 6.483219294750597e-07, "compile_s": 0.1813007739983732, "inner": 1, "median_ms": 228.8201529881917, "q1_ms": 164.3341215021792, "q3_ms": 269.0073210033006, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 6.325635695247911e-07, "rel_err_parts": [ 6.325635695247911e-07 ], "rel_vs_flex_repeat": 3.371636410065548e-07, "compile_s": 0.01536162500269711, "inner": 3, "median_ms": 14.682677331923818, "q1_ms": 14.528824501515675, "q3_ms": 15.017228664267652, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 60.625694274902344, "xla_arg_mb": 16.5, "xla_out_mb": 8.500022888183594, "residual_mb": 16.625003814697266, "residual_shapes": [ "float32[1, 1, 32, 1024, 64]", "float32[1, 1, 1024, 64]", "float32[1, 1, 1024, 64]", "uint32[]", "float32[1, 1, 32, 1024, 64]", "float32[1, 32, 1024]" ], "rel_err": 8.212336979340762e-07, "rel_err_parts": [ 6.714254823236843e-07, 8.212336979340762e-07, 8.074522384049487e-07 ], "rel_vs_flex_repeat": 5.982662969472585e-07, "compile_s": 0.4176329000038095, "inner": 1, "median_ms": 191.28611599444412, "q1_ms": 188.33416650886647, "q3_ms": 197.16976899508154, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 100.00069427490234, "xla_arg_mb": 16.5, "xla_out_mb": 8.500022888183594, "residual_mb": 32.125003814697266, "residual_shapes": [ "float32[1, 32, 1, 1024, 64]", "float32[1, 32, 1024, 64]", "float32[1, 32, 1024, 64]", "uint32[]", "float32[1, 32, 1, 1024, 64]", "float32[1, 32, 1024]" ], "rel_err": 6.714254823236843e-07, "rel_err_parts": [ 6.714254823236843e-07, 6.6607213966563e-07, 6.595964805455878e-07 ], "compile_s": 0.3972462379897479, "inner": 1, "median_ms": 180.34770998929162, "q1_ms": 179.8218754993286, "q3_ms": 182.8414534975309, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 392.25, "xla_arg_mb": 16.5, "xla_out_mb": 8.500022888183594, "rel_err": 7.899494676166796e-07, "rel_err_parts": [ 1.4470307974079333e-07, 7.899494676166796e-07, 1.9945356655171054e-07 ], "rel_vs_flex_repeat": 6.971709467507026e-07, "compile_s": 0.31785253901034594, "inner": 1, "median_ms": 252.44761900103185, "q1_ms": 251.15607500629267, "q3_ms": 256.5795709961094, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 7.159800929912308e-07, "rel_err_parts": [ 6.692216061310319e-07, 7.159800929912308e-07, 7.092019131960114e-07 ], "rel_vs_flex_repeat": 4.438918779214873e-07, "compile_s": 0.12239335699996445, "inner": 1, "median_ms": 129.40391099255066, "q1_ms": 128.4948510074173, "q3_ms": 134.57941399974516, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 60.625694274902344, "xla_arg_mb": 16.5, "xla_out_mb": 8.500022888183594, "residual_mb": 16.625003814697266, "residual_shapes": [ "float32[1, 1, 32, 1024, 64]", "float32[1, 1, 1024, 64]", "float32[1, 1, 1024, 64]", "uint32[]", "float32[1, 1, 32, 1024, 64]", "float32[1, 32, 1024]" ], "rel_err": 8.212336979340762e-07, "rel_err_parts": [ 6.714254823236843e-07, 8.212336979340762e-07, 8.074522384049487e-07 ], "rel_vs_flex_repeat": 5.982662969472585e-07, "compile_s": 0.39342178200604394, "inner": 1, "median_ms": 194.44847998965997, "q1_ms": 192.2392404958373, "q3_ms": 195.1244485026109, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 100.00069427490234, "xla_arg_mb": 16.5, "xla_out_mb": 8.500022888183594, "residual_mb": 32.125003814697266, "residual_shapes": [ "float32[1, 32, 1, 1024, 64]", "float32[1, 32, 1024, 64]", "float32[1, 32, 1024, 64]", "uint32[]", "float32[1, 32, 1, 1024, 64]", "float32[1, 32, 1024]" ], "rel_err": 6.714254823236843e-07, "rel_err_parts": [ 6.714254823236843e-07, 6.6607213966563e-07, 6.595964805455878e-07 ], "compile_s": 0.39361857299809344, "inner": 1, "median_ms": 184.80228900443763, "q1_ms": 182.377405501029, "q3_ms": 189.0978509982233, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 392.25, "xla_arg_mb": 16.5, "xla_out_mb": 8.500022888183594, "rel_err": 7.899494676166796e-07, "rel_err_parts": [ 1.4470307974079333e-07, 7.899494676166796e-07, 1.9945356655171054e-07 ], "rel_vs_flex_repeat": 6.971709467507026e-07, "compile_s": 0.3074600459949579, "inner": 1, "median_ms": 251.67155000963248, "q1_ms": 248.22414200025378, "q3_ms": 252.3273854967556, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 7.159800929912308e-07, "rel_err_parts": [ 6.692216061310319e-07, 7.159800929912308e-07, 7.092019131960114e-07 ], "rel_vs_flex_repeat": 4.438918779214873e-07, "compile_s": 0.12516868799866643, "inner": 1, "median_ms": 129.997142998036, "q1_ms": 120.36952749622287, "q3_ms": 156.65582849032944, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 60.625694274902344, "xla_arg_mb": 16.5, "xla_out_mb": 8.500022888183594, "residual_mb": 16.625003814697266, "residual_shapes": [ "float32[1, 1, 32, 1024, 64]", "float32[1, 1, 1024, 64]", "float32[1, 1, 1024, 64]", "uint32[]", "float32[1, 1, 32, 1024, 64]", "float32[1, 32, 1024]" ], "rel_err": 8.212336979340762e-07, "rel_err_parts": [ 6.714254823236843e-07, 8.212336979340762e-07, 8.074522384049487e-07 ], "rel_vs_flex_repeat": 5.982662969472585e-07, "compile_s": 0.5940356659993995, "inner": 1, "median_ms": 200.38029400166124, "q1_ms": 195.9047205018578, "q3_ms": 229.78742350096582, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 100.00069427490234, "xla_arg_mb": 16.5, "xla_out_mb": 8.500022888183594, "residual_mb": 32.125003814697266, "residual_shapes": [ "float32[1, 32, 1, 1024, 64]", "float32[1, 32, 1024, 64]", "float32[1, 32, 1024, 64]", "uint32[]", "float32[1, 32, 1, 1024, 64]", "float32[1, 32, 1024]" ], "rel_err": 6.714254823236843e-07, "rel_err_parts": [ 6.714254823236843e-07, 6.6607213966563e-07, 6.595964805455878e-07 ], "compile_s": 0.3996945699909702, "inner": 1, "median_ms": 181.87033099820837, "q1_ms": 181.18320100620622, "q3_ms": 185.99546149926027, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 392.25, "xla_arg_mb": 16.5, "xla_out_mb": 8.500022888183594, "rel_err": 7.899494676166796e-07, "rel_err_parts": [ 1.4470307974079333e-07, 7.899494676166796e-07, 1.9945356655171054e-07 ], "rel_vs_flex_repeat": 6.971709467507026e-07, "compile_s": 0.3219547399930889, "inner": 1, "median_ms": 246.53953999222722, "q1_ms": 245.59573400620138, "q3_ms": 247.78843249805504, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 7.159800929912308e-07, "rel_err_parts": [ 6.692216061310319e-07, 7.159800929912308e-07, 7.092019131960114e-07 ], "rel_vs_flex_repeat": 4.438918779214873e-07, "compile_s": 0.12942335799743887, "inner": 1, "median_ms": 124.54441700538155, "q1_ms": 124.09338850557106, "q3_ms": 126.45769699156517, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 49.09417724609375, "xla_arg_mb": 24.0, "xla_out_mb": 16.0, "rel_err": 6.522740250147763e-07, "rel_err_parts": [ 6.522740250147763e-07 ], "rel_vs_flex_repeat": 0.0, "compile_s": 0.2201788230013335, "inner": 1, "median_ms": 71.50585399358533, "q1_ms": 70.8132785075577, "q3_ms": 75.44719649740728, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 76.09417724609375, "xla_arg_mb": 24.0, "xla_out_mb": 16.0, "rel_err": 6.522740250147763e-07, "rel_err_parts": [ 6.522740250147763e-07 ], "compile_s": 0.2396207090059761, "inner": 1, "median_ms": 90.48796999559272, "q1_ms": 88.92633499635849, "q3_ms": 90.80711349815829, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 388.0, "xla_arg_mb": 24.0, "xla_out_mb": 16.0, "rel_err": 1.495774739623812e-07, "rel_err_parts": [ 1.495774739623812e-07 ], "rel_vs_flex_repeat": 6.522727744595613e-07, "compile_s": 0.1940616880019661, "inner": 1, "median_ms": 155.22787200461607, "q1_ms": 153.0371214976185, "q3_ms": 157.30917500331998, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 6.322396188807033e-07, "rel_err_parts": [ 6.322396188807033e-07 ], "rel_vs_flex_repeat": 3.4751246857922524e-07, "compile_s": 0.026460754990694113, "inner": 2, "median_ms": 21.740724994742777, "q1_ms": 18.945996998809278, "q3_ms": 24.512141000741394, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 49.09417724609375, "xla_arg_mb": 24.0, "xla_out_mb": 16.0, "rel_err": 6.522740250147763e-07, "rel_err_parts": [ 6.522740250147763e-07 ], "rel_vs_flex_repeat": 0.0, "compile_s": 0.21745674399426207, "inner": 1, "median_ms": 73.00135299738031, "q1_ms": 72.71798799047247, "q3_ms": 73.77566350623965, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 76.09417724609375, "xla_arg_mb": 24.0, "xla_out_mb": 16.0, "rel_err": 6.522740250147763e-07, "rel_err_parts": [ 6.522740250147763e-07 ], "compile_s": 0.2216784310003277, "inner": 1, "median_ms": 87.85795299627353, "q1_ms": 87.03723599319346, "q3_ms": 88.15602699905867, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 388.0, "xla_arg_mb": 24.0, "xla_out_mb": 16.0, "rel_err": 1.495774739623812e-07, "rel_err_parts": [ 1.495774739623812e-07 ], "rel_vs_flex_repeat": 6.522727744595613e-07, "compile_s": 0.1889380040083779, "inner": 1, "median_ms": 154.3165369948838, "q1_ms": 150.6533364954521, "q3_ms": 155.90339500340633, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 6.322396188807033e-07, "rel_err_parts": [ 6.322396188807033e-07 ], "rel_vs_flex_repeat": 3.4751246857922524e-07, "compile_s": 0.02525260999391321, "inner": 1, "median_ms": 27.942870001425035, "q1_ms": 23.41282849374693, "q3_ms": 37.9009504977148, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 49.09417724609375, "xla_arg_mb": 24.0, "xla_out_mb": 16.0, "rel_err": 6.522740250147763e-07, "rel_err_parts": [ 6.522740250147763e-07 ], "rel_vs_flex_repeat": 0.0, "compile_s": 0.5599192659865366, "inner": 1, "median_ms": 119.37883299833629, "q1_ms": 89.85974600363988, "q3_ms": 213.2403110008454, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 76.09417724609375, "xla_arg_mb": 24.0, "xla_out_mb": 16.0, "rel_err": 6.522740250147763e-07, "rel_err_parts": [ 6.522740250147763e-07 ], "compile_s": 0.23522157499974128, "inner": 1, "median_ms": 95.4114140040474, "q1_ms": 92.82617900316836, "q3_ms": 111.22851200343575, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 388.0, "xla_arg_mb": 24.0, "xla_out_mb": 16.0, "rel_err": 1.495774739623812e-07, "rel_err_parts": [ 1.495774739623812e-07 ], "rel_vs_flex_repeat": 6.522727744595613e-07, "compile_s": 0.18643052999686915, "inner": 1, "median_ms": 225.47712799860165, "q1_ms": 194.86523349769413, "q3_ms": 243.08429649681784, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 6.322396188807033e-07, "rel_err_parts": [ 6.322396188807033e-07 ], "rel_vs_flex_repeat": 3.4751246857922524e-07, "compile_s": 0.02525110699934885, "inner": 2, "median_ms": 25.011213001562282, "q1_ms": 24.54691450111568, "q3_ms": 25.092681997193722, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 92.25069427490234, "xla_arg_mb": 40.0, "xla_out_mb": 24.000022888183594, "residual_mb": 40.125003814697266, "residual_shapes": [ "float32[1, 8, 4, 1024, 128]", "float32[1, 8, 1024, 128]", "float32[1, 8, 1024, 128]", "uint32[]", "float32[1, 8, 4, 1024, 128]", "float32[1, 32, 1024]" ], "rel_err": 8.302061473841604e-07, "rel_err_parts": [ 6.701181973767234e-07, 8.256178034571349e-07, 8.302061473841604e-07 ], "rel_vs_flex_repeat": 6.013449933561787e-07, "compile_s": 0.4947536580002634, "inner": 1, "median_ms": 266.3821329915663, "q1_ms": 264.6641679966706, "q3_ms": 274.06265049648937, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 152.25069427490234, "xla_arg_mb": 40.0, "xla_out_mb": 24.000022888183594, "residual_mb": 64.12500381469727, "residual_shapes": [ "float32[1, 32, 1, 1024, 128]", "float32[1, 32, 1024, 128]", "float32[1, 32, 1024, 128]", "uint32[]", "float32[1, 32, 1, 1024, 128]", "float32[1, 32, 1024]" ], "rel_err": 6.727988761667802e-07, "rel_err_parts": [ 6.701181973767234e-07, 6.679061357317551e-07, 6.727988761667802e-07 ], "compile_s": 0.49497123999753967, "inner": 1, "median_ms": 293.96345900022425, "q1_ms": 288.7210605040309, "q3_ms": 320.4143414986902, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 532.125, "xla_arg_mb": 40.0, "xla_out_mb": 24.000022888183594, "rel_err": 7.818254061930929e-07, "rel_err_parts": [ 2.0110512366500188e-07, 7.818254061930929e-07, 2.1924437021425547e-07 ], "rel_vs_flex_repeat": 6.864132728878758e-07, "compile_s": 0.3753763269924093, "inner": 1, "median_ms": 297.5173930026358, "q1_ms": 294.9058829981368, "q3_ms": 307.727220999368, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 8.813879048830131e-07, "rel_err_parts": [ 8.813879048830131e-07, 8.777155358075106e-07, 8.759137131164607e-07 ], "rel_vs_flex_repeat": 7.083086757120327e-07, "compile_s": 0.0877764190081507, "inner": 1, "median_ms": 88.33807200426236, "q1_ms": 86.84878400526941, "q3_ms": 96.5751525000087, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 92.25069427490234, "xla_arg_mb": 40.0, "xla_out_mb": 24.000022888183594, "residual_mb": 40.125003814697266, "residual_shapes": [ "float32[1, 8, 4, 1024, 128]", "float32[1, 8, 1024, 128]", "float32[1, 8, 1024, 128]", "uint32[]", "float32[1, 8, 4, 1024, 128]", "float32[1, 32, 1024]" ], "rel_err": 8.302061473841604e-07, "rel_err_parts": [ 6.701181973767234e-07, 8.256178034571349e-07, 8.302061473841604e-07 ], "rel_vs_flex_repeat": 6.013449933561787e-07, "compile_s": 0.45286036300240085, "inner": 1, "median_ms": 260.787437990075, "q1_ms": 259.9067815026501, "q3_ms": 262.4831009961781, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 152.25069427490234, "xla_arg_mb": 40.0, "xla_out_mb": 24.000022888183594, "residual_mb": 64.12500381469727, "residual_shapes": [ "float32[1, 32, 1, 1024, 128]", "float32[1, 32, 1024, 128]", "float32[1, 32, 1024, 128]", "uint32[]", "float32[1, 32, 1, 1024, 128]", "float32[1, 32, 1024]" ], "rel_err": 6.727988761667802e-07, "rel_err_parts": [ 6.701181973767234e-07, 6.679061357317551e-07, 6.727988761667802e-07 ], "compile_s": 0.47611759800929576, "inner": 1, "median_ms": 284.98988000501413, "q1_ms": 283.57625599892344, "q3_ms": 289.6539399953326, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 532.125, "xla_arg_mb": 40.0, "xla_out_mb": 24.000022888183594, "rel_err": 7.818254061930929e-07, "rel_err_parts": [ 2.0110512366500188e-07, 7.818254061930929e-07, 2.1924437021425547e-07 ], "rel_vs_flex_repeat": 6.864132728878758e-07, "compile_s": 0.43407154400483705, "inner": 1, "median_ms": 303.9220420032507, "q1_ms": 298.4627519908827, "q3_ms": 311.11532149952836, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 8.813879048830131e-07, "rel_err_parts": [ 8.813879048830131e-07, 8.777155358075106e-07, 8.759137131164607e-07 ], "rel_vs_flex_repeat": 7.083086757120327e-07, "compile_s": 0.09252919199934695, "inner": 1, "median_ms": 86.17069100728258, "q1_ms": 85.86567349993857, "q3_ms": 98.76038650691044, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 92.25069427490234, "xla_arg_mb": 40.0, "xla_out_mb": 24.000022888183594, "residual_mb": 40.125003814697266, "residual_shapes": [ "float32[1, 8, 4, 1024, 128]", "float32[1, 8, 1024, 128]", "float32[1, 8, 1024, 128]", "uint32[]", "float32[1, 8, 4, 1024, 128]", "float32[1, 32, 1024]" ], "rel_err": 8.302061473841604e-07, "rel_err_parts": [ 6.701181973767234e-07, 8.256178034571349e-07, 8.302061473841604e-07 ], "rel_vs_flex_repeat": 6.013449933561787e-07, "compile_s": 0.5134005590080051, "inner": 1, "median_ms": 272.02982900780626, "q1_ms": 266.619787493255, "q3_ms": 275.04277499974705, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 152.25069427490234, "xla_arg_mb": 40.0, "xla_out_mb": 24.000022888183594, "residual_mb": 64.12500381469727, "residual_shapes": [ "float32[1, 32, 1, 1024, 128]", "float32[1, 32, 1024, 128]", "float32[1, 32, 1024, 128]", "uint32[]", "float32[1, 32, 1, 1024, 128]", "float32[1, 32, 1024]" ], "rel_err": 6.727988761667802e-07, "rel_err_parts": [ 6.701181973767234e-07, 6.679061357317551e-07, 6.727988761667802e-07 ], "compile_s": 0.49199918600788806, "inner": 1, "median_ms": 283.15926699724514, "q1_ms": 280.5357430042932, "q3_ms": 285.1346495008329, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 532.125, "xla_arg_mb": 40.0, "xla_out_mb": 24.000022888183594, "rel_err": 7.818254061930929e-07, "rel_err_parts": [ 2.0110512366500188e-07, 7.818254061930929e-07, 2.1924437021425547e-07 ], "rel_vs_flex_repeat": 6.864132728878758e-07, "compile_s": 0.35610073700081557, "inner": 1, "median_ms": 292.78207400057, "q1_ms": 291.11094199470244, "q3_ms": 294.0642344983644, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 8.813879048830131e-07, "rel_err_parts": [ 8.813879048830131e-07, 8.777155358075106e-07, 8.759137131164607e-07 ], "rel_vs_flex_repeat": 7.083086757120327e-07, "compile_s": 0.08530022398917936, "inner": 1, "median_ms": 91.73125200322829, "q1_ms": 89.53366400237428, "q3_ms": 204.90619650081499, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 40.21917724609375, "xla_arg_mb": 17.0, "xla_out_mb": 16.0, "rel_err": 6.574946382897906e-07, "rel_err_parts": [ 6.574946382897906e-07 ], "rel_vs_flex_repeat": 0.0, "compile_s": 0.20783917600056157, "inner": 1, "median_ms": 66.8391839863034, "q1_ms": 65.27367249509552, "q3_ms": 68.40301299234852, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 84.09417724609375, "xla_arg_mb": 17.0, "xla_out_mb": 16.0, "rel_err": 6.574946382897906e-07, "rel_err_parts": [ 6.574946382897906e-07 ], "compile_s": 0.2286366760090459, "inner": 1, "median_ms": 87.9906540067168, "q1_ms": 87.179214504431, "q3_ms": 90.4829345017788, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 384.0, "xla_arg_mb": 17.0, "xla_out_mb": 16.0, "rel_err": 1.5086362736838055e-07, "rel_err_parts": [ 1.5086362736838055e-07 ], "rel_vs_flex_repeat": 6.574374538104166e-07, "compile_s": 0.1838927880016854, "inner": 1, "median_ms": 159.15134499664418, "q1_ms": 154.4912779936567, "q3_ms": 175.2735445043072, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 6.380882950907107e-07, "rel_err_parts": [ 6.380882950907107e-07 ], "rel_vs_flex_repeat": 3.4660251913010143e-07, "compile_s": 0.03069843399862293, "inner": 1, "median_ms": 28.814217992476188, "q1_ms": 28.532675991300493, "q3_ms": 29.38575000007404, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 40.21917724609375, "xla_arg_mb": 17.0, "xla_out_mb": 16.0, "rel_err": 6.574946382897906e-07, "rel_err_parts": [ 6.574946382897906e-07 ], "rel_vs_flex_repeat": 0.0, "compile_s": 0.21168116600892972, "inner": 1, "median_ms": 67.92779400711879, "q1_ms": 67.4860030048876, "q3_ms": 69.72070100164274, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 84.09417724609375, "xla_arg_mb": 17.0, "xla_out_mb": 16.0, "rel_err": 6.574946382897906e-07, "rel_err_parts": [ 6.574946382897906e-07 ], "compile_s": 0.2328384959982941, "inner": 1, "median_ms": 91.96036600042135, "q1_ms": 91.41637600259855, "q3_ms": 92.20096100034425, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 384.0, "xla_arg_mb": 17.0, "xla_out_mb": 16.0, "rel_err": 1.5086362736838055e-07, "rel_err_parts": [ 1.5086362736838055e-07 ], "rel_vs_flex_repeat": 6.574374538104166e-07, "compile_s": 0.1892853909957921, "inner": 1, "median_ms": 153.57745900109876, "q1_ms": 151.31025499431416, "q3_ms": 160.12028800469125, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 6.380882950907107e-07, "rel_err_parts": [ 6.380882950907107e-07 ], "rel_vs_flex_repeat": 3.4660251913010143e-07, "compile_s": 0.029151706999982707, "inner": 2, "median_ms": 24.077546499029268, "q1_ms": 23.98603824985912, "q3_ms": 24.281855246954365, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 40.21917724609375, "xla_arg_mb": 17.0, "xla_out_mb": 16.0, "rel_err": 6.574946382897906e-07, "rel_err_parts": [ 6.574946382897906e-07 ], "rel_vs_flex_repeat": 0.0, "compile_s": 0.20661021099658683, "inner": 1, "median_ms": 65.99072899552993, "q1_ms": 65.05039499461418, "q3_ms": 66.58449250244303, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 84.09417724609375, "xla_arg_mb": 17.0, "xla_out_mb": 16.0, "rel_err": 6.574946382897906e-07, "rel_err_parts": [ 6.574946382897906e-07 ], "compile_s": 0.219399980007438, "inner": 1, "median_ms": 87.61945199512411, "q1_ms": 86.332973507524, "q3_ms": 88.02277549693827, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 384.0, "xla_arg_mb": 17.0, "xla_out_mb": 16.0, "rel_err": 1.5086362736838055e-07, "rel_err_parts": [ 1.5086362736838055e-07 ], "rel_vs_flex_repeat": 6.574374538104166e-07, "compile_s": 0.18633560999296606, "inner": 1, "median_ms": 156.48874700127635, "q1_ms": 151.75163450476248, "q3_ms": 161.27698249329114, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 6.380882950907107e-07, "rel_err_parts": [ 6.380882950907107e-07 ], "rel_vs_flex_repeat": 3.4660251913010143e-07, "compile_s": 0.025126560998614877, "inner": 2, "median_ms": 26.022626501799095, "q1_ms": 22.360537503118394, "q3_ms": 40.288723750563804, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 81.12569427490234, "xla_arg_mb": 33.0, "xla_out_mb": 17.000022888183594, "residual_mb": 33.125003814697266, "residual_shapes": [ "float32[1, 1, 32, 1024, 128]", "float32[1, 1, 1024, 128]", "float32[1, 1, 1024, 128]", "uint32[]", "float32[1, 1, 32, 1024, 128]", "float32[1, 32, 1024]" ], "rel_err": 8.299512614939886e-07, "rel_err_parts": [ 6.68250436319795e-07, 8.299512614939886e-07, 8.272837135336886e-07 ], "rel_vs_flex_repeat": 6.049148169040564e-07, "compile_s": 0.48272679299407173, "inner": 1, "median_ms": 259.7562870068941, "q1_ms": 256.9520564939012, "q3_ms": 262.25929700012784, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 160.00069427490234, "xla_arg_mb": 33.0, "xla_out_mb": 17.000022888183594, "residual_mb": 64.12500381469727, "residual_shapes": [ "float32[1, 32, 1, 1024, 128]", "float32[1, 32, 1024, 128]", "float32[1, 32, 1024, 128]", "uint32[]", "float32[1, 32, 1, 1024, 128]", "float32[1, 32, 1024]" ], "rel_err": 6.743950962118106e-07, "rel_err_parts": [ 6.68250436319795e-07, 6.743950962118106e-07, 6.7233384015708e-07 ], "compile_s": 0.4753612969943788, "inner": 1, "median_ms": 491.8366890051402, "q1_ms": 382.94984299864154, "q3_ms": 980.9455765062012, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 400.25, "xla_arg_mb": 33.0, "xla_out_mb": 17.000022888183594, "rel_err": 8.026575528674584e-07, "rel_err_parts": [ 1.9989511201856658e-07, 8.026575528674584e-07, 2.4365110107282817e-07 ], "rel_vs_flex_repeat": 7.041246590233641e-07, "compile_s": 0.4486116540065268, "inner": 1, "median_ms": 281.53183800168335, "q1_ms": 274.7468900051899, "q3_ms": 286.2992825030233, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 9.174459592031781e-07, "rel_err_parts": [ 8.756990723668423e-07, 9.174459592031781e-07, 9.142901262748637e-07 ], "rel_vs_flex_repeat": 7.501449204028177e-07, "compile_s": 0.1576680389989633, "inner": 1, "median_ms": 169.0500550030265, "q1_ms": 166.30845700274222, "q3_ms": 169.90870350855403, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 81.12569427490234, "xla_arg_mb": 33.0, "xla_out_mb": 17.000022888183594, "residual_mb": 33.125003814697266, "residual_shapes": [ "float32[1, 1, 32, 1024, 128]", "float32[1, 1, 1024, 128]", "float32[1, 1, 1024, 128]", "uint32[]", "float32[1, 1, 32, 1024, 128]", "float32[1, 32, 1024]" ], "rel_err": 8.299512614939886e-07, "rel_err_parts": [ 6.68250436319795e-07, 8.299512614939886e-07, 8.272837135336886e-07 ], "rel_vs_flex_repeat": 6.049148169040564e-07, "compile_s": 0.593535890991916, "inner": 1, "median_ms": 269.19997500954196, "q1_ms": 266.42363700375427, "q3_ms": 272.9762755043339, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 160.00069427490234, "xla_arg_mb": 33.0, "xla_out_mb": 17.000022888183594, "residual_mb": 64.12500381469727, "residual_shapes": [ "float32[1, 32, 1, 1024, 128]", "float32[1, 32, 1024, 128]", "float32[1, 32, 1024, 128]", "uint32[]", "float32[1, 32, 1, 1024, 128]", "float32[1, 32, 1024]" ], "rel_err": 6.743950962118106e-07, "rel_err_parts": [ 6.68250436319795e-07, 6.743950962118106e-07, 6.7233384015708e-07 ], "compile_s": 0.5051733730069827, "inner": 1, "median_ms": 278.8922640029341, "q1_ms": 275.8698745019501, "q3_ms": 281.1737605079543, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 400.25, "xla_arg_mb": 33.0, "xla_out_mb": 17.000022888183594, "rel_err": 8.026575528674584e-07, "rel_err_parts": [ 1.9989511201856658e-07, 8.026575528674584e-07, 2.4365110107282817e-07 ], "rel_vs_flex_repeat": 7.041246590233641e-07, "compile_s": 0.34142896700359415, "inner": 1, "median_ms": 281.133113996475, "q1_ms": 278.12816300138365, "q3_ms": 282.82808999938425, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 9.174459592031781e-07, "rel_err_parts": [ 8.756990723668423e-07, 9.174459592031781e-07, 9.142901262748637e-07 ], "rel_vs_flex_repeat": 7.501449204028177e-07, "compile_s": 0.15692824999860022, "inner": 1, "median_ms": 149.44768500572536, "q1_ms": 148.21140650019515, "q3_ms": 159.4057110050926, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 81.12569427490234, "xla_arg_mb": 33.0, "xla_out_mb": 17.000022888183594, "residual_mb": 33.125003814697266, "residual_shapes": [ "float32[1, 1, 32, 1024, 128]", "float32[1, 1, 1024, 128]", "float32[1, 1, 1024, 128]", "uint32[]", "float32[1, 1, 32, 1024, 128]", "float32[1, 32, 1024]" ], "rel_err": 8.299512614939886e-07, "rel_err_parts": [ 6.68250436319795e-07, 8.299512614939886e-07, 8.272837135336886e-07 ], "rel_vs_flex_repeat": 6.049148169040564e-07, "compile_s": 0.46727516700048, "inner": 1, "median_ms": 263.7341230001766, "q1_ms": 262.3377214986249, "q3_ms": 265.100603501196, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 160.00069427490234, "xla_arg_mb": 33.0, "xla_out_mb": 17.000022888183594, "residual_mb": 64.12500381469727, "residual_shapes": [ "float32[1, 32, 1, 1024, 128]", "float32[1, 32, 1024, 128]", "float32[1, 32, 1024, 128]", "uint32[]", "float32[1, 32, 1, 1024, 128]", "float32[1, 32, 1024]" ], "rel_err": 6.743950962118106e-07, "rel_err_parts": [ 6.68250436319795e-07, 6.743950962118106e-07, 6.7233384015708e-07 ], "compile_s": 0.4912084809911903, "inner": 1, "median_ms": 272.6912610087311, "q1_ms": 270.63117249053903, "q3_ms": 273.5675489966525, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 400.25, "xla_arg_mb": 33.0, "xla_out_mb": 17.000022888183594, "rel_err": 8.026575528674584e-07, "rel_err_parts": [ 1.9989511201856658e-07, 8.026575528674584e-07, 2.4365110107282817e-07 ], "rel_vs_flex_repeat": 7.041246590233641e-07, "compile_s": 0.3479857229976915, "inner": 1, "median_ms": 276.9224670046242, "q1_ms": 273.1843235087581, "q3_ms": 281.4238584978739, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "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": 9.174459592031781e-07, "rel_err_parts": [ 8.756990723668423e-07, 9.174459592031781e-07, 9.142901262748637e-07 ], "rel_vs_flex_repeat": 7.501449204028177e-07, "compile_s": 0.1805291339987889, "inner": 1, "median_ms": 166.57680500065908, "q1_ms": 162.33630199712934, "q3_ms": 173.98962499282788, "tag": null, "device_arg": "cpu", "run": "cpu_gqa" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "flex_native", "mode": "fwd", "round": 0, "xla_temp_mb": 40.25, "xla_arg_mb": 32.015625, "xla_out_mb": 0.015625, "rel_err": 7.022809427326138e-07, "rel_err_parts": [ 7.022809427326138e-07 ], "rel_vs_jax_xla": 8.209686939153471e-07, "compile_s": 0.11634879201301374, "inner": 3, "median_ms": 15.666078664556455, "q1_ms": 15.37810899996354, "q3_ms": 15.950418334493103, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "flex_repeat", "mode": "fwd", "round": 0, "xla_temp_mb": 136.25, "xla_arg_mb": 32.015625, "xla_out_mb": 0.015625, "rel_err": 7.140428124330356e-07, "rel_err_parts": [ 7.140428124330356e-07 ], "rel_vs_jax_xla": 5.656530106534774e-07, "compile_s": 0.14999222800543066, "inner": 1, "median_ms": 57.028629002161324, "q1_ms": 56.41991200536722, "q3_ms": 57.519257003150415, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "jax_xla", "mode": "fwd", "round": 0, "xla_temp_mb": 32.5, "xla_arg_mb": 32.015625, "xla_out_mb": 0.015625, "rel_err": 5.075322064840293e-07, "rel_err_parts": [ 5.075322064840293e-07 ], "compile_s": 0.06349994899937883, "inner": 4, "median_ms": 9.850096248555928, "q1_ms": 9.811878999244072, "q3_ms": 9.936489750543842, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "torch_sdpa", "mode": "fwd", "round": 0, "rel_err": 1.2908420785606722e-06, "rel_err_parts": [ 1.2908420785606722e-06 ], "rel_vs_jax_xla": 1.2430168681021314e-06, "compile_s": 0.25540087699482683, "inner": 15, "median_ms": 0.626806933238792, "q1_ms": 0.6251167326505916, "q3_ms": 0.6289559663855471, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "flex_native", "mode": "fwd", "round": 1, "xla_temp_mb": 40.25, "xla_arg_mb": 32.015625, "xla_out_mb": 0.015625, "rel_err": 7.022809427326138e-07, "rel_err_parts": [ 7.022809427326138e-07 ], "rel_vs_jax_xla": 8.209686939153471e-07, "compile_s": 0.11406678499770351, "inner": 3, "median_ms": 15.863645666589338, "q1_ms": 15.666375998989679, "q3_ms": 16.08291383308824, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "flex_repeat", "mode": "fwd", "round": 1, "xla_temp_mb": 136.25, "xla_arg_mb": 32.015625, "xla_out_mb": 0.015625, "rel_err": 7.140428124330356e-07, "rel_err_parts": [ 7.140428124330356e-07 ], "rel_vs_jax_xla": 5.656530106534774e-07, "compile_s": 0.14445350100868382, "inner": 1, "median_ms": 58.60169399238657, "q1_ms": 58.08279800112359, "q3_ms": 59.1242574955686, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "jax_xla", "mode": "fwd", "round": 1, "xla_temp_mb": 32.5, "xla_arg_mb": 32.015625, "xla_out_mb": 0.015625, "rel_err": 5.075322064840293e-07, "rel_err_parts": [ 5.075322064840293e-07 ], "compile_s": 0.0658561460004421, "inner": 4, "median_ms": 9.916847498971038, "q1_ms": 9.884977374895243, "q3_ms": 10.037604748504236, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "torch_sdpa", "mode": "fwd", "round": 1, "rel_err": 1.2908420785606722e-06, "rel_err_parts": [ 1.2908420785606722e-06 ], "rel_vs_jax_xla": 1.2430168681021314e-06, "compile_s": 0.0034407280036248267, "inner": 47, "median_ms": 0.6154505105595045, "q1_ms": 0.6134544786851002, "q3_ms": 0.6351342127210104, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "flex_native", "mode": "fwd", "round": 2, "xla_temp_mb": 40.25, "xla_arg_mb": 32.015625, "xla_out_mb": 0.015625, "rel_err": 7.022809427326138e-07, "rel_err_parts": [ 7.022809427326138e-07 ], "rel_vs_jax_xla": 8.209686939153471e-07, "compile_s": 0.1188869929901557, "inner": 3, "median_ms": 15.650900003189843, "q1_ms": 15.482032334451407, "q3_ms": 15.837713498816203, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "flex_repeat", "mode": "fwd", "round": 2, "xla_temp_mb": 136.25, "xla_arg_mb": 32.015625, "xla_out_mb": 0.015625, "rel_err": 7.140428124330356e-07, "rel_err_parts": [ 7.140428124330356e-07 ], "rel_vs_jax_xla": 5.656530106534774e-07, "compile_s": 0.15147599500778597, "inner": 1, "median_ms": 57.85837899020407, "q1_ms": 57.16694049624493, "q3_ms": 58.508427500783, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "jax_xla", "mode": "fwd", "round": 2, "xla_temp_mb": 32.5, "xla_arg_mb": 32.015625, "xla_out_mb": 0.015625, "rel_err": 5.075322064840293e-07, "rel_err_parts": [ 5.075322064840293e-07 ], "compile_s": 0.06637006800156087, "inner": 4, "median_ms": 9.946906749974005, "q1_ms": 9.732079375680769, "q3_ms": 10.076665001179208, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "torch_sdpa", "mode": "fwd", "round": 2, "rel_err": 1.2908420785606722e-06, "rel_err_parts": [ 1.2908420785606722e-06 ], "rel_vs_jax_xla": 1.2430168681021314e-06, "compile_s": 0.00495009298902005, "inner": 50, "median_ms": 0.564429999794811, "q1_ms": 0.5578231101389974, "q3_ms": 0.571644949959591, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "flex_native", "mode": "fwd", "round": 0, "xla_temp_mb": 40.25, "xla_arg_mb": 32.0625, "xla_out_mb": 0.0625, "rel_err": 7.969633770699147e-07, "rel_err_parts": [ 7.969633770699147e-07 ], "rel_vs_jax_xla": 9.091361334867543e-07, "compile_s": 0.1124783310078783, "inner": 3, "median_ms": 15.795052662724629, "q1_ms": 15.743493500243252, "q3_ms": 16.22589616696738, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "flex_repeat", "mode": "fwd", "round": 0, "xla_temp_mb": 136.25, "xla_arg_mb": 32.0625, "xla_out_mb": 0.0625, "rel_err": 5.635644129142747e-07, "rel_err_parts": [ 5.635644129142747e-07 ], "rel_vs_jax_xla": 8.475861363876902e-07, "compile_s": 0.15231890899303835, "inner": 1, "median_ms": 58.56954399496317, "q1_ms": 57.01278399646981, "q3_ms": 59.42576799134258, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "jax_xla", "mode": "fwd", "round": 0, "xla_temp_mb": 34.0, "xla_arg_mb": 32.0625, "xla_out_mb": 0.0625, "rel_err": 8.981768928606471e-07, "rel_err_parts": [ 8.981768928606471e-07 ], "compile_s": 0.07195632399816532, "inner": 4, "median_ms": 10.893582501012133, "q1_ms": 10.431911878185929, "q3_ms": 11.064161501053604, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "torch_sdpa", "mode": "fwd", "round": 0, "rel_err": 6.447160103562055e-07, "rel_err_parts": [ 6.447160103562055e-07 ], "rel_vs_jax_xla": 9.550253707857337e-07, "compile_s": 0.003335558998514898, "inner": 40, "median_ms": 0.7004969749687007, "q1_ms": 0.6972224126002402, "q3_ms": 0.7014286249614088, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "flex_native", "mode": "fwd", "round": 1, "xla_temp_mb": 40.25, "xla_arg_mb": 32.0625, "xla_out_mb": 0.0625, "rel_err": 7.969633770699147e-07, "rel_err_parts": [ 7.969633770699147e-07 ], "rel_vs_jax_xla": 9.091361334867543e-07, "compile_s": 0.1161446650075959, "inner": 3, "median_ms": 15.853730335948057, "q1_ms": 15.655336833636586, "q3_ms": 16.159665664114677, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "flex_repeat", "mode": "fwd", "round": 1, "xla_temp_mb": 136.25, "xla_arg_mb": 32.0625, "xla_out_mb": 0.0625, "rel_err": 5.635644129142747e-07, "rel_err_parts": [ 5.635644129142747e-07 ], "rel_vs_jax_xla": 8.475861363876902e-07, "compile_s": 0.1485176389978733, "inner": 1, "median_ms": 58.66151700320188, "q1_ms": 57.427748492045794, "q3_ms": 58.94006349990377, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "jax_xla", "mode": "fwd", "round": 1, "xla_temp_mb": 34.0, "xla_arg_mb": 32.0625, "xla_out_mb": 0.0625, "rel_err": 8.981768928606471e-07, "rel_err_parts": [ 8.981768928606471e-07 ], "compile_s": 0.06927654599712696, "inner": 4, "median_ms": 10.572157250862801, "q1_ms": 10.422690500490717, "q3_ms": 10.691743374991347, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "torch_sdpa", "mode": "fwd", "round": 1, "rel_err": 6.447160103562055e-07, "rel_err_parts": [ 6.447160103562055e-07 ], "rel_vs_jax_xla": 9.550253707857337e-07, "compile_s": 0.003346931000123732, "inner": 46, "median_ms": 0.7330728913075291, "q1_ms": 0.7203303151903168, "q3_ms": 0.8705391195215507, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "flex_native", "mode": "fwd", "round": 2, "xla_temp_mb": 40.25, "xla_arg_mb": 32.0625, "xla_out_mb": 0.0625, "rel_err": 7.969633770699147e-07, "rel_err_parts": [ 7.969633770699147e-07 ], "rel_vs_jax_xla": 9.091361334867543e-07, "compile_s": 0.11330924299545586, "inner": 2, "median_ms": 16.58337700064294, "q1_ms": 15.579881252051564, "q3_ms": 16.594821496255463, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "flex_repeat", "mode": "fwd", "round": 2, "xla_temp_mb": 136.25, "xla_arg_mb": 32.0625, "xla_out_mb": 0.0625, "rel_err": 5.635644129142747e-07, "rel_err_parts": [ 5.635644129142747e-07 ], "rel_vs_jax_xla": 8.475861363876902e-07, "compile_s": 0.14773671199509408, "inner": 1, "median_ms": 58.4554980014218, "q1_ms": 57.672306000313256, "q3_ms": 58.75362150254659, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "jax_xla", "mode": "fwd", "round": 2, "xla_temp_mb": 34.0, "xla_arg_mb": 32.0625, "xla_out_mb": 0.0625, "rel_err": 8.981768928606471e-07, "rel_err_parts": [ 8.981768928606471e-07 ], "compile_s": 0.0687990119913593, "inner": 4, "median_ms": 10.749332250270527, "q1_ms": 10.453130624227924, "q3_ms": 10.878371251237695, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "torch_sdpa", "mode": "fwd", "round": 2, "rel_err": 6.447160103562055e-07, "rel_err_parts": [ 6.447160103562055e-07 ], "rel_vs_jax_xla": 9.550253707857337e-07, "compile_s": 0.0033726890105754137, "inner": 41, "median_ms": 0.7631285612971136, "q1_ms": 0.7268514025604307, "q3_ms": 0.7882979879490803, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "flex_native", "mode": "fwd", "round": 0, "xla_temp_mb": 40.0, "xla_arg_mb": 32.25, "xla_out_mb": 0.25, "rel_err": 9.015959108182869e-07, "rel_err_parts": [ 9.015959108182869e-07 ], "rel_vs_jax_xla": 1.4201448550466012e-07, "compile_s": 0.10687658500683028, "inner": 3, "median_ms": 16.093911336308036, "q1_ms": 16.035040001346108, "q3_ms": 16.303228832839523, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "flex_repeat", "mode": "fwd", "round": 0, "xla_temp_mb": 136.0, "xla_arg_mb": 32.25, "xla_out_mb": 0.25, "rel_err": 8.49978107453353e-07, "rel_err_parts": [ 8.49978107453353e-07 ], "rel_vs_jax_xla": 7.010599460954836e-07, "compile_s": 0.14371641700563487, "inner": 1, "median_ms": 57.47414099460002, "q1_ms": 56.636382498254534, "q3_ms": 58.01792450074572, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "jax_xla", "mode": "fwd", "round": 0, "xla_temp_mb": 40.0, "xla_arg_mb": 32.25, "xla_out_mb": 0.25, "rel_err": 9.026235829878715e-07, "rel_err_parts": [ 9.026235829878715e-07 ], "compile_s": 0.0832201589946635, "inner": 3, "median_ms": 15.282924330676906, "q1_ms": 15.069854334190799, "q3_ms": 15.855388664931525, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "torch_sdpa", "mode": "fwd", "round": 0, "rel_err": 9.473249065194977e-07, "rel_err_parts": [ 9.473249065194977e-07 ], "rel_vs_jax_xla": 6.377138674906746e-07, "compile_s": 0.0040302319976035506, "inner": 22, "median_ms": 1.6600062272945333, "q1_ms": 1.6490896814502776, "q3_ms": 1.6789910001865462, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "flex_native", "mode": "fwd", "round": 1, "xla_temp_mb": 40.0, "xla_arg_mb": 32.25, "xla_out_mb": 0.25, "rel_err": 9.015959108182869e-07, "rel_err_parts": [ 9.015959108182869e-07 ], "rel_vs_jax_xla": 1.4201448550466012e-07, "compile_s": 0.10450895701069385, "inner": 3, "median_ms": 15.602161331723133, "q1_ms": 15.537206665612757, "q3_ms": 15.95036649814574, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "flex_repeat", "mode": "fwd", "round": 1, "xla_temp_mb": 136.0, "xla_arg_mb": 32.25, "xla_out_mb": 0.25, "rel_err": 8.49978107453353e-07, "rel_err_parts": [ 8.49978107453353e-07 ], "rel_vs_jax_xla": 7.010599460954836e-07, "compile_s": 0.14321832499990705, "inner": 1, "median_ms": 58.59291699016467, "q1_ms": 57.34163050510688, "q3_ms": 58.92338800185826, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "jax_xla", "mode": "fwd", "round": 1, "xla_temp_mb": 40.0, "xla_arg_mb": 32.25, "xla_out_mb": 0.25, "rel_err": 9.026235829878715e-07, "rel_err_parts": [ 9.026235829878715e-07 ], "compile_s": 0.07972650200827047, "inner": 3, "median_ms": 15.402339665646045, "q1_ms": 15.143740999822814, "q3_ms": 16.11863150174031, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "torch_sdpa", "mode": "fwd", "round": 1, "rel_err": 9.473249065194977e-07, "rel_err_parts": [ 9.473249065194977e-07 ], "rel_vs_jax_xla": 6.377138674906746e-07, "compile_s": 0.0038435710011981428, "inner": 22, "median_ms": 1.6554521816115912, "q1_ms": 1.6482738180705134, "q3_ms": 1.6615184090782846, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "flex_native", "mode": "fwd", "round": 2, "xla_temp_mb": 40.0, "xla_arg_mb": 32.25, "xla_out_mb": 0.25, "rel_err": 9.015959108182869e-07, "rel_err_parts": [ 9.015959108182869e-07 ], "rel_vs_jax_xla": 1.4201448550466012e-07, "compile_s": 0.10897296099574305, "inner": 2, "median_ms": 16.119380001327954, "q1_ms": 15.723114749562228, "q3_ms": 16.240265751548577, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "flex_repeat", "mode": "fwd", "round": 2, "xla_temp_mb": 136.0, "xla_arg_mb": 32.25, "xla_out_mb": 0.25, "rel_err": 8.49978107453353e-07, "rel_err_parts": [ 8.49978107453353e-07 ], "rel_vs_jax_xla": 7.010599460954836e-07, "compile_s": 0.14344089599035215, "inner": 1, "median_ms": 58.8080150046153, "q1_ms": 58.2393785007298, "q3_ms": 59.45659599819919, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "jax_xla", "mode": "fwd", "round": 2, "xla_temp_mb": 40.0, "xla_arg_mb": 32.25, "xla_out_mb": 0.25, "rel_err": 9.026235829878715e-07, "rel_err_parts": [ 9.026235829878715e-07 ], "compile_s": 0.08355704600398894, "inner": 3, "median_ms": 15.597041667206213, "q1_ms": 15.476391500366542, "q3_ms": 15.642674666499563, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "torch_sdpa", "mode": "fwd", "round": 2, "rel_err": 9.473249065194977e-07, "rel_err_parts": [ 9.473249065194977e-07 ], "rel_vs_jax_xla": 6.377138674906746e-07, "compile_s": 0.0038072720053605735, "inner": 21, "median_ms": 1.7025641905742563, "q1_ms": 1.655780238035645, "q3_ms": 1.790513166703195, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "flex_native", "mode": "fwd", "round": 0, "xla_temp_mb": 160.25, "xla_arg_mb": 128.015625, "xla_out_mb": 0.015625, "rel_err": 9.933024784913869e-07, "rel_err_parts": [ 9.933024784913869e-07 ], "rel_vs_jax_xla": 8.086619800451444e-07, "compile_s": 0.1673035410058219, "inner": 1, "median_ms": 70.36510399484541, "q1_ms": 67.63408350525424, "q3_ms": 74.04528450570069, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "flex_repeat", "mode": "fwd", "round": 0, "xla_temp_mb": 544.25, "xla_arg_mb": 128.015625, "xla_out_mb": 0.015625, "rel_err": 1.0012694247052423e-06, "rel_err_parts": [ 1.0012694247052423e-06 ], "rel_vs_jax_xla": 5.717565727536567e-07, "compile_s": 0.3059512920008274, "inner": 1, "median_ms": 239.07510600110982, "q1_ms": 234.62789399491157, "q3_ms": 241.9540569972014, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "jax_xla", "mode": "fwd", "round": 0, "xla_temp_mb": 130.0, "xla_arg_mb": 128.015625, "xla_out_mb": 0.015625, "rel_err": 9.497228461441409e-07, "rel_err_parts": [ 9.497228461441409e-07 ], "compile_s": 0.08801158200367354, "inner": 1, "median_ms": 40.49710898834746, "q1_ms": 39.658519002841786, "q3_ms": 41.596097995352466, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "torch_sdpa", "mode": "fwd", "round": 0, "rel_err": 2.4573278096795548e-06, "rel_err_parts": [ 2.4573278096795548e-06 ], "rel_vs_jax_xla": 2.321938609384233e-06, "compile_s": 0.01452435299870558, "inner": 3, "median_ms": 12.109498333302327, "q1_ms": 11.742649333124671, "q3_ms": 12.27342833347696, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "flex_native", "mode": "fwd", "round": 1, "xla_temp_mb": 160.25, "xla_arg_mb": 128.015625, "xla_out_mb": 0.015625, "rel_err": 9.933024784913869e-07, "rel_err_parts": [ 9.933024784913869e-07 ], "rel_vs_jax_xla": 8.086619800451444e-07, "compile_s": 0.16276304300117772, "inner": 1, "median_ms": 61.153731992817484, "q1_ms": 60.29824799770722, "q3_ms": 63.59791000431869, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "flex_repeat", "mode": "fwd", "round": 1, "xla_temp_mb": 544.25, "xla_arg_mb": 128.015625, "xla_out_mb": 0.015625, "rel_err": 1.0012694247052423e-06, "rel_err_parts": [ 1.0012694247052423e-06 ], "rel_vs_jax_xla": 5.717565727536567e-07, "compile_s": 0.3083021979982732, "inner": 1, "median_ms": 239.90454598970246, "q1_ms": 237.04576800082577, "q3_ms": 242.3771024987218, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "jax_xla", "mode": "fwd", "round": 1, "xla_temp_mb": 130.0, "xla_arg_mb": 128.015625, "xla_out_mb": 0.015625, "rel_err": 9.497228461441409e-07, "rel_err_parts": [ 9.497228461441409e-07 ], "compile_s": 0.09217956599604804, "inner": 1, "median_ms": 40.33981199609116, "q1_ms": 40.08563650131691, "q3_ms": 42.224251505103894, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "torch_sdpa", "mode": "fwd", "round": 1, "rel_err": 2.4573278096795548e-06, "rel_err_parts": [ 2.4573278096795548e-06 ], "rel_vs_jax_xla": 2.321938609384233e-06, "compile_s": 0.015131811989704147, "inner": 3, "median_ms": 12.681199664560458, "q1_ms": 12.563957832753658, "q3_ms": 12.69316233325905, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "flex_native", "mode": "fwd", "round": 2, "xla_temp_mb": 160.25, "xla_arg_mb": 128.015625, "xla_out_mb": 0.015625, "rel_err": 9.933024784913869e-07, "rel_err_parts": [ 9.933024784913869e-07 ], "rel_vs_jax_xla": 8.086619800451444e-07, "compile_s": 0.16430155398848, "inner": 1, "median_ms": 61.505806996137835, "q1_ms": 58.23072249768302, "q3_ms": 63.86915300390683, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "flex_repeat", "mode": "fwd", "round": 2, "xla_temp_mb": 544.25, "xla_arg_mb": 128.015625, "xla_out_mb": 0.015625, "rel_err": 1.0012694247052423e-06, "rel_err_parts": [ 1.0012694247052423e-06 ], "rel_vs_jax_xla": 5.717565727536567e-07, "compile_s": 0.3148990159970708, "inner": 1, "median_ms": 237.84407800121699, "q1_ms": 236.19321550359018, "q3_ms": 238.69009900226956, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "jax_xla", "mode": "fwd", "round": 2, "xla_temp_mb": 130.0, "xla_arg_mb": 128.015625, "xla_out_mb": 0.015625, "rel_err": 9.497228461441409e-07, "rel_err_parts": [ 9.497228461441409e-07 ], "compile_s": 0.09117830199829768, "inner": 1, "median_ms": 42.12993799592368, "q1_ms": 40.06719100289047, "q3_ms": 43.53065248869825, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "torch_sdpa", "mode": "fwd", "round": 2, "rel_err": 2.4573278096795548e-06, "rel_err_parts": [ 2.4573278096795548e-06 ], "rel_vs_jax_xla": 2.321938609384233e-06, "compile_s": 0.01581008500943426, "inner": 3, "median_ms": 11.62968399876263, "q1_ms": 11.548991499391075, "q3_ms": 12.068081666560223, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "flex_native", "mode": "fwd", "round": 0, "xla_temp_mb": 160.25, "xla_arg_mb": 128.0625, "xla_out_mb": 0.0625, "rel_err": 8.277044116766774e-07, "rel_err_parts": [ 8.277044116766774e-07 ], "rel_vs_jax_xla": 1.083634742826689e-06, "compile_s": 0.16828467800223734, "inner": 1, "median_ms": 65.39605600119103, "q1_ms": 64.37450849625748, "q3_ms": 74.10645649360958, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "flex_repeat", "mode": "fwd", "round": 0, "xla_temp_mb": 544.25, "xla_arg_mb": 128.0625, "xla_out_mb": 0.0625, "rel_err": 5.828962912346469e-07, "rel_err_parts": [ 5.828962912346469e-07 ], "rel_vs_jax_xla": 9.058765613190189e-07, "compile_s": 0.3019515960040735, "inner": 1, "median_ms": 235.7190889888443, "q1_ms": 233.34941700159106, "q3_ms": 239.9749839969445, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "jax_xla", "mode": "fwd", "round": 0, "xla_temp_mb": 136.0, "xla_arg_mb": 128.0625, "xla_out_mb": 0.0625, "rel_err": 8.982898975773423e-07, "rel_err_parts": [ 8.982898975773423e-07 ], "compile_s": 0.09874577599111944, "inner": 1, "median_ms": 42.63761801121291, "q1_ms": 42.5277299946174, "q3_ms": 43.04563049663557, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "torch_sdpa", "mode": "fwd", "round": 0, "rel_err": 6.749351086909883e-07, "rel_err_parts": [ 6.749351086909883e-07 ], "rel_vs_jax_xla": 9.691883633422549e-07, "compile_s": 0.016999333995045163, "inner": 3, "median_ms": 13.632699333053703, "q1_ms": 13.5957611661676, "q3_ms": 13.78052883471052, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "flex_native", "mode": "fwd", "round": 1, "xla_temp_mb": 160.25, "xla_arg_mb": 128.0625, "xla_out_mb": 0.0625, "rel_err": 8.277044116766774e-07, "rel_err_parts": [ 8.277044116766774e-07 ], "rel_vs_jax_xla": 1.083634742826689e-06, "compile_s": 0.1634872740105493, "inner": 1, "median_ms": 63.32906600437127, "q1_ms": 61.11619000148494, "q3_ms": 68.03532350022579, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "flex_repeat", "mode": "fwd", "round": 1, "xla_temp_mb": 544.25, "xla_arg_mb": 128.0625, "xla_out_mb": 0.0625, "rel_err": 5.828962912346469e-07, "rel_err_parts": [ 5.828962912346469e-07 ], "rel_vs_jax_xla": 9.058765613190189e-07, "compile_s": 0.3099737399898004, "inner": 1, "median_ms": 239.2573110118974, "q1_ms": 238.45122649800032, "q3_ms": 240.05418849992566, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "jax_xla", "mode": "fwd", "round": 1, "xla_temp_mb": 136.0, "xla_arg_mb": 128.0625, "xla_out_mb": 0.0625, "rel_err": 8.982898975773423e-07, "rel_err_parts": [ 8.982898975773423e-07 ], "compile_s": 0.0966246639873134, "inner": 1, "median_ms": 44.69286599487532, "q1_ms": 43.140012996445876, "q3_ms": 46.781195494986605, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "torch_sdpa", "mode": "fwd", "round": 1, "rel_err": 6.749351086909883e-07, "rel_err_parts": [ 6.749351086909883e-07 ], "rel_vs_jax_xla": 9.691883633422549e-07, "compile_s": 0.015549832009128295, "inner": 3, "median_ms": 14.991699999275928, "q1_ms": 14.926354500251666, "q3_ms": 15.632557166100014, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "flex_native", "mode": "fwd", "round": 2, "xla_temp_mb": 160.25, "xla_arg_mb": 128.0625, "xla_out_mb": 0.0625, "rel_err": 8.277044116766774e-07, "rel_err_parts": [ 8.277044116766774e-07 ], "rel_vs_jax_xla": 1.083634742826689e-06, "compile_s": 0.16650919998937752, "inner": 1, "median_ms": 61.54356899787672, "q1_ms": 59.3725170037942, "q3_ms": 63.0200115047046, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "flex_repeat", "mode": "fwd", "round": 2, "xla_temp_mb": 544.25, "xla_arg_mb": 128.0625, "xla_out_mb": 0.0625, "rel_err": 5.828962912346469e-07, "rel_err_parts": [ 5.828962912346469e-07 ], "rel_vs_jax_xla": 9.058765613190189e-07, "compile_s": 0.3116299829998752, "inner": 1, "median_ms": 236.86029799864627, "q1_ms": 232.01775750203524, "q3_ms": 241.51934999099467, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "jax_xla", "mode": "fwd", "round": 2, "xla_temp_mb": 136.0, "xla_arg_mb": 128.0625, "xla_out_mb": 0.0625, "rel_err": 8.982898975773423e-07, "rel_err_parts": [ 8.982898975773423e-07 ], "compile_s": 0.0936081879917765, "inner": 1, "median_ms": 45.24596199917141, "q1_ms": 44.75497349631041, "q3_ms": 45.71581600612262, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "torch_sdpa", "mode": "fwd", "round": 2, "rel_err": 6.749351086909883e-07, "rel_err_parts": [ 6.749351086909883e-07 ], "rel_vs_jax_xla": 9.691883633422549e-07, "compile_s": 0.017241743000340648, "inner": 3, "median_ms": 14.54355233484724, "q1_ms": 14.463010001539562, "q3_ms": 14.619303668344703, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "flex_native", "mode": "fwd", "round": 0, "xla_temp_mb": 160.0, "xla_arg_mb": 128.25, "xla_out_mb": 0.25, "rel_err": 1.0815224413818214e-06, "rel_err_parts": [ 1.0815224413818214e-06 ], "rel_vs_jax_xla": 8.167247642631992e-07, "compile_s": 0.15780609799548984, "inner": 1, "median_ms": 69.05963500321377, "q1_ms": 61.58849850180559, "q3_ms": 85.23594751022756, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "flex_repeat", "mode": "fwd", "round": 0, "xla_temp_mb": 544.0, "xla_arg_mb": 128.25, "xla_out_mb": 0.25, "rel_err": 9.17728186777822e-07, "rel_err_parts": [ 9.17728186777822e-07 ], "rel_vs_jax_xla": 8.236931421379268e-07, "compile_s": 0.3012962069915375, "inner": 1, "median_ms": 238.18471300182864, "q1_ms": 236.5872709924588, "q3_ms": 244.30061100429157, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "jax_xla", "mode": "fwd", "round": 0, "xla_temp_mb": 160.0, "xla_arg_mb": 128.25, "xla_out_mb": 0.25, "rel_err": 9.112583256865037e-07, "rel_err_parts": [ 9.112583256865037e-07 ], "compile_s": 0.1184284250048222, "inner": 1, "median_ms": 62.79788099345751, "q1_ms": 61.252553008671384, "q3_ms": 64.31906850775704, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "torch_sdpa", "mode": "fwd", "round": 0, "rel_err": 9.711479833640624e-07, "rel_err_parts": [ 9.711479833640624e-07 ], "rel_vs_jax_xla": 6.703700137222768e-07, "compile_s": 0.01684025299618952, "inner": 3, "median_ms": 14.890671998728067, "q1_ms": 14.61736183409812, "q3_ms": 16.2636104990573, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "flex_native", "mode": "fwd", "round": 1, "xla_temp_mb": 160.0, "xla_arg_mb": 128.25, "xla_out_mb": 0.25, "rel_err": 1.0815224413818214e-06, "rel_err_parts": [ 1.0815224413818214e-06 ], "rel_vs_jax_xla": 8.167247642631992e-07, "compile_s": 0.1618441949976841, "inner": 1, "median_ms": 64.12507199274842, "q1_ms": 63.055443002667744, "q3_ms": 64.54249099624576, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "flex_repeat", "mode": "fwd", "round": 1, "xla_temp_mb": 544.0, "xla_arg_mb": 128.25, "xla_out_mb": 0.25, "rel_err": 9.17728186777822e-07, "rel_err_parts": [ 9.17728186777822e-07 ], "rel_vs_jax_xla": 8.236931421379268e-07, "compile_s": 0.296773572001257, "inner": 1, "median_ms": 241.44501000409946, "q1_ms": 241.16263500764035, "q3_ms": 243.25728700205218, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "jax_xla", "mode": "fwd", "round": 1, "xla_temp_mb": 160.0, "xla_arg_mb": 128.25, "xla_out_mb": 0.25, "rel_err": 9.112583256865037e-07, "rel_err_parts": [ 9.112583256865037e-07 ], "compile_s": 0.11990570000489242, "inner": 1, "median_ms": 62.11051100399345, "q1_ms": 61.030708500766195, "q3_ms": 63.922538502083626, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "torch_sdpa", "mode": "fwd", "round": 1, "rel_err": 9.711479833640624e-07, "rel_err_parts": [ 9.711479833640624e-07 ], "rel_vs_jax_xla": 6.703700137222768e-07, "compile_s": 0.017704237005091272, "inner": 2, "median_ms": 16.071414000180084, "q1_ms": 15.909284502413357, "q3_ms": 16.30318450042978, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "flex_native", "mode": "fwd", "round": 2, "xla_temp_mb": 160.0, "xla_arg_mb": 128.25, "xla_out_mb": 0.25, "rel_err": 1.0815224413818214e-06, "rel_err_parts": [ 1.0815224413818214e-06 ], "rel_vs_jax_xla": 8.167247642631992e-07, "compile_s": 0.15607200699741952, "inner": 1, "median_ms": 73.95989399810787, "q1_ms": 68.84077149879886, "q3_ms": 86.34675400389824, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "flex_repeat", "mode": "fwd", "round": 2, "xla_temp_mb": 544.0, "xla_arg_mb": 128.25, "xla_out_mb": 0.25, "rel_err": 9.17728186777822e-07, "rel_err_parts": [ 9.17728186777822e-07 ], "rel_vs_jax_xla": 8.236931421379268e-07, "compile_s": 0.3028772180114174, "inner": 1, "median_ms": 243.41059700236656, "q1_ms": 237.42807150119916, "q3_ms": 248.50423649331788, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "jax_xla", "mode": "fwd", "round": 2, "xla_temp_mb": 160.0, "xla_arg_mb": 128.25, "xla_out_mb": 0.25, "rel_err": 9.112583256865037e-07, "rel_err_parts": [ 9.112583256865037e-07 ], "compile_s": 0.12190259600174613, "inner": 1, "median_ms": 62.28206600644626, "q1_ms": 61.16190698958235, "q3_ms": 64.2102174970205, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "torch_sdpa", "mode": "fwd", "round": 2, "rel_err": 9.711479833640624e-07, "rel_err_parts": [ 9.711479833640624e-07 ], "rel_vs_jax_xla": 6.703700137222768e-07, "compile_s": 0.0173950330063235, "inner": 3, "median_ms": 16.21365766429032, "q1_ms": 16.043374164534423, "q3_ms": 16.583076501168158, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "flex_native", "mode": "fwd", "round": 0, "xla_temp_mb": 768.5031204223633, "xla_arg_mb": 512.015625, "xla_out_mb": 0.015625, "rel_err": 1.7036660437952378e-06, "rel_err_parts": [ 1.7036660437952378e-06 ], "rel_vs_jax_xla": 7.564663633274904e-07, "compile_s": 0.49702596600400284, "inner": 1, "median_ms": 379.5294130104594, "q1_ms": 369.173941500776, "q3_ms": 383.1234175013378, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "flex_repeat", "mode": "fwd", "round": 0, "xla_temp_mb": 3072.5031204223633, "xla_arg_mb": 512.015625, "xla_out_mb": 0.015625, "rel_err": 1.7036660437952378e-06, "rel_err_parts": [ 1.7036660437952378e-06 ], "rel_vs_jax_xla": 7.564663633274904e-07, "compile_s": 1.0416322399978526, "inner": 1, "median_ms": 1125.9331439941889, "q1_ms": 1120.9458699959214, "q3_ms": 1127.5559784990037, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "jax_xla", "mode": "fwd", "round": 0, "xla_temp_mb": 520.0, "xla_arg_mb": 512.015625, "xla_out_mb": 0.015625, "rel_err": 1.7627528450248064e-06, "rel_err_parts": [ 1.7627528450248064e-06 ], "compile_s": 0.2822977420000825, "inner": 1, "median_ms": 171.22177200508304, "q1_ms": 169.64845699840225, "q3_ms": 172.1847144945059, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "torch_sdpa", "mode": "fwd", "round": 0, "rel_err": 4.8285537559422664e-06, "rel_err_parts": [ 4.8285537559422664e-06 ], "rel_vs_jax_xla": 4.543243449006695e-06, "compile_s": 0.06789087400829885, "inner": 1, "median_ms": 68.99374000204261, "q1_ms": 66.49536050099414, "q3_ms": 69.74887349497294, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "flex_native", "mode": "fwd", "round": 1, "xla_temp_mb": 768.5031204223633, "xla_arg_mb": 512.015625, "xla_out_mb": 0.015625, "rel_err": 1.7036660437952378e-06, "rel_err_parts": [ 1.7036660437952378e-06 ], "rel_vs_jax_xla": 7.564663633274904e-07, "compile_s": 0.4964666390005732, "inner": 1, "median_ms": 377.4747770075919, "q1_ms": 368.96936949779047, "q3_ms": 383.5960454962333, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "flex_repeat", "mode": "fwd", "round": 1, "xla_temp_mb": 3072.5031204223633, "xla_arg_mb": 512.015625, "xla_out_mb": 0.015625, "rel_err": 1.7036660437952378e-06, "rel_err_parts": [ 1.7036660437952378e-06 ], "rel_vs_jax_xla": 7.564663633274904e-07, "compile_s": 1.023276521009393, "inner": 1, "median_ms": 1119.6795349969761, "q1_ms": 1113.0009679909563, "q3_ms": 1128.347941499669, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "jax_xla", "mode": "fwd", "round": 1, "xla_temp_mb": 520.0, "xla_arg_mb": 512.015625, "xla_out_mb": 0.015625, "rel_err": 1.7627528450248064e-06, "rel_err_parts": [ 1.7627528450248064e-06 ], "compile_s": 0.29480008900281973, "inner": 1, "median_ms": 168.6413720017299, "q1_ms": 165.26249649905367, "q3_ms": 169.08551300002728, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "torch_sdpa", "mode": "fwd", "round": 1, "rel_err": 4.8285537559422664e-06, "rel_err_parts": [ 4.8285537559422664e-06 ], "rel_vs_jax_xla": 4.543243449006695e-06, "compile_s": 0.06872569400002249, "inner": 1, "median_ms": 66.22450299619231, "q1_ms": 65.53112050460186, "q3_ms": 68.69328200264135, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "flex_native", "mode": "fwd", "round": 2, "xla_temp_mb": 768.5031204223633, "xla_arg_mb": 512.015625, "xla_out_mb": 0.015625, "rel_err": 1.7036660437952378e-06, "rel_err_parts": [ 1.7036660437952378e-06 ], "rel_vs_jax_xla": 7.564663633274904e-07, "compile_s": 0.5068995730107417, "inner": 1, "median_ms": 369.7557369887363, "q1_ms": 367.1667794988025, "q3_ms": 372.7924909981084, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "flex_repeat", "mode": "fwd", "round": 2, "xla_temp_mb": 3072.5031204223633, "xla_arg_mb": 512.015625, "xla_out_mb": 0.015625, "rel_err": 1.7036660437952378e-06, "rel_err_parts": [ 1.7036660437952378e-06 ], "rel_vs_jax_xla": 7.564663633274904e-07, "compile_s": 1.0447687629930442, "inner": 1, "median_ms": 1126.4115669910097, "q1_ms": 1125.122535493574, "q3_ms": 1127.8336019968265, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "jax_xla", "mode": "fwd", "round": 2, "xla_temp_mb": 520.0, "xla_arg_mb": 512.015625, "xla_out_mb": 0.015625, "rel_err": 1.7627528450248064e-06, "rel_err_parts": [ 1.7627528450248064e-06 ], "compile_s": 0.27455685898894444, "inner": 1, "median_ms": 167.4963630066486, "q1_ms": 166.76388750056503, "q3_ms": 169.0197905045352, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "torch_sdpa", "mode": "fwd", "round": 2, "rel_err": 4.8285537559422664e-06, "rel_err_parts": [ 4.8285537559422664e-06 ], "rel_vs_jax_xla": 4.543243449006695e-06, "compile_s": 0.06993453700852115, "inner": 1, "median_ms": 62.86811899917666, "q1_ms": 62.36188599723391, "q3_ms": 63.13230849627871, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "flex_native", "mode": "fwd", "round": 0, "xla_temp_mb": 768.5031204223633, "xla_arg_mb": 512.0625, "xla_out_mb": 0.0625, "rel_err": 7.563425015177927e-07, "rel_err_parts": [ 7.563425015177927e-07 ], "rel_vs_jax_xla": 1.0326772326152422e-06, "compile_s": 0.5119682530057617, "inner": 1, "median_ms": 375.0324470020132, "q1_ms": 372.2075920086354, "q3_ms": 381.85588600026676, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "flex_repeat", "mode": "fwd", "round": 0, "xla_temp_mb": 3072.5031204223633, "xla_arg_mb": 512.0625, "xla_out_mb": 0.0625, "rel_err": 7.563425015177927e-07, "rel_err_parts": [ 7.563425015177927e-07 ], "rel_vs_jax_xla": 1.0326772326152422e-06, "compile_s": 1.0486346219986444, "inner": 1, "median_ms": 1122.4370669951895, "q1_ms": 1119.4563385070069, "q3_ms": 1130.3695044916822, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "jax_xla", "mode": "fwd", "round": 0, "xla_temp_mb": 544.0, "xla_arg_mb": 512.0625, "xla_out_mb": 0.0625, "rel_err": 9.016676472128893e-07, "rel_err_parts": [ 9.016676472128893e-07 ], "compile_s": 0.3085496889980277, "inner": 1, "median_ms": 183.65462200017646, "q1_ms": 180.38754149893066, "q3_ms": 185.75773199700052, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "torch_sdpa", "mode": "fwd", "round": 0, "rel_err": 7.516701998611097e-07, "rel_err_parts": [ 7.516701998611097e-07 ], "rel_vs_jax_xla": 1.0395217486802721e-06, "compile_s": 0.0714910829992732, "inner": 1, "median_ms": 67.94636699487455, "q1_ms": 67.02667149511399, "q3_ms": 68.29284699779237, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "flex_native", "mode": "fwd", "round": 1, "xla_temp_mb": 768.5031204223633, "xla_arg_mb": 512.0625, "xla_out_mb": 0.0625, "rel_err": 7.563425015177927e-07, "rel_err_parts": [ 7.563425015177927e-07 ], "rel_vs_jax_xla": 1.0326772326152422e-06, "compile_s": 0.5012487429921748, "inner": 1, "median_ms": 368.2383440027479, "q1_ms": 365.26220299856504, "q3_ms": 370.56388249766314, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "flex_repeat", "mode": "fwd", "round": 1, "xla_temp_mb": 3072.5031204223633, "xla_arg_mb": 512.0625, "xla_out_mb": 0.0625, "rel_err": 7.563425015177927e-07, "rel_err_parts": [ 7.563425015177927e-07 ], "rel_vs_jax_xla": 1.0326772326152422e-06, "compile_s": 1.0476559910021024, "inner": 1, "median_ms": 1128.3020330010913, "q1_ms": 1125.825913492008, "q3_ms": 1133.4746444990742, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "jax_xla", "mode": "fwd", "round": 1, "xla_temp_mb": 544.0, "xla_arg_mb": 512.0625, "xla_out_mb": 0.0625, "rel_err": 9.016676472128893e-07, "rel_err_parts": [ 9.016676472128893e-07 ], "compile_s": 0.31089757100562565, "inner": 1, "median_ms": 181.17774699931033, "q1_ms": 179.6276089953608, "q3_ms": 184.65332250343636, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "torch_sdpa", "mode": "fwd", "round": 1, "rel_err": 7.516701998611097e-07, "rel_err_parts": [ 7.516701998611097e-07 ], "rel_vs_jax_xla": 1.0395217486802721e-06, "compile_s": 0.07148434000555426, "inner": 1, "median_ms": 69.01253400428686, "q1_ms": 68.82458800100721, "q3_ms": 69.87167400075123, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "flex_native", "mode": "fwd", "round": 2, "xla_temp_mb": 768.5031204223633, "xla_arg_mb": 512.0625, "xla_out_mb": 0.0625, "rel_err": 7.563425015177927e-07, "rel_err_parts": [ 7.563425015177927e-07 ], "rel_vs_jax_xla": 1.0326772326152422e-06, "compile_s": 0.5006791859923396, "inner": 1, "median_ms": 382.5608490005834, "q1_ms": 377.0433910030988, "q3_ms": 383.3996664980077, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "flex_repeat", "mode": "fwd", "round": 2, "xla_temp_mb": 3072.5031204223633, "xla_arg_mb": 512.0625, "xla_out_mb": 0.0625, "rel_err": 7.563425015177927e-07, "rel_err_parts": [ 7.563425015177927e-07 ], "rel_vs_jax_xla": 1.0326772326152422e-06, "compile_s": 1.0539933510008268, "inner": 1, "median_ms": 1124.8639080004068, "q1_ms": 1122.7036765048979, "q3_ms": 1127.3863910028012, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "jax_xla", "mode": "fwd", "round": 2, "xla_temp_mb": 544.0, "xla_arg_mb": 512.0625, "xla_out_mb": 0.0625, "rel_err": 9.016676472128893e-07, "rel_err_parts": [ 9.016676472128893e-07 ], "compile_s": 0.29547211799945217, "inner": 1, "median_ms": 178.4511890000431, "q1_ms": 176.1526599948411, "q3_ms": 181.30890449538128, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "torch_sdpa", "mode": "fwd", "round": 2, "rel_err": 7.516701998611097e-07, "rel_err_parts": [ 7.516701998611097e-07 ], "rel_vs_jax_xla": 1.0395217486802721e-06, "compile_s": 0.07280593999894336, "inner": 1, "median_ms": 69.8102519963868, "q1_ms": 69.54998050059658, "q3_ms": 75.78557200031355, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "flex_native", "mode": "fwd", "round": 0, "xla_temp_mb": 768.2550048828125, "xla_arg_mb": 512.25, "xla_out_mb": 0.25, "rel_err": 1.0375609917900874e-06, "rel_err_parts": [ 1.0375609917900874e-06 ], "rel_vs_jax_xla": 7.614783612552856e-07, "compile_s": 0.49374731100397184, "inner": 1, "median_ms": 381.0336510068737, "q1_ms": 378.5119245003443, "q3_ms": 381.3116665041889, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "flex_repeat", "mode": "fwd", "round": 0, "xla_temp_mb": 3072.5050048828125, "xla_arg_mb": 512.25, "xla_out_mb": 0.25, "rel_err": 1.0375609917900874e-06, "rel_err_parts": [ 1.0375609917900874e-06 ], "rel_vs_jax_xla": 7.614783612552856e-07, "compile_s": 1.0064208690018859, "inner": 1, "median_ms": 1128.6939249985153, "q1_ms": 1123.3439479983645, "q3_ms": 1131.3346754977829, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "jax_xla", "mode": "fwd", "round": 0, "xla_temp_mb": 640.0, "xla_arg_mb": 512.25, "xla_out_mb": 0.25, "rel_err": 9.302755756834813e-07, "rel_err_parts": [ 9.302755756834813e-07 ], "compile_s": 0.35889272300119046, "inner": 1, "median_ms": 254.8527710023336, "q1_ms": 248.1254890008131, "q3_ms": 256.3966715024435, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "torch_sdpa", "mode": "fwd", "round": 0, "rel_err": 1.0325542234568275e-06, "rel_err_parts": [ 1.0325542234568275e-06 ], "rel_vs_jax_xla": 7.602030223097245e-07, "compile_s": 0.07439455599524081, "inner": 1, "median_ms": 76.52378399507143, "q1_ms": 74.805723503232, "q3_ms": 77.5579605033272, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "flex_native", "mode": "fwd", "round": 1, "xla_temp_mb": 768.2550048828125, "xla_arg_mb": 512.25, "xla_out_mb": 0.25, "rel_err": 1.0375609917900874e-06, "rel_err_parts": [ 1.0375609917900874e-06 ], "rel_vs_jax_xla": 7.614783612552856e-07, "compile_s": 0.4988867840002058, "inner": 1, "median_ms": 378.4538199979579, "q1_ms": 375.8994969975902, "q3_ms": 383.31404449854745, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "flex_repeat", "mode": "fwd", "round": 1, "xla_temp_mb": 3072.5050048828125, "xla_arg_mb": 512.25, "xla_out_mb": 0.25, "rel_err": 1.0375609917900874e-06, "rel_err_parts": [ 1.0375609917900874e-06 ], "rel_vs_jax_xla": 7.614783612552856e-07, "compile_s": 1.0025307089963462, "inner": 1, "median_ms": 1127.7054239908466, "q1_ms": 1123.3131300032255, "q3_ms": 1130.4898825037526, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "jax_xla", "mode": "fwd", "round": 1, "xla_temp_mb": 640.0, "xla_arg_mb": 512.25, "xla_out_mb": 0.25, "rel_err": 9.302755756834813e-07, "rel_err_parts": [ 9.302755756834813e-07 ], "compile_s": 0.34445870500348974, "inner": 1, "median_ms": 252.36599599884357, "q1_ms": 250.31766200845595, "q3_ms": 257.59787000424694, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "torch_sdpa", "mode": "fwd", "round": 1, "rel_err": 1.0325542234568275e-06, "rel_err_parts": [ 1.0325542234568275e-06 ], "rel_vs_jax_xla": 7.602030223097245e-07, "compile_s": 0.07496436299697962, "inner": 1, "median_ms": 76.92647598742042, "q1_ms": 76.52286200755043, "q3_ms": 77.69134300178848, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "flex_native", "mode": "fwd", "round": 2, "xla_temp_mb": 768.2550048828125, "xla_arg_mb": 512.25, "xla_out_mb": 0.25, "rel_err": 1.0375609917900874e-06, "rel_err_parts": [ 1.0375609917900874e-06 ], "rel_vs_jax_xla": 7.614783612552856e-07, "compile_s": 0.48401454799750354, "inner": 1, "median_ms": 384.015150993946, "q1_ms": 378.1035764986882, "q3_ms": 389.70703900122317, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "flex_repeat", "mode": "fwd", "round": 2, "xla_temp_mb": 3072.5050048828125, "xla_arg_mb": 512.25, "xla_out_mb": 0.25, "rel_err": 1.0375609917900874e-06, "rel_err_parts": [ 1.0375609917900874e-06 ], "rel_vs_jax_xla": 7.614783612552856e-07, "compile_s": 0.9913672489929013, "inner": 1, "median_ms": 1133.9798150002025, "q1_ms": 1124.2465100003756, "q3_ms": 1156.8393120032852, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "jax_xla", "mode": "fwd", "round": 2, "xla_temp_mb": 640.0, "xla_arg_mb": 512.25, "xla_out_mb": 0.25, "rel_err": 9.302755756834813e-07, "rel_err_parts": [ 9.302755756834813e-07 ], "compile_s": 0.34707664800225757, "inner": 1, "median_ms": 249.82008000370115, "q1_ms": 248.77245749667054, "q3_ms": 251.22256300528534, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "torch_sdpa", "mode": "fwd", "round": 2, "rel_err": 1.0325542234568275e-06, "rel_err_parts": [ 1.0325542234568275e-06 ], "rel_vs_jax_xla": 7.602030223097245e-07, "compile_s": 0.07807492799474858, "inner": 1, "median_ms": 76.18312000704464, "q1_ms": 73.25865200255066, "q3_ms": 77.16002750385087, "tag": "before", "device_arg": "cpu", "run": "cpu_decode_before", "note": "before: flex_attention padded queries to 16 rows on CPU (commit 788a603)" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "flex_native", "mode": "fwd", "round": 0, "xla_temp_mb": 32.5, "xla_arg_mb": 32.015625, "xla_out_mb": 0.015625, "rel_err": 7.140428124330356e-07, "rel_err_parts": [ 7.140428124330356e-07 ], "rel_vs_jax_xla": 5.656530106534774e-07, "compile_s": 0.06503144199086819, "inner": 5, "median_ms": 9.73496740043629, "q1_ms": 9.549060699646361, "q3_ms": 9.83841409906745, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "flex_repeat", "mode": "fwd", "round": 0, "xla_temp_mb": 128.5, "xla_arg_mb": 32.015625, "xla_out_mb": 0.015625, "rel_err": 1.1998938589385943e-06, "rel_err_parts": [ 1.1998938589385943e-06 ], "rel_vs_jax_xla": 1.26787563203834e-06, "compile_s": 0.10114790999796242, "inner": 1, "median_ms": 55.54588600352872, "q1_ms": 54.66133749723667, "q3_ms": 56.13782099680975, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "jax_xla", "mode": "fwd", "round": 0, "xla_temp_mb": 32.5, "xla_arg_mb": 32.015625, "xla_out_mb": 0.015625, "rel_err": 5.075322064840293e-07, "rel_err_parts": [ 5.075322064840293e-07 ], "compile_s": 0.06744262299616821, "inner": 4, "median_ms": 9.880620251351502, "q1_ms": 9.488562624028418, "q3_ms": 10.231320124148624, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "torch_sdpa", "mode": "fwd", "round": 0, "rel_err": 1.2908420785606722e-06, "rel_err_parts": [ 1.2908420785606722e-06 ], "rel_vs_jax_xla": 1.2430168681021314e-06, "compile_s": 0.25310317100957036, "inner": 10, "median_ms": 0.8893986989278346, "q1_ms": 0.8159677003277466, "q3_ms": 0.9876444994006306, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "flex_native", "mode": "fwd", "round": 1, "xla_temp_mb": 32.5, "xla_arg_mb": 32.015625, "xla_out_mb": 0.015625, "rel_err": 7.140428124330356e-07, "rel_err_parts": [ 7.140428124330356e-07 ], "rel_vs_jax_xla": 5.656530106534774e-07, "compile_s": 0.06453013299324084, "inner": 5, "median_ms": 9.66270859935321, "q1_ms": 9.354496499872766, "q3_ms": 9.795967899844982, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "flex_repeat", "mode": "fwd", "round": 1, "xla_temp_mb": 128.5, "xla_arg_mb": 32.015625, "xla_out_mb": 0.015625, "rel_err": 1.1998938589385943e-06, "rel_err_parts": [ 1.1998938589385943e-06 ], "rel_vs_jax_xla": 1.26787563203834e-06, "compile_s": 0.10065142999519594, "inner": 1, "median_ms": 56.418556996504776, "q1_ms": 55.651626498729456, "q3_ms": 56.48144600127125, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "jax_xla", "mode": "fwd", "round": 1, "xla_temp_mb": 32.5, "xla_arg_mb": 32.015625, "xla_out_mb": 0.015625, "rel_err": 5.075322064840293e-07, "rel_err_parts": [ 5.075322064840293e-07 ], "compile_s": 0.06682421400910243, "inner": 4, "median_ms": 9.230896252120147, "q1_ms": 9.131728626016411, "q3_ms": 10.272321624142933, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "torch_sdpa", "mode": "fwd", "round": 1, "rel_err": 1.2908420785606722e-06, "rel_err_parts": [ 1.2908420785606722e-06 ], "rel_vs_jax_xla": 1.2430168681021314e-06, "compile_s": 0.003907682999852113, "inner": 42, "median_ms": 0.7003864763343396, "q1_ms": 0.6913724168303557, "q3_ms": 0.7083928094722242, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "flex_native", "mode": "fwd", "round": 2, "xla_temp_mb": 32.5, "xla_arg_mb": 32.015625, "xla_out_mb": 0.015625, "rel_err": 7.140428124330356e-07, "rel_err_parts": [ 7.140428124330356e-07 ], "rel_vs_jax_xla": 5.656530106534774e-07, "compile_s": 0.06297673500375822, "inner": 5, "median_ms": 9.897821399499662, "q1_ms": 9.1489486003411, "q3_ms": 10.008077099337243, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "flex_repeat", "mode": "fwd", "round": 2, "xla_temp_mb": 128.5, "xla_arg_mb": 32.015625, "xla_out_mb": 0.015625, "rel_err": 1.1998938589385943e-06, "rel_err_parts": [ 1.1998938589385943e-06 ], "rel_vs_jax_xla": 1.26787563203834e-06, "compile_s": 0.09999289500410669, "inner": 1, "median_ms": 53.21812200418208, "q1_ms": 52.79530250118114, "q3_ms": 54.35884050530149, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "jax_xla", "mode": "fwd", "round": 2, "xla_temp_mb": 32.5, "xla_arg_mb": 32.015625, "xla_out_mb": 0.015625, "rel_err": 5.075322064840293e-07, "rel_err_parts": [ 5.075322064840293e-07 ], "compile_s": 0.06554156700440217, "inner": 4, "median_ms": 9.28370625115349, "q1_ms": 8.892236750398297, "q3_ms": 9.77359550051915, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "torch_sdpa", "mode": "fwd", "round": 2, "rel_err": 1.2908420785606722e-06, "rel_err_parts": [ 1.2908420785606722e-06 ], "rel_vs_jax_xla": 1.2430168681021314e-06, "compile_s": 0.0039263689977815375, "inner": 43, "median_ms": 0.6984499767747555, "q1_ms": 0.6926650813236551, "q3_ms": 0.7060282790240138, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "flex_native", "mode": "fwd", "round": 0, "xla_temp_mb": 34.0, "xla_arg_mb": 32.0625, "xla_out_mb": 0.0625, "rel_err": 5.635644129142747e-07, "rel_err_parts": [ 5.635644129142747e-07 ], "rel_vs_jax_xla": 8.475861363876902e-07, "compile_s": 0.07138857900281437, "inner": 4, "median_ms": 10.589646997686941, "q1_ms": 10.53042225066747, "q3_ms": 10.957083248285926, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "flex_repeat", "mode": "fwd", "round": 0, "xla_temp_mb": 130.0, "xla_arg_mb": 32.0625, "xla_out_mb": 0.0625, "rel_err": 5.635644129142747e-07, "rel_err_parts": [ 5.635644129142747e-07 ], "rel_vs_jax_xla": 8.475861363876902e-07, "compile_s": 0.1039525170053821, "inner": 1, "median_ms": 54.545473991311155, "q1_ms": 53.555485494143795, "q3_ms": 55.21621249499731, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "jax_xla", "mode": "fwd", "round": 0, "xla_temp_mb": 34.0, "xla_arg_mb": 32.0625, "xla_out_mb": 0.0625, "rel_err": 8.981768928606471e-07, "rel_err_parts": [ 8.981768928606471e-07 ], "compile_s": 0.07294020398694556, "inner": 4, "median_ms": 10.24303099984536, "q1_ms": 9.916037624861929, "q3_ms": 10.373882374551613, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "torch_sdpa", "mode": "fwd", "round": 0, "rel_err": 6.447160103562055e-07, "rel_err_parts": [ 6.447160103562055e-07 ], "rel_vs_jax_xla": 9.550253707857337e-07, "compile_s": 0.004236414999468252, "inner": 33, "median_ms": 1.1705004845597695, "q1_ms": 1.1535804695802545, "q3_ms": 1.1764789393937187, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "flex_native", "mode": "fwd", "round": 1, "xla_temp_mb": 34.0, "xla_arg_mb": 32.0625, "xla_out_mb": 0.0625, "rel_err": 5.635644129142747e-07, "rel_err_parts": [ 5.635644129142747e-07 ], "rel_vs_jax_xla": 8.475861363876902e-07, "compile_s": 0.07220437200157903, "inner": 4, "median_ms": 10.98834100048407, "q1_ms": 10.78495999900042, "q3_ms": 11.079687748861033, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "flex_repeat", "mode": "fwd", "round": 1, "xla_temp_mb": 130.0, "xla_arg_mb": 32.0625, "xla_out_mb": 0.0625, "rel_err": 5.635644129142747e-07, "rel_err_parts": [ 5.635644129142747e-07 ], "rel_vs_jax_xla": 8.475861363876902e-07, "compile_s": 0.10686310200253502, "inner": 1, "median_ms": 54.56438999681268, "q1_ms": 53.711590997409075, "q3_ms": 55.25732500245795, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "jax_xla", "mode": "fwd", "round": 1, "xla_temp_mb": 34.0, "xla_arg_mb": 32.0625, "xla_out_mb": 0.0625, "rel_err": 8.981768928606471e-07, "rel_err_parts": [ 8.981768928606471e-07 ], "compile_s": 0.07230799799435772, "inner": 4, "median_ms": 10.694562748540193, "q1_ms": 10.372725375418668, "q3_ms": 10.79276474956714, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "torch_sdpa", "mode": "fwd", "round": 1, "rel_err": 6.447160103562055e-07, "rel_err_parts": [ 6.447160103562055e-07 ], "rel_vs_jax_xla": 9.550253707857337e-07, "compile_s": 0.004051265001180582, "inner": 33, "median_ms": 1.173615788116656, "q1_ms": 1.1608035153282494, "q3_ms": 1.1907466819552197, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "flex_native", "mode": "fwd", "round": 2, "xla_temp_mb": 34.0, "xla_arg_mb": 32.0625, "xla_out_mb": 0.0625, "rel_err": 5.635644129142747e-07, "rel_err_parts": [ 5.635644129142747e-07 ], "rel_vs_jax_xla": 8.475861363876902e-07, "compile_s": 0.07040656100434717, "inner": 4, "median_ms": 10.904056500294246, "q1_ms": 10.676651252651936, "q3_ms": 11.131126248073997, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "flex_repeat", "mode": "fwd", "round": 2, "xla_temp_mb": 130.0, "xla_arg_mb": 32.0625, "xla_out_mb": 0.0625, "rel_err": 5.635644129142747e-07, "rel_err_parts": [ 5.635644129142747e-07 ], "rel_vs_jax_xla": 8.475861363876902e-07, "compile_s": 0.10893081399262883, "inner": 1, "median_ms": 55.291213997406885, "q1_ms": 54.900454997550696, "q3_ms": 56.03847400197992, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "jax_xla", "mode": "fwd", "round": 2, "xla_temp_mb": 34.0, "xla_arg_mb": 32.0625, "xla_out_mb": 0.0625, "rel_err": 8.981768928606471e-07, "rel_err_parts": [ 8.981768928606471e-07 ], "compile_s": 0.06835466899792664, "inner": 4, "median_ms": 10.316500251065008, "q1_ms": 10.137726749235298, "q3_ms": 10.491222001292044, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "torch_sdpa", "mode": "fwd", "round": 2, "rel_err": 6.447160103562055e-07, "rel_err_parts": [ 6.447160103562055e-07 ], "rel_vs_jax_xla": 9.550253707857337e-07, "compile_s": 0.0042464439902687445, "inner": 33, "median_ms": 1.1650617573953543, "q1_ms": 1.1616557577243922, "q3_ms": 1.1739924088804843, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "flex_native", "mode": "fwd", "round": 0, "xla_temp_mb": 40.0, "xla_arg_mb": 32.25, "xla_out_mb": 0.25, "rel_err": 9.015959108182869e-07, "rel_err_parts": [ 9.015959108182869e-07 ], "rel_vs_jax_xla": 1.4201448550466012e-07, "compile_s": 0.1068031989998417, "inner": 3, "median_ms": 15.92591999603125, "q1_ms": 15.630965332093183, "q3_ms": 16.009285832600046, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "flex_repeat", "mode": "fwd", "round": 0, "xla_temp_mb": 136.0, "xla_arg_mb": 32.25, "xla_out_mb": 0.25, "rel_err": 8.49978107453353e-07, "rel_err_parts": [ 8.49978107453353e-07 ], "rel_vs_jax_xla": 7.010599460954836e-07, "compile_s": 0.1433675640000729, "inner": 1, "median_ms": 57.644076994620264, "q1_ms": 57.11405249167001, "q3_ms": 59.16533950221492, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "jax_xla", "mode": "fwd", "round": 0, "xla_temp_mb": 40.0, "xla_arg_mb": 32.25, "xla_out_mb": 0.25, "rel_err": 9.026235829878715e-07, "rel_err_parts": [ 9.026235829878715e-07 ], "compile_s": 0.08515808900119737, "inner": 3, "median_ms": 15.498654999343367, "q1_ms": 15.379753332429875, "q3_ms": 15.886452000510568, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "torch_sdpa", "mode": "fwd", "round": 0, "rel_err": 9.473249065194977e-07, "rel_err_parts": [ 9.473249065194977e-07 ], "rel_vs_jax_xla": 6.377138674906746e-07, "compile_s": 0.004875814993283711, "inner": 17, "median_ms": 2.198824647005976, "q1_ms": 2.1775372939266484, "q3_ms": 2.2174875591778855, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "flex_native", "mode": "fwd", "round": 1, "xla_temp_mb": 40.0, "xla_arg_mb": 32.25, "xla_out_mb": 0.25, "rel_err": 9.015959108182869e-07, "rel_err_parts": [ 9.015959108182869e-07 ], "rel_vs_jax_xla": 1.4201448550466012e-07, "compile_s": 0.11092269299842883, "inner": 2, "median_ms": 15.80876349908067, "q1_ms": 15.502568996453192, "q3_ms": 16.168321002623998, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "flex_repeat", "mode": "fwd", "round": 1, "xla_temp_mb": 136.0, "xla_arg_mb": 32.25, "xla_out_mb": 0.25, "rel_err": 8.49978107453353e-07, "rel_err_parts": [ 8.49978107453353e-07 ], "rel_vs_jax_xla": 7.010599460954836e-07, "compile_s": 0.14401964799617417, "inner": 1, "median_ms": 59.10712300101295, "q1_ms": 58.4928764947108, "q3_ms": 59.39286950160749, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "jax_xla", "mode": "fwd", "round": 1, "xla_temp_mb": 40.0, "xla_arg_mb": 32.25, "xla_out_mb": 0.25, "rel_err": 9.026235829878715e-07, "rel_err_parts": [ 9.026235829878715e-07 ], "compile_s": 0.08770008799911011, "inner": 3, "median_ms": 15.307323332914772, "q1_ms": 15.108575331396423, "q3_ms": 16.087528834759723, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "torch_sdpa", "mode": "fwd", "round": 1, "rel_err": 9.473249065194977e-07, "rel_err_parts": [ 9.473249065194977e-07 ], "rel_vs_jax_xla": 6.377138674906746e-07, "compile_s": 0.004983867998817004, "inner": 19, "median_ms": 2.177010789685147, "q1_ms": 2.1674969734464726, "q3_ms": 2.188634578870781, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "flex_native", "mode": "fwd", "round": 2, "xla_temp_mb": 40.0, "xla_arg_mb": 32.25, "xla_out_mb": 0.25, "rel_err": 9.015959108182869e-07, "rel_err_parts": [ 9.015959108182869e-07 ], "rel_vs_jax_xla": 1.4201448550466012e-07, "compile_s": 0.10445129899017047, "inner": 3, "median_ms": 16.099876001438435, "q1_ms": 15.811451667104848, "q3_ms": 16.160997831320856, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "flex_repeat", "mode": "fwd", "round": 2, "xla_temp_mb": 136.0, "xla_arg_mb": 32.25, "xla_out_mb": 0.25, "rel_err": 8.49978107453353e-07, "rel_err_parts": [ 8.49978107453353e-07 ], "rel_vs_jax_xla": 7.010599460954836e-07, "compile_s": 0.14239102599094622, "inner": 1, "median_ms": 58.599994008545764, "q1_ms": 58.13169800239848, "q3_ms": 58.88998799491674, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "jax_xla", "mode": "fwd", "round": 2, "xla_temp_mb": 40.0, "xla_arg_mb": 32.25, "xla_out_mb": 0.25, "rel_err": 9.026235829878715e-07, "rel_err_parts": [ 9.026235829878715e-07 ], "compile_s": 0.08542637599748559, "inner": 3, "median_ms": 15.308949334818559, "q1_ms": 14.97062600052838, "q3_ms": 16.14546383158692, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 4096, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 32.0, "impl": "torch_sdpa", "mode": "fwd", "round": 2, "rel_err": 9.473249065194977e-07, "rel_err_parts": [ 9.473249065194977e-07 ], "rel_vs_jax_xla": 6.377138674906746e-07, "compile_s": 0.004851097008213401, "inner": 19, "median_ms": 2.2610829473705962, "q1_ms": 2.206180684305237, "q3_ms": 2.323644763428571, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "flex_native", "mode": "fwd", "round": 0, "xla_temp_mb": 130.0, "xla_arg_mb": 128.015625, "xla_out_mb": 0.015625, "rel_err": 9.940765721694333e-07, "rel_err_parts": [ 9.940765721694333e-07 ], "rel_vs_jax_xla": 5.61795616249583e-07, "compile_s": 0.08737722798832692, "inner": 1, "median_ms": 42.530764010734856, "q1_ms": 41.54909150383901, "q3_ms": 42.93745750328526, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 16384, "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.015625, "xla_out_mb": 0.015625, "rel_err": 2.443350012981682e-06, "rel_err_parts": [ 2.443350012981682e-06 ], "rel_vs_jax_xla": 2.3731429337203735e-06, "compile_s": 0.2417350739997346, "inner": 1, "median_ms": 222.78245800407603, "q1_ms": 221.0084164980799, "q3_ms": 224.96521250286605, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "jax_xla", "mode": "fwd", "round": 0, "xla_temp_mb": 130.0, "xla_arg_mb": 128.015625, "xla_out_mb": 0.015625, "rel_err": 9.497228461441409e-07, "rel_err_parts": [ 9.497228461441409e-07 ], "compile_s": 0.08888411799853202, "inner": 1, "median_ms": 40.41954901185818, "q1_ms": 39.755883495672606, "q3_ms": 42.66170550545212, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "torch_sdpa", "mode": "fwd", "round": 0, "rel_err": 2.4573278096795548e-06, "rel_err_parts": [ 2.4573278096795548e-06 ], "rel_vs_jax_xla": 2.321938609384233e-06, "compile_s": 0.01980533001187723, "inner": 3, "median_ms": 13.316810664643222, "q1_ms": 13.310025834167998, "q3_ms": 13.59718100138707, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "flex_native", "mode": "fwd", "round": 1, "xla_temp_mb": 130.0, "xla_arg_mb": 128.015625, "xla_out_mb": 0.015625, "rel_err": 9.940765721694333e-07, "rel_err_parts": [ 9.940765721694333e-07 ], "rel_vs_jax_xla": 5.61795616249583e-07, "compile_s": 0.09018482899409719, "inner": 1, "median_ms": 42.706916996394284, "q1_ms": 40.27131900511449, "q3_ms": 43.06318000453757, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 16384, "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.015625, "xla_out_mb": 0.015625, "rel_err": 2.443350012981682e-06, "rel_err_parts": [ 2.443350012981682e-06 ], "rel_vs_jax_xla": 2.3731429337203735e-06, "compile_s": 0.24075527100649197, "inner": 1, "median_ms": 219.98298200196587, "q1_ms": 218.51202900143107, "q3_ms": 222.1971855033189, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "jax_xla", "mode": "fwd", "round": 1, "xla_temp_mb": 130.0, "xla_arg_mb": 128.015625, "xla_out_mb": 0.015625, "rel_err": 9.497228461441409e-07, "rel_err_parts": [ 9.497228461441409e-07 ], "compile_s": 0.09452165399852674, "inner": 1, "median_ms": 40.36454499873798, "q1_ms": 40.25327449926408, "q3_ms": 42.17540050012758, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "torch_sdpa", "mode": "fwd", "round": 1, "rel_err": 2.4573278096795548e-06, "rel_err_parts": [ 2.4573278096795548e-06 ], "rel_vs_jax_xla": 2.321938609384233e-06, "compile_s": 0.01898622899898328, "inner": 3, "median_ms": 13.5445713288694, "q1_ms": 13.380000166459165, "q3_ms": 13.601984998482902, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "flex_native", "mode": "fwd", "round": 2, "xla_temp_mb": 130.0, "xla_arg_mb": 128.015625, "xla_out_mb": 0.015625, "rel_err": 9.940765721694333e-07, "rel_err_parts": [ 9.940765721694333e-07 ], "rel_vs_jax_xla": 5.61795616249583e-07, "compile_s": 0.09000827500130981, "inner": 1, "median_ms": 41.66895199159626, "q1_ms": 41.24910399696091, "q3_ms": 42.23195299709914, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 16384, "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.015625, "xla_out_mb": 0.015625, "rel_err": 2.443350012981682e-06, "rel_err_parts": [ 2.443350012981682e-06 ], "rel_vs_jax_xla": 2.3731429337203735e-06, "compile_s": 0.2370055959909223, "inner": 1, "median_ms": 221.52766400540713, "q1_ms": 220.17771950049791, "q3_ms": 224.8493234947091, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "jax_xla", "mode": "fwd", "round": 2, "xla_temp_mb": 130.0, "xla_arg_mb": 128.015625, "xla_out_mb": 0.015625, "rel_err": 9.497228461441409e-07, "rel_err_parts": [ 9.497228461441409e-07 ], "compile_s": 0.08963854500325397, "inner": 1, "median_ms": 41.26336499757599, "q1_ms": 41.18039299646625, "q3_ms": 42.356859499705024, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "torch_sdpa", "mode": "fwd", "round": 2, "rel_err": 2.4573278096795548e-06, "rel_err_parts": [ 2.4573278096795548e-06 ], "rel_vs_jax_xla": 2.321938609384233e-06, "compile_s": 0.01886035098868888, "inner": 3, "median_ms": 13.44345332957649, "q1_ms": 13.409150167717598, "q3_ms": 13.499863499115843, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "flex_native", "mode": "fwd", "round": 0, "xla_temp_mb": 136.0, "xla_arg_mb": 128.0625, "xla_out_mb": 0.0625, "rel_err": 5.67826589303877e-07, "rel_err_parts": [ 5.67826589303877e-07 ], "rel_vs_jax_xla": 8.519709240317752e-07, "compile_s": 0.09718500199960545, "inner": 1, "median_ms": 44.76126399822533, "q1_ms": 43.997909502650145, "q3_ms": 45.82937400118681, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "flex_repeat", "mode": "fwd", "round": 0, "xla_temp_mb": 520.0, "xla_arg_mb": 128.0625, "xla_out_mb": 0.0625, "rel_err": 5.67826589303877e-07, "rel_err_parts": [ 5.67826589303877e-07 ], "rel_vs_jax_xla": 8.519709240317752e-07, "compile_s": 0.24318020799546503, "inner": 1, "median_ms": 225.84469099820126, "q1_ms": 223.60249949997524, "q3_ms": 226.43703700305196, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "jax_xla", "mode": "fwd", "round": 0, "xla_temp_mb": 136.0, "xla_arg_mb": 128.0625, "xla_out_mb": 0.0625, "rel_err": 8.982898975773423e-07, "rel_err_parts": [ 8.982898975773423e-07 ], "compile_s": 0.09549293101008516, "inner": 1, "median_ms": 44.255486005567946, "q1_ms": 43.21456199977547, "q3_ms": 45.28879650024464, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "torch_sdpa", "mode": "fwd", "round": 0, "rel_err": 6.749351086909883e-07, "rel_err_parts": [ 6.749351086909883e-07 ], "rel_vs_jax_xla": 9.691883633422549e-07, "compile_s": 0.019728564002434723, "inner": 3, "median_ms": 15.172551332701309, "q1_ms": 14.913574833675133, "q3_ms": 15.582049498334527, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "flex_native", "mode": "fwd", "round": 1, "xla_temp_mb": 136.0, "xla_arg_mb": 128.0625, "xla_out_mb": 0.0625, "rel_err": 5.67826589303877e-07, "rel_err_parts": [ 5.67826589303877e-07 ], "rel_vs_jax_xla": 8.519709240317752e-07, "compile_s": 0.09401654799876269, "inner": 1, "median_ms": 43.59551800007466, "q1_ms": 42.502233998675365, "q3_ms": 44.17094049858861, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "flex_repeat", "mode": "fwd", "round": 1, "xla_temp_mb": 520.0, "xla_arg_mb": 128.0625, "xla_out_mb": 0.0625, "rel_err": 5.67826589303877e-07, "rel_err_parts": [ 5.67826589303877e-07 ], "rel_vs_jax_xla": 8.519709240317752e-07, "compile_s": 0.23180936599965207, "inner": 1, "median_ms": 228.08291600085795, "q1_ms": 226.32716400403297, "q3_ms": 230.7573760044761, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "jax_xla", "mode": "fwd", "round": 1, "xla_temp_mb": 136.0, "xla_arg_mb": 128.0625, "xla_out_mb": 0.0625, "rel_err": 8.982898975773423e-07, "rel_err_parts": [ 8.982898975773423e-07 ], "compile_s": 0.09724468500644434, "inner": 1, "median_ms": 45.89364500134252, "q1_ms": 44.594272498216014, "q3_ms": 46.131860493915156, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "torch_sdpa", "mode": "fwd", "round": 1, "rel_err": 6.749351086909883e-07, "rel_err_parts": [ 6.749351086909883e-07 ], "rel_vs_jax_xla": 9.691883633422549e-07, "compile_s": 0.017824823007686064, "inner": 3, "median_ms": 14.603872002529291, "q1_ms": 14.514776834403165, "q3_ms": 14.788695332147956, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "flex_native", "mode": "fwd", "round": 2, "xla_temp_mb": 136.0, "xla_arg_mb": 128.0625, "xla_out_mb": 0.0625, "rel_err": 5.67826589303877e-07, "rel_err_parts": [ 5.67826589303877e-07 ], "rel_vs_jax_xla": 8.519709240317752e-07, "compile_s": 0.09315649101336021, "inner": 1, "median_ms": 45.98346499551553, "q1_ms": 44.19141499965917, "q3_ms": 46.24210450128885, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "flex_repeat", "mode": "fwd", "round": 2, "xla_temp_mb": 520.0, "xla_arg_mb": 128.0625, "xla_out_mb": 0.0625, "rel_err": 5.67826589303877e-07, "rel_err_parts": [ 5.67826589303877e-07 ], "rel_vs_jax_xla": 8.519709240317752e-07, "compile_s": 0.2542350430012448, "inner": 1, "median_ms": 228.20573698845692, "q1_ms": 224.7080305023701, "q3_ms": 229.71688200777862, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "jax_xla", "mode": "fwd", "round": 2, "xla_temp_mb": 136.0, "xla_arg_mb": 128.0625, "xla_out_mb": 0.0625, "rel_err": 8.982898975773423e-07, "rel_err_parts": [ 8.982898975773423e-07 ], "compile_s": 0.09621693099325057, "inner": 1, "median_ms": 43.89719900791533, "q1_ms": 43.07214250729885, "q3_ms": 45.69060599897057, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "torch_sdpa", "mode": "fwd", "round": 2, "rel_err": 6.749351086909883e-07, "rel_err_parts": [ 6.749351086909883e-07 ], "rel_vs_jax_xla": 9.691883633422549e-07, "compile_s": 0.019257362990174443, "inner": 3, "median_ms": 14.51461000154571, "q1_ms": 14.287850834080018, "q3_ms": 14.780579830888502, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "flex_native", "mode": "fwd", "round": 0, "xla_temp_mb": 160.0, "xla_arg_mb": 128.25, "xla_out_mb": 0.25, "rel_err": 1.0815224413818214e-06, "rel_err_parts": [ 1.0815224413818214e-06 ], "rel_vs_jax_xla": 8.167247642631992e-07, "compile_s": 0.16557138199277688, "inner": 1, "median_ms": 64.69300200114958, "q1_ms": 62.74712549929973, "q3_ms": 68.7805749985273, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "flex_repeat", "mode": "fwd", "round": 0, "xla_temp_mb": 544.0, "xla_arg_mb": 128.25, "xla_out_mb": 0.25, "rel_err": 9.17728186777822e-07, "rel_err_parts": [ 9.17728186777822e-07 ], "rel_vs_jax_xla": 8.236931421379268e-07, "compile_s": 0.302745758992387, "inner": 1, "median_ms": 239.01981799281202, "q1_ms": 236.49778649996733, "q3_ms": 242.66805950173875, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "jax_xla", "mode": "fwd", "round": 0, "xla_temp_mb": 160.0, "xla_arg_mb": 128.25, "xla_out_mb": 0.25, "rel_err": 9.112583256865037e-07, "rel_err_parts": [ 9.112583256865037e-07 ], "compile_s": 0.12071588000981137, "inner": 1, "median_ms": 61.75192800583318, "q1_ms": 60.35778050136287, "q3_ms": 64.07022350322222, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "torch_sdpa", "mode": "fwd", "round": 0, "rel_err": 9.711479833640624e-07, "rel_err_parts": [ 9.711479833640624e-07 ], "rel_vs_jax_xla": 6.703700137222768e-07, "compile_s": 0.021022683009505272, "inner": 2, "median_ms": 16.369659497286193, "q1_ms": 16.309623253619066, "q3_ms": 17.14537674706662, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "flex_native", "mode": "fwd", "round": 1, "xla_temp_mb": 160.0, "xla_arg_mb": 128.25, "xla_out_mb": 0.25, "rel_err": 1.0815224413818214e-06, "rel_err_parts": [ 1.0815224413818214e-06 ], "rel_vs_jax_xla": 8.167247642631992e-07, "compile_s": 0.15706276400305796, "inner": 1, "median_ms": 73.26528900011908, "q1_ms": 70.15865650464548, "q3_ms": 77.00438299798407, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "flex_repeat", "mode": "fwd", "round": 1, "xla_temp_mb": 544.0, "xla_arg_mb": 128.25, "xla_out_mb": 0.25, "rel_err": 9.17728186777822e-07, "rel_err_parts": [ 9.17728186777822e-07 ], "rel_vs_jax_xla": 8.236931421379268e-07, "compile_s": 0.3019039169885218, "inner": 1, "median_ms": 241.3717169984011, "q1_ms": 236.84603450237773, "q3_ms": 244.66249249962857, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "jax_xla", "mode": "fwd", "round": 1, "xla_temp_mb": 160.0, "xla_arg_mb": 128.25, "xla_out_mb": 0.25, "rel_err": 9.112583256865037e-07, "rel_err_parts": [ 9.112583256865037e-07 ], "compile_s": 0.12634847599838395, "inner": 1, "median_ms": 62.76337000599597, "q1_ms": 60.90529599896399, "q3_ms": 63.90413399640238, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "torch_sdpa", "mode": "fwd", "round": 1, "rel_err": 9.711479833640624e-07, "rel_err_parts": [ 9.711479833640624e-07 ], "rel_vs_jax_xla": 6.703700137222768e-07, "compile_s": 0.020750958006829023, "inner": 2, "median_ms": 18.118733001756482, "q1_ms": 18.05361799779348, "q3_ms": 18.300912997801788, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "flex_native", "mode": "fwd", "round": 2, "xla_temp_mb": 160.0, "xla_arg_mb": 128.25, "xla_out_mb": 0.25, "rel_err": 1.0815224413818214e-06, "rel_err_parts": [ 1.0815224413818214e-06 ], "rel_vs_jax_xla": 8.167247642631992e-07, "compile_s": 0.15117111799190752, "inner": 1, "median_ms": 64.65727399336174, "q1_ms": 63.744577499164734, "q3_ms": 73.67056649673032, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "flex_repeat", "mode": "fwd", "round": 2, "xla_temp_mb": 544.0, "xla_arg_mb": 128.25, "xla_out_mb": 0.25, "rel_err": 9.17728186777822e-07, "rel_err_parts": [ 9.17728186777822e-07 ], "rel_vs_jax_xla": 8.236931421379268e-07, "compile_s": 0.3017561260057846, "inner": 1, "median_ms": 243.3331970096333, "q1_ms": 238.8413645021501, "q3_ms": 244.43864900968038, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "jax_xla", "mode": "fwd", "round": 2, "xla_temp_mb": 160.0, "xla_arg_mb": 128.25, "xla_out_mb": 0.25, "rel_err": 9.112583256865037e-07, "rel_err_parts": [ 9.112583256865037e-07 ], "compile_s": 0.1193194079969544, "inner": 1, "median_ms": 64.42965299356729, "q1_ms": 62.942815507994965, "q3_ms": 65.67044549592538, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 16384, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 128.0, "impl": "torch_sdpa", "mode": "fwd", "round": 2, "rel_err": 9.711479833640624e-07, "rel_err_parts": [ 9.711479833640624e-07 ], "rel_vs_jax_xla": 6.703700137222768e-07, "compile_s": 0.020525140993413515, "inner": 2, "median_ms": 16.119110499857925, "q1_ms": 15.943655998853501, "q3_ms": 16.24489075038582, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "flex_native", "mode": "fwd", "round": 0, "xla_temp_mb": 520.0, "xla_arg_mb": 512.015625, "xla_out_mb": 0.015625, "rel_err": 1.7800150544644566e-06, "rel_err_parts": [ 1.7800150544644566e-06 ], "rel_vs_jax_xla": 5.917842145208851e-07, "compile_s": 0.17982558900257573, "inner": 1, "median_ms": 168.4642540058121, "q1_ms": 163.46748099749675, "q3_ms": 170.06187150400365, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "flex_repeat", "mode": "fwd", "round": 0, "xla_temp_mb": 2056.0, "xla_arg_mb": 512.015625, "xla_out_mb": 0.015625, "rel_err": 4.832889771932969e-06, "rel_err_parts": [ 4.832889771932969e-06 ], "rel_vs_jax_xla": 4.6196241783036385e-06, "compile_s": 1.4444295290013542, "inner": 1, "median_ms": 1524.2878710123478, "q1_ms": 1482.3466034940793, "q3_ms": 1550.5527920031454, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "jax_xla", "mode": "fwd", "round": 0, "xla_temp_mb": 520.0, "xla_arg_mb": 512.015625, "xla_out_mb": 0.015625, "rel_err": 1.7627528450248064e-06, "rel_err_parts": [ 1.7627528450248064e-06 ], "compile_s": 0.20441718300571665, "inner": 1, "median_ms": 166.4549929992063, "q1_ms": 163.56166949117323, "q3_ms": 167.8723550066934, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "torch_sdpa", "mode": "fwd", "round": 0, "rel_err": 4.8285537559422664e-06, "rel_err_parts": [ 4.8285537559422664e-06 ], "rel_vs_jax_xla": 4.543243449006695e-06, "compile_s": 0.0743813709996175, "inner": 1, "median_ms": 68.98263700713869, "q1_ms": 67.30199750745669, "q3_ms": 71.35124350315891, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "flex_native", "mode": "fwd", "round": 1, "xla_temp_mb": 520.0, "xla_arg_mb": 512.015625, "xla_out_mb": 0.015625, "rel_err": 1.7800150544644566e-06, "rel_err_parts": [ 1.7800150544644566e-06 ], "rel_vs_jax_xla": 5.917842145208851e-07, "compile_s": 0.17498443899967242, "inner": 1, "median_ms": 167.8171410021605, "q1_ms": 167.51809499692172, "q3_ms": 169.24285700224573, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "flex_repeat", "mode": "fwd", "round": 1, "xla_temp_mb": 2056.0, "xla_arg_mb": 512.015625, "xla_out_mb": 0.015625, "rel_err": 4.832889771932969e-06, "rel_err_parts": [ 4.832889771932969e-06 ], "rel_vs_jax_xla": 4.6196241783036385e-06, "compile_s": 1.0620956589991692, "inner": 1, "median_ms": 1515.473646009923, "q1_ms": 1474.5898784967721, "q3_ms": 1554.7050020031747, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "jax_xla", "mode": "fwd", "round": 1, "xla_temp_mb": 520.0, "xla_arg_mb": 512.015625, "xla_out_mb": 0.015625, "rel_err": 1.7627528450248064e-06, "rel_err_parts": [ 1.7627528450248064e-06 ], "compile_s": 0.20383631400181912, "inner": 1, "median_ms": 163.8692909909878, "q1_ms": 162.55970350903226, "q3_ms": 168.68349399010185, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "torch_sdpa", "mode": "fwd", "round": 1, "rel_err": 4.8285537559422664e-06, "rel_err_parts": [ 4.8285537559422664e-06 ], "rel_vs_jax_xla": 4.543243449006695e-06, "compile_s": 0.0698897929978557, "inner": 1, "median_ms": 68.32004299212713, "q1_ms": 67.28243450197624, "q3_ms": 69.386832001328, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "flex_native", "mode": "fwd", "round": 2, "xla_temp_mb": 520.0, "xla_arg_mb": 512.015625, "xla_out_mb": 0.015625, "rel_err": 1.7800150544644566e-06, "rel_err_parts": [ 1.7800150544644566e-06 ], "rel_vs_jax_xla": 5.917842145208851e-07, "compile_s": 0.17785943900526036, "inner": 1, "median_ms": 168.74355300387833, "q1_ms": 168.24098699726164, "q3_ms": 169.6212535025552, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "flex_repeat", "mode": "fwd", "round": 2, "xla_temp_mb": 2056.0, "xla_arg_mb": 512.015625, "xla_out_mb": 0.015625, "rel_err": 4.832889771932969e-06, "rel_err_parts": [ 4.832889771932969e-06 ], "rel_vs_jax_xla": 4.6196241783036385e-06, "compile_s": 1.3316501639928902, "inner": 1, "median_ms": 1510.8915870077908, "q1_ms": 1495.2728185016895, "q3_ms": 1530.9557460059295, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "jax_xla", "mode": "fwd", "round": 2, "xla_temp_mb": 520.0, "xla_arg_mb": 512.015625, "xla_out_mb": 0.015625, "rel_err": 1.7627528450248064e-06, "rel_err_parts": [ 1.7627528450248064e-06 ], "compile_s": 0.20711118000326678, "inner": 1, "median_ms": 164.6163149998756, "q1_ms": 163.98236050008563, "q3_ms": 168.51002600742504, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 1, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "torch_sdpa", "mode": "fwd", "round": 2, "rel_err": 4.8285537559422664e-06, "rel_err_parts": [ 4.8285537559422664e-06 ], "rel_vs_jax_xla": 4.543243449006695e-06, "compile_s": 0.0686095910059521, "inner": 1, "median_ms": 69.91917800041847, "q1_ms": 69.1231675009476, "q3_ms": 70.53132200235268, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "flex_native", "mode": "fwd", "round": 0, "xla_temp_mb": 544.0, "xla_arg_mb": 512.0625, "xla_out_mb": 0.0625, "rel_err": 6.2710927295484e-07, "rel_err_parts": [ 6.2710927295484e-07 ], "rel_vs_jax_xla": 9.403548233422043e-07, "compile_s": 0.20448517199838534, "inner": 1, "median_ms": 181.3990460068453, "q1_ms": 176.00504000438377, "q3_ms": 182.99936849507503, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "flex_repeat", "mode": "fwd", "round": 0, "xla_temp_mb": 2080.0, "xla_arg_mb": 512.0625, "xla_out_mb": 0.0625, "rel_err": 6.2710927295484e-07, "rel_err_parts": [ 6.2710927295484e-07 ], "rel_vs_jax_xla": 9.403548233422043e-07, "compile_s": 1.356972991998191, "inner": 1, "median_ms": 1435.235315002501, "q1_ms": 1382.2644800020498, "q3_ms": 1490.8771374975913, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "jax_xla", "mode": "fwd", "round": 0, "xla_temp_mb": 544.0, "xla_arg_mb": 512.0625, "xla_out_mb": 0.0625, "rel_err": 9.016676472128893e-07, "rel_err_parts": [ 9.016676472128893e-07 ], "compile_s": 0.22453979299461935, "inner": 1, "median_ms": 184.46845200378448, "q1_ms": 182.23920550371986, "q3_ms": 185.02649800211657, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "torch_sdpa", "mode": "fwd", "round": 0, "rel_err": 7.516701998611097e-07, "rel_err_parts": [ 7.516701998611097e-07 ], "rel_vs_jax_xla": 1.0395217486802721e-06, "compile_s": 0.07346170100208838, "inner": 1, "median_ms": 73.82093099295162, "q1_ms": 70.8226235001348, "q3_ms": 76.8036635054159, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "flex_native", "mode": "fwd", "round": 1, "xla_temp_mb": 544.0, "xla_arg_mb": 512.0625, "xla_out_mb": 0.0625, "rel_err": 6.2710927295484e-07, "rel_err_parts": [ 6.2710927295484e-07 ], "rel_vs_jax_xla": 9.403548233422043e-07, "compile_s": 0.21190101199317724, "inner": 1, "median_ms": 182.17473301046994, "q1_ms": 179.04832349449862, "q3_ms": 184.87382450257428, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "flex_repeat", "mode": "fwd", "round": 1, "xla_temp_mb": 2080.0, "xla_arg_mb": 512.0625, "xla_out_mb": 0.0625, "rel_err": 6.2710927295484e-07, "rel_err_parts": [ 6.2710927295484e-07 ], "rel_vs_jax_xla": 9.403548233422043e-07, "compile_s": 1.4029119119950337, "inner": 1, "median_ms": 1373.3721559983678, "q1_ms": 1318.8574870000593, "q3_ms": 1492.6123350014677, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "jax_xla", "mode": "fwd", "round": 1, "xla_temp_mb": 544.0, "xla_arg_mb": 512.0625, "xla_out_mb": 0.0625, "rel_err": 9.016676472128893e-07, "rel_err_parts": [ 9.016676472128893e-07 ], "compile_s": 0.2275886610004818, "inner": 1, "median_ms": 179.23403000168037, "q1_ms": 178.09807000594446, "q3_ms": 183.14399150403915, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "torch_sdpa", "mode": "fwd", "round": 1, "rel_err": 7.516701998611097e-07, "rel_err_parts": [ 7.516701998611097e-07 ], "rel_vs_jax_xla": 1.0395217486802721e-06, "compile_s": 0.06821129799936898, "inner": 1, "median_ms": 65.99769000604283, "q1_ms": 65.39821100159315, "q3_ms": 68.74264299403876, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "flex_native", "mode": "fwd", "round": 2, "xla_temp_mb": 544.0, "xla_arg_mb": 512.0625, "xla_out_mb": 0.0625, "rel_err": 6.2710927295484e-07, "rel_err_parts": [ 6.2710927295484e-07 ], "rel_vs_jax_xla": 9.403548233422043e-07, "compile_s": 0.20846295700175688, "inner": 1, "median_ms": 176.28873699868564, "q1_ms": 175.7466164926882, "q3_ms": 179.93942400062224, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "flex_repeat", "mode": "fwd", "round": 2, "xla_temp_mb": 2080.0, "xla_arg_mb": 512.0625, "xla_out_mb": 0.0625, "rel_err": 6.2710927295484e-07, "rel_err_parts": [ 6.2710927295484e-07 ], "rel_vs_jax_xla": 9.403548233422043e-07, "compile_s": 1.3620280450122664, "inner": 1, "median_ms": 1496.4008329989156, "q1_ms": 1419.4272275053663, "q3_ms": 1497.6983970045694, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "jax_xla", "mode": "fwd", "round": 2, "xla_temp_mb": 544.0, "xla_arg_mb": 512.0625, "xla_out_mb": 0.0625, "rel_err": 9.016676472128893e-07, "rel_err_parts": [ 9.016676472128893e-07 ], "compile_s": 0.21493793699482922, "inner": 1, "median_ms": 183.0273660016246, "q1_ms": 180.3963184938766, "q3_ms": 183.7643909966573, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 4, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "torch_sdpa", "mode": "fwd", "round": 2, "rel_err": 7.516701998611097e-07, "rel_err_parts": [ 7.516701998611097e-07 ], "rel_vs_jax_xla": 1.0395217486802721e-06, "compile_s": 0.06723339800373651, "inner": 1, "median_ms": 66.0066470009042, "q1_ms": 65.79318299191073, "q3_ms": 66.06824749906082, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "flex_native", "mode": "fwd", "round": 0, "xla_temp_mb": 768.2550048828125, "xla_arg_mb": 512.25, "xla_out_mb": 0.25, "rel_err": 1.0375609917900874e-06, "rel_err_parts": [ 1.0375609917900874e-06 ], "rel_vs_jax_xla": 7.614783612552856e-07, "compile_s": 0.49601687499671243, "inner": 1, "median_ms": 378.7586870021187, "q1_ms": 375.51953299407614, "q3_ms": 386.7806130001554, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "flex_repeat", "mode": "fwd", "round": 0, "xla_temp_mb": 3072.5050048828125, "xla_arg_mb": 512.25, "xla_out_mb": 0.25, "rel_err": 1.0375609917900874e-06, "rel_err_parts": [ 1.0375609917900874e-06 ], "rel_vs_jax_xla": 7.614783612552856e-07, "compile_s": 1.0072035389894154, "inner": 1, "median_ms": 1120.8426719967974, "q1_ms": 1112.5766929981182, "q3_ms": 1124.0469585027313, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "jax_xla", "mode": "fwd", "round": 0, "xla_temp_mb": 640.0, "xla_arg_mb": 512.25, "xla_out_mb": 0.25, "rel_err": 9.302755756834813e-07, "rel_err_parts": [ 9.302755756834813e-07 ], "compile_s": 0.36814145000244025, "inner": 1, "median_ms": 251.09095399966463, "q1_ms": 249.54909249208868, "q3_ms": 256.4131129984162, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "torch_sdpa", "mode": "fwd", "round": 0, "rel_err": 1.0325542234568275e-06, "rel_err_parts": [ 1.0325542234568275e-06 ], "rel_vs_jax_xla": 7.602030223097245e-07, "compile_s": 0.06743392700445838, "inner": 1, "median_ms": 67.31949999812059, "q1_ms": 67.16488800157094, "q3_ms": 67.5716879995889, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "flex_native", "mode": "fwd", "round": 1, "xla_temp_mb": 768.2550048828125, "xla_arg_mb": 512.25, "xla_out_mb": 0.25, "rel_err": 1.0375609917900874e-06, "rel_err_parts": [ 1.0375609917900874e-06 ], "rel_vs_jax_xla": 7.614783612552856e-07, "compile_s": 0.4935250810085563, "inner": 1, "median_ms": 381.72629999462515, "q1_ms": 376.3281680003274, "q3_ms": 382.5681890011765, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "flex_repeat", "mode": "fwd", "round": 1, "xla_temp_mb": 3072.5050048828125, "xla_arg_mb": 512.25, "xla_out_mb": 0.25, "rel_err": 1.0375609917900874e-06, "rel_err_parts": [ 1.0375609917900874e-06 ], "rel_vs_jax_xla": 7.614783612552856e-07, "compile_s": 1.006779616989661, "inner": 1, "median_ms": 1121.3658610067796, "q1_ms": 1114.7164610010805, "q3_ms": 1129.0352904979954, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "jax_xla", "mode": "fwd", "round": 1, "xla_temp_mb": 640.0, "xla_arg_mb": 512.25, "xla_out_mb": 0.25, "rel_err": 9.302755756834813e-07, "rel_err_parts": [ 9.302755756834813e-07 ], "compile_s": 0.3642599160084501, "inner": 1, "median_ms": 249.93661099870224, "q1_ms": 248.9837475004606, "q3_ms": 253.1012874969747, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "torch_sdpa", "mode": "fwd", "round": 1, "rel_err": 1.0325542234568275e-06, "rel_err_parts": [ 1.0325542234568275e-06 ], "rel_vs_jax_xla": 7.602030223097245e-07, "compile_s": 0.07423433299118187, "inner": 1, "median_ms": 70.49615000141785, "q1_ms": 66.8609530039248, "q3_ms": 71.49818100151606, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "flex_native", "mode": "fwd", "round": 2, "xla_temp_mb": 768.2550048828125, "xla_arg_mb": 512.25, "xla_out_mb": 0.25, "rel_err": 1.0375609917900874e-06, "rel_err_parts": [ 1.0375609917900874e-06 ], "rel_vs_jax_xla": 7.614783612552856e-07, "compile_s": 0.519770293001784, "inner": 1, "median_ms": 373.39752400293946, "q1_ms": 370.5533324973658, "q3_ms": 374.75092399836285, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "flex_repeat", "mode": "fwd", "round": 2, "xla_temp_mb": 3072.5050048828125, "xla_arg_mb": 512.25, "xla_out_mb": 0.25, "rel_err": 1.0375609917900874e-06, "rel_err_parts": [ 1.0375609917900874e-06 ], "rel_vs_jax_xla": 7.614783612552856e-07, "compile_s": 1.0029599410045194, "inner": 1, "median_ms": 1121.7593150067842, "q1_ms": 1118.4053944962216, "q3_ms": 1127.3024675028864, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "jax_xla", "mode": "fwd", "round": 2, "xla_temp_mb": 640.0, "xla_arg_mb": 512.25, "xla_out_mb": 0.25, "rel_err": 9.302755756834813e-07, "rel_err_parts": [ 9.302755756834813e-07 ], "compile_s": 0.3687738260050537, "inner": 1, "median_ms": 250.38133200723678, "q1_ms": 245.43331600580132, "q3_ms": 254.1742375033209, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" }, { "suite": "decode", "dtype": "float32", "batch": 1, "lq": 16, "lkv": 65536, "hq": 32, "hkv": 8, "head_dim": 128, "kv_mb": 512.0, "impl": "torch_sdpa", "mode": "fwd", "round": 2, "rel_err": 1.0325542234568275e-06, "rel_err_parts": [ 1.0325542234568275e-06 ], "rel_vs_jax_xla": 7.602030223097245e-07, "compile_s": 0.07030487799784169, "inner": 1, "median_ms": 69.60352101305034, "q1_ms": 69.1383970042807, "q3_ms": 70.55199149908731, "tag": "after", "device_arg": "cpu", "run": "cpu_decode_after", "note": "after: no tile padding on the CPU XLA path" } ] }