{ "device": { "jax": "0.11.2", "backend": "gpu", "devices": [ "cuda:0" ], "python": "3.13.18", "device_kind": "NVIDIA H100" }, "config": { "output": "/home/tu/tu_tu/tu_iivgl01/tabdiffpfn/experiments/results/h100_branch/attention.json", "batch": 4, "heads": 8, "head_dims": [ 64, 128 ], "seq_lens": [ 128, 512, 1024, 2048, 4096, 8192 ], "ragged_seq_lens": [ 600, 1500, 3000 ], "dtypes": [ "bfloat16", "float32" ], "modes": [ "fwd", "fwd_bwd" ], "causal": [ 0, 1 ], "impls": [ "probjax_mha", "probjax_flex", "probjax_default", "jax_xla", "jax_cudnn", "torch_flash", "torch_efficient", "torch_math" ], "append": false, "tag": "h100_nvl", "torch_note": "torch rows: torch 2.14.1+cu130, separate process run with XLA_PYTHON_CLIENT_PREALLOCATE=false and the platform allocator (JAX only builds inputs and references)" }, "results": [ { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 64, "seq": 128, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0021501779556274414, "rel_err": 0.002294705016538501, "xla_temp_mb": 0.0, "xla_arg_mb": 1.5, "xla_out_mb": 0.5, "compile_s": 0.58245162316598, "inner": 200, "median_ms": 0.032379410695284605, "q1_ms": 0.03235680749639869, "q3_ms": 0.03259640536271036, "tflops": 4.145156601616163 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 64, "seq": 128, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0021501779556274414, "rel_err": 0.002294705016538501, "xla_temp_mb": 0.0, "xla_arg_mb": 1.5, "xla_out_mb": 0.5, "compile_s": 0.3656611740589142, "inner": 200, "median_ms": 0.03260710509493947, "q1_ms": 0.0323298602597788, "q3_ms": 0.032728229998610914, "tflops": 4.11621110212664 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 64, "seq": 128, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0021501779556274414, "rel_err": 0.002294705016538501, "xla_temp_mb": 0.0, "xla_arg_mb": 1.5, "xla_out_mb": 0.5, "compile_s": 0.3736522209364921, "inner": 200, "median_ms": 0.024089179933071136, "q1_ms": 0.023997932439669967, "q3_ms": 0.024731452576816082, "tflops": 5.571701833474932 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 64, "seq": 128, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0021272897720336914, "rel_err": 0.002270278288051486, "xla_temp_mb": 3.0, "xla_arg_mb": 1.5, "xla_out_mb": 0.5, "compile_s": 0.21017099916934967, "inner": 200, "median_ms": 0.04397152573801577, "q1_ms": 0.04381847509648651, "q3_ms": 0.04421169694978744, "tflops": 3.052378232214978 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 64, "seq": 128, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0021501779556274414, "rel_err": 0.002294705016538501, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 1.5, "xla_out_mb": 0.5, "compile_s": 1.8839773400686681, "inner": 200, "median_ms": 0.0483914150390774, "q1_ms": 0.0482587848091498, "q3_ms": 0.057820758083835244, "tflops": 2.773585519076379 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 64, "seq": 128, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0032709836959838867, "rel_err": 0.0035336874425411224, "xla_temp_mb": 1.031494140625, "xla_arg_mb": 2.5, "xla_out_mb": 1.5000228881835938, "compile_s": 0.49407712602987885, "inner": 200, "median_ms": 0.03703694441355765, "q1_ms": 0.03648591984529048, "q3_ms": 0.043835199903696775, "tflops": 12.683607015594951 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 64, "seq": 128, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0032709836959838867, "rel_err": 0.0035336874425411224, "xla_temp_mb": 0.53173828125, "xla_arg_mb": 2.5, "xla_out_mb": 1.5000228881835938, "compile_s": 0.5104248500429094, "inner": 200, "median_ms": 0.041541935643181205, "q1_ms": 0.04106380802113563, "q3_ms": 0.04165493242908269, "tflops": 11.308140574742522 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 64, "seq": 128, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0032709836959838867, "rel_err": 0.0035336874425411224, "xla_temp_mb": 0.53173828125, "xla_arg_mb": 2.5, "xla_out_mb": 1.5000228881835938, "compile_s": 0.48620230006054044, "inner": 200, "median_ms": 0.04087468492798507, "q1_ms": 0.04078475700225681, "q3_ms": 0.04143867990933359, "tflops": 11.49273808049282 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 64, "seq": 128, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00531381368637085, "rel_err": 0.005376721732318401, "xla_temp_mb": 5.0, "xla_arg_mb": 2.5, "xla_out_mb": 1.5000228881835938, "compile_s": 0.03203871799632907, "inner": 200, "median_ms": 0.05848812987096608, "q1_ms": 0.0583962828386575, "q3_ms": 0.058601482305675745, "tflops": 8.031750186514225 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 64, "seq": 128, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0032709836959838867, "rel_err": 0.0035336874425411224, "xla_temp_mb": 2.031982421875, "xla_arg_mb": 2.5, "xla_out_mb": 1.5000228881835938, "compile_s": 1.385311946971342, "inner": 200, "median_ms": 0.09394488530233502, "q1_ms": 0.09378393238876015, "q3_ms": 0.0957827502861619, "tflops": 5.000400463401533 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 64, "seq": 128, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.005635261535644531, "rel_err": 0.0019494958687573671, "xla_temp_mb": 0.0, "xla_arg_mb": 1.5, "xla_out_mb": 0.5, "compile_s": 0.5003675778862089, "inner": 200, "median_ms": 0.03290495602414012, "q1_ms": 0.032707182108424604, "q3_ms": 0.033099629799835384, "tflops": 2.039475875633045 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 64, "seq": 128, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.005635261535644531, "rel_err": 0.0019494958687573671, "xla_temp_mb": 0.0, "xla_arg_mb": 1.5, "xla_out_mb": 0.5, "compile_s": 0.4955592919141054, "inner": 200, "median_ms": 0.03272150410339236, "q1_ms": 0.029213064699433744, "q3_ms": 0.03321520518511534, "tflops": 2.050910122833949 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 64, "seq": 128, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.005635261535644531, "rel_err": 0.0019494958687573671, "xla_temp_mb": 0.0, "xla_arg_mb": 1.5, "xla_out_mb": 0.5, "compile_s": 0.49652559007517993, "inner": 200, "median_ms": 0.03315305453725159, "q1_ms": 0.033122305176220834, "q3_ms": 0.03340060240589082, "tflops": 2.024213603744862 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 64, "seq": 128, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.007589459419250488, "rel_err": 0.0026255426928400993, "xla_temp_mb": 3.0, "xla_arg_mb": 1.5, "xla_out_mb": 0.5, "compile_s": 0.26251308782957494, "inner": 200, "median_ms": 0.034733699867501855, "q1_ms": 0.03456257516518235, "q3_ms": 0.039627409423701465, "tflops": 1.932096616715156 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 64, "seq": 128, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.005635261535644531, "rel_err": 0.0019494958687573671, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 1.5, "xla_out_mb": 0.5, "compile_s": 0.35177920293062925, "inner": 200, "median_ms": 0.06298706983216107, "q1_ms": 0.05812208284623921, "q3_ms": 0.06320609478279948, "tflops": 1.0654387349470629 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 64, "seq": 128, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.009996175765991211, "rel_err": 0.005375280044972897, "xla_temp_mb": 1.031494140625, "xla_arg_mb": 2.5, "xla_out_mb": 1.5000228881835938, "compile_s": 0.6479109970387071, "inner": 200, "median_ms": 0.04939231090247631, "q1_ms": 0.049186357646249235, "q3_ms": 0.04964848281815648, "tflops": 4.7554167786108605 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 64, "seq": 128, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.009996175765991211, "rel_err": 0.005375280044972897, "xla_temp_mb": 0.53173828125, "xla_arg_mb": 2.5, "xla_out_mb": 1.5000228881835938, "compile_s": 0.6658460050821304, "inner": 200, "median_ms": 0.042945980094373226, "q1_ms": 0.04292097466532141, "q3_ms": 0.043132202699780464, "tflops": 5.469220250273763 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 64, "seq": 128, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.009996175765991211, "rel_err": 0.005375280044972897, "xla_temp_mb": 0.53173828125, "xla_arg_mb": 2.5, "xla_out_mb": 1.5000228881835938, "compile_s": 0.7093678100500256, "inner": 200, "median_ms": 0.05475329468026757, "q1_ms": 0.05173700279556215, "q3_ms": 0.05489754257723689, "tflops": 4.289806218449321 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 64, "seq": 128, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.011534929275512695, "rel_err": 0.005375280044972897, "xla_temp_mb": 5.0, "xla_arg_mb": 2.5, "xla_out_mb": 1.5000228881835938, "compile_s": 0.8127072351053357, "inner": 200, "median_ms": 0.06344031542539597, "q1_ms": 0.06328529270831496, "q3_ms": 0.06953857489861548, "tflops": 3.702393697525251 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 64, "seq": 128, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.009996175765991211, "rel_err": 0.005375280044972897, "xla_temp_mb": 2.031982421875, "xla_arg_mb": 2.5, "xla_out_mb": 1.5000228881835938, "compile_s": 1.1307117850519717, "inner": 200, "median_ms": 0.10058536543510854, "q1_ms": 0.09757464285939932, "q3_ms": 0.10108178772497922, "tflops": 2.3351411309583674 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 64, "seq": 512, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00117570161819458, "rel_err": 0.002601401647552848, "xla_temp_mb": 0.0, "xla_arg_mb": 6.0, "xla_out_mb": 2.0, "compile_s": 0.42613860103301704, "inner": 200, "median_ms": 0.03812893992289901, "q1_ms": 0.037962092901580036, "q3_ms": 0.03829291497822851, "tflops": 56.32161954521821 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 64, "seq": 512, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00117570161819458, "rel_err": 0.002601401647552848, "xla_temp_mb": 0.0, "xla_arg_mb": 6.0, "xla_out_mb": 2.0, "compile_s": 0.427638057153672, "inner": 200, "median_ms": 0.022344185272231698, "q1_ms": 0.02198373491410166, "q3_ms": 0.02245718496851623, "tflops": 96.10928399653002 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 64, "seq": 512, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00117570161819458, "rel_err": 0.002601401647552848, "xla_temp_mb": 0.0, "xla_arg_mb": 6.0, "xla_out_mb": 2.0, "compile_s": 0.42889614310115576, "inner": 200, "median_ms": 0.022598389768972993, "q1_ms": 0.022451785043813288, "q3_ms": 0.0228771852562204, "tflops": 95.02817103139091 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 64, "seq": 512, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.001338571310043335, "rel_err": 0.002961773192510009, "xla_temp_mb": 48.0, "xla_arg_mb": 6.0, "xla_out_mb": 2.0, "compile_s": 0.2180318171158433, "inner": 200, "median_ms": 0.047950310399755836, "q1_ms": 0.04793259024154395, "q3_ms": 0.04803713760338724, "tflops": 44.7856047249057 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 64, "seq": 512, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00117570161819458, "rel_err": 0.002601401647552848, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 6.0, "xla_out_mb": 2.0, "compile_s": 0.30465239798650146, "inner": 200, "median_ms": 0.06459241500124335, "q1_ms": 0.06438524054829031, "q3_ms": 0.06464693986345083, "tflops": 33.246684582991094 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 64, "seq": 512, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.001962810754776001, "rel_err": 0.005092823877930641, "xla_temp_mb": 4.125244140625, "xla_arg_mb": 10.0, "xla_out_mb": 6.000022888183594, "compile_s": 0.5757848720531911, "inner": 200, "median_ms": 0.044875419698655605, "q1_ms": 0.04485132289119065, "q3_ms": 0.04511879466008395, "tflops": 167.49019437527787 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 64, "seq": 512, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.001962810754776001, "rel_err": 0.005092823877930641, "xla_temp_mb": 2.12548828125, "xla_arg_mb": 10.0, "xla_out_mb": 6.000022888183594, "compile_s": 0.5716010550968349, "inner": 200, "median_ms": 0.04710156470537186, "q1_ms": 0.04706729028839618, "q3_ms": 0.04752036300487816, "tflops": 159.5741630881063 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 64, "seq": 512, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.001962810754776001, "rel_err": 0.005092823877930641, "xla_temp_mb": 2.12548828125, "xla_arg_mb": 10.0, "xla_out_mb": 6.000022888183594, "compile_s": 0.5824770429171622, "inner": 200, "median_ms": 0.05467874580062926, "q1_ms": 0.049840257270261645, "q3_ms": 0.05551491514779627, "tflops": 137.46095778066476 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 64, "seq": 512, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0027485042810440063, "rel_err": 0.0069497511722147465, "xla_temp_mb": 80.0, "xla_arg_mb": 10.0, "xla_out_mb": 6.000022888183594, "compile_s": 0.03245029808022082, "inner": 181, "median_ms": 0.1510475242500147, "q1_ms": 0.15102050560449368, "q3_ms": 0.15114412709905956, "tflops": 49.76044993335446 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 64, "seq": 512, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.001962810754776001, "rel_err": 0.005092823877930641, "xla_temp_mb": 8.125732421875, "xla_arg_mb": 10.0, "xla_out_mb": 6.000022888183594, "compile_s": 0.1330546799581498, "inner": 200, "median_ms": 0.10925773996859789, "q1_ms": 0.10908590978942811, "q3_ms": 0.10932043776847422, "tflops": 68.79322938732078 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 64, "seq": 512, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.007361173629760742, "rel_err": 0.0025465681683272123, "xla_temp_mb": 0.0, "xla_arg_mb": 6.0, "xla_out_mb": 2.0, "compile_s": 0.5092254630289972, "inner": 200, "median_ms": 0.02354988013394177, "q1_ms": 0.022647185251116753, "q3_ms": 0.031813884852454066, "tflops": 45.59436472258076 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 64, "seq": 512, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.007361173629760742, "rel_err": 0.0025465681683272123, "xla_temp_mb": 0.0, "xla_arg_mb": 6.0, "xla_out_mb": 2.0, "compile_s": 0.5063336009625345, "inner": 200, "median_ms": 0.02224993542768061, "q1_ms": 0.02213678730186075, "q3_ms": 0.02237271226476878, "tflops": 48.258199557028085 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 64, "seq": 512, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.007361173629760742, "rel_err": 0.0025465681683272123, "xla_temp_mb": 0.0, "xla_arg_mb": 6.0, "xla_out_mb": 2.0, "compile_s": 0.5078775370493531, "inner": 200, "median_ms": 0.022002740297466516, "q1_ms": 0.021819387329742312, "q3_ms": 0.022213064949028194, "tflops": 48.800368021597514 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 64, "seq": 512, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.007361173629760742, "rel_err": 0.0025465681683272123, "xla_temp_mb": 48.0, "xla_arg_mb": 6.0, "xla_out_mb": 2.0, "compile_s": 0.711287965066731, "inner": 200, "median_ms": 0.07699678069911897, "q1_ms": 0.07692502986174077, "q3_ms": 0.07703394978307188, "tflops": 13.945282052711669 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 64, "seq": 512, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.007361173629760742, "rel_err": 0.0025465681683272123, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 6.0, "xla_out_mb": 2.0, "compile_s": 0.05738696618936956, "inner": 200, "median_ms": 0.07047769497148693, "q1_ms": 0.07027137442491949, "q3_ms": 0.07067804457619786, "tflops": 15.235200646593256 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 64, "seq": 512, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.01617288589477539, "rel_err": 0.0037439705338329077, "xla_temp_mb": 4.125244140625, "xla_arg_mb": 10.0, "xla_out_mb": 6.000022888183594, "compile_s": 0.6280911031644791, "inner": 200, "median_ms": 0.04170603002421558, "q1_ms": 0.04164363257586956, "q3_ms": 0.04175345762632787, "tflops": 90.10918521417535 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 64, "seq": 512, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.01617288589477539, "rel_err": 0.0037439705338329077, "xla_temp_mb": 2.12548828125, "xla_arg_mb": 10.0, "xla_out_mb": 6.000022888183594, "compile_s": 0.6857734378427267, "inner": 200, "median_ms": 0.05390154430642724, "q1_ms": 0.05258607503492385, "q3_ms": 0.05449037009384483, "tflops": 69.7214974516395 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 64, "seq": 512, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.01617288589477539, "rel_err": 0.0037439705338329077, "xla_temp_mb": 2.12548828125, "xla_arg_mb": 10.0, "xla_out_mb": 6.000022888183594, "compile_s": 0.6475662270095199, "inner": 200, "median_ms": 0.049328404711559415, "q1_ms": 0.04931290808599442, "q3_ms": 0.04933960735797882, "tflops": 76.18524065343112 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 64, "seq": 512, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.01617288589477539, "rel_err": 0.00455421581864357, "xla_temp_mb": 80.0, "xla_arg_mb": 10.0, "xla_out_mb": 6.000022888183594, "compile_s": 0.049216388957574964, "inner": 143, "median_ms": 0.2504382998897479, "q1_ms": 0.2503502980473783, "q3_ms": 0.250633968808851, "tflops": 15.00607688861668 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 64, "seq": 512, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.01617288589477539, "rel_err": 0.0037439705338329077, "xla_temp_mb": 8.125732421875, "xla_arg_mb": 10.0, "xla_out_mb": 6.000022888183594, "compile_s": 0.3084611869417131, "inner": 200, "median_ms": 0.09652367560192943, "q1_ms": 0.09640619973652065, "q3_ms": 0.09660287236329168, "tflops": 38.93445168311513 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 64, "seq": 600, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0010502040386199951, "rel_err": 0.00308937462978065, "xla_temp_mb": 10.0, "xla_arg_mb": 7.03125, "xla_out_mb": 2.34375, "compile_s": 0.503125251038, "inner": 200, "median_ms": 0.04834071500226855, "q1_ms": 0.047993610496632755, "q3_ms": 0.049567632377147675, "tflops": 61.006958624041914 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 64, "seq": 600, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0010502040386199951, "rel_err": 0.00308937462978065, "xla_temp_mb": 10.0, "xla_arg_mb": 7.03125, "xla_out_mb": 2.34375, "compile_s": 0.49826927413232625, "inner": 200, "median_ms": 0.05013680434785783, "q1_ms": 0.04949445778038353, "q3_ms": 0.05104913027025759, "tflops": 58.821459372210775 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 64, "seq": 600, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0010502040386199951, "rel_err": 0.00308937462978065, "xla_temp_mb": 10.0, "xla_arg_mb": 7.03125, "xla_out_mb": 2.34375, "compile_s": 0.5027708020061255, "inner": 200, "median_ms": 0.050257855327799916, "q1_ms": 0.05003805737942457, "q3_ms": 0.050424630171619356, "tflops": 58.67978211097096 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 64, "seq": 600, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0012940764427185059, "rel_err": 0.0038067714776843786, "xla_temp_mb": 65.91798400878906, "xla_arg_mb": 7.03125, "xla_out_mb": 2.34375, "compile_s": 0.21607089205645025, "inner": 184, "median_ms": 0.08133889692228125, "q1_ms": 0.08121718786677103, "q3_ms": 0.08137841106635396, "tflops": 36.257191965829875 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 64, "seq": 600, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0010502040386199951, "rel_err": 0.00308937462978065, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 7.03125, "xla_out_mb": 2.34375, "compile_s": 0.40210924996063113, "inner": 200, "median_ms": 0.06317637045867741, "q1_ms": 0.06310004217084497, "q3_ms": 0.07228302012663335, "tflops": 46.68074437623746 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 64, "seq": 600, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.002047628164291382, "rel_err": 0.004772087559103966, "xla_temp_mb": 20.00048828125, "xla_arg_mb": 11.71875, "xla_out_mb": 7.031272888183594, "compile_s": 0.609733286080882, "inner": 200, "median_ms": 0.11402752483263612, "q1_ms": 0.11399537499528378, "q3_ms": 0.11458014778327197, "tflops": 90.52130189750234 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 64, "seq": 600, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.002047628164291382, "rel_err": 0.004772087559103966, "xla_temp_mb": 15.078651428222656, "xla_arg_mb": 11.71875, "xla_out_mb": 7.031272888183594, "compile_s": 0.649293112102896, "inner": 200, "median_ms": 0.10831573978066444, "q1_ms": 0.10830096493009478, "q3_ms": 0.10862863739021122, "tflops": 95.2947375967844 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 64, "seq": 600, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.002047628164291382, "rel_err": 0.004772087559103966, "xla_temp_mb": 15.078651428222656, "xla_arg_mb": 11.71875, "xla_out_mb": 7.031272888183594, "compile_s": 0.6556209339760244, "inner": 198, "median_ms": 0.10859125265569398, "q1_ms": 0.10856236342923961, "q3_ms": 0.10859784364436914, "tflops": 95.05296004575345 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 64, "seq": 600, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0025545358657836914, "rel_err": 0.005635968875139952, "xla_temp_mb": 109.93653869628906, "xla_arg_mb": 11.71875, "xla_out_mb": 7.031272888183594, "compile_s": 0.03584496700204909, "inner": 122, "median_ms": 0.24762199084717232, "q1_ms": 0.24629798698712324, "q3_ms": 0.24771384062764584, "tflops": 41.684181460161575 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 64, "seq": 600, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.002047628164291382, "rel_err": 0.004772087559103966, "xla_temp_mb": 9.522216796875, "xla_arg_mb": 11.71875, "xla_out_mb": 7.031272888183594, "compile_s": 0.7258995128795505, "inner": 200, "median_ms": 0.0958570244256407, "q1_ms": 0.09581279999110848, "q3_ms": 0.09613394970074296, "tflops": 107.68037148917591 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 64, "seq": 600, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.004250764846801758, "rel_err": 0.0016897449968382716, "xla_temp_mb": 10.0, "xla_arg_mb": 7.03125, "xla_out_mb": 2.34375, "compile_s": 0.5423730579204857, "inner": 200, "median_ms": 0.036055295495316386, "q1_ms": 0.03599955001845956, "q3_ms": 0.036076094838790596, "tflops": 40.8971825010156 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 64, "seq": 600, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.004250764846801758, "rel_err": 0.0016897449968382716, "xla_temp_mb": 10.0, "xla_arg_mb": 7.03125, "xla_out_mb": 2.34375, "compile_s": 0.5810728068463504, "inner": 200, "median_ms": 0.036241349298506975, "q1_ms": 0.03616817237343639, "q3_ms": 0.03630174498539418, "tflops": 40.687226842868874 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 64, "seq": 600, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.004250764846801758, "rel_err": 0.0016897449968382716, "xla_temp_mb": 10.0, "xla_arg_mb": 7.03125, "xla_out_mb": 2.34375, "compile_s": 0.5769043900072575, "inner": 200, "median_ms": 0.03850923967547715, "q1_ms": 0.03742836939636618, "q3_ms": 0.0386758369859308, "tflops": 38.29107020617201 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 64, "seq": 600, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.008065938949584961, "rel_err": 0.0032063359394669533, "xla_temp_mb": 65.91798400878906, "xla_arg_mb": 7.03125, "xla_out_mb": 2.34375, "compile_s": 0.6546666361391544, "inner": 200, "median_ms": 0.12268134974874556, "q1_ms": 0.1224144222214818, "q3_ms": 0.12274559703655541, "tflops": 12.019430850898978 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 64, "seq": 600, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.004250764846801758, "rel_err": 0.0016897449968382716, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 7.03125, "xla_out_mb": 2.34375, "compile_s": 0.36299878102727234, "inner": 200, "median_ms": 0.05076125496998429, "q1_ms": 0.050716979312710464, "q3_ms": 0.05140920169651508, "tflops": 29.048927195986863 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 64, "seq": 600, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.014807701110839844, "rel_err": 0.0035223369486629963, "xla_temp_mb": 20.00048828125, "xla_arg_mb": 11.71875, "xla_out_mb": 7.031272888183594, "compile_s": 0.6480019560549408, "inner": 200, "median_ms": 0.08669144939631224, "q1_ms": 0.08664187742397189, "q3_ms": 0.08729544992092997, "tflops": 59.53251486668005 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 64, "seq": 600, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.014807701110839844, "rel_err": 0.0035223369486629963, "xla_temp_mb": 15.07861328125, "xla_arg_mb": 11.71875, "xla_out_mb": 7.031272888183594, "compile_s": 0.6623648749664426, "inner": 200, "median_ms": 0.08538305526599288, "q1_ms": 0.08501575444824994, "q3_ms": 0.08573622966650873, "tflops": 60.44478010211884 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 64, "seq": 600, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.014807701110839844, "rel_err": 0.0035223369486629963, "xla_temp_mb": 15.07861328125, "xla_arg_mb": 11.71875, "xla_out_mb": 7.031272888183594, "compile_s": 0.6582728859502822, "inner": 200, "median_ms": 0.0857161555904895, "q1_ms": 0.08485900762025267, "q3_ms": 0.08582645503338426, "tflops": 60.2098865079365 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 64, "seq": 600, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.014807701110839844, "rel_err": 0.0036034160293638706, "xla_temp_mb": 109.93653869628906, "xla_arg_mb": 11.71875, "xla_out_mb": 7.031272888183594, "compile_s": 0.03911998798139393, "inner": 115, "median_ms": 0.3133883213867312, "q1_ms": 0.31334793033159297, "q3_ms": 0.31434879132101073, "tflops": 16.468258859050497 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 64, "seq": 600, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.014807701110839844, "rel_err": 0.0035223369486629963, "xla_temp_mb": 9.522216796875, "xla_arg_mb": 11.71875, "xla_out_mb": 7.031272888183594, "compile_s": 0.7875301770400256, "inner": 200, "median_ms": 0.09768972522579134, "q1_ms": 0.09721432346850634, "q3_ms": 0.10652009514160454, "tflops": 52.83012095766895 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 64, "seq": 1024, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0011154413223266602, "rel_err": 0.0039638252928853035, "xla_temp_mb": 0.0, "xla_arg_mb": 12.0, "xla_out_mb": 4.0, "compile_s": 0.4229674700181931, "inner": 200, "median_ms": 0.034384349128231406, "q1_ms": 0.03428167721722275, "q3_ms": 0.034651875612325966, "tflops": 249.82106131964557 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 64, "seq": 1024, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0011154413223266602, "rel_err": 0.0039638252928853035, "xla_temp_mb": 0.0, "xla_arg_mb": 12.0, "xla_out_mb": 4.0, "compile_s": 0.4270467881578952, "inner": 200, "median_ms": 0.034298700047656894, "q1_ms": 0.03426480514463037, "q3_ms": 0.03472172247711569, "tflops": 250.44490257836515 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 64, "seq": 1024, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0011154413223266602, "rel_err": 0.0039638252928853035, "xla_temp_mb": 0.0, "xla_arg_mb": 12.0, "xla_out_mb": 4.0, "compile_s": 0.4258241110946983, "inner": 200, "median_ms": 0.034288050374016166, "q1_ms": 0.03414000500924885, "q3_ms": 0.034786424366757274, "tflops": 250.52268934221877 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 64, "seq": 1024, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0012263059616088867, "rel_err": 0.004357793368399143, "xla_temp_mb": 192.0, "xla_arg_mb": 12.0, "xla_out_mb": 4.0, "compile_s": 0.21097166603431106, "inner": 153, "median_ms": 0.2303781639175867, "q1_ms": 0.2295634286246853, "q3_ms": 0.23038362402444573, "tflops": 37.28623601268426 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 64, "seq": 1024, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0011154413223266602, "rel_err": 0.0039638252928853035, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 12.0, "xla_out_mb": 4.0, "compile_s": 0.05436793388798833, "inner": 200, "median_ms": 0.05153810023330152, "q1_ms": 0.05014340975321829, "q3_ms": 0.05232307536061853, "tflops": 166.67154111454002 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 64, "seq": 1024, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0015159845352172852, "rel_err": 0.004941033665090799, "xla_temp_mb": 8.250244140625, "xla_arg_mb": 20.0, "xla_out_mb": 12.000022888183594, "compile_s": 0.5621474818326533, "inner": 188, "median_ms": 0.15017005337521117, "q1_ms": 0.13700141472385285, "q3_ms": 0.15314807186517468, "tflops": 200.20483709145998 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 64, "seq": 1024, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0015159845352172852, "rel_err": 0.004941033665090799, "xla_temp_mb": 4.25048828125, "xla_arg_mb": 20.0, "xla_out_mb": 12.000022888183594, "compile_s": 0.5501581460703164, "inner": 200, "median_ms": 0.14255574555136263, "q1_ms": 0.14131499454379082, "q3_ms": 0.1491851225728169, "tflops": 210.89834685875712 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 64, "seq": 1024, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0015159845352172852, "rel_err": 0.004941033665090799, "xla_temp_mb": 4.25048828125, "xla_arg_mb": 20.0, "xla_out_mb": 12.000022888183594, "compile_s": 0.5571218759287149, "inner": 200, "median_ms": 0.1466561306733638, "q1_ms": 0.14220574230421335, "q3_ms": 0.14889519778080285, "tflops": 205.00180206554754 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 64, "seq": 1024, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0015907883644104004, "rel_err": 0.004941033665090799, "xla_temp_mb": 320.0, "xla_arg_mb": 20.0, "xla_out_mb": 12.000022888183594, "compile_s": 0.0332969450391829, "inner": 69, "median_ms": 0.5886829559407804, "q1_ms": 0.5868160000502848, "q3_ms": 0.5889148974413241, "tflops": 51.071244323615886 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 64, "seq": 1024, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0015159845352172852, "rel_err": 0.004941033665090799, "xla_temp_mb": 16.250732421875, "xla_arg_mb": 20.0, "xla_out_mb": 12.000022888183594, "compile_s": 0.13883083290420473, "inner": 190, "median_ms": 0.13411509020155982, "q1_ms": 0.13080346826954106, "q3_ms": 0.13486335029531465, "tflops": 224.1714263981484 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 64, "seq": 1024, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.006038188934326172, "rel_err": 0.002593584591522813, "xla_temp_mb": 0.0, "xla_arg_mb": 12.0, "xla_out_mb": 4.0, "compile_s": 0.4977843570522964, "inner": 200, "median_ms": 0.033763599349185824, "q1_ms": 0.03373642801307142, "q3_ms": 0.034146204707212746, "tflops": 127.20703298191367 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 64, "seq": 1024, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.006038188934326172, "rel_err": 0.002593584591522813, "xla_temp_mb": 0.0, "xla_arg_mb": 12.0, "xla_out_mb": 4.0, "compile_s": 0.49566772300750017, "inner": 200, "median_ms": 0.03413255559280515, "q1_ms": 0.03412047983147204, "q3_ms": 0.03416717750951648, "tflops": 125.83198712801168 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 64, "seq": 1024, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.006038188934326172, "rel_err": 0.002593584591522813, "xla_temp_mb": 0.0, "xla_arg_mb": 12.0, "xla_out_mb": 4.0, "compile_s": 0.4977842268999666, "inner": 200, "median_ms": 0.033911250066012144, "q1_ms": 0.033852802007459104, "q3_ms": 0.033949530334211886, "tflops": 126.65316930633206 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 64, "seq": 1024, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.008179187774658203, "rel_err": 0.0035132081247866154, "xla_temp_mb": 192.0, "xla_arg_mb": 12.0, "xla_out_mb": 4.0, "compile_s": 0.6520703558344394, "inner": 126, "median_ms": 0.2869752934202552, "q1_ms": 0.28627580641547135, "q3_ms": 0.28770711524264203, "tflops": 14.966331229463444 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 64, "seq": 1024, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.006038188934326172, "rel_err": 0.002593584591522813, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 12.0, "xla_out_mb": 4.0, "compile_s": 0.05847171205095947, "inner": 200, "median_ms": 0.0652286154218018, "q1_ms": 0.05483569228090346, "q3_ms": 0.06536951230373234, "tflops": 65.84483310317918 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 64, "seq": 1024, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.012666940689086914, "rel_err": 0.0043540229089558125, "xla_temp_mb": 8.250244140625, "xla_arg_mb": 20.0, "xla_out_mb": 12.000022888183594, "compile_s": 0.57534185401164, "inner": 200, "median_ms": 0.10937329032458365, "q1_ms": 0.10769379267003387, "q3_ms": 0.11170250421855599, "tflops": 137.44110185758208 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 64, "seq": 1024, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.012666940689086914, "rel_err": 0.0043540229089558125, "xla_temp_mb": 4.25048828125, "xla_arg_mb": 20.0, "xla_out_mb": 12.000022888183594, "compile_s": 0.58005625102669, "inner": 200, "median_ms": 0.11379657546058297, "q1_ms": 0.10999165999237448, "q3_ms": 0.11405807570554316, "tflops": 132.09875143568746 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 64, "seq": 1024, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.012666940689086914, "rel_err": 0.0043540229089558125, "xla_temp_mb": 4.25048828125, "xla_arg_mb": 20.0, "xla_out_mb": 12.000022888183594, "compile_s": 0.5905174519866705, "inner": 200, "median_ms": 0.11245882487855852, "q1_ms": 0.10979018697980791, "q3_ms": 0.112516077933833, "tflops": 133.67012817565094 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 64, "seq": 1024, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.012666940689086914, "rel_err": 0.005755797028541565, "xla_temp_mb": 320.0, "xla_arg_mb": 20.0, "xla_out_mb": 12.000022888183594, "compile_s": 0.039001098833978176, "inner": 53, "median_ms": 0.7965000003646567, "q1_ms": 0.7957084706264004, "q3_ms": 0.7976975362255888, "tflops": 18.873051511761222 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 64, "seq": 1024, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.012666940689086914, "rel_err": 0.0043540229089558125, "xla_temp_mb": 16.250732421875, "xla_arg_mb": 20.0, "xla_out_mb": 12.000022888183594, "compile_s": 0.1318127829581499, "inner": 189, "median_ms": 0.12007378233921906, "q1_ms": 0.11998685397128894, "q3_ms": 0.12041484165897287, "tflops": 125.19290425558663 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 64, "seq": 1500, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0006500333547592163, "rel_err": 0.0023746348451822996, "xla_temp_mb": 24.0, "xla_arg_mb": 17.578125, "xla_out_mb": 5.859375, "compile_s": 0.49730531801469624, "inner": 200, "median_ms": 0.1280789845623076, "q1_ms": 0.12762803758960217, "q3_ms": 0.1285910332808271, "tflops": 143.91119716469362 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 64, "seq": 1500, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0006500333547592163, "rel_err": 0.0023746348451822996, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 17.578125, "xla_out_mb": 5.859375, "compile_s": 0.054749823873862624, "inner": 200, "median_ms": 0.066941810073331, "q1_ms": 0.06308451935183257, "q3_ms": 0.06697715725749731, "tflops": 275.3436152952658 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 64, "seq": 1500, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0006500333547592163, "rel_err": 0.0023746348451822996, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 17.578125, "xla_out_mb": 5.859375, "compile_s": 0.05098122381605208, "inner": 200, "median_ms": 0.06661640945822, "q1_ms": 0.06378039252012968, "q3_ms": 0.06879970489535481, "tflops": 276.6885839375664 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 64, "seq": 1500, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.000764697790145874, "rel_err": 0.0027935151010751724, "xla_temp_mb": 411.98731994628906, "xla_arg_mb": 17.578125, "xla_out_mb": 5.859375, "compile_s": 0.03004060499370098, "inner": 79, "median_ms": 0.500483002183558, "q1_ms": 0.4967441458160742, "q3_ms": 0.5005446448332712, "tflops": 36.828423581985795 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 64, "seq": 1500, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0006500333547592163, "rel_err": 0.0023746348451822996, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 17.578125, "xla_out_mb": 5.859375, "compile_s": 0.05626824009232223, "inner": 200, "median_ms": 0.06640736013650894, "q1_ms": 0.06370099203195423, "q3_ms": 0.06645096000283957, "tflops": 277.55959523327886 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 64, "seq": 1500, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0012814700603485107, "rel_err": 0.003724043257534504, "xla_temp_mb": 48.00048828125, "xla_arg_mb": 29.296875, "xla_out_mb": 17.578147888183594, "compile_s": 0.5683768140152097, "inner": 98, "median_ms": 0.42179358971058106, "q1_ms": 0.41543401549665293, "q3_ms": 0.4281139283973191, "tflops": 152.94684787472877 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 64, "seq": 1500, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0012814700603485107, "rel_err": 0.003724043257534504, "xla_temp_mb": 23.804443359375, "xla_arg_mb": 29.296875, "xla_out_mb": 17.578147888183594, "compile_s": 0.1321250929031521, "inner": 133, "median_ms": 0.2650814218432608, "q1_ms": 0.2586698234851535, "q3_ms": 0.26891825114958046, "tflops": 243.36673445996948 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 64, "seq": 1500, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0012814700603485107, "rel_err": 0.003724043257534504, "xla_temp_mb": 23.804443359375, "xla_arg_mb": 29.296875, "xla_out_mb": 17.578147888183594, "compile_s": 0.13054726598784328, "inner": 129, "median_ms": 0.2667661323103794, "q1_ms": 0.25601132660959813, "q3_ms": 0.26734857711681104, "tflops": 241.82979841286982 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 64, "seq": 1500, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0013823211193084717, "rel_err": 0.004368611145764589, "xla_temp_mb": 686.82861328125, "xla_arg_mb": 29.296875, "xla_out_mb": 17.578147888183594, "compile_s": 0.03328278404660523, "inner": 33, "median_ms": 1.3868718154051087, "q1_ms": 1.3862089112852558, "q3_ms": 1.3886479992475926, "tflops": 46.516195140324406 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 64, "seq": 1500, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0012814700603485107, "rel_err": 0.003724043257534504, "xla_temp_mb": 23.804443359375, "xla_arg_mb": 29.296875, "xla_out_mb": 17.578147888183594, "compile_s": 0.1302407579496503, "inner": 133, "median_ms": 0.26493367591970846, "q1_ms": 0.2624478910636521, "q3_ms": 0.2668110786804131, "tflops": 243.50245311793125 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 64, "seq": 1500, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.007173061370849609, "rel_err": 0.0024814915377646685, "xla_temp_mb": 24.0, "xla_arg_mb": 17.578125, "xla_out_mb": 5.859375, "compile_s": 0.5361827171873301, "inner": 200, "median_ms": 0.09341313038021326, "q1_ms": 0.09326638479251415, "q3_ms": 0.09386808262206614, "tflops": 98.6585072407779 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 64, "seq": 1500, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.007173061370849609, "rel_err": 0.0024814915377646685, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 17.578125, "xla_out_mb": 5.859375, "compile_s": 0.05724893603473902, "inner": 200, "median_ms": 0.06410026922821999, "q1_ms": 0.06400751473847777, "q3_ms": 0.06459881435148418, "tflops": 143.77474714166533 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 64, "seq": 1500, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.007173061370849609, "rel_err": 0.0024814915377646685, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 17.578125, "xla_out_mb": 5.859375, "compile_s": 0.052480840124189854, "inner": 200, "median_ms": 0.06029703072272241, "q1_ms": 0.05997252941597253, "q3_ms": 0.06055142788682133, "tflops": 152.84334716878573 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 64, "seq": 1500, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.011619091033935547, "rel_err": 0.004019577521830797, "xla_temp_mb": 411.98731994628906, "xla_arg_mb": 17.578125, "xla_out_mb": 5.859375, "compile_s": 0.45758337201550603, "inner": 83, "median_ms": 0.5046537465763739, "q1_ms": 0.5034393440169025, "q3_ms": 0.5048790438863707, "tflops": 18.262026315116753 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 64, "seq": 1500, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.007173061370849609, "rel_err": 0.0024814915377646685, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 17.578125, "xla_out_mb": 5.859375, "compile_s": 0.054428393952548504, "inner": 200, "median_ms": 0.04865206079557538, "q1_ms": 0.04860226239543408, "q3_ms": 0.04885073518380523, "tflops": 189.42671388008586 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 64, "seq": 1500, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.014818668365478516, "rel_err": 0.005191786680370569, "xla_temp_mb": 48.00048828125, "xla_arg_mb": 29.296875, "xla_out_mb": 17.578147888183594, "compile_s": 0.6177412141114473, "inner": 132, "median_ms": 0.2680770368456389, "q1_ms": 0.25970740049766994, "q3_ms": 0.27272350041929516, "tflops": 120.32362181984759 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 64, "seq": 1500, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.014818668365478516, "rel_err": 0.005191786680370569, "xla_temp_mb": 23.804443359375, "xla_arg_mb": 29.296875, "xla_out_mb": 17.578147888183594, "compile_s": 0.1332472988869995, "inner": 146, "median_ms": 0.20483091684083826, "q1_ms": 0.20193576107873287, "q3_ms": 0.20545992477197353, "tflops": 157.4762272097048 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 64, "seq": 1500, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.014818668365478516, "rel_err": 0.005191786680370569, "xla_temp_mb": 23.804443359375, "xla_arg_mb": 29.296875, "xla_out_mb": 17.578147888183594, "compile_s": 0.13333006901666522, "inner": 144, "median_ms": 0.20380261088980156, "q1_ms": 0.2010240803226932, "q3_ms": 0.20481858276374018, "tflops": 158.27078887346144 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 64, "seq": 1500, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.014818668365478516, "rel_err": 0.004564434289932251, "xla_temp_mb": 686.82861328125, "xla_arg_mb": 29.296875, "xla_out_mb": 17.578147888183594, "compile_s": 0.040124944876879454, "inner": 29, "median_ms": 1.5571338349375232, "q1_ms": 1.5524936727537164, "q3_ms": 1.5574439814121561, "tflops": 20.714982409520506 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 64, "seq": 1500, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.014818668365478516, "rel_err": 0.005191786680370569, "xla_temp_mb": 23.804443359375, "xla_arg_mb": 29.296875, "xla_out_mb": 17.578147888183594, "compile_s": 0.1321197219658643, "inner": 139, "median_ms": 0.20165668411863794, "q1_ms": 0.1986099711934225, "q3_ms": 0.20707498150442144, "tflops": 159.95502524986114 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 64, "seq": 2048, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.000682830810546875, "rel_err": 0.0027175149880349636, "xla_temp_mb": 0.0, "xla_arg_mb": 24.0, "xla_out_mb": 8.0, "compile_s": 0.42018455802462995, "inner": 200, "median_ms": 0.1296285295393318, "q1_ms": 0.12453486793674529, "q3_ms": 0.13064770260825753, "tflops": 265.0630882731304 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 64, "seq": 2048, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.000682830810546875, "rel_err": 0.0027175149880349636, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 24.0, "xla_out_mb": 8.0, "compile_s": 0.05532688181847334, "inner": 200, "median_ms": 0.10517940041609108, "q1_ms": 0.09672415035311133, "q3_ms": 0.10835523717105389, "tflops": 326.6774504520127 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 64, "seq": 2048, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.000682830810546875, "rel_err": 0.0027175149880349636, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 24.0, "xla_out_mb": 8.0, "compile_s": 0.05246266000904143, "inner": 200, "median_ms": 0.10416459990665317, "q1_ms": 0.10078108753077686, "q3_ms": 0.10712809220422059, "tflops": 329.86003305145306 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 64, "seq": 2048, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.000711590051651001, "rel_err": 0.002831970574334264, "xla_temp_mb": 768.0, "xla_arg_mb": 24.0, "xla_out_mb": 8.0, "compile_s": 0.22422304889187217, "inner": 55, "median_ms": 0.8099185980179093, "q1_ms": 0.8093511452898383, "q3_ms": 0.810960953293199, "tflops": 42.42369350708529 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 64, "seq": 2048, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.000682830810546875, "rel_err": 0.0027175149880349636, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 24.0, "xla_out_mb": 8.0, "compile_s": 0.053120998200029135, "inner": 200, "median_ms": 0.09637472569011152, "q1_ms": 0.09627612482290715, "q3_ms": 0.10356118029449135, "tflops": 356.52229484400453 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 64, "seq": 2048, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0015923678874969482, "rel_err": 0.00450897915288806, "xla_temp_mb": 16.500244140625, "xla_arg_mb": 40.0, "xla_out_mb": 24.000022888183594, "compile_s": 0.5320834890007973, "inner": 78, "median_ms": 0.5679877462964028, "q1_ms": 0.5647563390457668, "q3_ms": 0.5826235954793026, "tflops": 211.72830764775534 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 64, "seq": 2048, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0015923678874969482, "rel_err": 0.00450897915288806, "xla_temp_mb": 32.500732421875, "xla_arg_mb": 40.0, "xla_out_mb": 24.000022888183594, "compile_s": 0.13277606084011495, "inner": 97, "median_ms": 0.40614306350651475, "q1_ms": 0.4027564217787736, "q3_ms": 0.41790818435514404, "tflops": 296.1003033013045 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 64, "seq": 2048, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0015923678874969482, "rel_err": 0.00450897915288806, "xla_temp_mb": 32.500732421875, "xla_arg_mb": 40.0, "xla_out_mb": 24.000022888183594, "compile_s": 0.13106274488382041, "inner": 96, "median_ms": 0.4091828959644772, "q1_ms": 0.40263431340766453, "q3_ms": 0.4126208538461166, "tflops": 293.9005649406229 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 64, "seq": 2048, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0015923678874969482, "rel_err": 0.004729839973151684, "xla_temp_mb": 1280.0, "xla_arg_mb": 40.0, "xla_out_mb": 24.000022888183594, "compile_s": 0.03454238106496632, "inner": 22, "median_ms": 2.1262457276779143, "q1_ms": 2.1235434793528505, "q3_ms": 2.1268321116539566, "tflops": 56.559353757919446 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 64, "seq": 2048, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0015923678874969482, "rel_err": 0.00450897915288806, "xla_temp_mb": 32.500732421875, "xla_arg_mb": 40.0, "xla_out_mb": 24.000022888183594, "compile_s": 0.13114561489783227, "inner": 99, "median_ms": 0.4156270920477732, "q1_ms": 0.4049195011964802, "q3_ms": 0.4189729436554692, "tflops": 289.34370879311473 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 64, "seq": 2048, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.008076190948486328, "rel_err": 0.002793925581499934, "xla_temp_mb": 0.0, "xla_arg_mb": 24.0, "xla_out_mb": 8.0, "compile_s": 0.5005266789812595, "inner": 200, "median_ms": 0.09738532011397183, "q1_ms": 0.09663472534157336, "q3_ms": 0.09914409252814949, "tflops": 176.4112821510889 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 64, "seq": 2048, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.008076190948486328, "rel_err": 0.002793925581499934, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 24.0, "xla_out_mb": 8.0, "compile_s": 0.05727774603292346, "inner": 200, "median_ms": 0.07088224519975483, "q1_ms": 0.06726518564391881, "q3_ms": 0.0735415902454406, "tflops": 242.37196685270098 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 64, "seq": 2048, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.008076190948486328, "rel_err": 0.002793925581499934, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 24.0, "xla_out_mb": 8.0, "compile_s": 0.05293866898864508, "inner": 200, "median_ms": 0.06708091008476913, "q1_ms": 0.06705885753035545, "q3_ms": 0.07131869497243315, "tflops": 256.10668016116745 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 64, "seq": 2048, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.008959054946899414, "rel_err": 0.0030993486288934946, "xla_temp_mb": 768.0, "xla_arg_mb": 24.0, "xla_out_mb": 8.0, "compile_s": 0.5502286360133439, "inner": 48, "median_ms": 0.9505414588299269, "q1_ms": 0.9456578142514143, "q3_ms": 0.9525287290064929, "tflops": 18.073771558736254 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 64, "seq": 2048, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.008076190948486328, "rel_err": 0.002793925581499934, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 24.0, "xla_out_mb": 8.0, "compile_s": 0.05629051895812154, "inner": 200, "median_ms": 0.07195798913016915, "q1_ms": 0.06813655490987003, "q3_ms": 0.07495368947274983, "tflops": 238.74860028289973 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 64, "seq": 2048, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.016491413116455078, "rel_err": 0.008016865700483322, "xla_temp_mb": 16.500244140625, "xla_arg_mb": 40.0, "xla_out_mb": 24.000022888183594, "compile_s": 0.5763023220933974, "inner": 106, "median_ms": 0.38035617007891526, "q1_ms": 0.3691689486896514, "q3_ms": 0.38737554722553436, "tflops": 158.08746347278787 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 64, "seq": 2048, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.016491413116455078, "rel_err": 0.008016865700483322, "xla_temp_mb": 32.500732421875, "xla_arg_mb": 40.0, "xla_out_mb": 24.000022888183594, "compile_s": 0.13139121397398412, "inner": 119, "median_ms": 0.30662382350248446, "q1_ms": 0.300894769206017, "q3_ms": 0.30723302977300493, "tflops": 196.10199056666838 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 64, "seq": 2048, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.016491413116455078, "rel_err": 0.008016865700483322, "xla_temp_mb": 32.500732421875, "xla_arg_mb": 40.0, "xla_out_mb": 24.000022888183594, "compile_s": 0.13364165811799467, "inner": 117, "median_ms": 0.30671630818874407, "q1_ms": 0.30212259739955777, "q3_ms": 0.3081777523088659, "tflops": 196.04285960235956 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 64, "seq": 2048, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.01198720932006836, "rel_err": 0.0051156082190573215, "xla_temp_mb": 1280.0, "xla_arg_mb": 40.0, "xla_out_mb": 24.000022888183594, "compile_s": 0.040311614982783794, "inner": 17, "median_ms": 2.703152884564856, "q1_ms": 2.6948849382974647, "q3_ms": 2.7035952011561566, "tflops": 22.244225432953804 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 64, "seq": 2048, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.016491413116455078, "rel_err": 0.008016865700483322, "xla_temp_mb": 32.500732421875, "xla_arg_mb": 40.0, "xla_out_mb": 24.000022888183594, "compile_s": 0.1314633151050657, "inner": 110, "median_ms": 0.30844275500964036, "q1_ms": 0.29865378217602323, "q3_ms": 0.30926029113205994, "tflops": 194.94554878463802 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 64, "seq": 3000, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0005548298358917236, "rel_err": 0.0030985635239630938, "xla_temp_mb": 48.0, "xla_arg_mb": 35.15625, "xla_out_mb": 11.71875, "compile_s": 0.49461298482492566, "inner": 109, "median_ms": 0.36728299231430805, "q1_ms": 0.36582005009431084, "q3_ms": 0.369445597026668, "tflops": 200.73894392829965 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 64, "seq": 3000, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0005548298358917236, "rel_err": 0.0030985635239630938, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 35.15625, "xla_out_mb": 11.71875, "compile_s": 0.056779767852276564, "inner": 136, "median_ms": 0.2517208974014091, "q1_ms": 0.236822445431779, "q3_ms": 0.2527194592403248, "tflops": 292.8958253411474 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 64, "seq": 3000, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0005548298358917236, "rel_err": 0.0030985635239630938, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 35.15625, "xla_out_mb": 11.71875, "compile_s": 0.05281553999520838, "inner": 143, "median_ms": 0.24073413358284876, "q1_ms": 0.2396636849295322, "q3_ms": 0.2472949193113229, "tflops": 306.26317465955225 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 64, "seq": 3000, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0006076693534851074, "rel_err": 0.0033936568070203066, "xla_temp_mb": 1647.949234008789, "xla_arg_mb": 35.15625, "xla_out_mb": 11.71875, "compile_s": 0.01861283602192998, "inner": 27, "median_ms": 1.715882888270749, "q1_ms": 1.712611035650803, "q3_ms": 1.7192015614084624, "tflops": 42.96796739683231 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 64, "seq": 3000, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0005548298358917236, "rel_err": 0.0030985635239630938, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 35.15625, "xla_out_mb": 11.71875, "compile_s": 0.056368298130109906, "inner": 140, "median_ms": 0.2430879494308361, "q1_ms": 0.23922017509383814, "q3_ms": 0.25274434592574835, "tflops": 303.2976343443847 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 64, "seq": 3000, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0008320212364196777, "rel_err": 0.003266180632635951, "xla_temp_mb": 96.00048828125, "xla_arg_mb": 58.59375, "xla_out_mb": 35.156272888183594, "compile_s": 0.5794230629689991, "inner": 32, "median_ms": 1.5003231892478652, "q1_ms": 1.4548216095136013, "q3_ms": 1.5247103292495012, "tflops": 171.99494205602684 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 64, "seq": 3000, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0008320212364196777, "rel_err": 0.003266180632635951, "xla_temp_mb": 47.608154296875, "xla_arg_mb": 58.59375, "xla_out_mb": 35.156272888183594, "compile_s": 0.13375873700715601, "inner": 53, "median_ms": 0.8858961474923593, "q1_ms": 0.8671596988965318, "q3_ms": 0.8994389329773357, "tflops": 291.28470727684885 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 64, "seq": 3000, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0008320212364196777, "rel_err": 0.003266180632635951, "xla_temp_mb": 47.608154296875, "xla_arg_mb": 58.59375, "xla_out_mb": 35.156272888183594, "compile_s": 0.1344597260467708, "inner": 53, "median_ms": 0.8772912628526958, "q1_ms": 0.8748784432855418, "q3_ms": 0.8884376499204423, "tflops": 294.1417644590498 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 64, "seq": 3000, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0011211037635803223, "rel_err": 0.004401002544909716, "xla_temp_mb": 2746.9482421875, "xla_arg_mb": 58.59375, "xla_out_mb": 35.156272888183594, "compile_s": 0.03855330985970795, "inner": 9, "median_ms": 5.248559457767342, "q1_ms": 5.244629509333107, "q3_ms": 5.251248897467223, "tflops": 49.16549046960206 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 64, "seq": 3000, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0008320212364196777, "rel_err": 0.003266180632635951, "xla_temp_mb": 47.608154296875, "xla_arg_mb": 58.59375, "xla_out_mb": 35.156272888183594, "compile_s": 0.1338913170620799, "inner": 52, "median_ms": 0.8690788661344693, "q1_ms": 0.8684478077786759, "q3_ms": 0.8880751534329299, "tflops": 296.9212692373458 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 64, "seq": 3000, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.006508588790893555, "rel_err": 0.002251619938760996, "xla_temp_mb": 48.0, "xla_arg_mb": 35.15625, "xla_out_mb": 11.71875, "compile_s": 0.5371573241427541, "inner": 141, "median_ms": 0.2554050144732844, "q1_ms": 0.2529778218900798, "q3_ms": 0.2605811707635826, "tflops": 144.33545901995598 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 64, "seq": 3000, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.006508588790893555, "rel_err": 0.002251619938760996, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 35.15625, "xla_out_mb": 11.71875, "compile_s": 0.05634375894442201, "inner": 172, "median_ms": 0.14301424380391836, "q1_ms": 0.13791384895570405, "q3_ms": 0.14608949463446297, "tflops": 257.76453463294814 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 64, "seq": 3000, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.006508588790893555, "rel_err": 0.002251619938760996, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 35.15625, "xla_out_mb": 11.71875, "compile_s": 0.05419873516075313, "inner": 186, "median_ms": 0.14471184332124007, "q1_ms": 0.14013919656375246, "q3_ms": 0.14607148978518703, "tflops": 254.74072580339583 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 64, "seq": 3000, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.007374167442321777, "rel_err": 0.0025510634295642376, "xla_temp_mb": 1647.949234008789, "xla_arg_mb": 35.15625, "xla_out_mb": 11.71875, "compile_s": 0.4470303209964186, "inner": 26, "median_ms": 1.7746041546790647, "q1_ms": 1.7740036128088832, "q3_ms": 1.7754066771326158, "tflops": 20.77308333962895 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 64, "seq": 3000, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.006508588790893555, "rel_err": 0.002251619938760996, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 35.15625, "xla_out_mb": 11.71875, "compile_s": 0.05737573583610356, "inner": 179, "median_ms": 0.14345103864396752, "q1_ms": 0.1382520090653327, "q3_ms": 0.14591681856828526, "tflops": 256.97966601338527 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 64, "seq": 3000, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.01353597640991211, "rel_err": 0.0039185951463878155, "xla_temp_mb": 96.00048828125, "xla_arg_mb": 58.59375, "xla_out_mb": 35.156272888183594, "compile_s": 0.6216016530524939, "inner": 52, "median_ms": 0.8971941737959591, "q1_ms": 0.8688017394608603, "q3_ms": 0.9028062721391996, "tflops": 143.8083346597197 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 64, "seq": 3000, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.01353597640991211, "rel_err": 0.0039185951463878155, "xla_temp_mb": 47.608154296875, "xla_arg_mb": 58.59375, "xla_out_mb": 35.156272888183594, "compile_s": 0.1344268750399351, "inner": 73, "median_ms": 0.5742794521510193, "q1_ms": 0.5581366162976785, "q3_ms": 0.5748158205963978, "tflops": 224.6711065783881 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 64, "seq": 3000, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.01353597640991211, "rel_err": 0.0039185951463878155, "xla_temp_mb": 47.608154296875, "xla_arg_mb": 58.59375, "xla_out_mb": 35.156272888183594, "compile_s": 0.1347675141878426, "inner": 71, "median_ms": 0.5738483512328124, "q1_ms": 0.5626939565048251, "q3_ms": 0.5740876552242209, "tflops": 224.83988970747166 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 64, "seq": 3000, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.01353597640991211, "rel_err": 0.00468168780207634, "xla_temp_mb": 2746.9482421875, "xla_arg_mb": 58.59375, "xla_out_mb": 35.156272888183594, "compile_s": 0.04391791392117739, "inner": 8, "median_ms": 5.535922886338085, "q1_ms": 5.529904199647717, "q3_ms": 5.542023514863104, "tflops": 23.30668303173332 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 64, "seq": 3000, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.01353597640991211, "rel_err": 0.0039185951463878155, "xla_temp_mb": 47.608154296875, "xla_arg_mb": 58.59375, "xla_out_mb": 35.156272888183594, "compile_s": 0.13477759389206767, "inner": 72, "median_ms": 0.5726807220425042, "q1_ms": 0.5555184163515352, "q3_ms": 0.5763149513061055, "tflops": 225.29831201551056 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 64, "seq": 4096, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00045752525329589844, "rel_err": 0.0028358343988656998, "xla_temp_mb": 0.0, "xla_arg_mb": 48.0, "xla_out_mb": 16.0, "compile_s": 0.420315558090806, "inner": 91, "median_ms": 0.4968916495832113, "q1_ms": 0.4957884788042405, "q3_ms": 0.49957972541860823, "tflops": 276.597430420259 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 64, "seq": 4096, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00045752525329589844, "rel_err": 0.0028358343988656998, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 48.0, "xla_out_mb": 16.0, "compile_s": 0.05306143802590668, "inner": 102, "median_ms": 0.3784900878136065, "q1_ms": 0.36657199049003275, "q3_ms": 0.39375569059641335, "tflops": 363.12431394421094 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 64, "seq": 4096, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00045752525329589844, "rel_err": 0.0028358343988656998, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 48.0, "xla_out_mb": 16.0, "compile_s": 0.052137140883132815, "inner": 99, "median_ms": 0.37978185020913985, "q1_ms": 0.3696534890832022, "q3_ms": 0.39633148881333946, "tflops": 361.88920928241976 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 64, "seq": 4096, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0006755292415618896, "rel_err": 0.0041870675049722195, "xla_temp_mb": 3072.0, "xla_arg_mb": 48.0, "xla_out_mb": 16.0, "compile_s": 0.2624216100666672, "inner": 16, "median_ms": 3.0240288615459576, "q1_ms": 3.0234004370868206, "q3_ms": 3.02438728249399, "tflops": 45.44895560346532 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 64, "seq": 4096, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00045752525329589844, "rel_err": 0.0028358343988656998, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 48.0, "xla_out_mb": 16.0, "compile_s": 0.05437371414154768, "inner": 102, "median_ms": 0.39318759752181814, "q1_ms": 0.37977523421503456, "q3_ms": 0.39761195612121736, "tflops": 349.5505818043344 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 64, "seq": 4096, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0007363855838775635, "rel_err": 0.003625117940828204, "xla_temp_mb": 33.000244140625, "xla_arg_mb": 80.0, "xla_out_mb": 48.000022888183594, "compile_s": 0.5460895490832627, "inner": 23, "median_ms": 2.1974519603764233, "q1_ms": 2.172768996704532, "q3_ms": 2.253900940084587, "tflops": 218.9064178993922 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 64, "seq": 4096, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0007363855838775635, "rel_err": 0.003625117940828204, "xla_temp_mb": 65.000732421875, "xla_arg_mb": 80.0, "xla_out_mb": 48.000022888183594, "compile_s": 0.13070700597018003, "inner": 33, "median_ms": 1.433954087340019, "q1_ms": 1.414836123303482, "q3_ms": 1.4551521779856447, "tflops": 335.46146379367076 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 64, "seq": 4096, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0007363855838775635, "rel_err": 0.003625117940828204, "xla_temp_mb": 65.000732421875, "xla_arg_mb": 80.0, "xla_out_mb": 48.000022888183594, "compile_s": 0.1317147130612284, "inner": 33, "median_ms": 1.4528113092775599, "q1_ms": 1.413215043474779, "q3_ms": 1.47344578748964, "tflops": 331.10723607403986 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 64, "seq": 4096, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0008107870817184448, "rel_err": 0.0038248642813414335, "xla_temp_mb": 5120.000015258789, "xla_arg_mb": 80.0, "xla_out_mb": 48.000022888183594, "compile_s": 0.040669484063982964, "inner": 6, "median_ms": 8.282971180354556, "q1_ms": 8.280047254326444, "q3_ms": 8.29048043427368, "tflops": 58.07533633497552 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 64, "seq": 4096, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0007363855838775635, "rel_err": 0.003625117940828204, "xla_temp_mb": 65.000732421875, "xla_arg_mb": 80.0, "xla_out_mb": 48.000022888183594, "compile_s": 0.13484337413683534, "inner": 33, "median_ms": 1.4123326365017528, "q1_ms": 1.4084309840461975, "q3_ms": 1.4174863816041388, "tflops": 340.59705533923835 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 64, "seq": 4096, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.004129528999328613, "rel_err": 0.0016415518475696445, "xla_temp_mb": 0.0, "xla_arg_mb": 48.0, "xla_out_mb": 16.0, "compile_s": 0.5003275491762906, "inner": 125, "median_ms": 0.3316164966672659, "q1_ms": 0.3294507758691907, "q3_ms": 0.33533396013081074, "tflops": 207.22574849752147 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 64, "seq": 4096, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.004129528999328613, "rel_err": 0.0016415518475696445, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 48.0, "xla_out_mb": 16.0, "compile_s": 0.054862982826307416, "inner": 138, "median_ms": 0.23780765173875767, "q1_ms": 0.22488656922823924, "q3_ms": 0.23931667016333213, "tflops": 288.9708393886813 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 64, "seq": 4096, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.004129528999328613, "rel_err": 0.0016415518475696445, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 48.0, "xla_out_mb": 16.0, "compile_s": 0.05494763283059001, "inner": 147, "median_ms": 0.2320342393097829, "q1_ms": 0.22583016983176374, "q3_ms": 0.2368578598612831, "tflops": 296.16093271585845 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 64, "seq": 4096, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0068624019622802734, "rel_err": 0.002727911341935396, "xla_temp_mb": 3072.0, "xla_arg_mb": 48.0, "xla_out_mb": 16.0, "compile_s": 0.5487603510264307, "inner": 13, "median_ms": 3.736681609342878, "q1_ms": 3.72361780430835, "q3_ms": 3.7486861871842008, "tflops": 18.390508991769522 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 64, "seq": 4096, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.004129528999328613, "rel_err": 0.0016415518475696445, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 48.0, "xla_out_mb": 16.0, "compile_s": 0.05748815601691604, "inner": 157, "median_ms": 0.23530524929949811, "q1_ms": 0.22553203176635844, "q3_ms": 0.2380772834869137, "tflops": 292.04395966761194 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 64, "seq": 4096, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.015132904052734375, "rel_err": 0.003882892895489931, "xla_temp_mb": 33.000244140625, "xla_arg_mb": 80.0, "xla_out_mb": 48.000022888183594, "compile_s": 0.5757806529290974, "inner": 35, "median_ms": 1.3794780575803347, "q1_ms": 1.3759300711431672, "q3_ms": 1.3852450298145413, "tflops": 174.35447215295287 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 64, "seq": 4096, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.015132904052734375, "rel_err": 0.003882892895489931, "xla_temp_mb": 65.000732421875, "xla_arg_mb": 80.0, "xla_out_mb": 48.000022888183594, "compile_s": 0.13595590111799538, "inner": 50, "median_ms": 0.9253255603834987, "q1_ms": 0.9119962994009256, "q3_ms": 0.9301450382918119, "tflops": 259.9281581244961 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 64, "seq": 4096, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.015132904052734375, "rel_err": 0.003882892895489931, "xla_temp_mb": 65.000732421875, "xla_arg_mb": 80.0, "xla_out_mb": 48.000022888183594, "compile_s": 0.13564767292700708, "inner": 50, "median_ms": 0.9105873806402087, "q1_ms": 0.9003061288967729, "q3_ms": 0.9305737400427461, "tflops": 264.13518756091077 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 64, "seq": 4096, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.015132904052734375, "rel_err": 0.003882892895489931, "xla_temp_mb": 5120.000015258789, "xla_arg_mb": 80.0, "xla_out_mb": 48.000022888183594, "compile_s": 0.049433848122134805, "inner": 4, "median_ms": 9.989916521590203, "q1_ms": 9.976121509680524, "q3_ms": 9.996584005421028, "tflops": 24.07609393493852 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 64, "seq": 4096, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.015132904052734375, "rel_err": 0.003882892895489931, "xla_temp_mb": 65.000732421875, "xla_arg_mb": 80.0, "xla_out_mb": 48.000022888183594, "compile_s": 0.13809308595955372, "inner": 50, "median_ms": 0.9262433601543307, "q1_ms": 0.9076946973800659, "q3_ms": 0.9314271295443177, "tflops": 259.6705994587911 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 64, "seq": 8192, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00031141191720962524, "rel_err": 0.0025766564067453146, "xla_temp_mb": 0.0, "xla_arg_mb": 96.0, "xla_out_mb": 32.0, "compile_s": 0.42445946601219475, "inner": 25, "median_ms": 1.919437712058425, "q1_ms": 1.916142520494759, "q3_ms": 1.9296598620712757, "tflops": 286.41503208688977 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 64, "seq": 8192, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00031141191720962524, "rel_err": 0.0025766564067453146, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 96.0, "xla_out_mb": 32.0, "compile_s": 0.05602489993907511, "inner": 33, "median_ms": 1.5016108765407945, "q1_ms": 1.474079270282704, "q3_ms": 1.5182970338642146, "tflops": 366.11070316329364 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 64, "seq": 8192, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00031141191720962524, "rel_err": 0.0025766564067453146, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 96.0, "xla_out_mb": 32.0, "compile_s": 0.05599656002596021, "inner": 33, "median_ms": 1.491209061703447, "q1_ms": 1.4875948316220082, "q3_ms": 1.4913999651925582, "tflops": 368.6644803915016 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 64, "seq": 8192, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0003448277711868286, "rel_err": 0.0028531427960842848, "xla_temp_mb": 12288.0, "xla_arg_mb": 96.0, "xla_out_mb": 32.0, "compile_s": 0.029638994950801134, "inner": 4, "median_ms": 11.722026509232819, "q1_ms": 11.720011505531147, "q3_ms": 11.722750234184787, "tflops": 46.89938326406159 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 64, "seq": 8192, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00031141191720962524, "rel_err": 0.0025766564067453146, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 96.0, "xla_out_mb": 32.0, "compile_s": 0.057911094976589084, "inner": 32, "median_ms": 1.4877476278343238, "q1_ms": 1.4452755349338986, "q3_ms": 1.567193499795394, "tflops": 369.5222251426241 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 64, "seq": 8192, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0006060153245925903, "rel_err": 0.0036204783245921135, "xla_temp_mb": 66.000244140625, "xla_arg_mb": 160.0, "xla_out_mb": 96.0000228881836, "compile_s": 0.5423598398920149, "inner": 6, "median_ms": 8.490913004303971, "q1_ms": 8.4445207612589, "q3_ms": 8.744112836817902, "tflops": 226.61230277976787 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 64, "seq": 8192, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0006060153245925903, "rel_err": 0.0036204783245921135, "xla_temp_mb": 130.000732421875, "xla_arg_mb": 160.0, "xla_out_mb": 96.0000228881836, "compile_s": 0.13651272095739841, "inner": 9, "median_ms": 5.49245910305116, "q1_ms": 5.434397724457085, "q3_ms": 5.577927274215552, "tflops": 350.32492960013167 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 64, "seq": 8192, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0006060153245925903, "rel_err": 0.0036204783245921135, "xla_temp_mb": 130.000732421875, "xla_arg_mb": 160.0, "xla_out_mb": 96.0000228881836, "compile_s": 0.13655729009769857, "inner": 9, "median_ms": 5.460833664983511, "q1_ms": 5.262114896645976, "q3_ms": 5.476949762346017, "tflops": 352.35377355406223 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 64, "seq": 8192, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0006311088800430298, "rel_err": 0.004330759402364492, "xla_temp_mb": 20480.0, "xla_arg_mb": 160.0, "xla_out_mb": 96.0000228881836, "compile_s": 0.06460637110285461, "inner": 1, "median_ms": 32.366849947720766, "q1_ms": 32.34734549187124, "q3_ms": 32.38391503691673, "tflops": 59.448026351526245 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 64, "seq": 8192, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0006060153245925903, "rel_err": 0.0036204783245921135, "xla_temp_mb": 130.000732421875, "xla_arg_mb": 160.0, "xla_out_mb": 96.0000228881836, "compile_s": 0.1376003078185022, "inner": 9, "median_ms": 5.3235096598251, "q1_ms": 5.272343677158156, "q3_ms": 5.573592886018257, "tflops": 361.443008759604 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 64, "seq": 8192, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.007561683654785156, "rel_err": 0.003005886683240533, "xla_temp_mb": 0.0, "xla_arg_mb": 96.0, "xla_out_mb": 32.0, "compile_s": 0.5019453940913081, "inner": 43, "median_ms": 1.114728768507755, "q1_ms": 1.1055585576307982, "q3_ms": 1.1352169554853857, "tflops": 246.58725486377156 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 64, "seq": 8192, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.007561683654785156, "rel_err": 0.003005886683240533, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 96.0, "xla_out_mb": 32.0, "compile_s": 0.05523768602870405, "inner": 56, "median_ms": 0.7995604662158128, "q1_ms": 0.7923019043768623, "q3_ms": 0.8228587061499378, "tflops": 343.78626577793625 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 64, "seq": 8192, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.007561683654785156, "rel_err": 0.003005886683240533, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 96.0, "xla_out_mb": 32.0, "compile_s": 0.0552148160059005, "inner": 57, "median_ms": 0.8255010348205504, "q1_ms": 0.8082373941872726, "q3_ms": 0.8505399220443347, "tflops": 332.9831161311066 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 64, "seq": 8192, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.008063316345214844, "rel_err": 0.003205293556675315, "xla_temp_mb": 12288.0, "xla_arg_mb": 96.0, "xla_out_mb": 32.0, "compile_s": 0.4955595429055393, "inner": 3, "median_ms": 14.991421640540162, "q1_ms": 14.970621831404667, "q3_ms": 15.000418371831376, "tflops": 18.335679799750853 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 64, "seq": 8192, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.007561683654785156, "rel_err": 0.003005886683240533, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 96.0, "xla_out_mb": 32.0, "compile_s": 0.05901136598549783, "inner": 56, "median_ms": 0.8316203768897269, "q1_ms": 0.7954952847545168, "q3_ms": 0.8471008685384211, "tflops": 330.53291451569237 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 64, "seq": 8192, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.01398777961730957, "rel_err": 0.0038633455988019705, "xla_temp_mb": 66.000244140625, "xla_arg_mb": 160.0, "xla_out_mb": 96.0000228881836, "compile_s": 0.6113709609489888, "inner": 9, "median_ms": 5.241883246021138, "q1_ms": 5.23975377695428, "q3_ms": 5.260532062190274, "tflops": 183.53569302297282 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 64, "seq": 8192, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.01398777961730957, "rel_err": 0.0038633455988019705, "xla_temp_mb": 130.000732421875, "xla_arg_mb": 160.0, "xla_out_mb": 96.0000228881836, "compile_s": 0.1386431148275733, "inner": 16, "median_ms": 3.027764134458266, "q1_ms": 3.0261634310591035, "q3_ms": 3.1469928435399197, "tflops": 317.75020496308775 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 64, "seq": 8192, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.01398777961730957, "rel_err": 0.0038633455988019705, "xla_temp_mb": 130.000732421875, "xla_arg_mb": 160.0, "xla_out_mb": 96.0000228881836, "compile_s": 0.13939219200983644, "inner": 16, "median_ms": 3.00752287148498, "q1_ms": 3.0017306853551418, "q3_ms": 3.1535934103885666, "tflops": 319.8887308308221 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 64, "seq": 8192, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.01398777961730957, "rel_err": 0.004312852397561073, "xla_temp_mb": 20480.0, "xla_arg_mb": 160.0, "xla_out_mb": 96.0000228881836, "compile_s": 0.0776635438669473, "inner": 1, "median_ms": 38.71144214645028, "q1_ms": 38.58794807456434, "q3_ms": 38.741967058740556, "tflops": 24.8524110950028 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 64, "seq": 8192, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.01398777961730957, "rel_err": 0.0038633455988019705, "xla_temp_mb": 130.000732421875, "xla_arg_mb": 160.0, "xla_out_mb": 96.0000228881836, "compile_s": 0.14467269787564874, "inner": 16, "median_ms": 3.0280853679869324, "q1_ms": 3.0055788083700463, "q3_ms": 3.0974554683780298, "tflops": 317.7164965278323 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 128, "seq": 128, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0027510523796081543, "rel_err": 0.0036955196410417557, "xla_temp_mb": 0.0, "xla_arg_mb": 3.0, "xla_out_mb": 1.0, "compile_s": 0.27958805300295353, "inner": 200, "median_ms": 0.037999694468453526, "q1_ms": 0.03786759451031685, "q3_ms": 0.038041192456148565, "tflops": 7.064147745262766 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 128, "seq": 128, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0027510523796081543, "rel_err": 0.0036955196410417557, "xla_temp_mb": 0.0, "xla_arg_mb": 3.0, "xla_out_mb": 1.0, "compile_s": 0.28169999714009464, "inner": 200, "median_ms": 0.025178685318678617, "q1_ms": 0.024908980121836066, "q3_ms": 0.025400230078957975, "tflops": 10.661218113753666 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 128, "seq": 128, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0027510523796081543, "rel_err": 0.0036955196410417557, "xla_temp_mb": 0.0, "xla_arg_mb": 3.0, "xla_out_mb": 1.0, "compile_s": 0.2824166249483824, "inner": 200, "median_ms": 0.025109935086220503, "q1_ms": 0.024999277666211128, "q3_ms": 0.025545277749188244, "tflops": 10.690408202102777 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 128, "seq": 128, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0032942891120910645, "rel_err": 0.004425255581736565, "xla_temp_mb": 3.0, "xla_arg_mb": 3.0, "xla_out_mb": 1.0, "compile_s": 0.17832715390250087, "inner": 200, "median_ms": 0.03465610556304455, "q1_ms": 0.034621680388227105, "q3_ms": 0.0347095524193719, "tflops": 7.745690164513046 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 128, "seq": 128, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0027510523796081543, "rel_err": 0.0036955196410417557, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 3.0, "xla_out_mb": 1.0, "compile_s": 0.3834223640151322, "inner": 200, "median_ms": 0.05230060080066323, "q1_ms": 0.05147501069586724, "q3_ms": 0.052455299883149564, "tflops": 5.1325501407355905 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 128, "seq": 128, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.003998756408691406, "rel_err": 0.004262624774128199, "xla_temp_mb": 2.031494140625, "xla_arg_mb": 5.0, "xla_out_mb": 3.0000228881835938, "compile_s": 0.6263463089708239, "inner": 200, "median_ms": 0.0365924509242177, "q1_ms": 0.03651264763902873, "q3_ms": 0.03693357284646481, "tflops": 25.67535303786394 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 128, "seq": 128, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.003998756408691406, "rel_err": 0.004262624774128199, "xla_temp_mb": 1.03173828125, "xla_arg_mb": 5.0, "xla_out_mb": 3.0000228881835938, "compile_s": 0.4525914320256561, "inner": 200, "median_ms": 0.03957533976063132, "q1_ms": 0.03941594040952623, "q3_ms": 0.039755639736540616, "tflops": 23.740139735568814 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 128, "seq": 128, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.003998756408691406, "rel_err": 0.004262624774128199, "xla_temp_mb": 1.03173828125, "xla_arg_mb": 5.0, "xla_out_mb": 3.0000228881835938, "compile_s": 0.4463101190049201, "inner": 200, "median_ms": 0.03865198930725455, "q1_ms": 0.038511594757437706, "q3_ms": 0.038985792780295014, "tflops": 24.307263683933126 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 128, "seq": 128, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.004075884819030762, "rel_err": 0.004262624774128199, "xla_temp_mb": 5.0, "xla_arg_mb": 5.0, "xla_out_mb": 3.0000228881835938, "compile_s": 0.03179212100803852, "inner": 200, "median_ms": 0.05856673466041684, "q1_ms": 0.058350362232886255, "q3_ms": 0.05869013548363, "tflops": 16.041940897807827 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 128, "seq": 128, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.003998756408691406, "rel_err": 0.004262624774128199, "xla_temp_mb": 4.031982421875, "xla_arg_mb": 5.0, "xla_out_mb": 3.0000228881835938, "compile_s": 0.9537154780700803, "inner": 200, "median_ms": 0.1113095402251929, "q1_ms": 0.11128184269182384, "q3_ms": 0.1114249147940427, "tflops": 8.440643040113427 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 128, "seq": 128, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.005361318588256836, "rel_err": 0.001854726462624967, "xla_temp_mb": 0.0, "xla_arg_mb": 3.0, "xla_out_mb": 1.0, "compile_s": 0.36887998413294554, "inner": 200, "median_ms": 0.038452944718301296, "q1_ms": 0.038310542004182935, "q3_ms": 0.03866141720209271, "tflops": 3.4904408227575976 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 128, "seq": 128, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.005361318588256836, "rel_err": 0.001854726462624967, "xla_temp_mb": 0.0, "xla_arg_mb": 3.0, "xla_out_mb": 1.0, "compile_s": 0.36772252805531025, "inner": 200, "median_ms": 0.026051325257867575, "q1_ms": 0.02570465498138219, "q3_ms": 0.026344250072725117, "tflops": 5.1520499119124805 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 128, "seq": 128, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.005361318588256836, "rel_err": 0.001854726462624967, "xla_temp_mb": 0.0, "xla_arg_mb": 3.0, "xla_out_mb": 1.0, "compile_s": 0.36704474897123873, "inner": 200, "median_ms": 0.03338115522637963, "q1_ms": 0.033101479639299214, "q3_ms": 0.03342098440043628, "tflops": 4.0207634244465495 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 128, "seq": 128, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.006750822067260742, "rel_err": 0.002335419412702322, "xla_temp_mb": 3.0, "xla_arg_mb": 3.0, "xla_out_mb": 1.0, "compile_s": 0.2825764052104205, "inner": 200, "median_ms": 0.034562755608931184, "q1_ms": 0.03448815492447466, "q3_ms": 0.03474157827440649, "tflops": 3.8833051831468404 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 128, "seq": 128, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.005361318588256836, "rel_err": 0.001854726462624967, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 3.0, "xla_out_mb": 1.0, "compile_s": 0.4309009318239987, "inner": 200, "median_ms": 0.06660821498371661, "q1_ms": 0.06548994511831552, "q3_ms": 0.06696883530821651, "tflops": 2.0150326507445304 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 128, "seq": 128, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.01621723175048828, "rel_err": 0.00388345867395401, "xla_temp_mb": 2.031494140625, "xla_arg_mb": 5.0, "xla_out_mb": 3.0000228881835938, "compile_s": 0.784612599061802, "inner": 200, "median_ms": 0.04989226465113461, "q1_ms": 0.04976418800652027, "q3_ms": 0.05032873013988137, "tflops": 9.415528665310186 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 128, "seq": 128, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.01621723175048828, "rel_err": 0.00388345867395401, "xla_temp_mb": 1.03173828125, "xla_arg_mb": 5.0, "xla_out_mb": 3.0000228881835938, "compile_s": 0.46985696302726865, "inner": 200, "median_ms": 0.04019833984784782, "q1_ms": 0.04003636247944087, "q3_ms": 0.04034166282508522, "tflops": 11.686105689390818 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 128, "seq": 128, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.01621723175048828, "rel_err": 0.00388345867395401, "xla_temp_mb": 1.03173828125, "xla_arg_mb": 5.0, "xla_out_mb": 3.0000228881835938, "compile_s": 0.47549624810926616, "inner": 200, "median_ms": 0.06039123050868511, "q1_ms": 0.04796329187229276, "q3_ms": 0.06064690765924752, "tflops": 7.778646734685122 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 128, "seq": 128, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.01621723175048828, "rel_err": 0.005187384784221649, "xla_temp_mb": 5.0, "xla_arg_mb": 5.0, "xla_out_mb": 3.0000228881835938, "compile_s": 0.7761802719905972, "inner": 200, "median_ms": 0.07628843537531793, "q1_ms": 0.07619636482559144, "q3_ms": 0.07634163717739284, "tflops": 6.157709824417044 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 128, "seq": 128, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.01621723175048828, "rel_err": 0.00388345867395401, "xla_temp_mb": 4.031982421875, "xla_arg_mb": 5.0, "xla_out_mb": 3.0000228881835938, "compile_s": 0.02151871006935835, "inner": 200, "median_ms": 0.11332123540341854, "q1_ms": 0.11305733525659889, "q3_ms": 0.11358155985362828, "tflops": 4.145401753940186 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 128, "seq": 512, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0011283457279205322, "rel_err": 0.0027316263876855373, "xla_temp_mb": 0.0, "xla_arg_mb": 12.0, "xla_out_mb": 4.0, "compile_s": 0.3224048928823322, "inner": 200, "median_ms": 0.02197809051722288, "q1_ms": 0.021757989889010787, "q3_ms": 0.02222591545432806, "tflops": 195.4204025429005 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 128, "seq": 512, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0011283457279205322, "rel_err": 0.0027316263876855373, "xla_temp_mb": 0.0, "xla_arg_mb": 12.0, "xla_out_mb": 4.0, "compile_s": 0.3263449619989842, "inner": 200, "median_ms": 0.022545489482581615, "q1_ms": 0.022341214935295284, "q3_ms": 0.022670990438200533, "tflops": 190.50228646923998 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 128, "seq": 512, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0011283457279205322, "rel_err": 0.0027316263876855373, "xla_temp_mb": 0.0, "xla_arg_mb": 12.0, "xla_out_mb": 4.0, "compile_s": 0.34105254197493196, "inner": 200, "median_ms": 0.022214240161702037, "q1_ms": 0.022196187637746334, "q3_ms": 0.022351534571498632, "tflops": 193.34297570999715 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 128, "seq": 512, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0012100040912628174, "rel_err": 0.002929314039647579, "xla_temp_mb": 52.00001525878906, "xla_arg_mb": 12.0, "xla_out_mb": 4.0, "compile_s": 0.018514588940888643, "inner": 200, "median_ms": 0.05962362978607416, "q1_ms": 0.05949671263806522, "q3_ms": 0.05990566045511514, "tflops": 72.03464987640089 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 128, "seq": 512, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0011283457279205322, "rel_err": 0.0027316263876855373, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 12.0, "xla_out_mb": 4.0, "compile_s": 0.05701475217938423, "inner": 200, "median_ms": 0.0509457103908062, "q1_ms": 0.050903654191643, "q3_ms": 0.05115668289363384, "tflops": 84.30478764655878 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 128, "seq": 512, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0017972588539123535, "rel_err": 0.003949459176510572, "xla_temp_mb": 8.125244140625, "xla_arg_mb": 20.0, "xla_out_mb": 12.000022888183594, "compile_s": 0.6129108360037208, "inner": 200, "median_ms": 0.07436739513650537, "q1_ms": 0.07160822744481266, "q3_ms": 0.07553141273092479, "tflops": 202.1367765861268 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 128, "seq": 512, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0017972588539123535, "rel_err": 0.003949459176510572, "xla_temp_mb": 4.12548828125, "xla_arg_mb": 20.0, "xla_out_mb": 12.000022888183594, "compile_s": 0.4400263971183449, "inner": 200, "median_ms": 0.07323745056055486, "q1_ms": 0.06977405748330057, "q3_ms": 0.07493106706533581, "tflops": 205.25544541683334 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 128, "seq": 512, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0017972588539123535, "rel_err": 0.003949459176510572, "xla_temp_mb": 4.12548828125, "xla_arg_mb": 20.0, "xla_out_mb": 12.000022888183594, "compile_s": 0.44478632300160825, "inner": 200, "median_ms": 0.07253379444591701, "q1_ms": 0.06943180749658495, "q3_ms": 0.07270892267115414, "tflops": 207.24664483406443 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 128, "seq": 512, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0020554959774017334, "rel_err": 0.004599467851221561, "xla_temp_mb": 80.00001525878906, "xla_arg_mb": 20.0, "xla_out_mb": 12.000022888183594, "compile_s": 0.03494559298269451, "inner": 160, "median_ms": 0.1749440751154907, "q1_ms": 0.1748879840306472, "q3_ms": 0.17500988760730252, "tflops": 85.92680561531594 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 128, "seq": 512, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0017972588539123535, "rel_err": 0.003949459176510572, "xla_temp_mb": 16.125732421875, "xla_arg_mb": 20.0, "xla_out_mb": 12.000022888183594, "compile_s": 0.1340898780617863, "inner": 200, "median_ms": 0.09597973432391882, "q1_ms": 0.0958810601150617, "q3_ms": 0.09623625490348786, "tflops": 156.62041202643573 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 128, "seq": 512, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.006621599197387695, "rel_err": 0.002405368722975254, "xla_temp_mb": 0.0, "xla_arg_mb": 12.0, "xla_out_mb": 4.0, "compile_s": 0.37553706602193415, "inner": 200, "median_ms": 0.022253035567700863, "q1_ms": 0.02220946247689426, "q3_ms": 0.02229266450740397, "tflops": 96.50295311247164 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 128, "seq": 512, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.006621599197387695, "rel_err": 0.002405368722975254, "xla_temp_mb": 0.0, "xla_arg_mb": 12.0, "xla_out_mb": 4.0, "compile_s": 0.3727867540437728, "inner": 200, "median_ms": 0.030558659927919507, "q1_ms": 0.030349865555763245, "q3_ms": 0.030654040165245533, "tflops": 70.27414333826793 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 128, "seq": 512, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.006621599197387695, "rel_err": 0.002405368722975254, "xla_temp_mb": 0.0, "xla_arg_mb": 12.0, "xla_out_mb": 4.0, "compile_s": 0.37328379205428064, "inner": 200, "median_ms": 0.03056391491554678, "q1_ms": 0.030374042107723653, "q3_ms": 0.030696315225213766, "tflops": 70.26206079731138 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 128, "seq": 512, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.01019430160522461, "rel_err": 0.0037031921092420816, "xla_temp_mb": 52.00001525878906, "xla_arg_mb": 12.0, "xla_out_mb": 4.0, "compile_s": 0.7401493121869862, "inner": 200, "median_ms": 0.0892982049845159, "q1_ms": 0.08923475223127753, "q3_ms": 0.0894355500349775, "tflops": 24.048452579448472 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 128, "seq": 512, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.006621599197387695, "rel_err": 0.002405368722975254, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 12.0, "xla_out_mb": 4.0, "compile_s": 0.05873396713286638, "inner": 200, "median_ms": 0.06720691453665495, "q1_ms": 0.06673776195384562, "q3_ms": 0.06740931246895343, "tflops": 31.953314072003007 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 128, "seq": 512, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.013897895812988281, "rel_err": 0.006165427155792713, "xla_temp_mb": 8.125244140625, "xla_arg_mb": 20.0, "xla_out_mb": 12.000022888183594, "compile_s": 0.763766375835985, "inner": 200, "median_ms": 0.05852438509464264, "q1_ms": 0.058058887952938676, "q3_ms": 0.059711160720326006, "tflops": 128.4283936660795 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 128, "seq": 512, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.013897895812988281, "rel_err": 0.006165427155792713, "xla_temp_mb": 4.12548828125, "xla_arg_mb": 20.0, "xla_out_mb": 12.000022888183594, "compile_s": 0.45436320709995925, "inner": 200, "median_ms": 0.06612751516513526, "q1_ms": 0.06599019223358482, "q3_ms": 0.06626828981097788, "tflops": 113.66210796262912 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 128, "seq": 512, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.013897895812988281, "rel_err": 0.006165427155792713, "xla_temp_mb": 4.12548828125, "xla_arg_mb": 20.0, "xla_out_mb": 12.000022888183594, "compile_s": 0.4727944650221616, "inner": 200, "median_ms": 0.06621401524171233, "q1_ms": 0.06617324019316584, "q3_ms": 0.06628029281273484, "tflops": 113.51362306850531 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 128, "seq": 512, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.013897895812988281, "rel_err": 0.006165427155792713, "xla_temp_mb": 80.00001525878906, "xla_arg_mb": 20.0, "xla_out_mb": 12.000022888183594, "compile_s": 0.03929809993132949, "inner": 126, "median_ms": 0.26944568306798977, "q1_ms": 0.2692013143724392, "q3_ms": 0.27062404727090206, "tflops": 27.895020185212704 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 128, "seq": 512, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.013897895812988281, "rel_err": 0.006165427155792713, "xla_temp_mb": 16.125732421875, "xla_arg_mb": 20.0, "xla_out_mb": 12.000022888183594, "compile_s": 1.0341696948744357, "inner": 200, "median_ms": 0.09695098036900163, "q1_ms": 0.09692902967799455, "q3_ms": 0.09706963028293103, "tflops": 77.52570153899309 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 128, "seq": 600, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0011000633239746094, "rel_err": 0.002210719743743539, "xla_temp_mb": 20.0, "xla_arg_mb": 14.0625, "xla_out_mb": 4.6875, "compile_s": 0.37245421507395804, "inner": 200, "median_ms": 0.059816784923896194, "q1_ms": 0.05956918466836214, "q3_ms": 0.05984425777569413, "tflops": 98.6050990119951 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 128, "seq": 600, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0011000633239746094, "rel_err": 0.002210719743743539, "xla_temp_mb": 20.0, "xla_arg_mb": 14.0625, "xla_out_mb": 4.6875, "compile_s": 0.3748027279507369, "inner": 200, "median_ms": 0.05981593509204686, "q1_ms": 0.05969343241304159, "q3_ms": 0.05993193248286843, "tflops": 98.60649993891396 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 128, "seq": 600, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0011000633239746094, "rel_err": 0.002210719743743539, "xla_temp_mb": 20.0, "xla_arg_mb": 14.0625, "xla_out_mb": 4.6875, "compile_s": 0.37193914502859116, "inner": 200, "median_ms": 0.06005578557960689, "q1_ms": 0.05974023544695228, "q3_ms": 0.060088507598266006, "tflops": 98.21268580662547 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 128, "seq": 600, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.002395838499069214, "rel_err": 0.004814747720956802, "xla_temp_mb": 65.91798400878906, "xla_arg_mb": 14.0625, "xla_out_mb": 4.6875, "compile_s": 0.018599048955366015, "inner": 198, "median_ms": 0.09621937895626431, "q1_ms": 0.09399746436475202, "q3_ms": 0.096800237320213, "tflops": 61.29991758397229 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 128, "seq": 600, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0011000633239746094, "rel_err": 0.002210719743743539, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 14.0625, "xla_out_mb": 4.6875, "compile_s": 0.4061016011983156, "inner": 200, "median_ms": 0.06456141942180693, "q1_ms": 0.06449947250075638, "q3_ms": 0.06483854493126273, "tflops": 91.35858617767238 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 128, "seq": 600, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.002751857042312622, "rel_err": 0.0050126416608691216, "xla_temp_mb": 40.00048828125, "xla_arg_mb": 23.4375, "xla_out_mb": 14.062522888183594, "compile_s": 0.7448338789399713, "inner": 168, "median_ms": 0.17960128579510465, "q1_ms": 0.17912090143432752, "q3_ms": 0.18058015155561624, "tflops": 114.9426069452042 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 128, "seq": 600, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.002751857042312622, "rel_err": 0.0050126416608691216, "xla_temp_mb": 30.078651428222656, "xla_arg_mb": 23.4375, "xla_out_mb": 14.062522888183594, "compile_s": 0.5032411299180239, "inner": 157, "median_ms": 0.17533117199352213, "q1_ms": 0.1752242537891599, "q3_ms": 0.1767288372929024, "tflops": 117.74198372872746 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 128, "seq": 600, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.002751857042312622, "rel_err": 0.0050126416608691216, "xla_temp_mb": 30.078651428222656, "xla_arg_mb": 23.4375, "xla_out_mb": 14.062522888183594, "compile_s": 0.5026589019689709, "inner": 159, "median_ms": 0.17737547732280487, "q1_ms": 0.17625252487805654, "q3_ms": 0.17767912948375228, "tflops": 116.3849722159189 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 128, "seq": 600, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0027889013290405273, "rel_err": 0.0050126416608691216, "xla_temp_mb": 109.93678283691406, "xla_arg_mb": 23.4375, "xla_out_mb": 14.062522888183594, "compile_s": 0.03658725810237229, "inner": 120, "median_ms": 0.2917953592259437, "q1_ms": 0.2899518158907691, "q3_ms": 0.2921739340915034, "tflops": 70.74766389281403 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 128, "seq": 600, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.002751857042312622, "rel_err": 0.0050126416608691216, "xla_temp_mb": 18.897216796875, "xla_arg_mb": 23.4375, "xla_out_mb": 14.062522888183594, "compile_s": 0.9846655121073127, "inner": 170, "median_ms": 0.13148004654794931, "q1_ms": 0.13095178257893114, "q3_ms": 0.1315901323002489, "tflops": 157.0112008780847 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 128, "seq": 600, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.006784915924072266, "rel_err": 0.0023472141474485397, "xla_temp_mb": 20.0, "xla_arg_mb": 14.0625, "xla_out_mb": 4.6875, "compile_s": 0.3982698619365692, "inner": 200, "median_ms": 0.04916146514005959, "q1_ms": 0.049103485071100295, "q3_ms": 0.04921846557408571, "tflops": 59.98844809848613 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 128, "seq": 600, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.006784915924072266, "rel_err": 0.0023472141474485397, "xla_temp_mb": 20.0, "xla_arg_mb": 14.0625, "xla_out_mb": 4.6875, "compile_s": 0.40467746485956013, "inner": 200, "median_ms": 0.04923961008898914, "q1_ms": 0.049221289809793234, "q3_ms": 0.04926063760649413, "tflops": 59.89324437521238 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 128, "seq": 600, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.006784915924072266, "rel_err": 0.0023472141474485397, "xla_temp_mb": 20.0, "xla_arg_mb": 14.0625, "xla_out_mb": 4.6875, "compile_s": 0.3992117098532617, "inner": 200, "median_ms": 0.04921286483295262, "q1_ms": 0.049173664883710444, "q3_ms": 0.04923788772430271, "tflops": 59.92579399737135 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 128, "seq": 600, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.007506132125854492, "rel_err": 0.0025967159308493137, "xla_temp_mb": 65.91798400878906, "xla_arg_mb": 14.0625, "xla_out_mb": 4.6875, "compile_s": 0.6863689008168876, "inner": 200, "median_ms": 0.13532397453673184, "q1_ms": 0.13507205236237496, "q3_ms": 0.13542567321565002, "tflops": 21.79303416187722 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 128, "seq": 600, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.006784915924072266, "rel_err": 0.0023472141474485397, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 14.0625, "xla_out_mb": 4.6875, "compile_s": 0.44837676291354, "inner": 200, "median_ms": 0.05382339935749769, "q1_ms": 0.053558400250039995, "q3_ms": 0.053998049697838724, "tflops": 54.79252583828455 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 128, "seq": 600, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.012672662734985352, "rel_err": 0.004155972506850958, "xla_temp_mb": 40.00048828125, "xla_arg_mb": 23.4375, "xla_out_mb": 14.062522888183594, "compile_s": 0.7845551690552384, "inner": 200, "median_ms": 0.12939748936332762, "q1_ms": 0.1293225446715951, "q3_ms": 0.12990485993213952, "tflops": 79.76909019476943 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 128, "seq": 600, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.012672662734985352, "rel_err": 0.004155972506850958, "xla_temp_mb": 30.07861328125, "xla_arg_mb": 23.4375, "xla_out_mb": 14.062522888183594, "compile_s": 0.5035284098703414, "inner": 200, "median_ms": 0.12481459998525679, "q1_ms": 0.12449557718355209, "q3_ms": 0.12534635490737855, "tflops": 82.69801770962079 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 128, "seq": 600, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.012672662734985352, "rel_err": 0.004155972506850958, "xla_temp_mb": 30.07861328125, "xla_arg_mb": 23.4375, "xla_out_mb": 14.062522888183594, "compile_s": 0.5097135729156435, "inner": 200, "median_ms": 0.12495310511440039, "q1_ms": 0.12471385241951793, "q3_ms": 0.12510890024714172, "tflops": 82.60635052286057 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 128, "seq": 600, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.012672662734985352, "rel_err": 0.00421340623870492, "xla_temp_mb": 109.93678283691406, "xla_arg_mb": 23.4375, "xla_out_mb": 14.062522888183594, "compile_s": 0.03925111098214984, "inner": 105, "median_ms": 0.3520147342767034, "q1_ms": 0.3512800671160221, "q3_ms": 0.35264078039853347, "tflops": 29.322408964524733 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 128, "seq": 600, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.012672662734985352, "rel_err": 0.004155972506850958, "xla_temp_mb": 18.897216796875, "xla_arg_mb": 23.4375, "xla_out_mb": 14.062522888183594, "compile_s": 0.02264803694561124, "inner": 172, "median_ms": 0.11191928358531968, "q1_ms": 0.11184742149008914, "q3_ms": 0.11208730541321255, "tflops": 92.22646597921856 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 128, "seq": 1024, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0007505267858505249, "rel_err": 0.002722856355831027, "xla_temp_mb": 0.0, "xla_arg_mb": 24.0, "xla_out_mb": 8.0, "compile_s": 0.3217971450649202, "inner": 200, "median_ms": 0.059329235227778554, "q1_ms": 0.05913146014790982, "q3_ms": 0.05936431000009179, "tflops": 289.5683572869689 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 128, "seq": 1024, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0007505267858505249, "rel_err": 0.002722856355831027, "xla_temp_mb": 0.0, "xla_arg_mb": 24.0, "xla_out_mb": 8.0, "compile_s": 0.32511990517377853, "inner": 200, "median_ms": 0.062470374396070845, "q1_ms": 0.05964471260085702, "q3_ms": 0.06266957789193839, "tflops": 275.00826351827584 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 128, "seq": 1024, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0007505267858505249, "rel_err": 0.002722856355831027, "xla_temp_mb": 0.0, "xla_arg_mb": 24.0, "xla_out_mb": 8.0, "compile_s": 0.3270147400908172, "inner": 200, "median_ms": 0.06181302480399608, "q1_ms": 0.05992871010676026, "q3_ms": 0.06397602264769375, "tflops": 277.93283435774134 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 128, "seq": 1024, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0009716004133224487, "rel_err": 0.003524895291775465, "xla_temp_mb": 192.00001525878906, "xla_arg_mb": 24.0, "xla_out_mb": 8.0, "compile_s": 0.019132476998493075, "inner": 138, "median_ms": 0.25218733977796376, "q1_ms": 0.2518176333978772, "q3_ms": 0.2524019413006803, "tflops": 68.12344029294204 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 128, "seq": 1024, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0007505267858505249, "rel_err": 0.002722856355831027, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 24.0, "xla_out_mb": 8.0, "compile_s": 0.056836241856217384, "inner": 200, "median_ms": 0.04725011996924877, "q1_ms": 0.04720576689578593, "q3_ms": 0.04728476982563734, "tflops": 363.5941918281047 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 128, "seq": 1024, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0015575289726257324, "rel_err": 0.0038550400640815496, "xla_temp_mb": 16.250244140625, "xla_arg_mb": 40.0, "xla_out_mb": 24.000022888183594, "compile_s": 0.6156550680752844, "inner": 134, "median_ms": 0.2654745597948334, "q1_ms": 0.2593301009011469, "q3_ms": 0.26804891783299284, "tflops": 226.49832130984564 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 128, "seq": 1024, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0015575289726257324, "rel_err": 0.0038550400640815496, "xla_temp_mb": 8.25048828125, "xla_arg_mb": 40.0, "xla_out_mb": 24.000022888183594, "compile_s": 0.4295217760372907, "inner": 150, "median_ms": 0.25617195448527735, "q1_ms": 0.24674168011794487, "q3_ms": 0.2599996468052268, "tflops": 234.72336097375464 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 128, "seq": 1024, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0015575289726257324, "rel_err": 0.0038550400640815496, "xla_temp_mb": 8.25048828125, "xla_arg_mb": 40.0, "xla_out_mb": 24.000022888183594, "compile_s": 0.42133169807493687, "inner": 148, "median_ms": 0.25591476367333454, "q1_ms": 0.25554851320496685, "q3_ms": 0.26349105708603116, "tflops": 234.95925471792265 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 128, "seq": 1024, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0019018948078155518, "rel_err": 0.003942885436117649, "xla_temp_mb": 320.00001525878906, "xla_arg_mb": 40.0, "xla_out_mb": 24.000022888183594, "compile_s": 0.03595892991870642, "inner": 64, "median_ms": 0.6579028595297132, "q1_ms": 0.6558339537150459, "q3_ms": 0.6581563757208642, "tflops": 91.39577564229198 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 128, "seq": 1024, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0015575289726257324, "rel_err": 0.0038550400640815496, "xla_temp_mb": 32.250732421875, "xla_arg_mb": 40.0, "xla_out_mb": 24.000022888183594, "compile_s": 0.1531360240187496, "inner": 144, "median_ms": 0.2273536746441904, "q1_ms": 0.21365454530395153, "q3_ms": 0.22846530191600323, "tflops": 264.47578750641713 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 128, "seq": 1024, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.005032062530517578, "rel_err": 0.0017408216372132301, "xla_temp_mb": 0.0, "xla_arg_mb": 24.0, "xla_out_mb": 8.0, "compile_s": 0.3756416959222406, "inner": 200, "median_ms": 0.05456989980302751, "q1_ms": 0.0544668227666989, "q3_ms": 0.0545738241635263, "tflops": 157.41158812835928 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 128, "seq": 1024, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.005032062530517578, "rel_err": 0.0017408216372132301, "xla_temp_mb": 0.0, "xla_arg_mb": 24.0, "xla_out_mb": 8.0, "compile_s": 0.3736600410193205, "inner": 200, "median_ms": 0.05411020014435053, "q1_ms": 0.054038624512031674, "q3_ms": 0.054383547976613045, "tflops": 158.748897048699 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 128, "seq": 1024, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.005032062530517578, "rel_err": 0.0017408216372132301, "xla_temp_mb": 0.0, "xla_arg_mb": 24.0, "xla_out_mb": 8.0, "compile_s": 0.3742310998495668, "inner": 200, "median_ms": 0.05444959970191121, "q1_ms": 0.0544349750271067, "q3_ms": 0.05448004754725844, "tflops": 157.75937084985566 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 128, "seq": 1024, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.005884408950805664, "rel_err": 0.002035687444731593, "xla_temp_mb": 192.00001525878906, "xla_arg_mb": 24.0, "xla_out_mb": 8.0, "compile_s": 0.6201782459393144, "inner": 115, "median_ms": 0.3114912254006966, "q1_ms": 0.3109534433030564, "q3_ms": 0.3118960014746889, "tflops": 27.576810810481305 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 128, "seq": 1024, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.005032062530517578, "rel_err": 0.0017408216372132301, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 24.0, "xla_out_mb": 8.0, "compile_s": 0.05923623498529196, "inner": 200, "median_ms": 0.06679031415842474, "q1_ms": 0.050929537974298, "q3_ms": 0.06684969004709274, "tflops": 128.6104834246612 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 128, "seq": 1024, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.013173818588256836, "rel_err": 0.003089054487645626, "xla_temp_mb": 16.250244140625, "xla_arg_mb": 40.0, "xla_out_mb": 24.000022888183594, "compile_s": 0.7463977299630642, "inner": 179, "median_ms": 0.18800322913341017, "q1_ms": 0.18350905010362242, "q3_ms": 0.18824465304883664, "tflops": 159.9162483037222 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 128, "seq": 1024, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.013173818588256836, "rel_err": 0.003089054487645626, "xla_temp_mb": 8.25048828125, "xla_arg_mb": 40.0, "xla_out_mb": 24.000022888183594, "compile_s": 0.48212246084585786, "inner": 166, "median_ms": 0.1943730244255928, "q1_ms": 0.19007144041404306, "q3_ms": 0.19473169652565595, "tflops": 154.67563547383594 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 128, "seq": 1024, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.013173818588256836, "rel_err": 0.003089054487645626, "xla_temp_mb": 8.25048828125, "xla_arg_mb": 40.0, "xla_out_mb": 24.000022888183594, "compile_s": 0.47479035099968314, "inner": 174, "median_ms": 0.19419337381962043, "q1_ms": 0.19076335626579394, "q3_ms": 0.19442570133498688, "tflops": 154.81872774879605 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 128, "seq": 1024, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.013173818588256836, "rel_err": 0.003089054487645626, "xla_temp_mb": 320.00001525878906, "xla_arg_mb": 40.0, "xla_out_mb": 24.000022888183594, "compile_s": 0.04005301999859512, "inner": 50, "median_ms": 0.8498550811782479, "q1_ms": 0.8494162815622985, "q3_ms": 0.852910962421447, "tflops": 35.376350318830696 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 128, "seq": 1024, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.013173818588256836, "rel_err": 0.003089054487645626, "xla_temp_mb": 32.250732421875, "xla_arg_mb": 40.0, "xla_out_mb": 24.000022888183594, "compile_s": 0.1360454980749637, "inner": 152, "median_ms": 0.19544242777077384, "q1_ms": 0.19310980018130258, "q3_ms": 0.1964791307201315, "tflops": 153.82929599739572 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 128, "seq": 1500, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0007283985614776611, "rel_err": 0.0028605633415281773, "xla_temp_mb": 48.0, "xla_arg_mb": 35.15625, "xla_out_mb": 11.71875, "compile_s": 0.37466551503166556, "inner": 170, "median_ms": 0.18842336447799907, "q1_ms": 0.1856142853605835, "q3_ms": 0.18889348173294873, "tflops": 195.64452689891525 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 128, "seq": 1500, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0007283985614776611, "rel_err": 0.0028605633415281773, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 35.15625, "xla_out_mb": 11.71875, "compile_s": 0.05553427804261446, "inner": 200, "median_ms": 0.10026662494055927, "q1_ms": 0.09557444020174444, "q3_ms": 0.10180499753914773, "tflops": 367.6597274701723 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 128, "seq": 1500, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0007283985614776611, "rel_err": 0.0028605633415281773, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 35.15625, "xla_out_mb": 11.71875, "compile_s": 0.05269227595999837, "inner": 200, "median_ms": 0.0992582249455154, "q1_ms": 0.09644718433264643, "q3_ms": 0.10063837748020887, "tflops": 371.394914831847 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 128, "seq": 1500, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0008879303932189941, "rel_err": 0.0034870761446654797, "xla_temp_mb": 411.98731994628906, "xla_arg_mb": 35.15625, "xla_out_mb": 11.71875, "compile_s": 0.021270492114126682, "inner": 73, "median_ms": 0.5621536705030562, "q1_ms": 0.560012438704502, "q3_ms": 0.562361068385717, "tflops": 65.57637516982038 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 128, "seq": 1500, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0007283985614776611, "rel_err": 0.0028605633415281773, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 35.15625, "xla_out_mb": 11.71875, "compile_s": 0.057312542805448174, "inner": 200, "median_ms": 0.0936925399582833, "q1_ms": 0.0936227454803884, "q3_ms": 0.10051379969809204, "tflops": 393.45715268700934 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 128, "seq": 1500, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0016910731792449951, "rel_err": 0.004783402197062969, "xla_temp_mb": 96.00048828125, "xla_arg_mb": 58.59375, "xla_out_mb": 35.156272888183594, "compile_s": 0.7539875479415059, "inner": 70, "median_ms": 0.6342162710747549, "q1_ms": 0.6105350496779595, "q3_ms": 0.6420282430813781, "tflops": 203.43848917870474 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 128, "seq": 1500, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0016910731792449951, "rel_err": 0.004783402197062969, "xla_temp_mb": 47.241943359375, "xla_arg_mb": 58.59375, "xla_out_mb": 35.156272888183594, "compile_s": 0.13748215418308973, "inner": 89, "median_ms": 0.42832984453004397, "q1_ms": 0.4248539550920551, "q3_ms": 0.42880798841753365, "tflops": 301.22579980753585 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 128, "seq": 1500, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0016910731792449951, "rel_err": 0.004783402197062969, "xla_temp_mb": 47.241943359375, "xla_arg_mb": 58.59375, "xla_out_mb": 35.156272888183594, "compile_s": 0.13438790291547775, "inner": 91, "median_ms": 0.4265183956454416, "q1_ms": 0.4178842409401313, "q3_ms": 0.4293552592171581, "tflops": 302.5051236178233 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 128, "seq": 1500, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0016910731792449951, "rel_err": 0.004783402197062969, "xla_temp_mb": 686.8286285400391, "xla_arg_mb": 58.59375, "xla_out_mb": 35.156272888183594, "compile_s": 0.037204637890681624, "inner": 30, "median_ms": 1.5478514367714524, "q1_ms": 1.5448987794419131, "q3_ms": 1.5524460817687213, "tflops": 83.35683705480258 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 128, "seq": 1500, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0016910731792449951, "rel_err": 0.004783402197062969, "xla_temp_mb": 47.241943359375, "xla_arg_mb": 58.59375, "xla_out_mb": 35.156272888183594, "compile_s": 0.13624040805734694, "inner": 95, "median_ms": 0.42740377891612685, "q1_ms": 0.4174838001888833, "q3_ms": 0.42818193201367793, "tflops": 301.87847268734487 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 128, "seq": 1500, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0072650909423828125, "rel_err": 0.0025133287999778986, "xla_temp_mb": 48.0, "xla_arg_mb": 35.15625, "xla_out_mb": 11.71875, "compile_s": 0.3997378770727664, "inner": 184, "median_ms": 0.1512858915426161, "q1_ms": 0.1508136085011875, "q3_ms": 0.15130170314754968, "tflops": 121.83555129995611 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 128, "seq": 1500, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0072650909423828125, "rel_err": 0.0025133287999778986, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 35.15625, "xla_out_mb": 11.71875, "compile_s": 0.05643430491909385, "inner": 200, "median_ms": 0.07312340079806745, "q1_ms": 0.06688439229037613, "q3_ms": 0.0742910971166566, "tflops": 252.06705102379667 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 128, "seq": 1500, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0072650909423828125, "rel_err": 0.0025133287999778986, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 35.15625, "xla_out_mb": 11.71875, "compile_s": 0.05389420292340219, "inner": 200, "median_ms": 0.0714100559707731, "q1_ms": 0.06693886767607182, "q3_ms": 0.07317447452805936, "tflops": 258.1149076195081 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 128, "seq": 1500, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.009048938751220703, "rel_err": 0.003130443627014756, "xla_temp_mb": 411.98731994628906, "xla_arg_mb": 35.15625, "xla_out_mb": 11.71875, "compile_s": 0.467581151984632, "inner": 74, "median_ms": 0.5684839718899614, "q1_ms": 0.5669600135218855, "q3_ms": 0.5692276282462518, "tflops": 32.423077714436936 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 128, "seq": 1500, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0072650909423828125, "rel_err": 0.0025133287999778986, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 35.15625, "xla_out_mb": 11.71875, "compile_s": 0.05976511584594846, "inner": 200, "median_ms": 0.07443269481882453, "q1_ms": 0.07051678199786693, "q3_ms": 0.0747494725510478, "tflops": 247.63311398122892 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 128, "seq": 1500, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.01283717155456543, "rel_err": 0.004822417162358761, "xla_temp_mb": 96.00048828125, "xla_arg_mb": 58.59375, "xla_out_mb": 35.156272888183594, "compile_s": 0.7824872699566185, "inner": 97, "median_ms": 0.43716200041709485, "q1_ms": 0.4351092670496899, "q3_ms": 0.4373207105054837, "tflops": 147.5700082313863 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 128, "seq": 1500, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.01283717155456543, "rel_err": 0.004822417162358761, "xla_temp_mb": 47.241943359375, "xla_arg_mb": 58.59375, "xla_out_mb": 35.156272888183594, "compile_s": 1.0657460750080645, "inner": 108, "median_ms": 0.32576150907617474, "q1_ms": 0.32298884320038335, "q3_ms": 0.33164219887651225, "tflops": 198.03444606745967 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 128, "seq": 1500, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.01283717155456543, "rel_err": 0.004822417162358761, "xla_temp_mb": 47.241943359375, "xla_arg_mb": 58.59375, "xla_out_mb": 35.156272888183594, "compile_s": 0.02214453904889524, "inner": 104, "median_ms": 0.3341112012724177, "q1_ms": 0.328795063139226, "q3_ms": 0.33433591418613035, "tflops": 193.08541513817764 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 128, "seq": 1500, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.01283717155456543, "rel_err": 0.004428453277796507, "xla_temp_mb": 686.8286285400391, "xla_arg_mb": 58.59375, "xla_out_mb": 35.156272888183594, "compile_s": 0.04123534681275487, "inner": 27, "median_ms": 1.7220604802585309, "q1_ms": 1.7146330706223294, "q3_ms": 1.723667313517244, "tflops": 37.46209888651233 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 128, "seq": 1500, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.01283717155456543, "rel_err": 0.004822417162358761, "xla_temp_mb": 47.241943359375, "xla_arg_mb": 58.59375, "xla_out_mb": 35.156272888183594, "compile_s": 0.021121161989867687, "inner": 110, "median_ms": 0.3342668991535902, "q1_ms": 0.33090427750721574, "q3_ms": 0.33438158623704856, "tflops": 192.99547805467205 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 128, "seq": 2048, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0005683600902557373, "rel_err": 0.0023447908461093903, "xla_temp_mb": 0.0, "xla_arg_mb": 48.0, "xla_out_mb": 16.0, "compile_s": 0.32379062403924763, "inner": 145, "median_ms": 0.23869859407944927, "q1_ms": 0.23673066211029375, "q3_ms": 0.23993458598852158, "tflops": 287.89225592643066 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 128, "seq": 2048, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0005683600902557373, "rel_err": 0.0023447908461093903, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 48.0, "xla_out_mb": 16.0, "compile_s": 0.054182460997253656, "inner": 184, "median_ms": 0.15216474370706987, "q1_ms": 0.14822250016985217, "q3_ms": 0.1626568365558658, "tflops": 451.6123450271166 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 128, "seq": 2048, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0005683600902557373, "rel_err": 0.0023447908461093903, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 48.0, "xla_out_mb": 16.0, "compile_s": 0.05436867079697549, "inner": 185, "median_ms": 0.16231960491151423, "q1_ms": 0.15541308660160852, "q3_ms": 0.16355109463974432, "tflops": 423.3590685084605 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 128, "seq": 2048, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0006561279296875, "rel_err": 0.002706880448386073, "xla_temp_mb": 768.0, "xla_arg_mb": 48.0, "xla_out_mb": 16.0, "compile_s": 0.017817370826378465, "inner": 49, "median_ms": 0.8741931239028975, "q1_ms": 0.8731820995976427, "q3_ms": 0.8772356644728962, "tflops": 78.60903369863743 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 128, "seq": 2048, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0005683600902557373, "rel_err": 0.0023447908461093903, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 48.0, "xla_out_mb": 16.0, "compile_s": 0.05587040679529309, "inner": 200, "median_ms": 0.1582311699166894, "q1_ms": 0.15545489732176065, "q3_ms": 0.16213527997024357, "tflops": 434.2979753747736 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 128, "seq": 2048, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0010292530059814453, "rel_err": 0.0036079285200685263, "xla_temp_mb": 32.500244140625, "xla_arg_mb": 80.0, "xla_out_mb": 48.000022888183594, "compile_s": 0.605848893057555, "inner": 51, "median_ms": 0.915514154578833, "q1_ms": 0.904058617557965, "q3_ms": 0.9411097773551649, "tflops": 262.7137629419246 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 128, "seq": 2048, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0010292530059814453, "rel_err": 0.0036079285200685263, "xla_temp_mb": 64.500732421875, "xla_arg_mb": 80.0, "xla_out_mb": 48.000022888183594, "compile_s": 0.13250372698530555, "inner": 62, "median_ms": 0.700946322702352, "q1_ms": 0.6934764040183397, "q3_ms": 0.7017693715170026, "tflops": 343.1335050717328 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 128, "seq": 2048, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0010292530059814453, "rel_err": 0.0036079285200685263, "xla_temp_mb": 64.500732421875, "xla_arg_mb": 80.0, "xla_out_mb": 48.000022888183594, "compile_s": 0.13262801803648472, "inner": 64, "median_ms": 0.7012852656771429, "q1_ms": 0.6916910788277164, "q3_ms": 0.702684561474598, "tflops": 342.9676628722006 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 128, "seq": 2048, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.000937044620513916, "rel_err": 0.0036079285200685263, "xla_temp_mb": 1280.000015258789, "xla_arg_mb": 80.0, "xla_out_mb": 48.000022888183594, "compile_s": 0.03664973005652428, "inner": 20, "median_ms": 2.2997527034021914, "q1_ms": 2.296073955949396, "q3_ms": 2.3004309739917517, "tflops": 104.58436170992819 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 128, "seq": 2048, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0010292530059814453, "rel_err": 0.0036079285200685263, "xla_temp_mb": 64.500732421875, "xla_arg_mb": 80.0, "xla_out_mb": 48.000022888183594, "compile_s": 0.1327230369206518, "inner": 64, "median_ms": 0.699947624525521, "q1_ms": 0.6848161792731844, "q3_ms": 0.7006825308053521, "tflops": 343.6230942837215 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 128, "seq": 2048, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.007552385330200195, "rel_err": 0.0026127169840037823, "xla_temp_mb": 0.0, "xla_arg_mb": 48.0, "xla_out_mb": 16.0, "compile_s": 0.37431973591446877, "inner": 194, "median_ms": 0.16859742769606642, "q1_ms": 0.16342743825570671, "q3_ms": 0.1691343837667295, "tflops": 203.79752430114715 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 128, "seq": 2048, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.007552385330200195, "rel_err": 0.0026127169840037823, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 48.0, "xla_out_mb": 16.0, "compile_s": 0.05492996913380921, "inner": 192, "median_ms": 0.11417354107834399, "q1_ms": 0.10406211125276363, "q3_ms": 0.11724480949245239, "tflops": 300.94309104788925 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 128, "seq": 2048, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.007552385330200195, "rel_err": 0.0026127169840037823, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 48.0, "xla_out_mb": 16.0, "compile_s": 0.053986432030797005, "inner": 198, "median_ms": 0.1154826104998438, "q1_ms": 0.10950477779701803, "q3_ms": 0.1163451369108651, "tflops": 297.5317081877576 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 128, "seq": 2048, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.008072614669799805, "rel_err": 0.0027926883194595575, "xla_temp_mb": 768.0, "xla_arg_mb": 48.0, "xla_out_mb": 16.0, "compile_s": 0.5443709208630025, "inner": 44, "median_ms": 1.0204847309399736, "q1_ms": 1.011431886581704, "q3_ms": 1.0212043628350576, "tflops": 33.670017126420966 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 128, "seq": 2048, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.007552385330200195, "rel_err": 0.0026127169840037823, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 48.0, "xla_out_mb": 16.0, "compile_s": 0.05896112788468599, "inner": 200, "median_ms": 0.11307169101200998, "q1_ms": 0.10838635265827179, "q3_ms": 0.11517153528984636, "tflops": 303.8756921425227 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 128, "seq": 2048, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.013039588928222656, "rel_err": 0.00369092496111989, "xla_temp_mb": 32.500244140625, "xla_arg_mb": 80.0, "xla_out_mb": 48.000022888183594, "compile_s": 0.765973465051502, "inner": 74, "median_ms": 0.5975255264422378, "q1_ms": 0.566169196374815, "q3_ms": 0.5980682295017146, "tflops": 201.2618356307583 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 128, "seq": 2048, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.013039588928222656, "rel_err": 0.00369092496111989, "xla_temp_mb": 64.500732421875, "xla_arg_mb": 80.0, "xla_out_mb": 48.000022888183594, "compile_s": 0.13500462099909782, "inner": 78, "median_ms": 0.5059511794780309, "q1_ms": 0.4987567752742996, "q3_ms": 0.5067985501880639, "tflops": 237.68910749860567 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 128, "seq": 2048, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.013039588928222656, "rel_err": 0.00369092496111989, "xla_temp_mb": 64.500732421875, "xla_arg_mb": 80.0, "xla_out_mb": 48.000022888183594, "compile_s": 0.13602404785342515, "inner": 82, "median_ms": 0.5052490983294641, "q1_ms": 0.5034285671327535, "q3_ms": 0.5069598967277604, "tflops": 238.0193941673918 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 128, "seq": 2048, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.013039588928222656, "rel_err": 0.00369092496111989, "xla_temp_mb": 1280.000015258789, "xla_arg_mb": 80.0, "xla_out_mb": 48.000022888183594, "compile_s": 0.04162506596185267, "inner": 16, "median_ms": 2.9172820650273934, "q1_ms": 2.8988289341214113, "q3_ms": 2.923328247561585, "tflops": 41.22298824980805 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 128, "seq": 2048, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.013039588928222656, "rel_err": 0.00369092496111989, "xla_temp_mb": 64.500732421875, "xla_arg_mb": 80.0, "xla_out_mb": 48.000022888183594, "compile_s": 0.13756865402683616, "inner": 80, "median_ms": 0.504381739301607, "q1_ms": 0.5012396315578371, "q3_ms": 0.5072774249128997, "tflops": 238.42870373244864 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 128, "seq": 3000, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0005130171775817871, "rel_err": 0.0019032340496778488, "xla_temp_mb": 96.0, "xla_arg_mb": 70.3125, "xla_out_mb": 23.4375, "compile_s": 0.3661382079590112, "inner": 69, "median_ms": 0.6224103339448356, "q1_ms": 0.6204923543998081, "q3_ms": 0.6286884283484973, "tflops": 236.9112335674508 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 128, "seq": 3000, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0005130171775817871, "rel_err": 0.0019032340496778488, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 70.3125, "xla_out_mb": 23.4375, "compile_s": 0.05497039994224906, "inner": 108, "median_ms": 0.3607254906523007, "q1_ms": 0.35757861095825555, "q3_ms": 0.3652338141196028, "tflops": 408.77621299607904 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 128, "seq": 3000, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0005130171775817871, "rel_err": 0.0019032340496778488, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 70.3125, "xla_out_mb": 23.4375, "compile_s": 0.053051005117595196, "inner": 107, "median_ms": 0.36748487798270774, "q1_ms": 0.35185497210613476, "q3_ms": 0.3733720660871156, "tflops": 401.25732740202346 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 128, "seq": 3000, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0006972551345825195, "rel_err": 0.002586735412478447, "xla_temp_mb": 1647.949234008789, "xla_arg_mb": 70.3125, "xla_out_mb": 23.4375, "compile_s": 0.019150747917592525, "inner": 25, "median_ms": 1.8437785562127829, "q1_ms": 1.8248236225917935, "q3_ms": 1.8520873179659247, "tflops": 79.97489693278692 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 128, "seq": 3000, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0005130171775817871, "rel_err": 0.0019032340496778488, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 70.3125, "xla_out_mb": 23.4375, "compile_s": 0.058127901051193476, "inner": 113, "median_ms": 0.35925786653780833, "q1_ms": 0.3489287347590502, "q3_ms": 0.36146547398076645, "tflops": 410.4461272373606 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 128, "seq": 3000, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0009522885084152222, "rel_err": 0.003991737496107817, "xla_temp_mb": 192.00048828125, "xla_arg_mb": 117.1875, "xla_out_mb": 70.3125228881836, "compile_s": 0.7494788798503578, "inner": 22, "median_ms": 2.217314862222834, "q1_ms": 2.1549752564169466, "q3_ms": 2.2786730921573257, "tflops": 232.75720051892827 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 128, "seq": 3000, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0009522885084152222, "rel_err": 0.003991737496107817, "xla_temp_mb": 94.483154296875, "xla_arg_mb": 117.1875, "xla_out_mb": 70.3125228881836, "compile_s": 0.1348008611239493, "inner": 33, "median_ms": 1.3871158783634503, "q1_ms": 1.3564087125952495, "q3_ms": 1.41508218296098, "tflops": 372.06408494790026 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 128, "seq": 3000, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0009522885084152222, "rel_err": 0.003991737496107817, "xla_temp_mb": 94.483154296875, "xla_arg_mb": 117.1875, "xla_out_mb": 70.3125228881836, "compile_s": 0.13496070099063218, "inner": 33, "median_ms": 1.3785150367766619, "q1_ms": 1.370454439893365, "q3_ms": 1.4422642257572575, "tflops": 374.38547003939175 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 128, "seq": 3000, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0009522885084152222, "rel_err": 0.003991737496107817, "xla_temp_mb": 2746.948257446289, "xla_arg_mb": 117.1875, "xla_out_mb": 70.3125228881836, "compile_s": 0.03949791193008423, "inner": 9, "median_ms": 5.522362680898772, "q1_ms": 5.510699883517292, "q3_ms": 5.5299548872022175, "tflops": 93.45565110837751 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 128, "seq": 3000, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0009522885084152222, "rel_err": 0.003991737496107817, "xla_temp_mb": 94.483154296875, "xla_arg_mb": 117.1875, "xla_out_mb": 70.3125228881836, "compile_s": 0.13363748393021524, "inner": 34, "median_ms": 1.3839300338398008, "q1_ms": 1.3746528660275918, "q3_ms": 1.3899216325679684, "tflops": 372.92058657622977 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 128, "seq": 3000, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0070171356201171875, "rel_err": 0.0024275495670735836, "xla_temp_mb": 96.0, "xla_arg_mb": 70.3125, "xla_out_mb": 23.4375, "compile_s": 0.39901529904454947, "inner": 95, "median_ms": 0.427930726130542, "q1_ms": 0.42130084169146265, "q3_ms": 0.43685643713137035, "tflops": 172.289568142646 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 128, "seq": 3000, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0070171356201171875, "rel_err": 0.0024275495670735836, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 70.3125, "xla_out_mb": 23.4375, "compile_s": 0.05727446312084794, "inner": 142, "median_ms": 0.2222515062206018, "q1_ms": 0.21441913091800582, "q3_ms": 0.2228316086882227, "tflops": 331.73228498536815 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 128, "seq": 3000, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0070171356201171875, "rel_err": 0.0024275495670735836, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 70.3125, "xla_out_mb": 23.4375, "compile_s": 0.054976239101961255, "inner": 140, "median_ms": 0.2226214639709464, "q1_ms": 0.2173129078333399, "q3_ms": 0.22268003168781952, "tflops": 331.18100422527993 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 128, "seq": 3000, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0070171356201171875, "rel_err": 0.0024275495670735836, "xla_temp_mb": 1647.949234008789, "xla_arg_mb": 70.3125, "xla_out_mb": 23.4375, "compile_s": 0.43164451001212, "inner": 25, "median_ms": 1.8820304423570633, "q1_ms": 1.8785502621904016, "q3_ms": 1.8848444428294897, "tflops": 39.174711705333905 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 128, "seq": 3000, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0070171356201171875, "rel_err": 0.0024275495670735836, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 70.3125, "xla_out_mb": 23.4375, "compile_s": 0.059482666896656156, "inner": 147, "median_ms": 0.22191013729035045, "q1_ms": 0.21516672750542154, "q3_ms": 0.22226789820508486, "tflops": 332.2425955851364 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 128, "seq": 3000, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.011957883834838867, "rel_err": 0.004573028534650803, "xla_temp_mb": 192.00048828125, "xla_arg_mb": 117.1875, "xla_out_mb": 70.3125228881836, "compile_s": 0.7846149348188192, "inner": 35, "median_ms": 1.3715828037155526, "q1_ms": 1.322959253697523, "q3_ms": 1.3800546393862794, "tflops": 188.13884170970957 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 128, "seq": 3000, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.011957883834838867, "rel_err": 0.004573028534650803, "xla_temp_mb": 94.483154296875, "xla_arg_mb": 117.1875, "xla_out_mb": 70.3125228881836, "compile_s": 0.13818088197149336, "inner": 48, "median_ms": 0.9411184582859278, "q1_ms": 0.9243816263430441, "q3_ms": 0.9423480538922983, "tflops": 274.1929007215377 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 128, "seq": 3000, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.011957883834838867, "rel_err": 0.004573028534650803, "xla_temp_mb": 94.483154296875, "xla_arg_mb": 117.1875, "xla_out_mb": 70.3125228881836, "compile_s": 0.13848496112041175, "inner": 47, "median_ms": 0.9352802150347765, "q1_ms": 0.917511318433792, "q3_ms": 0.9396689265292693, "tflops": 275.9044785208088 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 128, "seq": 3000, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.011933326721191406, "rel_err": 0.004184063524007797, "xla_temp_mb": 2746.948257446289, "xla_arg_mb": 117.1875, "xla_out_mb": 70.3125228881836, "compile_s": 0.04498838703148067, "inner": 8, "median_ms": 5.854256625752896, "q1_ms": 5.828710942296311, "q3_ms": 5.8592465065885335, "tflops": 44.07869632240683 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 128, "seq": 3000, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.011957883834838867, "rel_err": 0.004573028534650803, "xla_temp_mb": 94.483154296875, "xla_arg_mb": 117.1875, "xla_out_mb": 70.3125228881836, "compile_s": 0.13841020199470222, "inner": 48, "median_ms": 0.9416644995023186, "q1_ms": 0.929530479576594, "q3_ms": 0.9422885608122064, "tflops": 274.0339050016026 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 128, "seq": 4096, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.000495031476020813, "rel_err": 0.003562052734196186, "xla_temp_mb": 0.0, "xla_arg_mb": 96.0, "xla_out_mb": 32.0, "compile_s": 0.32331868587061763, "inner": 52, "median_ms": 0.884610098392631, "q1_ms": 0.8774993363588762, "q3_ms": 0.8934252590943987, "tflops": 310.73340383911875 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 128, "seq": 4096, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.000495031476020813, "rel_err": 0.003562052734196186, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 96.0, "xla_out_mb": 32.0, "compile_s": 0.05337549396790564, "inner": 75, "median_ms": 0.5764762933055559, "q1_ms": 0.5733494320884347, "q3_ms": 0.5886249182124933, "tflops": 476.8243033340202 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 128, "seq": 4096, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.000495031476020813, "rel_err": 0.003562052734196186, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 96.0, "xla_out_mb": 32.0, "compile_s": 0.05411775200627744, "inner": 76, "median_ms": 0.5865222100462569, "q1_ms": 0.5760730342253259, "q3_ms": 0.5977890855559197, "tflops": 468.6572856675305 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 128, "seq": 4096, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.000495031476020813, "rel_err": 0.003562052734196186, "xla_temp_mb": 3072.0, "xla_arg_mb": 96.0, "xla_out_mb": 32.0, "compile_s": 0.020255984971299767, "inner": 15, "median_ms": 3.219590460260709, "q1_ms": 3.215200236688058, "q3_ms": 3.222222168308993, "tflops": 85.37666834860157 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 128, "seq": 4096, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.000495031476020813, "rel_err": 0.003562052734196186, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 96.0, "xla_out_mb": 32.0, "compile_s": 0.05811002105474472, "inner": 77, "median_ms": 0.5781998428479804, "q1_ms": 0.5720283718199118, "q3_ms": 0.5861100326473643, "tflops": 475.40294302063745 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 128, "seq": 4096, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.000783964991569519, "rel_err": 0.004214399494230747, "xla_temp_mb": 65.000244140625, "xla_arg_mb": 160.0, "xla_out_mb": 96.0000228881836, "compile_s": 0.6291449191048741, "inner": 14, "median_ms": 3.516918214570199, "q1_ms": 3.4785101472932314, "q3_ms": 3.5399403256763304, "tflops": 273.55560055910325 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 128, "seq": 4096, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.000783964991569519, "rel_err": 0.004214399494230747, "xla_temp_mb": 129.000732421875, "xla_arg_mb": 160.0, "xla_out_mb": 96.0000228881836, "compile_s": 0.1351878708228469, "inner": 20, "median_ms": 2.3853269522078335, "q1_ms": 2.3363581218291074, "q3_ms": 2.4176076462026685, "tflops": 403.32947792063294 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 128, "seq": 4096, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.000783964991569519, "rel_err": 0.004214399494230747, "xla_temp_mb": 129.000732421875, "xla_arg_mb": 160.0, "xla_out_mb": 96.0000228881836, "compile_s": 0.13440350210294127, "inner": 20, "median_ms": 2.301324694417417, "q1_ms": 2.2917957277968526, "q3_ms": 2.4183601257391274, "tflops": 418.0516885069752 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 128, "seq": 4096, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0010238289833068848, "rel_err": 0.004597060848027468, "xla_temp_mb": 5120.000015258789, "xla_arg_mb": 160.0, "xla_out_mb": 96.0000228881836, "compile_s": 0.04208004497922957, "inner": 5, "median_ms": 8.734190184623003, "q1_ms": 8.732028002850711, "q3_ms": 8.737523993477225, "tflops": 110.15018610400527 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 128, "seq": 4096, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.000783964991569519, "rel_err": 0.004214399494230747, "xla_temp_mb": 129.000732421875, "xla_arg_mb": 160.0, "xla_out_mb": 96.0000228881836, "compile_s": 0.13510232092812657, "inner": 20, "median_ms": 2.370409050490707, "q1_ms": 2.3349906085059047, "q3_ms": 2.4243179010227323, "tflops": 405.8677864501225 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 128, "seq": 4096, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.006810665130615234, "rel_err": 0.002356122015044093, "xla_temp_mb": 0.0, "xla_arg_mb": 96.0, "xla_out_mb": 32.0, "compile_s": 0.376660960027948, "inner": 83, "median_ms": 0.5327075809599405, "q1_ms": 0.5278992831303053, "q3_ms": 0.5469304938654763, "tflops": 258.00074634630624 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 128, "seq": 4096, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.006810665130615234, "rel_err": 0.002356122015044093, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 96.0, "xla_out_mb": 32.0, "compile_s": 0.05828284099698067, "inner": 105, "median_ms": 0.3637194855227357, "q1_ms": 0.3552797904592895, "q3_ms": 0.36757628189488534, "tflops": 377.87074639257634 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 128, "seq": 4096, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.006810665130615234, "rel_err": 0.002356122015044093, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 96.0, "xla_out_mb": 32.0, "compile_s": 0.05639642593450844, "inner": 111, "median_ms": 0.3646044060587883, "q1_ms": 0.3550488969546046, "q3_ms": 0.36748615722801237, "tflops": 376.9536275155148 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 128, "seq": 4096, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.006975531578063965, "rel_err": 0.002413156908005476, "xla_temp_mb": 3072.0, "xla_arg_mb": 96.0, "xla_out_mb": 32.0, "compile_s": 0.5593992099165916, "inner": 12, "median_ms": 3.9074518329774337, "q1_ms": 3.8999089641341316, "q3_ms": 3.9157867819691696, "tflops": 35.17355026927436 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 128, "seq": 4096, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.006810665130615234, "rel_err": 0.002356122015044093, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 96.0, "xla_out_mb": 32.0, "compile_s": 0.05873015895485878, "inner": 108, "median_ms": 0.35995836973328277, "q1_ms": 0.3510655186363254, "q3_ms": 0.3622200747486204, "tflops": 381.81902416614923 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 128, "seq": 4096, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.015851497650146484, "rel_err": 0.003916752990335226, "xla_temp_mb": 65.000244140625, "xla_arg_mb": 160.0, "xla_out_mb": 96.0000228881836, "compile_s": 0.7542233171407133, "inner": 25, "median_ms": 1.9766114000231028, "q1_ms": 1.9506106618791819, "q3_ms": 2.0357394218444824, "tflops": 243.36414185731076 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 128, "seq": 4096, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.015851497650146484, "rel_err": 0.003916752990335226, "xla_temp_mb": 129.000732421875, "xla_arg_mb": 160.0, "xla_out_mb": 96.0000228881836, "compile_s": 0.13595719798468053, "inner": 31, "median_ms": 1.5126749077030728, "q1_ms": 1.5069573381615262, "q3_ms": 1.5131395813378115, "tflops": 318.0037790687171 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 128, "seq": 4096, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.015851497650146484, "rel_err": 0.003916752990335226, "xla_temp_mb": 129.000732421875, "xla_arg_mb": 160.0, "xla_out_mb": 96.0000228881836, "compile_s": 0.13677743612788618, "inner": 32, "median_ms": 1.5053289971547201, "q1_ms": 1.4700105202791747, "q3_ms": 1.5220020795823075, "tflops": 319.5556174505541 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 128, "seq": 4096, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.015851497650146484, "rel_err": 0.004538213834166527, "xla_temp_mb": 5152.000015258789, "xla_arg_mb": 160.0, "xla_out_mb": 96.0000228881836, "compile_s": 0.04875202709808946, "inner": 4, "median_ms": 10.54205879336223, "q1_ms": 10.525441117351875, "q3_ms": 10.55330625968054, "tflops": 45.630208157716105 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 128, "seq": 4096, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.015851497650146484, "rel_err": 0.003916752990335226, "xla_temp_mb": 129.000732421875, "xla_arg_mb": 160.0, "xla_out_mb": 96.0000228881836, "compile_s": 0.1372115151025355, "inner": 31, "median_ms": 1.508261355000638, "q1_ms": 1.4803365789233676, "q3_ms": 1.5142207116549533, "tflops": 318.93433824126356 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 128, "seq": 8192, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0003512650728225708, "rel_err": 0.003266321960836649, "xla_temp_mb": 0.0, "xla_arg_mb": 192.0, "xla_out_mb": 64.0, "compile_s": 0.32673251698724926, "inner": 15, "median_ms": 3.355496795848012, "q1_ms": 3.347341188540061, "q3_ms": 3.3813527319580317, "tflops": 327.67476611406744 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 128, "seq": 8192, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0003512650728225708, "rel_err": 0.003266321960836649, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 192.0, "xla_out_mb": 64.0, "compile_s": 0.05560096795670688, "inner": 22, "median_ms": 2.2701883095909245, "q1_ms": 2.236415044700896, "q3_ms": 2.3073680032129316, "tflops": 484.32617819890277 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 128, "seq": 8192, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0003512650728225708, "rel_err": 0.003266321960836649, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 192.0, "xla_out_mb": 64.0, "compile_s": 0.05486185988411307, "inner": 22, "median_ms": 2.2735128560188143, "q1_ms": 2.2694910664788703, "q3_ms": 2.3095568652603435, "tflops": 483.61795046163627 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 128, "seq": 8192, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.000379122793674469, "rel_err": 0.0035253637470304966, "xla_temp_mb": 12288.000015258789, "xla_arg_mb": 192.0, "xla_out_mb": 64.0, "compile_s": 0.029108510119840503, "inner": 4, "median_ms": 12.416483485139906, "q1_ms": 12.413351010764018, "q3_ms": 12.43111613439396, "tflops": 88.55257844073965 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 128, "seq": 8192, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0003512650728225708, "rel_err": 0.003266321960836649, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 192.0, "xla_out_mb": 64.0, "compile_s": 0.05757425189949572, "inner": 22, "median_ms": 2.277453768659722, "q1_ms": 2.2418611333705485, "q3_ms": 2.3196902475319803, "tflops": 482.78109655023235 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 128, "seq": 8192, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0006157159805297852, "rel_err": 0.005359946750104427, "xla_temp_mb": 130.000244140625, "xla_arg_mb": 320.0, "xla_out_mb": 192.0000228881836, "compile_s": 0.6228717670310289, "inner": 3, "median_ms": 13.586166392390927, "q1_ms": 13.454311682532232, "q3_ms": 13.929998657355707, "tflops": 283.2506673384536 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 128, "seq": 8192, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0006157159805297852, "rel_err": 0.005359946750104427, "xla_temp_mb": 258.000732421875, "xla_arg_mb": 320.0, "xla_out_mb": 192.0000228881836, "compile_s": 0.1421828221064061, "inner": 5, "median_ms": 8.8779057841748, "q1_ms": 8.85526379570365, "q3_ms": 9.086013119667768, "tflops": 433.4682965520677 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 128, "seq": 8192, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0006157159805297852, "rel_err": 0.005359946750104427, "xla_temp_mb": 258.000732421875, "xla_arg_mb": 320.0, "xla_out_mb": 192.0000228881836, "compile_s": 0.14296656800433993, "inner": 5, "median_ms": 8.79976600408554, "q1_ms": 8.60941840801388, "q3_ms": 9.044585190713406, "tflops": 437.31738951118945 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 128, "seq": 8192, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0006582215428352356, "rel_err": 0.005359946750104427, "xla_temp_mb": 20480.00001525879, "xla_arg_mb": 320.0, "xla_out_mb": 192.0000228881836, "compile_s": 0.06629871902987361, "inner": 1, "median_ms": 33.73347781598568, "q1_ms": 33.67433499079198, "q3_ms": 33.79080246668309, "tflops": 114.07927514050643 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 128, "seq": 8192, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0006157159805297852, "rel_err": 0.005359946750104427, "xla_temp_mb": 258.000732421875, "xla_arg_mb": 320.0, "xla_out_mb": 192.0000228881836, "compile_s": 0.14206700399518013, "inner": 5, "median_ms": 8.864306379109621, "q1_ms": 8.80007571540773, "q3_ms": 8.880838518962264, "tflops": 434.13331315862564 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 128, "seq": 8192, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.004915475845336914, "rel_err": 0.0017004889668896794, "xla_temp_mb": 0.0, "xla_arg_mb": 192.0, "xla_out_mb": 64.0, "compile_s": 0.3745421590283513, "inner": 26, "median_ms": 1.8891427003277035, "q1_ms": 1.8793088652623389, "q3_ms": 1.9298939240308335, "tflops": 291.0080926086926 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 128, "seq": 8192, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.004915475845336914, "rel_err": 0.0017004889668896794, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 192.0, "xla_out_mb": 64.0, "compile_s": 0.05522048403508961, "inner": 38, "median_ms": 1.2741321575288709, "q1_ms": 1.2361847483062822, "q3_ms": 1.2820064868336838, "tflops": 431.474718410867 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 128, "seq": 8192, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.004915475845336914, "rel_err": 0.0017004889668896794, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 192.0, "xla_out_mb": 64.0, "compile_s": 0.054860214004293084, "inner": 38, "median_ms": 1.2749687364128859, "q1_ms": 1.2504540684044754, "q3_ms": 1.3021243166325516, "tflops": 431.1916035170662 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 128, "seq": 8192, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.011793375015258789, "rel_err": 0.004079870413988829, "xla_temp_mb": 12288.000015258789, "xla_arg_mb": 192.0, "xla_out_mb": 64.0, "compile_s": 0.514892578125, "inner": 3, "median_ms": 15.848161264633138, "q1_ms": 15.724598507707316, "q3_ms": 15.879462977560857, "tflops": 34.688933606123676 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 128, "seq": 8192, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.004915475845336914, "rel_err": 0.0017004889668896794, "xla_temp_mb": 0.00048828125, "xla_arg_mb": 192.0, "xla_out_mb": 64.0, "compile_s": 0.0586824850179255, "inner": 37, "median_ms": 1.2912784602392364, "q1_ms": 1.2390187444366716, "q3_ms": 1.3006019745827526, "tflops": 425.7453607536722 }, { "impl": "probjax_mha", "dtype": "bfloat16", "head_dim": 128, "seq": 8192, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.01760244369506836, "rel_err": 0.005058928392827511, "xla_temp_mb": 130.000244140625, "xla_arg_mb": 320.0, "xla_out_mb": 192.0000228881836, "compile_s": 0.7510640430264175, "inner": 7, "median_ms": 7.358489125700934, "q1_ms": 7.281649990805558, "q3_ms": 7.520703632118446, "tflops": 261.4864703526643 }, { "impl": "probjax_flex", "dtype": "bfloat16", "head_dim": 128, "seq": 8192, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.01760244369506836, "rel_err": 0.005058928392827511, "xla_temp_mb": 258.000732421875, "xla_arg_mb": 320.0, "xla_out_mb": 192.0000228881836, "compile_s": 0.13877141289412975, "inner": 10, "median_ms": 5.011330684646964, "q1_ms": 4.986435291357338, "q3_ms": 5.088557046838105, "tflops": 383.95896612908336 }, { "impl": "probjax_default", "dtype": "bfloat16", "head_dim": 128, "seq": 8192, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.01760244369506836, "rel_err": 0.005058928392827511, "xla_temp_mb": 258.000732421875, "xla_arg_mb": 320.0, "xla_out_mb": 192.0000228881836, "compile_s": 0.14005058002658188, "inner": 10, "median_ms": 4.889067099429667, "q1_ms": 4.881027084775269, "q3_ms": 5.062533658929169, "tflops": 393.5608388014271 }, { "impl": "jax_xla", "dtype": "bfloat16", "head_dim": 128, "seq": 8192, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.01760244369506836, "rel_err": 0.00520186685025692, "xla_temp_mb": 20480.00001525879, "xla_arg_mb": 320.0, "xla_out_mb": 192.0000228881836, "compile_s": 0.0792325499933213, "inner": 1, "median_ms": 40.3525831643492, "q1_ms": 40.21145903971046, "q3_ms": 40.41439283173531, "tflops": 47.68332527241896 }, { "impl": "jax_cudnn", "dtype": "bfloat16", "head_dim": 128, "seq": 8192, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.01760244369506836, "rel_err": 0.005058928392827511, "xla_temp_mb": 258.000732421875, "xla_arg_mb": 320.0, "xla_out_mb": 192.0000228881836, "compile_s": 0.1422828531358391, "inner": 10, "median_ms": 5.043369601480663, "q1_ms": 5.0140852108597755, "q3_ms": 5.083770107012242, "tflops": 381.5197974074908 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 64, "seq": 128, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0004020603373646736, "rel_err": 0.00043028933578170836, "xla_temp_mb": 0.0, "xla_arg_mb": 3.0, "xla_out_mb": 1.0, "compile_s": 0.6630640050861984, "inner": 200, "median_ms": 0.03833124530501664, "q1_ms": 0.03803140018135309, "q3_ms": 0.038388899993151426, "tflops": 3.5015227637917135 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 64, "seq": 128, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0004020603373646736, "rel_err": 0.00043028933578170836, "xla_temp_mb": 0.0, "xla_arg_mb": 3.0, "xla_out_mb": 1.0, "compile_s": 0.6649310300126672, "inner": 200, "median_ms": 0.038431850261986256, "q1_ms": 0.038255450199358165, "q3_ms": 0.038559322129003704, "tflops": 3.492356654312778 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 64, "seq": 128, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0004020603373646736, "rel_err": 0.00043028933578170836, "xla_temp_mb": 0.0, "xla_arg_mb": 3.0, "xla_out_mb": 1.0, "compile_s": 0.6666091170627624, "inner": 200, "median_ms": 0.02383923507295549, "q1_ms": 0.023670162190683186, "q3_ms": 0.023991313064470887, "tflops": 5.630118902274 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 64, "seq": 128, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00037729740142822266, "rel_err": 0.0004037877661176026, "xla_temp_mb": 4.0, "xla_arg_mb": 3.0, "xla_out_mb": 1.0, "compile_s": 0.17435061908327043, "inner": 200, "median_ms": 0.03456091042608023, "q1_ms": 0.03437363542616367, "q3_ms": 0.034847757196985185, "tflops": 3.8835125100962937 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 64, "seq": 128, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0005728304386138916, "rel_err": 0.0007045819074846804, "xla_temp_mb": 1.031494140625, "xla_arg_mb": 4.0, "xla_out_mb": 3.0000228881835938, "compile_s": 0.8486492349766195, "inner": 200, "median_ms": 0.0343070097733289, "q1_ms": 0.03427083254791796, "q3_ms": 0.034395832917653024, "tflops": 13.692888162033999 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 64, "seq": 128, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0005728304386138916, "rel_err": 0.0007045819074846804, "xla_temp_mb": 0.03173828125, "xla_arg_mb": 4.0, "xla_out_mb": 3.0000228881835938, "compile_s": 0.8691200409084558, "inner": 200, "median_ms": 0.036857599625363946, "q1_ms": 0.036656277370639145, "q3_ms": 0.03716780222021043, "tflops": 12.745323970493411 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 64, "seq": 128, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0005728304386138916, "rel_err": 0.0007045819074846804, "xla_temp_mb": 0.03173828125, "xla_arg_mb": 4.0, "xla_out_mb": 3.0000228881835938, "compile_s": 0.8487190550658852, "inner": 200, "median_ms": 0.037257850635796785, "q1_ms": 0.03701087727677077, "q3_ms": 0.03728932701051235, "tflops": 12.608404402927626 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 64, "seq": 128, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0005763769149780273, "rel_err": 0.0007045819074846804, "xla_temp_mb": 6.0, "xla_arg_mb": 4.0, "xla_out_mb": 3.0000228881835938, "compile_s": 0.03122011781670153, "inner": 200, "median_ms": 0.05862234509550035, "q1_ms": 0.05853809474501759, "q3_ms": 0.058677069609984756, "tflops": 8.013361581402469 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 64, "seq": 128, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0009530782699584961, "rel_err": 0.00031446287175640464, "xla_temp_mb": 0.0, "xla_arg_mb": 3.0, "xla_out_mb": 1.0, "compile_s": 1.100318348966539, "inner": 200, "median_ms": 0.023094939533621073, "q1_ms": 0.022867265506647527, "q3_ms": 0.023132787318900228, "tflops": 2.9057821910425217 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 64, "seq": 128, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0009530782699584961, "rel_err": 0.00031446287175640464, "xla_temp_mb": 0.0, "xla_arg_mb": 3.0, "xla_out_mb": 1.0, "compile_s": 0.03865455696359277, "inner": 200, "median_ms": 0.024037935072556138, "q1_ms": 0.023910337476991117, "q3_ms": 0.024571759859099984, "tflops": 2.79178988533909 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 64, "seq": 128, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0009530782699584961, "rel_err": 0.00031446287175640464, "xla_temp_mb": 0.0, "xla_arg_mb": 3.0, "xla_out_mb": 1.0, "compile_s": 0.03872656705789268, "inner": 200, "median_ms": 0.023288684897124767, "q1_ms": 0.023195089888758957, "q3_ms": 0.023563890135847032, "tflops": 2.881608141311805 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 64, "seq": 128, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0013991594314575195, "rel_err": 0.0004616449004970491, "xla_temp_mb": 4.0, "xla_arg_mb": 3.0, "xla_out_mb": 1.0, "compile_s": 0.282288053072989, "inner": 200, "median_ms": 0.03435685997828841, "q1_ms": 0.03411031037103385, "q3_ms": 0.03443255787715316, "tflops": 1.9532886312197622 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 64, "seq": 128, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0016770362854003906, "rel_err": 0.0007809422677382827, "xla_temp_mb": 1.031494140625, "xla_arg_mb": 4.0, "xla_out_mb": 3.0000228881835938, "compile_s": 1.1731115658767521, "inner": 200, "median_ms": 0.03499935497529805, "q1_ms": 0.03494198259431869, "q3_ms": 0.035437155165709555, "tflops": 6.71101007906503 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 64, "seq": 128, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0016770362854003906, "rel_err": 0.0007809422677382827, "xla_temp_mb": 0.03173828125, "xla_arg_mb": 4.0, "xla_out_mb": 3.0000228881835938, "compile_s": 1.1963365951087326, "inner": 200, "median_ms": 0.05995368934236467, "q1_ms": 0.03955436986871064, "q3_ms": 0.06020159227773547, "tflops": 3.9177075935846974 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 64, "seq": 128, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0016770362854003906, "rel_err": 0.0007809422677382827, "xla_temp_mb": 0.03173828125, "xla_arg_mb": 4.0, "xla_out_mb": 3.0000228881835938, "compile_s": 1.2055549500510097, "inner": 200, "median_ms": 0.045800129882991314, "q1_ms": 0.04555818042717874, "q3_ms": 0.046112252748571336, "tflops": 5.128392094085027 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 64, "seq": 128, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0016134977340698242, "rel_err": 0.0007285389583557844, "xla_temp_mb": 6.0, "xla_arg_mb": 4.0, "xla_out_mb": 3.0000228881835938, "compile_s": 0.7476080220658332, "inner": 200, "median_ms": 0.05851370049640536, "q1_ms": 0.058463921886868775, "q3_ms": 0.05875304515939206, "tflops": 4.014120146348107 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 64, "seq": 512, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0003509223461151123, "rel_err": 0.0008548463811166584, "xla_temp_mb": 0.0, "xla_arg_mb": 12.0, "xla_out_mb": 4.0, "compile_s": 0.9365642420016229, "inner": 200, "median_ms": 0.04150834050960839, "q1_ms": 0.04150351509451866, "q3_ms": 0.04151873989030719, "tflops": 51.73619618695425 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 64, "seq": 512, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0003509223461151123, "rel_err": 0.0008548463811166584, "xla_temp_mb": 0.0, "xla_arg_mb": 12.0, "xla_out_mb": 4.0, "compile_s": 0.9402725519612432, "inner": 200, "median_ms": 0.04162663943134248, "q1_ms": 0.04138051706831902, "q3_ms": 0.04166546277701855, "tflops": 51.589166873342826 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 64, "seq": 512, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0003509223461151123, "rel_err": 0.0008548463811166584, "xla_temp_mb": 0.0, "xla_arg_mb": 12.0, "xla_out_mb": 4.0, "compile_s": 0.944615620886907, "inner": 200, "median_ms": 0.04132283967919648, "q1_ms": 0.041316040442325175, "q3_ms": 0.04161346470937133, "tflops": 51.96844323070873 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 64, "seq": 512, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0004189014434814453, "rel_err": 0.0010204433929175138, "xla_temp_mb": 68.00001525878906, "xla_arg_mb": 12.0, "xla_out_mb": 4.0, "compile_s": 0.02075071493163705, "inner": 200, "median_ms": 0.07451220415532589, "q1_ms": 0.07386907760519534, "q3_ms": 0.0751804024912417, "tflops": 28.820562649353665 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 64, "seq": 512, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00036704540252685547, "rel_err": 0.0007074366440065205, "xla_temp_mb": 4.125244140625, "xla_arg_mb": 16.0, "xla_out_mb": 12.000022888183594, "compile_s": 1.0255694061052054, "inner": 200, "median_ms": 0.13854203512892127, "q1_ms": 0.13662021548952907, "q3_ms": 0.1391879323637113, "tflops": 54.2520741882112 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 64, "seq": 512, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00036704540252685547, "rel_err": 0.0007074366440065205, "xla_temp_mb": 0.12548828125, "xla_arg_mb": 16.0, "xla_out_mb": 12.000022888183594, "compile_s": 1.0183339659124613, "inner": 200, "median_ms": 0.13984228018671274, "q1_ms": 0.13973818044178188, "q3_ms": 0.14189262525178492, "tflops": 53.7476416857951 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 64, "seq": 512, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00036704540252685547, "rel_err": 0.0007074366440065205, "xla_temp_mb": 0.12548828125, "xla_arg_mb": 16.0, "xla_out_mb": 12.000022888183594, "compile_s": 1.0164034711197019, "inner": 200, "median_ms": 0.13871543458662927, "q1_ms": 0.1386265578912571, "q3_ms": 0.13935368217062205, "tflops": 54.184257075632466 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 64, "seq": 512, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.000362396240234375, "rel_err": 0.0007074366440065205, "xla_temp_mb": 96.0, "xla_arg_mb": 16.0, "xla_out_mb": 12.000022888183594, "compile_s": 0.03555755619890988, "inner": 156, "median_ms": 0.20126452615771157, "q1_ms": 0.19958626571445703, "q3_ms": 0.20153129410643417, "tflops": 37.344846165837914 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 64, "seq": 512, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0012368559837341309, "rel_err": 0.0004327311762608588, "xla_temp_mb": 0.0, "xla_arg_mb": 12.0, "xla_out_mb": 4.0, "compile_s": 1.1136323739774525, "inner": 200, "median_ms": 0.0403639895375818, "q1_ms": 0.04028999479487538, "q3_ms": 0.04046909511089325, "tflops": 26.601479098102242 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 64, "seq": 512, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0012368559837341309, "rel_err": 0.0004327311762608588, "xla_temp_mb": 0.0, "xla_arg_mb": 12.0, "xla_out_mb": 4.0, "compile_s": 0.038917966885492206, "inner": 200, "median_ms": 0.04022514563985169, "q1_ms": 0.040170992724597454, "q3_ms": 0.04025669244583696, "tflops": 26.69329860514481 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 64, "seq": 512, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0012368559837341309, "rel_err": 0.0004327311762608588, "xla_temp_mb": 0.0, "xla_arg_mb": 12.0, "xla_out_mb": 4.0, "compile_s": 0.03960388619452715, "inner": 200, "median_ms": 0.04018069477751851, "q1_ms": 0.04014434525743127, "q3_ms": 0.040189792634919286, "tflops": 26.72282871028823 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 64, "seq": 512, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0017076730728149414, "rel_err": 0.0005974530358798802, "xla_temp_mb": 68.00001525878906, "xla_arg_mb": 12.0, "xla_out_mb": 4.0, "compile_s": 0.6701728869229555, "inner": 200, "median_ms": 0.10426916996948421, "q1_ms": 0.1040872506564483, "q3_ms": 0.10429349960759282, "tflops": 10.29778816033776 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 64, "seq": 512, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0020322799682617188, "rel_err": 0.000902797793969512, "xla_temp_mb": 4.125244140625, "xla_arg_mb": 16.0, "xla_out_mb": 12.000022888183594, "compile_s": 1.1966907540336251, "inner": 200, "median_ms": 0.10613407008349895, "q1_ms": 0.10601942020002753, "q3_ms": 0.10715468961279839, "tflops": 35.408953798185536 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 64, "seq": 512, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0020322799682617188, "rel_err": 0.000902797793969512, "xla_temp_mb": 0.12548828125, "xla_arg_mb": 16.0, "xla_out_mb": 12.000022888183594, "compile_s": 1.212104853009805, "inner": 200, "median_ms": 0.12367747491225599, "q1_ms": 0.12147467990871519, "q3_ms": 0.12420694576576351, "tflops": 30.38626384203116 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 64, "seq": 512, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0020322799682617188, "rel_err": 0.000902797793969512, "xla_temp_mb": 0.12548828125, "xla_arg_mb": 16.0, "xla_out_mb": 12.000022888183594, "compile_s": 1.2122641720343381, "inner": 200, "median_ms": 0.12520771473646164, "q1_ms": 0.12278757523745298, "q3_ms": 0.12553969281725585, "tflops": 30.01489478431961 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 64, "seq": 512, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0012679100036621094, "rel_err": 0.0006988496752455831, "xla_temp_mb": 96.0, "xla_arg_mb": 16.0, "xla_out_mb": 12.000022888183594, "compile_s": 0.03832619916647673, "inner": 122, "median_ms": 0.2966234021530044, "q1_ms": 0.29659171428409264, "q3_ms": 0.29669200035850285, "tflops": 12.669588295199636 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 64, "seq": 600, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00024455785751342773, "rel_err": 0.0005569846834987402, "xla_temp_mb": 20.0, "xla_arg_mb": 14.0625, "xla_out_mb": 4.6875, "compile_s": 1.116810865001753, "inner": 200, "median_ms": 0.11393139488063753, "q1_ms": 0.1139017246896401, "q3_ms": 0.11440064990893006, "tflops": 25.885051289767002 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 64, "seq": 600, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00024455785751342773, "rel_err": 0.0005569846834987402, "xla_temp_mb": 20.0, "xla_arg_mb": 14.0625, "xla_out_mb": 4.6875, "compile_s": 1.1850355549249798, "inner": 200, "median_ms": 0.11446584947407246, "q1_ms": 0.11387902253773063, "q3_ms": 0.11449932237155735, "tflops": 25.76419092288309 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 64, "seq": 600, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00024455785751342773, "rel_err": 0.0005569846834987402, "xla_temp_mb": 20.0, "xla_arg_mb": 14.0625, "xla_out_mb": 4.6875, "compile_s": 0.04143515112809837, "inner": 200, "median_ms": 0.11355954455211759, "q1_ms": 0.11290294991340488, "q3_ms": 0.11376350012142211, "tflops": 25.96981179901189 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 64, "seq": 600, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0003241300582885742, "rel_err": 0.0007382117328234017, "xla_temp_mb": 87.89064025878906, "xla_arg_mb": 14.0625, "xla_out_mb": 4.6875, "compile_s": 0.021855521947145462, "inner": 197, "median_ms": 0.12191013210845478, "q1_ms": 0.12134473066455519, "q3_ms": 0.12203066746898109, "tflops": 24.19093433002252 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 64, "seq": 600, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.000510871410369873, "rel_err": 0.0005970043712295592, "xla_temp_mb": 40.00048828125, "xla_arg_mb": 18.75, "xla_out_mb": 14.062522888183594, "compile_s": 1.2177663380280137, "inner": 111, "median_ms": 0.3213767121705386, "q1_ms": 0.3195247384258085, "q3_ms": 0.3219791488755528, "tflops": 32.1178218866172 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 64, "seq": 600, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.000510871410369873, "rel_err": 0.0005970043712295592, "xla_temp_mb": 30.078651428222656, "xla_arg_mb": 18.75, "xla_out_mb": 14.062522888183594, "compile_s": 1.2376245758496225, "inner": 117, "median_ms": 0.3303502383840899, "q1_ms": 0.32704687104202235, "q3_ms": 0.33055467861425925, "tflops": 31.24538384016228 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 64, "seq": 600, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.000510871410369873, "rel_err": 0.0005970043712295592, "xla_temp_mb": 30.078651428222656, "xla_arg_mb": 18.75, "xla_out_mb": 14.062522888183594, "compile_s": 1.2615070620086044, "inner": 115, "median_ms": 0.32789034809431306, "q1_ms": 0.32695591368752974, "q3_ms": 0.33112643185355095, "tflops": 31.479792131700822 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 64, "seq": 600, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0005156397819519043, "rel_err": 0.0005970043712295592, "xla_temp_mb": 131.9091796875, "xla_arg_mb": 18.75, "xla_out_mb": 14.062522888183594, "compile_s": 0.031787126092240214, "inner": 114, "median_ms": 0.3374705087034064, "q1_ms": 0.33560963633486574, "q3_ms": 0.33878134216361666, "tflops": 30.586139333056963 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 64, "seq": 600, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0010406970977783203, "rel_err": 0.00044006036478094757, "xla_temp_mb": 20.0, "xla_arg_mb": 14.0625, "xla_out_mb": 4.6875, "compile_s": 1.1803848869167268, "inner": 200, "median_ms": 0.08504388039000332, "q1_ms": 0.08501277538016438, "q3_ms": 0.08564631978515536, "tflops": 17.33881371872738 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 64, "seq": 600, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0010406970977783203, "rel_err": 0.00044006036478094757, "xla_temp_mb": 20.0, "xla_arg_mb": 14.0625, "xla_out_mb": 4.6875, "compile_s": 1.1696537958923727, "inner": 200, "median_ms": 0.08326962939463556, "q1_ms": 0.08284187759272754, "q3_ms": 0.08328508003614843, "tflops": 17.708257028642365 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 64, "seq": 600, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0010406970977783203, "rel_err": 0.00044006036478094757, "xla_temp_mb": 20.0, "xla_arg_mb": 14.0625, "xla_out_mb": 4.6875, "compile_s": 0.042617907049134374, "inner": 200, "median_ms": 0.08258778019808233, "q1_ms": 0.0825522851664573, "q3_ms": 0.08261695504188538, "tflops": 17.85445735995504 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 64, "seq": 600, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.000980377197265625, "rel_err": 0.0004145540005993098, "xla_temp_mb": 87.89064025878906, "xla_arg_mb": 14.0625, "xla_out_mb": 4.6875, "compile_s": 0.49110779096372426, "inner": 200, "median_ms": 0.12559406459331512, "q1_ms": 0.12509521679021418, "q3_ms": 0.12577234243508428, "tflops": 11.740682211175805 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 64, "seq": 600, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0018483996391296387, "rel_err": 0.0007540053920820355, "xla_temp_mb": 40.00048828125, "xla_arg_mb": 18.75, "xla_out_mb": 14.062522888183594, "compile_s": 1.273083272157237, "inner": 155, "median_ms": 0.22333844474727108, "q1_ms": 0.2218573812335249, "q3_ms": 0.2234267288698785, "tflops": 23.1082472426103 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 64, "seq": 600, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0018483996391296387, "rel_err": 0.0007540053920820355, "xla_temp_mb": 30.07861328125, "xla_arg_mb": 18.75, "xla_out_mb": 14.062522888183594, "compile_s": 1.2725585428997874, "inner": 160, "median_ms": 0.2141683697118424, "q1_ms": 0.21384036881499924, "q3_ms": 0.2162483062420506, "tflops": 24.097676080477843 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 64, "seq": 600, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0018483996391296387, "rel_err": 0.0007540053920820355, "xla_temp_mb": 30.07861328125, "xla_arg_mb": 18.75, "xla_out_mb": 14.062522888183594, "compile_s": 1.2653799219988286, "inner": 148, "median_ms": 0.213988554802396, "q1_ms": 0.21381811883829133, "q3_ms": 0.2156423374922392, "tflops": 24.117925394495035 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 64, "seq": 600, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0014612674713134766, "rel_err": 0.0005960851558484137, "xla_temp_mb": 131.9091796875, "xla_arg_mb": 18.75, "xla_out_mb": 14.062522888183594, "compile_s": 0.03890814702026546, "inner": 98, "median_ms": 0.38319143789763355, "q1_ms": 0.38262168943349806, "q3_ms": 0.3840704092147703, "tflops": 13.468359387974395 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 64, "seq": 1024, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00022867321968078613, "rel_err": 0.0007638643146492541, "xla_temp_mb": 0.0, "xla_arg_mb": 24.0, "xla_out_mb": 8.0, "compile_s": 0.9381521979812533, "inner": 186, "median_ms": 0.14680767712253395, "q1_ms": 0.14617859125537896, "q3_ms": 0.14718396524067529, "tflops": 58.511480873240416 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 64, "seq": 1024, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00022867321968078613, "rel_err": 0.0007638643146492541, "xla_temp_mb": 0.0, "xla_arg_mb": 24.0, "xla_out_mb": 8.0, "compile_s": 0.9400831428356469, "inner": 190, "median_ms": 0.14688061109106793, "q1_ms": 0.14599669528634923, "q3_ms": 0.14722703184050165, "tflops": 58.4824268376316 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 64, "seq": 1024, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00022867321968078613, "rel_err": 0.0007638643146492541, "xla_temp_mb": 0.0, "xla_arg_mb": 24.0, "xla_out_mb": 8.0, "compile_s": 0.9422805269714445, "inner": 200, "median_ms": 0.1454237150028348, "q1_ms": 0.14517389179673046, "q3_ms": 0.14859495451673865, "tflops": 59.06832040312375 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 64, "seq": 1024, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0003216266632080078, "rel_err": 0.001074367668479681, "xla_temp_mb": 256.0, "xla_arg_mb": 24.0, "xla_out_mb": 8.0, "compile_s": 0.01739074499346316, "inner": 117, "median_ms": 0.3144517361831206, "q1_ms": 0.31360903850350624, "q3_ms": 0.31458023979336536, "tflops": 27.317179724514737 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 64, "seq": 1024, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00035509467124938965, "rel_err": 0.0007818220183253288, "xla_temp_mb": 8.250244140625, "xla_arg_mb": 32.0, "xla_out_mb": 24.000022888183594, "compile_s": 1.0087425210513175, "inner": 83, "median_ms": 0.4952661693096161, "q1_ms": 0.49490657921162357, "q3_ms": 0.49673114953764597, "tflops": 60.70426961306332 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 64, "seq": 1024, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00035509467124938965, "rel_err": 0.0007818220183253288, "xla_temp_mb": 0.25048828125, "xla_arg_mb": 32.0, "xla_out_mb": 24.000022888183594, "compile_s": 1.0378280910663307, "inner": 81, "median_ms": 0.5052046056974817, "q1_ms": 0.5033449815977136, "q3_ms": 0.5087558220336099, "tflops": 59.510089046976915 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 64, "seq": 1024, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00035509467124938965, "rel_err": 0.0007818220183253288, "xla_temp_mb": 0.25048828125, "xla_arg_mb": 32.0, "xla_out_mb": 24.000022888183594, "compile_s": 1.065715879900381, "inner": 79, "median_ms": 0.5053983046378516, "q1_ms": 0.5019302671492288, "q3_ms": 0.5085272080842642, "tflops": 59.48728121188144 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 64, "seq": 1024, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0003515779972076416, "rel_err": 0.0007826067158021033, "xla_temp_mb": 384.0, "xla_arg_mb": 32.0, "xla_out_mb": 24.000022888183594, "compile_s": 0.03235615510493517, "inner": 55, "median_ms": 0.7755083285949447, "q1_ms": 0.7735909640111707, "q3_ms": 0.7760625078596852, "tflops": 38.76782487490616 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 64, "seq": 1024, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0009534358978271484, "rel_err": 0.0003247723507229239, "xla_temp_mb": 0.0, "xla_arg_mb": 24.0, "xla_out_mb": 8.0, "compile_s": 1.1300215320661664, "inner": 200, "median_ms": 0.12583511997945607, "q1_ms": 0.1246572972740978, "q3_ms": 0.12597524444572628, "tflops": 34.13170581234555 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 64, "seq": 1024, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0009534358978271484, "rel_err": 0.0003247723507229239, "xla_temp_mb": 0.0, "xla_arg_mb": 24.0, "xla_out_mb": 8.0, "compile_s": 0.03835535002872348, "inner": 200, "median_ms": 0.12500616954639554, "q1_ms": 0.12377644947264344, "q3_ms": 0.12516892224084586, "tflops": 34.35804257969795 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 64, "seq": 1024, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0009534358978271484, "rel_err": 0.0003247723507229239, "xla_temp_mb": 0.0, "xla_arg_mb": 24.0, "xla_out_mb": 8.0, "compile_s": 0.03922384697943926, "inner": 200, "median_ms": 0.12457357021048666, "q1_ms": 0.12447982269804925, "q3_ms": 0.12531497224699706, "tflops": 34.47735574041088 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 64, "seq": 1024, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.001054137945175171, "rel_err": 0.00035907485289499164, "xla_temp_mb": 256.0, "xla_arg_mb": 24.0, "xla_out_mb": 8.0, "compile_s": 0.64617003290914, "inner": 112, "median_ms": 0.35881476651411504, "q1_ms": 0.35799691977445036, "q3_ms": 0.3590589684400973, "tflops": 11.969873307404825 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 64, "seq": 1024, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0013262033462524414, "rel_err": 0.000557274033781141, "xla_temp_mb": 8.250244140625, "xla_arg_mb": 32.0, "xla_out_mb": 24.000022888183594, "compile_s": 1.2062800908461213, "inner": 115, "median_ms": 0.3298445220064858, "q1_ms": 0.32748961821198463, "q3_ms": 0.3306143046559199, "tflops": 45.57415549773603 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 64, "seq": 1024, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0013262033462524414, "rel_err": 0.000557274033781141, "xla_temp_mb": 0.25048828125, "xla_arg_mb": 32.0, "xla_out_mb": 24.000022888183594, "compile_s": 1.211765965912491, "inner": 102, "median_ms": 0.35639523583299976, "q1_ms": 0.35314867729503735, "q3_ms": 0.3568156378106305, "tflops": 42.17897442109439 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 64, "seq": 1024, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0013262033462524414, "rel_err": 0.000557274033781141, "xla_temp_mb": 0.25048828125, "xla_arg_mb": 32.0, "xla_out_mb": 24.000022888183594, "compile_s": 1.207650267984718, "inner": 109, "median_ms": 0.35676300961421714, "q1_ms": 0.354681178978687, "q3_ms": 0.35865093611283316, "tflops": 42.135493677596095 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 64, "seq": 1024, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0013262033462524414, "rel_err": 0.0004974621115252376, "xla_temp_mb": 384.0, "xla_arg_mb": 32.0, "xla_out_mb": 24.000022888183594, "compile_s": 0.03916597692295909, "inner": 48, "median_ms": 0.9133138519246131, "q1_ms": 0.9105391776150402, "q3_ms": 0.9169282178239275, "tflops": 16.459167354488788 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 64, "seq": 1500, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00014570355415344238, "rel_err": 0.0005782486987300217, "xla_temp_mb": 48.0, "xla_arg_mb": 35.15625, "xla_out_mb": 11.71875, "compile_s": 1.1261195910628885, "inner": 101, "median_ms": 0.39063263361775635, "q1_ms": 0.38790734632328006, "q3_ms": 0.39196803031655236, "tflops": 47.18499790787107 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 64, "seq": 1500, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00014570355415344238, "rel_err": 0.0005782486987300217, "xla_temp_mb": 48.0, "xla_arg_mb": 35.15625, "xla_out_mb": 11.71875, "compile_s": 1.1263665209989995, "inner": 101, "median_ms": 0.39200293186056145, "q1_ms": 0.3912447672798344, "q3_ms": 0.39505772292613983, "tflops": 47.020056489160154 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 64, "seq": 1500, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00014570355415344238, "rel_err": 0.0005782486987300217, "xla_temp_mb": 48.0, "xla_arg_mb": 35.15625, "xla_out_mb": 11.71875, "compile_s": 0.04235825897194445, "inner": 93, "median_ms": 0.3944143417581756, "q1_ms": 0.39128294610168024, "q3_ms": 0.3949172421789137, "tflops": 46.73258055940846 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 64, "seq": 1500, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00019633769989013672, "rel_err": 0.0007791986572556198, "xla_temp_mb": 549.3166656494141, "xla_arg_mb": 35.15625, "xla_out_mb": 11.71875, "compile_s": 0.04736873600631952, "inner": 63, "median_ms": 0.6591928753233145, "q1_ms": 0.6563379595588361, "q3_ms": 0.6606104831758237, "tflops": 27.96146725790938 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 64, "seq": 1500, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00029903650283813477, "rel_err": 0.0005960106500424445, "xla_temp_mb": 96.00048828125, "xla_arg_mb": 46.875, "xla_out_mb": 35.156272888183594, "compile_s": 1.2216632801573724, "inner": 37, "median_ms": 1.2251492149221737, "q1_ms": 1.223477311558216, "q3_ms": 1.2260358371595677, "tflops": 52.65644316157691 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 64, "seq": 1500, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00029903650283813477, "rel_err": 0.0005960106500424445, "xla_temp_mb": 72.18807983398438, "xla_arg_mb": 46.875, "xla_out_mb": 35.156272888183594, "compile_s": 1.2511975830420852, "inner": 36, "median_ms": 1.2159162498493161, "q1_ms": 1.2143586063757539, "q3_ms": 1.2266249622270051, "tflops": 53.05628574993938 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 64, "seq": 1500, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00029903650283813477, "rel_err": 0.0005960106500424445, "xla_temp_mb": 72.18807983398438, "xla_arg_mb": 46.875, "xla_out_mb": 35.156272888183594, "compile_s": 1.2351302551105618, "inner": 37, "median_ms": 1.217568972827615, "q1_ms": 1.2144101878024034, "q3_ms": 1.2261479860171676, "tflops": 52.98426737187701 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 64, "seq": 1500, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00031238794326782227, "rel_err": 0.0006205877871252596, "xla_temp_mb": 824.1582183837891, "xla_arg_mb": 46.875, "xla_out_mb": 35.156272888183594, "compile_s": 0.0368137639015913, "inner": 26, "median_ms": 1.8087448462700615, "q1_ms": 1.799579114259149, "q3_ms": 1.810979075694027, "tflops": 35.66672222067954 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 64, "seq": 1500, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0013152360916137695, "rel_err": 0.00044277688721194863, "xla_temp_mb": 48.0, "xla_arg_mb": 35.15625, "xla_out_mb": 11.71875, "compile_s": 1.179918410954997, "inner": 129, "median_ms": 0.2898091627212689, "q1_ms": 0.2879549815976458, "q3_ms": 0.2900399727920114, "tflops": 31.80023679535528 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 64, "seq": 1500, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0013152360916137695, "rel_err": 0.00044277688721194863, "xla_temp_mb": 48.0, "xla_arg_mb": 35.15625, "xla_out_mb": 11.71875, "compile_s": 1.1692618278320879, "inner": 127, "median_ms": 0.28249130796963773, "q1_ms": 0.2813811494609503, "q3_ms": 0.2837302795735987, "tflops": 32.624012633303884 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 64, "seq": 1500, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0013152360916137695, "rel_err": 0.00044277688721194863, "xla_temp_mb": 48.0, "xla_arg_mb": 35.15625, "xla_out_mb": 11.71875, "compile_s": 0.04308899608440697, "inner": 125, "median_ms": 0.281408304348588, "q1_ms": 0.280608544126153, "q3_ms": 0.2830575406551361, "tflops": 32.74956658202913 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 64, "seq": 1500, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0013723373413085938, "rel_err": 0.000462000141851604, "xla_temp_mb": 549.3166656494141, "xla_arg_mb": 35.15625, "xla_out_mb": 11.71875, "compile_s": 0.4189130300655961, "inner": 68, "median_ms": 0.6554871044286034, "q1_ms": 0.6531650536036229, "q3_ms": 0.6575462860329186, "tflops": 14.059773163705037 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 64, "seq": 1500, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0014880597591400146, "rel_err": 0.0005150618962943554, "xla_temp_mb": 96.00048828125, "xla_arg_mb": 46.875, "xla_out_mb": 35.156272888183594, "compile_s": 1.2861243919469416, "inner": 56, "median_ms": 0.7579790918888258, "q1_ms": 0.7573556878404426, "q3_ms": 0.7630494281849158, "tflops": 42.55526352266593 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 64, "seq": 1500, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0014880597591400146, "rel_err": 0.0005150618962943554, "xla_temp_mb": 72.18798828125, "xla_arg_mb": 46.875, "xla_out_mb": 35.156272888183594, "compile_s": 1.2637640300672501, "inner": 56, "median_ms": 0.7757079646190894, "q1_ms": 0.7749184982718101, "q3_ms": 0.7794154454521569, "tflops": 41.582659288330596 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 64, "seq": 1500, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0014880597591400146, "rel_err": 0.0005150618962943554, "xla_temp_mb": 72.18798828125, "xla_arg_mb": 46.875, "xla_out_mb": 35.156272888183594, "compile_s": 1.2716774391010404, "inner": 56, "median_ms": 0.7772324315737933, "q1_ms": 0.7749827782390639, "q3_ms": 0.7797036793947751, "tflops": 41.50109888580672 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 64, "seq": 1500, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.001489013433456421, "rel_err": 0.0005425030831247568, "xla_temp_mb": 823.9751129150391, "xla_arg_mb": 46.875, "xla_out_mb": 35.156272888183594, "compile_s": 0.038827777840197086, "inner": 25, "median_ms": 1.8417587596923113, "q1_ms": 1.8377189617604017, "q3_ms": 1.8455433612689376, "tflops": 17.513694358857705 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 64, "seq": 2048, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00013641268014907837, "rel_err": 0.0006847892655059695, "xla_temp_mb": 0.0, "xla_arg_mb": 48.0, "xla_out_mb": 16.0, "compile_s": 0.9380263648927212, "inner": 73, "median_ms": 0.5519816983644276, "q1_ms": 0.5503730060634752, "q3_ms": 0.5541203977310494, "tflops": 62.24796668043716 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 64, "seq": 2048, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00013641268014907837, "rel_err": 0.0006847892655059695, "xla_temp_mb": 0.0, "xla_arg_mb": 48.0, "xla_out_mb": 16.0, "compile_s": 0.9424483641050756, "inner": 72, "median_ms": 0.5482632938461999, "q1_ms": 0.5469830761689487, "q3_ms": 0.5541414656262431, "tflops": 62.67014179803669 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 64, "seq": 2048, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00013641268014907837, "rel_err": 0.0006847892655059695, "xla_temp_mb": 0.0, "xla_arg_mb": 48.0, "xla_out_mb": 16.0, "compile_s": 0.9488540180027485, "inner": 74, "median_ms": 0.5503526095905014, "q1_ms": 0.547708959218014, "q3_ms": 0.5519480815761395, "tflops": 62.43222575716668 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 64, "seq": 2048, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00015877187252044678, "rel_err": 0.0007970320293679833, "xla_temp_mb": 1024.000015258789, "xla_arg_mb": 48.0, "xla_out_mb": 16.0, "compile_s": 0.018277692142874002, "inner": 40, "median_ms": 1.1213263031095266, "q1_ms": 1.1183290771441534, "q3_ms": 1.1217887979000807, "tflops": 30.642051535505523 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 64, "seq": 2048, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00025177001953125, "rel_err": 0.0006236736080609262, "xla_temp_mb": 16.500244140625, "xla_arg_mb": 64.0, "xla_out_mb": 48.000022888183594, "compile_s": 1.030411692801863, "inner": 24, "median_ms": 1.9225811993237585, "q1_ms": 1.9118624768452719, "q3_ms": 1.9610046195642403, "tflops": 62.55084795913924 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 64, "seq": 2048, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00025177001953125, "rel_err": 0.0006236736080609262, "xla_temp_mb": 0.50048828125, "xla_arg_mb": 64.0, "xla_out_mb": 48.000022888183594, "compile_s": 1.0493711731396616, "inner": 24, "median_ms": 1.9255037089654554, "q1_ms": 1.9144007916717478, "q3_ms": 1.951219436402122, "tflops": 62.45590892816998 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 64, "seq": 2048, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00025177001953125, "rel_err": 0.0006236736080609262, "xla_temp_mb": 0.50048828125, "xla_arg_mb": 64.0, "xla_out_mb": 48.000022888183594, "compile_s": 1.0464642709121108, "inner": 24, "median_ms": 1.9263020464374374, "q1_ms": 1.907082502536165, "q3_ms": 1.9435515811589237, "tflops": 62.4300246736543 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 64, "seq": 2048, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0002517402172088623, "rel_err": 0.000623599742539227, "xla_temp_mb": 1536.000503540039, "xla_arg_mb": 64.0, "xla_out_mb": 48.000022888183594, "compile_s": 0.0344984601251781, "inner": 16, "median_ms": 2.8391563100740314, "q1_ms": 2.8372200613375753, "q3_ms": 2.8397659698384814, "tflops": 42.35733124706481 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 64, "seq": 2048, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0008454322814941406, "rel_err": 0.00032341157202608883, "xla_temp_mb": 0.0, "xla_arg_mb": 48.0, "xla_out_mb": 16.0, "compile_s": 1.115656099980697, "inner": 104, "median_ms": 0.3621403929383422, "q1_ms": 0.36135539882637274, "q3_ms": 0.36836413290494907, "tflops": 47.43980378605551 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 64, "seq": 2048, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0008454322814941406, "rel_err": 0.00032341157202608883, "xla_temp_mb": 0.0, "xla_arg_mb": 48.0, "xla_out_mb": 16.0, "compile_s": 0.03909953706897795, "inner": 100, "median_ms": 0.3632518486119807, "q1_ms": 0.3605189046356827, "q3_ms": 0.36339359590783715, "tflops": 47.294650391032796 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 64, "seq": 2048, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0008454322814941406, "rel_err": 0.00032341157202608883, "xla_temp_mb": 0.0, "xla_arg_mb": 48.0, "xla_out_mb": 16.0, "compile_s": 0.03977192612364888, "inner": 103, "median_ms": 0.36118642734454093, "q1_ms": 0.3593928436070536, "q3_ms": 0.36272783402028014, "tflops": 47.56510179606465 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 64, "seq": 2048, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0013530254364013672, "rel_err": 0.0005175861879251897, "xla_temp_mb": 1024.000015258789, "xla_arg_mb": 48.0, "xla_out_mb": 16.0, "compile_s": 0.515314576216042, "inner": 36, "median_ms": 1.2737049708246355, "q1_ms": 1.2683070866235842, "q3_ms": 1.2764022071173207, "tflops": 13.488107197129981 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 64, "seq": 2048, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.001447916030883789, "rel_err": 0.0005423835827969015, "xla_temp_mb": 16.500244140625, "xla_arg_mb": 64.0, "xla_out_mb": 48.000022888183594, "compile_s": 1.200056267902255, "inner": 41, "median_ms": 1.1044327033365646, "q1_ms": 1.1042527082126314, "q3_ms": 1.1098507433993423, "tflops": 54.443826194520184 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 64, "seq": 2048, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.001447916030883789, "rel_err": 0.0005423835827969015, "xla_temp_mb": 0.50048828125, "xla_arg_mb": 64.0, "xla_out_mb": 48.000022888183594, "compile_s": 1.2020838619209826, "inner": 39, "median_ms": 1.1578631014204943, "q1_ms": 1.150603895672621, "q3_ms": 1.1611573326473055, "tflops": 51.93147797026405 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 64, "seq": 2048, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.001447916030883789, "rel_err": 0.0005423835827969015, "xla_temp_mb": 0.50048828125, "xla_arg_mb": 64.0, "xla_out_mb": 48.000022888183594, "compile_s": 1.2244977939408273, "inner": 39, "median_ms": 1.1573579783240955, "q1_ms": 1.1535173592468104, "q3_ms": 1.1633325768157077, "tflops": 51.95414320387732 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 64, "seq": 2048, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0014529973268508911, "rel_err": 0.000590299372561276, "xla_temp_mb": 1536.000503540039, "xla_arg_mb": 64.0, "xla_out_mb": 48.000022888183594, "compile_s": 0.040480653988197446, "inner": 14, "median_ms": 3.3214355020650794, "q1_ms": 3.2831159935865015, "q3_ms": 3.3234833889374773, "tflops": 18.103480289355275 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 64, "seq": 3000, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00011336803436279297, "rel_err": 0.0006243938696570694, "xla_temp_mb": 96.0, "xla_arg_mb": 70.3125, "xla_out_mb": 23.4375, "compile_s": 1.1201440969016403, "inner": 32, "median_ms": 1.4123147193458863, "q1_ms": 1.4118669205345213, "q3_ms": 1.4273757769842632, "tflops": 52.20366182556473 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 64, "seq": 3000, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00011336803436279297, "rel_err": 0.0006243938696570694, "xla_temp_mb": 96.0, "xla_arg_mb": 70.3125, "xla_out_mb": 23.4375, "compile_s": 1.1823693050537258, "inner": 30, "median_ms": 1.4621497985596457, "q1_ms": 1.4550148509442806, "q3_ms": 1.4708403148688376, "tflops": 50.42438201108325 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 64, "seq": 3000, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00011336803436279297, "rel_err": 0.0006243938696570694, "xla_temp_mb": 96.0, "xla_arg_mb": 70.3125, "xla_out_mb": 23.4375, "compile_s": 0.04339800588786602, "inner": 30, "median_ms": 1.4546838278571765, "q1_ms": 1.4537343328508237, "q3_ms": 1.4637591627736888, "tflops": 50.68317842551746 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 64, "seq": 3000, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00011709332466125488, "rel_err": 0.0006449115462601185, "xla_temp_mb": 2197.265884399414, "xla_arg_mb": 70.3125, "xla_out_mb": 23.4375, "compile_s": 0.019015519879758358, "inner": 19, "median_ms": 2.357126421932327, "q1_ms": 2.3557491061326705, "q3_ms": 2.3578208944711245, "tflops": 31.27876354614837 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 64, "seq": 3000, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00015710480511188507, "rel_err": 0.0007145414128899574, "xla_temp_mb": 192.00048828125, "xla_arg_mb": 93.75, "xla_out_mb": 70.3125228881836, "compile_s": 1.2253672908991575, "inner": 10, "median_ms": 4.571020998992026, "q1_ms": 4.570599994622171, "q3_ms": 4.586304456461221, "tflops": 56.45303315318464 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 64, "seq": 3000, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00015710480511188507, "rel_err": 0.0007145414128899574, "xla_temp_mb": 144.37548828125, "xla_arg_mb": 93.75, "xla_out_mb": 70.3125228881836, "compile_s": 1.2988526080735028, "inner": 10, "median_ms": 4.570539016276598, "q1_ms": 4.545201547443867, "q3_ms": 4.58229900104925, "tflops": 56.458986364855384 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 64, "seq": 3000, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00015710480511188507, "rel_err": 0.0007145414128899574, "xla_temp_mb": 144.37548828125, "xla_arg_mb": 93.75, "xla_out_mb": 70.3125228881836, "compile_s": 1.2937759410124272, "inner": 10, "median_ms": 4.55215300898999, "q1_ms": 4.53742160461843, "q3_ms": 4.579158953856677, "tflops": 56.687022490321446 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 64, "seq": 3000, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00015683472156524658, "rel_err": 0.0007214891957119107, "xla_temp_mb": 3296.264907836914, "xla_arg_mb": 93.75, "xla_out_mb": 70.3125228881836, "compile_s": 0.03869562805630267, "inner": 7, "median_ms": 6.525501409279449, "q1_ms": 6.522800706859146, "q3_ms": 6.544686338332083, "tflops": 39.544547432485174 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 64, "seq": 3000, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0009112358093261719, "rel_err": 0.00033203110797330737, "xla_temp_mb": 96.0, "xla_arg_mb": 70.3125, "xla_out_mb": 23.4375, "compile_s": 1.1822618381120265, "inner": 50, "median_ms": 0.8646249398589134, "q1_ms": 0.8562058582901955, "q3_ms": 0.8673237380571663, "tflops": 42.63582774516699 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 64, "seq": 3000, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0009112358093261719, "rel_err": 0.00033203110797330737, "xla_temp_mb": 96.0, "xla_arg_mb": 70.3125, "xla_out_mb": 23.4375, "compile_s": 1.1845574968028814, "inner": 52, "median_ms": 0.8406349224969745, "q1_ms": 0.8399765580319442, "q3_ms": 0.8500901816520267, "tflops": 43.8525678786949 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 64, "seq": 3000, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0009112358093261719, "rel_err": 0.00033203110797330737, "xla_temp_mb": 96.0, "xla_arg_mb": 70.3125, "xla_out_mb": 23.4375, "compile_s": 0.04346047714352608, "inner": 50, "median_ms": 0.8463439997285604, "q1_ms": 0.8388822129927576, "q3_ms": 0.8469141880050302, "tflops": 43.556757077291294 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 64, "seq": 3000, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.001251220703125, "rel_err": 0.00045591293019242585, "xla_temp_mb": 2197.265884399414, "xla_arg_mb": 70.3125, "xla_out_mb": 23.4375, "compile_s": 0.4325488011818379, "inner": 20, "median_ms": 2.3637103498913348, "q1_ms": 2.3627485788892955, "q3_ms": 2.36422031885013, "tflops": 15.595819513881946 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 64, "seq": 3000, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.001356363296508789, "rel_err": 0.0007203135173767805, "xla_temp_mb": 192.00048828125, "xla_arg_mb": 93.75, "xla_out_mb": 70.3125228881836, "compile_s": 1.291944975964725, "inner": 18, "median_ms": 2.547875551196436, "q1_ms": 2.5475806065110698, "q3_ms": 2.550076137089895, "tflops": 50.639835976059615 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 64, "seq": 3000, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.001356363296508789, "rel_err": 0.0007203135173767805, "xla_temp_mb": 144.37548828125, "xla_arg_mb": 93.75, "xla_out_mb": 70.3125228881836, "compile_s": 1.2975841220468283, "inner": 19, "median_ms": 2.511991371743773, "q1_ms": 2.5021073984374342, "q3_ms": 2.527755789311701, "tflops": 51.36323374806585 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 64, "seq": 3000, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.001356363296508789, "rel_err": 0.0007203135173767805, "xla_temp_mb": 144.37548828125, "xla_arg_mb": 93.75, "xla_out_mb": 70.3125228881836, "compile_s": 1.290293080965057, "inner": 19, "median_ms": 2.509715060066236, "q1_ms": 2.504787156603446, "q3_ms": 2.5222099717020203, "tflops": 51.40982020349148 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 64, "seq": 3000, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0013564825057983398, "rel_err": 0.0005977512337267399, "xla_temp_mb": 3296.264907836914, "xla_arg_mb": 93.75, "xla_out_mb": 70.3125228881836, "compile_s": 1.0678949519060552, "inner": 7, "median_ms": 6.674471144963588, "q1_ms": 6.647113344765135, "q3_ms": 6.681726131189082, "tflops": 19.330969779884168 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 64, "seq": 4096, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00011709332466125488, "rel_err": 0.0005204850458540022, "xla_temp_mb": 0.0, "xla_arg_mb": 96.0, "xla_out_mb": 32.0, "compile_s": 0.9394409470260143, "inner": 21, "median_ms": 2.111457853711077, "q1_ms": 2.1062871175152917, "q3_ms": 2.1180604096679456, "tflops": 65.0919710428691 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 64, "seq": 4096, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00011709332466125488, "rel_err": 0.0005204850458540022, "xla_temp_mb": 0.0, "xla_arg_mb": 96.0, "xla_out_mb": 32.0, "compile_s": 0.9413465620018542, "inner": 21, "median_ms": 2.107610659939902, "q1_ms": 2.1052983134896275, "q3_ms": 2.120858523994684, "tflops": 65.21078873074168 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 64, "seq": 4096, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00011709332466125488, "rel_err": 0.0005204850458540022, "xla_temp_mb": 0.0, "xla_arg_mb": 96.0, "xla_out_mb": 32.0, "compile_s": 0.9424188800621778, "inner": 22, "median_ms": 2.116953086277301, "q1_ms": 2.1093313169496306, "q3_ms": 2.121399458370764, "tflops": 64.92300389787513 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 64, "seq": 4096, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.000124320387840271, "rel_err": 0.0005526096792891622, "xla_temp_mb": 4096.0, "xla_arg_mb": 96.0, "xla_out_mb": 32.0, "compile_s": 0.020992256002500653, "inner": 11, "median_ms": 4.211733625693755, "q1_ms": 4.210951325314289, "q3_ms": 4.212942175482484, "tflops": 32.63239456397509 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 64, "seq": 4096, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00012206658720970154, "rel_err": 0.0006644873647019267, "xla_temp_mb": 33.000244140625, "xla_arg_mb": 128.0, "xla_out_mb": 96.0000228881836, "compile_s": 1.0275019768159837, "inner": 6, "median_ms": 7.579628494568169, "q1_ms": 7.539544487372041, "q3_ms": 7.599346844169001, "tflops": 63.4643686688243 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 64, "seq": 4096, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00012206658720970154, "rel_err": 0.0006644873647019267, "xla_temp_mb": 1.00048828125, "xla_arg_mb": 128.0, "xla_out_mb": 96.0000228881836, "compile_s": 1.059117894852534, "inner": 6, "median_ms": 7.560941856354475, "q1_ms": 7.482142905549457, "q3_ms": 7.620285168134918, "tflops": 63.62121892892491 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 64, "seq": 4096, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00012206658720970154, "rel_err": 0.0006644873647019267, "xla_temp_mb": 1.00048828125, "xla_arg_mb": 128.0, "xla_out_mb": 96.0000228881836, "compile_s": 1.052788021042943, "inner": 6, "median_ms": 7.512786968921621, "q1_ms": 7.450952233436206, "q3_ms": 7.66289490275085, "tflops": 64.0290133530896 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 64, "seq": 4096, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00012206658720970154, "rel_err": 0.0006644873647019267, "xla_temp_mb": 6144.0, "xla_arg_mb": 128.0, "xla_out_mb": 96.0000228881836, "compile_s": 0.04243549797683954, "inner": 4, "median_ms": 10.686997033189982, "q1_ms": 10.685394488973543, "q3_ms": 10.68781825597398, "tflops": 45.01136621055228 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 64, "seq": 4096, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0009582042694091797, "rel_err": 0.000337931647663936, "xla_temp_mb": 0.0, "xla_arg_mb": 96.0, "xla_out_mb": 32.0, "compile_s": 1.11674726405181, "inner": 37, "median_ms": 1.1921036497664612, "q1_ms": 1.1850801497899197, "q3_ms": 1.2008398362259203, "tflops": 57.64555519099574 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 64, "seq": 4096, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0009582042694091797, "rel_err": 0.000337931647663936, "xla_temp_mb": 0.0, "xla_arg_mb": 96.0, "xla_out_mb": 32.0, "compile_s": 0.04137500189244747, "inner": 37, "median_ms": 1.189507166477474, "q1_ms": 1.18709312487595, "q3_ms": 1.1918806609370418, "tflops": 57.771385219562156 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 64, "seq": 4096, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0009582042694091797, "rel_err": 0.000337931647663936, "xla_temp_mb": 0.0, "xla_arg_mb": 96.0, "xla_out_mb": 32.0, "compile_s": 0.04065299406647682, "inner": 37, "median_ms": 1.188954193698796, "q1_ms": 1.1887809614071974, "q3_ms": 1.196365526600464, "tflops": 57.79825421382808 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 64, "seq": 4096, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.001003861427307129, "rel_err": 0.0003540336329024285, "xla_temp_mb": 4096.0, "xla_arg_mb": 96.0, "xla_out_mb": 32.0, "compile_s": 0.5023693279363215, "inner": 10, "median_ms": 4.8450424103066325, "q1_ms": 4.831437394022942, "q3_ms": 4.854979296214879, "tflops": 14.183462375853773 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 64, "seq": 4096, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.001342177391052246, "rel_err": 0.0006075144046917558, "xla_temp_mb": 33.000244140625, "xla_arg_mb": 128.0, "xla_out_mb": 96.0000228881836, "compile_s": 1.2254456309601665, "inner": 12, "median_ms": 4.083973491409172, "q1_ms": 4.037983201366539, "q3_ms": 4.136401703969264, "tflops": 58.893175747085806 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 64, "seq": 4096, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.001342177391052246, "rel_err": 0.0006075144046917558, "xla_temp_mb": 1.00048828125, "xla_arg_mb": 128.0, "xla_out_mb": 96.0000228881836, "compile_s": 1.2117465469054878, "inner": 11, "median_ms": 4.125566539269957, "q1_ms": 4.120645600117065, "q3_ms": 4.175525878302076, "tflops": 58.29942779654237 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 64, "seq": 4096, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.001342177391052246, "rel_err": 0.0006075144046917558, "xla_temp_mb": 1.00048828125, "xla_arg_mb": 128.0, "xla_out_mb": 96.0000228881836, "compile_s": 1.213885940844193, "inner": 11, "median_ms": 4.169669997116381, "q1_ms": 4.131215592761608, "q3_ms": 4.174255459061401, "tflops": 57.682782748355415 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 64, "seq": 4096, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0013422966003417969, "rel_err": 0.0005387319251894951, "xla_temp_mb": 6144.0, "xla_arg_mb": 128.0, "xla_out_mb": 96.0000228881836, "compile_s": 0.049199221888557076, "inner": 4, "median_ms": 12.313162966165692, "q1_ms": 12.279666785616428, "q3_ms": 12.345293973339722, "tflops": 19.533418767939622 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 64, "seq": 8192, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 6.511062383651733e-05, "rel_err": 0.0005790345603600144, "xla_temp_mb": 0.0, "xla_arg_mb": 192.0, "xla_out_mb": 64.0, "compile_s": 0.9478051150217652, "inner": 5, "median_ms": 8.190814638510346, "q1_ms": 8.188112708739936, "q3_ms": 8.22895357850939, "tflops": 67.11857588660844 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 64, "seq": 8192, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 6.511062383651733e-05, "rel_err": 0.0005790345603600144, "xla_temp_mb": 0.0, "xla_arg_mb": 192.0, "xla_out_mb": 64.0, "compile_s": 0.9477268350310624, "inner": 5, "median_ms": 8.175568608567119, "q1_ms": 8.174620592035353, "q3_ms": 8.19998059887439, "tflops": 67.24374049187416 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 64, "seq": 8192, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 6.511062383651733e-05, "rel_err": 0.0005790345603600144, "xla_temp_mb": 0.0, "xla_arg_mb": 192.0, "xla_out_mb": 64.0, "compile_s": 0.9472633260302246, "inner": 5, "median_ms": 8.184004621580243, "q1_ms": 8.177696610800922, "q3_ms": 8.19216959644109, "tflops": 67.17442612854342 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 64, "seq": 8192, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 7.735192775726318e-05, "rel_err": 0.0006878975545987487, "xla_temp_mb": 16384.0, "xla_arg_mb": 192.0, "xla_out_mb": 64.0, "compile_s": 0.03300302312709391, "inner": 3, "median_ms": 16.30902065274616, "q1_ms": 16.307754131654896, "q3_ms": 16.31238750026872, "tflops": 33.70869567176804 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 64, "seq": 8192, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00010075420141220093, "rel_err": 0.0006305324495770037, "xla_temp_mb": 66.000244140625, "xla_arg_mb": 256.0, "xla_out_mb": 192.0000228881836, "compile_s": 1.0383139690384269, "inner": 1, "median_ms": 30.137700960040092, "q1_ms": 29.917001957073808, "q3_ms": 30.269065871834755, "tflops": 63.84512711036736 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 64, "seq": 8192, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00010075420141220093, "rel_err": 0.0006305324495770037, "xla_temp_mb": 2.00048828125, "xla_arg_mb": 256.0, "xla_out_mb": 192.0000228881836, "compile_s": 1.0529542409349233, "inner": 1, "median_ms": 30.289700953289866, "q1_ms": 29.793627094477415, "q3_ms": 30.61029000673443, "tflops": 63.52473903836981 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 64, "seq": 8192, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00010075420141220093, "rel_err": 0.0006305324495770037, "xla_temp_mb": 2.00048828125, "xla_arg_mb": 256.0, "xla_out_mb": 192.0000228881836, "compile_s": 1.0476899140048772, "inner": 1, "median_ms": 30.82504984922707, "q1_ms": 29.80337687768042, "q3_ms": 31.207983614876866, "tflops": 62.42148376140412 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 64, "seq": 8192, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00010014325380325317, "rel_err": 0.0006270954618230462, "xla_temp_mb": 24576.0, "xla_arg_mb": 256.0, "xla_out_mb": 192.0000228881836, "compile_s": 0.07256955094635487, "inner": 1, "median_ms": 41.41590092331171, "q1_ms": 41.33753199130297, "q3_ms": 41.429082048125565, "tflops": 46.45909676505332 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 64, "seq": 8192, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0007331967353820801, "rel_err": 0.0002779076457954943, "xla_temp_mb": 0.0, "xla_arg_mb": 192.0, "xla_out_mb": 64.0, "compile_s": 1.134107128949836, "inner": 11, "median_ms": 4.263250644064762, "q1_ms": 4.2498061877929345, "q3_ms": 4.281018310311166, "tflops": 64.47613098391979 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 64, "seq": 8192, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0007331967353820801, "rel_err": 0.0002779076457954943, "xla_temp_mb": 0.0, "xla_arg_mb": 192.0, "xla_out_mb": 64.0, "compile_s": 0.04359711613506079, "inner": 11, "median_ms": 4.280082542787899, "q1_ms": 4.2690088104626, "q3_ms": 4.291875594803555, "tflops": 64.22257145652007 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 64, "seq": 8192, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0007331967353820801, "rel_err": 0.0002779076457954943, "xla_temp_mb": 0.0, "xla_arg_mb": 192.0, "xla_out_mb": 64.0, "compile_s": 0.04397983499802649, "inner": 11, "median_ms": 4.266842543570833, "q1_ms": 4.254722036421299, "q3_ms": 4.288981552235782, "tflops": 64.42185389713498 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 64, "seq": 8192, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0008221864700317383, "rel_err": 0.00031163793755695224, "xla_temp_mb": 16384.0, "xla_arg_mb": 192.0, "xla_out_mb": 64.0, "compile_s": 0.49155392684042454, "inner": 2, "median_ms": 18.603841541334987, "q1_ms": 18.376746971625835, "q3_ms": 18.794523319229484, "tflops": 14.775330478560674 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 64, "seq": 8192, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0016331672668457031, "rel_err": 0.0008911209879443049, "xla_temp_mb": 66.000244140625, "xla_arg_mb": 256.0, "xla_out_mb": 192.0000228881836, "compile_s": 1.2756553289946169, "inner": 3, "median_ms": 15.362080031385025, "q1_ms": 15.35468501970172, "q3_ms": 15.621735791986186, "tflops": 62.62645893905428 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 64, "seq": 8192, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0016331672668457031, "rel_err": 0.0008911209879443049, "xla_temp_mb": 2.00048828125, "xla_arg_mb": 256.0, "xla_out_mb": 192.0000228881836, "compile_s": 1.2515066920313984, "inner": 3, "median_ms": 16.05522135893504, "q1_ms": 15.899320016615093, "q3_ms": 16.192234470508993, "tflops": 59.922728737002934 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 64, "seq": 8192, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0016331672668457031, "rel_err": 0.0008911209879443049, "xla_temp_mb": 2.00048828125, "xla_arg_mb": 256.0, "xla_out_mb": 192.0000228881836, "compile_s": 1.2347067662049085, "inner": 3, "median_ms": 16.037428053095937, "q1_ms": 15.96838499729832, "q3_ms": 16.08534968302896, "tflops": 59.9892121803332 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 64, "seq": 8192, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0014424324035644531, "rel_err": 0.0007870484841987491, "xla_temp_mb": 24576.0, "xla_arg_mb": 256.0, "xla_out_mb": 192.0000228881836, "compile_s": 1.065112448995933, "inner": 1, "median_ms": 47.10189695470035, "q1_ms": 46.83802754152566, "q3_ms": 47.49793093651533, "tflops": 20.425348797082655 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 128, "seq": 128, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0005143880844116211, "rel_err": 0.0006051270756870508, "xla_temp_mb": 0.0, "xla_arg_mb": 6.0, "xla_out_mb": 2.0, "compile_s": 0.5320823909714818, "inner": 200, "median_ms": 0.035517209907993674, "q1_ms": 0.03511785762384534, "q3_ms": 0.03573048219550401, "tflops": 7.557898176556504 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 128, "seq": 128, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0005143880844116211, "rel_err": 0.0006051270756870508, "xla_temp_mb": 0.0, "xla_arg_mb": 6.0, "xla_out_mb": 2.0, "compile_s": 0.5310147739946842, "inner": 200, "median_ms": 0.03687035525217652, "q1_ms": 0.03659368027001619, "q3_ms": 0.03694082726724446, "tflops": 7.280522635706196 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 128, "seq": 128, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0005143880844116211, "rel_err": 0.0006051270756870508, "xla_temp_mb": 0.0, "xla_arg_mb": 6.0, "xla_out_mb": 2.0, "compile_s": 0.5353186919819564, "inner": 200, "median_ms": 0.029342324705794454, "q1_ms": 0.02929292561020702, "q3_ms": 0.0295221246778965, "tflops": 9.148404521165633 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 128, "seq": 128, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0004310011863708496, "rel_err": 0.0005070305778644979, "xla_temp_mb": 4.0, "xla_arg_mb": 6.0, "xla_out_mb": 2.0, "compile_s": 0.18917072610929608, "inner": 200, "median_ms": 0.0343177595641464, "q1_ms": 0.03423910995479673, "q3_ms": 0.03449320967774838, "tflops": 7.822056550581142 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 128, "seq": 128, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.000563204288482666, "rel_err": 0.0008504130528308451, "xla_temp_mb": 2.031494140625, "xla_arg_mb": 8.0, "xla_out_mb": 6.000022888183594, "compile_s": 1.448246518848464, "inner": 200, "median_ms": 0.06538397981785238, "q1_ms": 0.06531252758577466, "q3_ms": 0.06539192982017994, "tflops": 14.369331732594736 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 128, "seq": 128, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.000563204288482666, "rel_err": 0.0008504130528308451, "xla_temp_mb": 0.03173828125, "xla_arg_mb": 8.0, "xla_out_mb": 6.000022888183594, "compile_s": 0.8985309940762818, "inner": 200, "median_ms": 0.06501063471660018, "q1_ms": 0.0647705077426508, "q3_ms": 0.0650205573765561, "tflops": 14.451852379163077 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 128, "seq": 128, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.000563204288482666, "rel_err": 0.0008504130528308451, "xla_temp_mb": 0.03173828125, "xla_arg_mb": 8.0, "xla_out_mb": 6.000022888183594, "compile_s": 0.8907003139611334, "inner": 200, "median_ms": 0.06503448006696999, "q1_ms": 0.0647990801371634, "q3_ms": 0.06504832999780774, "tflops": 14.44655350565599 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 128, "seq": 128, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0005519390106201172, "rel_err": 0.0008334030280821025, "xla_temp_mb": 4.0, "xla_arg_mb": 8.0, "xla_out_mb": 6.000022888183594, "compile_s": 0.03295764420181513, "inner": 200, "median_ms": 0.0582879944704473, "q1_ms": 0.058110199752263725, "q3_ms": 0.058367797173559666, "tflops": 16.118655385824773 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 128, "seq": 128, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0009992122650146484, "rel_err": 0.0003183129010722041, "xla_temp_mb": 0.0, "xla_arg_mb": 6.0, "xla_out_mb": 2.0, "compile_s": 0.615049023879692, "inner": 200, "median_ms": 0.022647290024906397, "q1_ms": 0.02234796469565481, "q3_ms": 0.022838115110062063, "tflops": 5.9264365781686825 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 128, "seq": 128, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0009992122650146484, "rel_err": 0.0003183129010722041, "xla_temp_mb": 0.0, "xla_arg_mb": 6.0, "xla_out_mb": 2.0, "compile_s": 0.610185147030279, "inner": 200, "median_ms": 0.02365258987993002, "q1_ms": 0.023381815408356488, "q3_ms": 0.02401703502982855, "tflops": 5.674546790915613 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 128, "seq": 128, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0009992122650146484, "rel_err": 0.0003183129010722041, "xla_temp_mb": 0.0, "xla_arg_mb": 6.0, "xla_out_mb": 2.0, "compile_s": 0.6128234299831092, "inner": 200, "median_ms": 0.023090639151632786, "q1_ms": 0.02242614224087447, "q3_ms": 0.023356362362392247, "tflops": 5.812646723142317 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 128, "seq": 128, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0011777877807617188, "rel_err": 0.0003752006159629673, "xla_temp_mb": 4.0, "xla_arg_mb": 6.0, "xla_out_mb": 2.0, "compile_s": 0.2589289741590619, "inner": 200, "median_ms": 0.03439466585405171, "q1_ms": 0.03431855991948396, "q3_ms": 0.03463678527623415, "tflops": 3.9022832368696814 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 128, "seq": 128, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.001806020736694336, "rel_err": 0.0005865312996320426, "xla_temp_mb": 2.031494140625, "xla_arg_mb": 8.0, "xla_out_mb": 6.000022888183594, "compile_s": 1.682232677936554, "inner": 200, "median_ms": 0.06050228956155479, "q1_ms": 0.06013761798385531, "q3_ms": 0.06066359463147819, "tflops": 7.764368115723388 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 128, "seq": 128, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.001806020736694336, "rel_err": 0.0005865312996320426, "xla_temp_mb": 0.03173828125, "xla_arg_mb": 8.0, "xla_out_mb": 6.000022888183594, "compile_s": 0.9933245091233402, "inner": 200, "median_ms": 0.06734602502547204, "q1_ms": 0.0667049502953887, "q3_ms": 0.06750129454303533, "tflops": 6.975349292290429 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 128, "seq": 128, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.001806020736694336, "rel_err": 0.0005865312996320426, "xla_temp_mb": 0.03173828125, "xla_arg_mb": 8.0, "xla_out_mb": 6.000022888183594, "compile_s": 1.0087986579164863, "inner": 200, "median_ms": 0.06707837455905974, "q1_ms": 0.06695002492051572, "q3_ms": 0.06760684540495276, "tflops": 7.003181742073877 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 128, "seq": 128, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.001806020736694336, "rel_err": 0.0006100231548771262, "xla_temp_mb": 4.0, "xla_arg_mb": 8.0, "xla_out_mb": 6.000022888183594, "compile_s": 0.7180497571825981, "inner": 200, "median_ms": 0.07078916067257524, "q1_ms": 0.0705362152075395, "q3_ms": 0.0709807372186333, "tflops": 6.636073143638681 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 128, "seq": 512, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00029471516609191895, "rel_err": 0.00043253504554741085, "xla_temp_mb": 0.0, "xla_arg_mb": 24.0, "xla_out_mb": 8.0, "compile_s": 0.5411123381927609, "inner": 200, "median_ms": 0.04924532026052475, "q1_ms": 0.049232947058044374, "q3_ms": 0.04927299451082945, "tflops": 87.21574503482036 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 128, "seq": 512, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00029471516609191895, "rel_err": 0.00043253504554741085, "xla_temp_mb": 0.0, "xla_arg_mb": 24.0, "xla_out_mb": 8.0, "compile_s": 0.5440514711663127, "inner": 200, "median_ms": 0.04921552026644349, "q1_ms": 0.04918652237392962, "q3_ms": 0.049231472657993436, "tflops": 87.26855416234272 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 128, "seq": 512, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00029471516609191895, "rel_err": 0.00043253504554741085, "xla_temp_mb": 0.0, "xla_arg_mb": 24.0, "xla_out_mb": 8.0, "compile_s": 0.5429858539719135, "inner": 200, "median_ms": 0.04922157037071884, "q1_ms": 0.049202697118744254, "q3_ms": 0.04922504769638181, "tflops": 87.2578274860367 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 128, "seq": 512, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0003890395164489746, "rel_err": 0.0005709690158255398, "xla_temp_mb": 72.00001525878906, "xla_arg_mb": 24.0, "xla_out_mb": 8.0, "compile_s": 0.018136641941964626, "inner": 200, "median_ms": 0.10086274007335305, "q1_ms": 0.09908323991112411, "q3_ms": 0.10194240952841938, "tflops": 42.58229840748386 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 128, "seq": 512, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0005009770393371582, "rel_err": 0.0006858734996058047, "xla_temp_mb": 8.125244140625, "xla_arg_mb": 32.0, "xla_out_mb": 24.000022888183594, "compile_s": 1.5380317179951817, "inner": 93, "median_ms": 0.42852006542185944, "q1_ms": 0.42767366719862787, "q3_ms": 0.4350721976289185, "tflops": 35.07977046816062 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 128, "seq": 512, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0005009770393371582, "rel_err": 0.0006858734996058047, "xla_temp_mb": 0.12548828125, "xla_arg_mb": 32.0, "xla_out_mb": 24.000022888183594, "compile_s": 0.9145427639596164, "inner": 96, "median_ms": 0.4053942490524302, "q1_ms": 0.404495667074419, "q3_ms": 0.410931953714074, "tflops": 37.08090475169972 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 128, "seq": 512, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0005009770393371582, "rel_err": 0.0006858734996058047, "xla_temp_mb": 0.12548828125, "xla_arg_mb": 32.0, "xla_out_mb": 24.000022888183594, "compile_s": 0.911350553156808, "inner": 96, "median_ms": 0.4055619695767139, "q1_ms": 0.404849880093631, "q3_ms": 0.411247463489417, "tflops": 37.06556990954881 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 128, "seq": 512, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0005003213882446289, "rel_err": 0.0006710723391734064, "xla_temp_mb": 96.0, "xla_arg_mb": 32.0, "xla_out_mb": 24.000022888183594, "compile_s": 0.031869968166574836, "inner": 126, "median_ms": 0.269158261029848, "q1_ms": 0.2657198217800922, "q3_ms": 0.26998215536069536, "tflops": 55.84961605296223 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 128, "seq": 512, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0009677410125732422, "rel_err": 0.0002774132299236953, "xla_temp_mb": 0.0, "xla_arg_mb": 24.0, "xla_out_mb": 8.0, "compile_s": 0.62613774696365, "inner": 200, "median_ms": 0.05144226481206715, "q1_ms": 0.0514163903426379, "q3_ms": 0.0514547893544659, "tflops": 41.7455113192499 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 128, "seq": 512, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0009677410125732422, "rel_err": 0.0002774132299236953, "xla_temp_mb": 0.0, "xla_arg_mb": 24.0, "xla_out_mb": 8.0, "compile_s": 0.6236200227867812, "inner": 200, "median_ms": 0.05121021531522274, "q1_ms": 0.05120106739923358, "q3_ms": 0.051241592736914754, "tflops": 41.934673283079114 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 128, "seq": 512, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0009677410125732422, "rel_err": 0.0002774132299236953, "xla_temp_mb": 0.0, "xla_arg_mb": 24.0, "xla_out_mb": 8.0, "compile_s": 0.6268553847912699, "inner": 200, "median_ms": 0.05125731462612748, "q1_ms": 0.05122066708281636, "q3_ms": 0.05127866461407393, "tflops": 41.8961403589676 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 128, "seq": 512, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00148773193359375, "rel_err": 0.0004264741437509656, "xla_temp_mb": 72.00001525878906, "xla_arg_mb": 24.0, "xla_out_mb": 8.0, "compile_s": 0.6891566319391131, "inner": 200, "median_ms": 0.1269210153259337, "q1_ms": 0.12685141758993268, "q3_ms": 0.12705079221632332, "tflops": 16.919842962847824 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 128, "seq": 512, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.001967310905456543, "rel_err": 0.000911962240934372, "xla_temp_mb": 8.125244140625, "xla_arg_mb": 32.0, "xla_out_mb": 24.000022888183594, "compile_s": 1.6707123520318419, "inner": 134, "median_ms": 0.2621230443439154, "q1_ms": 0.260744470609491, "q3_ms": 0.2636758388200802, "tflops": 28.674292208122182 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 128, "seq": 512, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.001967310905456543, "rel_err": 0.000911962240934372, "xla_temp_mb": 0.12548828125, "xla_arg_mb": 32.0, "xla_out_mb": 24.000022888183594, "compile_s": 0.9868635660968721, "inner": 124, "median_ms": 0.2928266598422441, "q1_ms": 0.2919381931273928, "q3_ms": 0.29503338267996665, "tflops": 25.66772018657466 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 128, "seq": 512, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.001967310905456543, "rel_err": 0.000911962240934372, "xla_temp_mb": 0.12548828125, "xla_arg_mb": 32.0, "xla_out_mb": 24.000022888183594, "compile_s": 0.9967790090013295, "inner": 123, "median_ms": 0.2941831386065095, "q1_ms": 0.291796430483521, "q3_ms": 0.2950760980734859, "tflops": 25.549366301558948 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 128, "seq": 512, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0019021034240722656, "rel_err": 0.0004917786573059857, "xla_temp_mb": 96.0, "xla_arg_mb": 32.0, "xla_out_mb": 24.000022888183594, "compile_s": 0.03864142904058099, "inner": 101, "median_ms": 0.35732540550946007, "q1_ms": 0.35683604966996624, "q3_ms": 0.3575419403133121, "tflops": 21.0345882271755 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 128, "seq": 600, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0002888590097427368, "rel_err": 0.0006460706354118884, "xla_temp_mb": 40.0, "xla_arg_mb": 28.125, "xla_out_mb": 9.375, "compile_s": 0.6413314370438457, "inner": 184, "median_ms": 0.15725567936897278, "q1_ms": 0.15722138050477952, "q3_ms": 0.15954086088838623, "tflops": 37.50732580004833 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 128, "seq": 600, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0002888590097427368, "rel_err": 0.0006460706354118884, "xla_temp_mb": 40.0, "xla_arg_mb": 28.125, "xla_out_mb": 9.375, "compile_s": 0.6343865250237286, "inner": 193, "median_ms": 0.1600620148635899, "q1_ms": 0.15979261352280105, "q3_ms": 0.16100414053879564, "tflops": 36.84971731129759 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 128, "seq": 600, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0002888590097427368, "rel_err": 0.0006460706354118884, "xla_temp_mb": 40.0, "xla_arg_mb": 28.125, "xla_out_mb": 9.375, "compile_s": 0.6366270000580698, "inner": 186, "median_ms": 0.16216591866787083, "q1_ms": 0.15952410013665275, "q3_ms": 0.1623284732622485, "tflops": 36.37163744670717 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 128, "seq": 600, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0003879666328430176, "rel_err": 0.0008677377481944859, "xla_temp_mb": 87.89088439941406, "xla_arg_mb": 28.125, "xla_out_mb": 9.375, "compile_s": 0.018378312001004815, "inner": 174, "median_ms": 0.15299528805089407, "q1_ms": 0.15198038262047475, "q3_ms": 0.1531849172362394, "tflops": 38.551775516367165 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 128, "seq": 600, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0003733038902282715, "rel_err": 0.0008604301838204265, "xla_temp_mb": 80.00048828125, "xla_arg_mb": 37.5, "xla_out_mb": 28.125022888183594, "compile_s": 1.7133067410904914, "inner": 51, "median_ms": 0.8294954667707868, "q1_ms": 0.8254291183844793, "q3_ms": 0.8354529227111853, "tflops": 24.887224616628895 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 128, "seq": 600, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0003733038902282715, "rel_err": 0.0008604301838204265, "xla_temp_mb": 60.078651428222656, "xla_arg_mb": 37.5, "xla_out_mb": 28.125022888183594, "compile_s": 0.9940975462086499, "inner": 54, "median_ms": 0.7959729258867877, "q1_ms": 0.7933816372382421, "q3_ms": 0.8026481847520227, "tflops": 25.935354493371552 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 128, "seq": 600, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0003733038902282715, "rel_err": 0.0008604301838204265, "xla_temp_mb": 60.078651428222656, "xla_arg_mb": 37.5, "xla_out_mb": 28.125022888183594, "compile_s": 0.9658215600065887, "inner": 54, "median_ms": 0.7983847624725766, "q1_ms": 0.7944429353638379, "q3_ms": 0.8032862420400811, "tflops": 25.85700650907536 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 128, "seq": 600, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0003719925880432129, "rel_err": 0.0008574077510274947, "xla_temp_mb": 131.90968322753906, "xla_arg_mb": 37.5, "xla_out_mb": 28.125022888183594, "compile_s": 0.03258794592693448, "inner": 94, "median_ms": 0.4258941081253455, "q1_ms": 0.42153150462446376, "q3_ms": 0.4307157397171126, "tflops": 48.47176705699878 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 128, "seq": 600, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0012070536613464355, "rel_err": 0.0003694213228300214, "xla_temp_mb": 40.0, "xla_arg_mb": 28.125, "xla_out_mb": 9.375, "compile_s": 0.6441521099768579, "inner": 200, "median_ms": 0.1263868692331016, "q1_ms": 0.1263076445320621, "q3_ms": 0.1272801449522376, "tflops": 23.334069574591574 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 128, "seq": 600, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0012070536613464355, "rel_err": 0.0003694213228300214, "xla_temp_mb": 40.0, "xla_arg_mb": 28.125, "xla_out_mb": 9.375, "compile_s": 0.6517216891516, "inner": 200, "median_ms": 0.12652197037823498, "q1_ms": 0.12646852468606085, "q3_ms": 0.12727859255392104, "tflops": 23.30915327340906 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 128, "seq": 600, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0012070536613464355, "rel_err": 0.0003694213228300214, "xla_temp_mb": 40.0, "xla_arg_mb": 28.125, "xla_out_mb": 9.375, "compile_s": 0.6466877029743046, "inner": 200, "median_ms": 0.12713732081465423, "q1_ms": 0.1264283445198089, "q3_ms": 0.1271568424999714, "tflops": 23.196335907529015 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 128, "seq": 600, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.001227259635925293, "rel_err": 0.00037560539203695953, "xla_temp_mb": 87.89088439941406, "xla_arg_mb": 28.125, "xla_out_mb": 9.375, "compile_s": 0.4684437080286443, "inner": 185, "median_ms": 0.15611467580940272, "q1_ms": 0.15544245913121346, "q3_ms": 0.15619291382766254, "tflops": 18.890728784528378 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 128, "seq": 600, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0019232630729675293, "rel_err": 0.0007195310317911208, "xla_temp_mb": 80.00048828125, "xla_arg_mb": 37.5, "xla_out_mb": 28.125022888183594, "compile_s": 1.7126034130342305, "inner": 81, "median_ms": 0.510692987361072, "q1_ms": 0.5097145367022833, "q3_ms": 0.5153301976805116, "tflops": 20.211595333111866 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 128, "seq": 600, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0019232630729675293, "rel_err": 0.0007195310317911208, "xla_temp_mb": 60.07861328125, "xla_arg_mb": 37.5, "xla_out_mb": 28.125022888183594, "compile_s": 0.9975453468505293, "inner": 81, "median_ms": 0.4871376174973485, "q1_ms": 0.4855932978292306, "q3_ms": 0.48998600722830976, "tflops": 21.18891998739182 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 128, "seq": 600, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0019232630729675293, "rel_err": 0.0007195310317911208, "xla_temp_mb": 60.07861328125, "xla_arg_mb": 37.5, "xla_out_mb": 28.125022888183594, "compile_s": 1.0215109349228442, "inner": 78, "median_ms": 0.48565898890583176, "q1_ms": 0.48510469824601055, "q3_ms": 0.493190000550105, "tflops": 21.253431390727123 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 128, "seq": 600, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0017398595809936523, "rel_err": 0.0005386921693570912, "xla_temp_mb": 131.90968322753906, "xla_arg_mb": 37.5, "xla_out_mb": 28.125022888183594, "compile_s": 0.037959781009703875, "inner": 84, "median_ms": 0.4663012860830696, "q1_ms": 0.4611332634729999, "q3_ms": 0.466713124686586, "tflops": 22.135731356660237 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 128, "seq": 1024, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00017982721328735352, "rel_err": 0.00048050330951809883, "xla_temp_mb": 0.0, "xla_arg_mb": 48.0, "xla_out_mb": 16.0, "compile_s": 0.5399880819022655, "inner": 160, "median_ms": 0.19744692399399355, "q1_ms": 0.1952195532794576, "q3_ms": 0.19903594729839824, "tflops": 87.01006243339917 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 128, "seq": 1024, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00017982721328735352, "rel_err": 0.00048050330951809883, "xla_temp_mb": 0.0, "xla_arg_mb": 48.0, "xla_out_mb": 16.0, "compile_s": 0.5412535180803388, "inner": 163, "median_ms": 0.19586880897040382, "q1_ms": 0.19538197509281657, "q3_ms": 0.19850251850517797, "tflops": 87.71110251962533 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 128, "seq": 1024, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00017982721328735352, "rel_err": 0.00048050330951809883, "xla_temp_mb": 0.0, "xla_arg_mb": 48.0, "xla_out_mb": 16.0, "compile_s": 0.5430864740628749, "inner": 168, "median_ms": 0.19528448908767176, "q1_ms": 0.1951789253369151, "q3_ms": 0.19948227938065038, "tflops": 87.97354702496214 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 128, "seq": 1024, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00016279518604278564, "rel_err": 0.00043499324237927794, "xla_temp_mb": 256.00025939941406, "xla_arg_mb": 48.0, "xla_out_mb": 16.0, "compile_s": 0.01998933800496161, "inner": 103, "median_ms": 0.37146048412377974, "q1_ms": 0.3711183067158009, "q3_ms": 0.3717715537754222, "tflops": 46.24952025388317 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 128, "seq": 1024, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0004494190216064453, "rel_err": 0.000607359572313726, "xla_temp_mb": 16.250244140625, "xla_arg_mb": 64.0, "xla_out_mb": 48.000022888183594, "compile_s": 1.5061539320740849, "inner": 28, "median_ms": 1.5852622870754982, "q1_ms": 1.5807351779325733, "q3_ms": 1.5978672683039414, "tflops": 37.93034290554363 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 128, "seq": 1024, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0004494190216064453, "rel_err": 0.000607359572313726, "xla_temp_mb": 0.25048828125, "xla_arg_mb": 64.0, "xla_out_mb": 48.000022888183594, "compile_s": 0.9464004219044, "inner": 30, "median_ms": 1.5140277333557606, "q1_ms": 1.5100858833951254, "q3_ms": 1.5222546954949698, "tflops": 39.71495423715002 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 128, "seq": 1024, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0004494190216064453, "rel_err": 0.000607359572313726, "xla_temp_mb": 0.25048828125, "xla_arg_mb": 64.0, "xla_out_mb": 48.000022888183594, "compile_s": 0.9360280279070139, "inner": 29, "median_ms": 1.510879859842103, "q1_ms": 1.5104512386987436, "q3_ms": 1.527976533719178, "tflops": 39.79769916999485 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 128, "seq": 1024, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0004520416259765625, "rel_err": 0.000610903836786747, "xla_temp_mb": 384.00001525878906, "xla_arg_mb": 64.0, "xla_out_mb": 48.000022888183594, "compile_s": 0.03676208411343396, "inner": 49, "median_ms": 0.9114404674619436, "q1_ms": 0.9093189574017817, "q3_ms": 0.9139945507239626, "tflops": 65.97199081081031 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 128, "seq": 1024, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0011320114135742188, "rel_err": 0.0003317581140436232, "xla_temp_mb": 0.0, "xla_arg_mb": 48.0, "xla_out_mb": 16.0, "compile_s": 0.6254824879579246, "inner": 166, "median_ms": 0.18116362611707076, "q1_ms": 0.17971594899286886, "q3_ms": 0.1817814211910927, "tflops": 47.415338145467736 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 128, "seq": 1024, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0011320114135742188, "rel_err": 0.0003317581140436232, "xla_temp_mb": 0.0, "xla_arg_mb": 48.0, "xla_out_mb": 16.0, "compile_s": 0.6296691379975528, "inner": 171, "median_ms": 0.18102280136568155, "q1_ms": 0.17830976884136773, "q3_ms": 0.1810721580581794, "tflops": 47.452224400436705 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 128, "seq": 1024, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0011320114135742188, "rel_err": 0.0003317581140436232, "xla_temp_mb": 0.0, "xla_arg_mb": 48.0, "xla_out_mb": 16.0, "compile_s": 0.6236564428545535, "inner": 171, "median_ms": 0.17983871321493422, "q1_ms": 0.17955228046388827, "q3_ms": 0.1818958127990975, "tflops": 47.76465777829349 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 128, "seq": 1024, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0013012886047363281, "rel_err": 0.0003813680959865451, "xla_temp_mb": 256.00025939941406, "xla_arg_mb": 48.0, "xla_out_mb": 16.0, "compile_s": 0.604048253968358, "inner": 94, "median_ms": 0.4167075323416514, "q1_ms": 0.4151914667516471, "q3_ms": 0.41713039985204, "tflops": 20.613821266271852 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 128, "seq": 1024, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0017702579498291016, "rel_err": 0.0008775067399255931, "xla_temp_mb": 16.250244140625, "xla_arg_mb": 64.0, "xla_out_mb": 48.000022888183594, "compile_s": 1.673870462924242, "inner": 48, "median_ms": 0.8939241233747452, "q1_ms": 0.8903127818484791, "q3_ms": 0.900009634885161, "tflops": 33.632352328181256 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 128, "seq": 1024, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0017702579498291016, "rel_err": 0.0008775067399255931, "xla_temp_mb": 0.25048828125, "xla_arg_mb": 64.0, "xla_out_mb": 48.000022888183594, "compile_s": 0.9561426560394466, "inner": 46, "median_ms": 0.9547311954362238, "q1_ms": 0.9522679459263126, "q3_ms": 0.9623649884420244, "tflops": 31.490299275560155 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 128, "seq": 1024, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0017702579498291016, "rel_err": 0.0008775067399255931, "xla_temp_mb": 0.25048828125, "xla_arg_mb": 64.0, "xla_out_mb": 48.000022888183594, "compile_s": 0.9723222241736948, "inner": 45, "median_ms": 0.9502319567319419, "q1_ms": 0.9499709768634703, "q3_ms": 0.960621054077314, "tflops": 31.639402210171298 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 128, "seq": 1024, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.001769423484802246, "rel_err": 0.0006764822173863649, "xla_temp_mb": 384.00001525878906, "xla_arg_mb": 64.0, "xla_out_mb": 48.000022888183594, "compile_s": 0.038582568988204, "inner": 43, "median_ms": 1.0515612084418535, "q1_ms": 1.0467261065144178, "q3_ms": 1.052384348694495, "tflops": 28.59060493164097 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 128, "seq": 1500, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00018177181482315063, "rel_err": 0.0008210560190491378, "xla_temp_mb": 96.0, "xla_arg_mb": 70.3125, "xla_out_mb": 23.4375, "compile_s": 0.642299324972555, "inner": 79, "median_ms": 0.5070223425856873, "q1_ms": 0.50341052480655, "q3_ms": 0.5106848732367912, "tflops": 72.70685511017682 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 128, "seq": 1500, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00018177181482315063, "rel_err": 0.0008210560190491378, "xla_temp_mb": 96.0, "xla_arg_mb": 70.3125, "xla_out_mb": 23.4375, "compile_s": 0.6399649509694427, "inner": 77, "median_ms": 0.5165661818214825, "q1_ms": 0.5143926170497358, "q3_ms": 0.5216098121403681, "tflops": 71.3635566889271 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 128, "seq": 1500, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00018177181482315063, "rel_err": 0.0008210560190491378, "xla_temp_mb": 96.0, "xla_arg_mb": 70.3125, "xla_out_mb": 23.4375, "compile_s": 0.6336204069666564, "inner": 77, "median_ms": 0.5167752879296804, "q1_ms": 0.5149114483894852, "q3_ms": 0.5248041487867374, "tflops": 71.33468039403661 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 128, "seq": 1500, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0001644045114517212, "rel_err": 0.0007426086231134832, "xla_temp_mb": 549.3166656494141, "xla_arg_mb": 70.3125, "xla_out_mb": 23.4375, "compile_s": 0.02184364292770624, "inner": 58, "median_ms": 0.743663757397183, "q1_ms": 0.7397966892940218, "q3_ms": 0.7454085761906002, "tflops": 49.570790069188924 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 128, "seq": 1500, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00042620301246643066, "rel_err": 0.0011019905796274543, "xla_temp_mb": 192.00048828125, "xla_arg_mb": 93.75, "xla_out_mb": 70.3125228881836, "compile_s": 1.7275641430169344, "inner": 13, "median_ms": 3.4482172296310845, "q1_ms": 3.43609988116301, "q3_ms": 3.4729247697843952, "tflops": 37.417596226617064 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 128, "seq": 1500, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00042620301246643066, "rel_err": 0.0011019905796274543, "xla_temp_mb": 144.18807983398438, "xla_arg_mb": 93.75, "xla_out_mb": 70.3125228881836, "compile_s": 0.9749578549526632, "inner": 14, "median_ms": 3.356789138966373, "q1_ms": 3.3556662882412116, "q3_ms": 3.389050493881639, "tflops": 38.43673065497621 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 128, "seq": 1500, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00042620301246643066, "rel_err": 0.0011019905796274543, "xla_temp_mb": 144.18807983398438, "xla_arg_mb": 93.75, "xla_out_mb": 70.3125228881836, "compile_s": 0.9701613748911768, "inner": 14, "median_ms": 3.371576279667871, "q1_ms": 3.3571630649800808, "q3_ms": 3.3893851380396103, "tflops": 38.268153913074144 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 128, "seq": 1500, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00042435526847839355, "rel_err": 0.001103577553294599, "xla_temp_mb": 824.1582183837891, "xla_arg_mb": 93.75, "xla_out_mb": 70.3125228881836, "compile_s": 0.03380066202953458, "inner": 22, "median_ms": 2.139868553389203, "q1_ms": 2.134396959180859, "q3_ms": 2.143557405691933, "tflops": 60.295292341974466 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 128, "seq": 1500, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0008329153060913086, "rel_err": 0.0003060622839257121, "xla_temp_mb": 96.0, "xla_arg_mb": 70.3125, "xla_out_mb": 23.4375, "compile_s": 0.6444049701094627, "inner": 93, "median_ms": 0.40858367868568946, "q1_ms": 0.4071648012325969, "q3_ms": 0.41196823635849583, "tflops": 45.11193413131697 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 128, "seq": 1500, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0008329153060913086, "rel_err": 0.0003060622839257121, "xla_temp_mb": 96.0, "xla_arg_mb": 70.3125, "xla_out_mb": 23.4375, "compile_s": 0.6506771328859031, "inner": 96, "median_ms": 0.4123775749273288, "q1_ms": 0.41075987489117927, "q3_ms": 0.41582131719527143, "tflops": 44.69690187020518 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 128, "seq": 1500, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0008329153060913086, "rel_err": 0.0003060622839257121, "xla_temp_mb": 96.0, "xla_arg_mb": 70.3125, "xla_out_mb": 23.4375, "compile_s": 0.6472583319991827, "inner": 96, "median_ms": 0.4130635109807675, "q1_ms": 0.41050575479554635, "q3_ms": 0.4143357497620551, "tflops": 44.622677893371716 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 128, "seq": 1500, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0012900829315185547, "rel_err": 0.00047405267832800746, "xla_temp_mb": 549.3166656494141, "xla_arg_mb": 70.3125, "xla_out_mb": 23.4375, "compile_s": 0.4298571851104498, "inner": 58, "median_ms": 0.7446408271789551, "q1_ms": 0.7423563551254058, "q3_ms": 0.7457641888698883, "tflops": 24.752873233971023 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 128, "seq": 1500, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.002327442169189453, "rel_err": 0.0008719622856006026, "xla_temp_mb": 192.00048828125, "xla_arg_mb": 93.75, "xla_out_mb": 70.3125228881836, "compile_s": 1.709251469001174, "inner": 22, "median_ms": 2.033117449503731, "q1_ms": 2.031155629083514, "q3_ms": 2.0458517511459915, "tflops": 31.730582026014726 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 128, "seq": 1500, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.002327442169189453, "rel_err": 0.0008719622856006026, "xla_temp_mb": 144.18798828125, "xla_arg_mb": 93.75, "xla_out_mb": 70.3125228881836, "compile_s": 0.9881407681386918, "inner": 22, "median_ms": 2.0411269612271674, "q1_ms": 2.0367674296721816, "q3_ms": 2.055592633868483, "tflops": 31.606069208558228 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 128, "seq": 1500, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.002327442169189453, "rel_err": 0.0008719622856006026, "xla_temp_mb": 144.18798828125, "xla_arg_mb": 93.75, "xla_out_mb": 70.3125228881836, "compile_s": 0.996055677998811, "inner": 22, "median_ms": 2.050735596144064, "q1_ms": 2.0334120014343746, "q3_ms": 2.0566094529137695, "tflops": 31.45798030779783 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 128, "seq": 1500, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0017006397247314453, "rel_err": 0.0006371345371007919, "xla_temp_mb": 823.9751129150391, "xla_arg_mb": 93.75, "xla_out_mb": 70.3125228881836, "compile_s": 0.9908906309865415, "inner": 22, "median_ms": 2.2131874522363595, "q1_ms": 2.1878422756510028, "q3_ms": 2.217357205650346, "tflops": 29.148909160321036 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 128, "seq": 2048, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00010614842176437378, "rel_err": 0.0005807357956655324, "xla_temp_mb": 0.0, "xla_arg_mb": 96.0, "xla_out_mb": 32.0, "compile_s": 0.541711196070537, "inner": 57, "median_ms": 0.7209067371788255, "q1_ms": 0.7204833060601039, "q3_ms": 0.7370017701759934, "tflops": 95.32367113799589 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 128, "seq": 2048, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00010614842176437378, "rel_err": 0.0005807357956655324, "xla_temp_mb": 0.0, "xla_arg_mb": 96.0, "xla_out_mb": 32.0, "compile_s": 0.541806774912402, "inner": 57, "median_ms": 0.7215965449352536, "q1_ms": 0.7212451503922541, "q3_ms": 0.7292763513856029, "tflops": 95.23254679963297 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 128, "seq": 2048, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00010614842176437378, "rel_err": 0.0005807357956655324, "xla_temp_mb": 0.0, "xla_arg_mb": 96.0, "xla_out_mb": 32.0, "compile_s": 0.5457619251683354, "inner": 57, "median_ms": 0.7258854896287646, "q1_ms": 0.7233572470252974, "q3_ms": 0.7305661759649713, "tflops": 94.66985870063995 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 128, "seq": 2048, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00011555105447769165, "rel_err": 0.0006321773980744183, "xla_temp_mb": 1024.000015258789, "xla_arg_mb": 96.0, "xla_out_mb": 32.0, "compile_s": 0.01813064212910831, "inner": 36, "median_ms": 1.2610192449452977, "q1_ms": 1.259196181005488, "q3_ms": 1.2615836781656575, "tflops": 54.49518475745468 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 128, "seq": 2048, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00021569430828094482, "rel_err": 0.0008874887716956437, "xla_temp_mb": 32.500244140625, "xla_arg_mb": 128.0, "xla_out_mb": 96.0000228881836, "compile_s": 1.5151424719952047, "inner": 7, "median_ms": 6.096101299460445, "q1_ms": 6.035498659392553, "q3_ms": 6.123837715546999, "tflops": 39.45442451838651 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 128, "seq": 2048, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00021569430828094482, "rel_err": 0.0008874887716956437, "xla_temp_mb": 0.50048828125, "xla_arg_mb": 128.0, "xla_out_mb": 96.0000228881836, "compile_s": 0.9489481698255986, "inner": 8, "median_ms": 5.799400008982047, "q1_ms": 5.78291185956914, "q3_ms": 5.841296064318158, "tflops": 41.47293999439391 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 128, "seq": 2048, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00021569430828094482, "rel_err": 0.0008874887716956437, "xla_temp_mb": 0.50048828125, "xla_arg_mb": 128.0, "xla_out_mb": 96.0000228881836, "compile_s": 0.947917551966384, "inner": 8, "median_ms": 5.797973746666685, "q1_ms": 5.779893748695031, "q3_ms": 5.839464807650074, "tflops": 41.483142057046464 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 128, "seq": 2048, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00021566450595855713, "rel_err": 0.0008873661863617599, "xla_temp_mb": 1568.000259399414, "xla_arg_mb": 128.0, "xla_out_mb": 96.0000228881836, "compile_s": 0.03408810100518167, "inner": 15, "median_ms": 3.178007062524557, "q1_ms": 3.166382131166756, "q3_ms": 3.180517434763412, "tflops": 75.68207491173297 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 128, "seq": 2048, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0008929818868637085, "rel_err": 0.00025550625286996365, "xla_temp_mb": 0.0, "xla_arg_mb": 96.0, "xla_out_mb": 32.0, "compile_s": 0.6289706500247121, "inner": 82, "median_ms": 0.5110776954807523, "q1_ms": 0.5075794746872129, "q3_ms": 0.5111616579020715, "tflops": 67.2299704562122 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 128, "seq": 2048, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0008929818868637085, "rel_err": 0.00025550625286996365, "xla_temp_mb": 0.0, "xla_arg_mb": 96.0, "xla_out_mb": 32.0, "compile_s": 0.623072185087949, "inner": 78, "median_ms": 0.5108693178790884, "q1_ms": 0.505269269267909, "q3_ms": 0.5125378330166523, "tflops": 67.25739277247455 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 128, "seq": 2048, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0008929818868637085, "rel_err": 0.00025550625286996365, "xla_temp_mb": 0.0, "xla_arg_mb": 96.0, "xla_out_mb": 32.0, "compile_s": 0.6236765128560364, "inner": 79, "median_ms": 0.5113124682367602, "q1_ms": 0.5076902659161936, "q3_ms": 0.512248037760205, "tflops": 67.19910133716888 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 128, "seq": 2048, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0011209994554519653, "rel_err": 0.0003207482513971627, "xla_temp_mb": 1024.000015258789, "xla_arg_mb": 96.0, "xla_out_mb": 32.0, "compile_s": 0.5312968520447612, "inner": 33, "median_ms": 1.4244108460843563, "q1_ms": 1.4201893478236864, "q3_ms": 1.4259752422345406, "tflops": 24.12207016146601 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 128, "seq": 2048, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.002022385597229004, "rel_err": 0.0009868371998891234, "xla_temp_mb": 32.500244140625, "xla_arg_mb": 128.0, "xla_out_mb": 96.0000228881836, "compile_s": 1.6596710770390928, "inner": 15, "median_ms": 3.132779197767377, "q1_ms": 3.1300122306371727, "q3_ms": 3.16173043102026, "tflops": 38.3873476859475 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 128, "seq": 2048, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.002022385597229004, "rel_err": 0.0009868371998891234, "xla_temp_mb": 0.50048828125, "xla_arg_mb": 128.0, "xla_out_mb": 96.0000228881836, "compile_s": 0.9611204979009926, "inner": 14, "median_ms": 3.2947429322770665, "q1_ms": 3.28376900870353, "q3_ms": 3.3202707168779204, "tflops": 36.50029357673935 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 128, "seq": 2048, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.002022385597229004, "rel_err": 0.0009868371998891234, "xla_temp_mb": 0.50048828125, "xla_arg_mb": 128.0, "xla_out_mb": 96.0000228881836, "compile_s": 0.9959395180922002, "inner": 14, "median_ms": 3.2808608541797315, "q1_ms": 3.2772111423712755, "q3_ms": 3.3177871456635852, "tflops": 36.65473472756184 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 128, "seq": 2048, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0013439655303955078, "rel_err": 0.0006167421233840287, "xla_temp_mb": 1568.000259399414, "xla_arg_mb": 128.0, "xla_out_mb": 96.0000228881836, "compile_s": 0.04078252403996885, "inner": 13, "median_ms": 3.6776505518131533, "q1_ms": 3.6767116085124703, "q3_ms": 3.6825804272666574, "tflops": 32.69997586603489 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 128, "seq": 3000, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 9.392201900482178e-05, "rel_err": 0.0005303416983224452, "xla_temp_mb": 192.0, "xla_arg_mb": 140.625, "xla_out_mb": 46.875, "compile_s": 0.6469288419466466, "inner": 25, "median_ms": 1.7681758385151625, "q1_ms": 1.7626650212332606, "q3_ms": 1.7878029588609934, "tflops": 83.39442084211893 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 128, "seq": 3000, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 9.392201900482178e-05, "rel_err": 0.0005303416983224452, "xla_temp_mb": 192.0, "xla_arg_mb": 140.625, "xla_out_mb": 46.875, "compile_s": 0.631896052043885, "inner": 25, "median_ms": 1.812836080789566, "q1_ms": 1.7977691581472754, "q3_ms": 1.823994885198772, "tflops": 81.33995211292174 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 128, "seq": 3000, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 9.392201900482178e-05, "rel_err": 0.0005303416983224452, "xla_temp_mb": 192.0, "xla_arg_mb": 140.625, "xla_out_mb": 46.875, "compile_s": 0.629877737024799, "inner": 25, "median_ms": 1.8066437169909477, "q1_ms": 1.7990473238751292, "q3_ms": 1.8264116765931249, "tflops": 81.61874896152469 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 128, "seq": 3000, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00012783706188201904, "rel_err": 0.0007218469399958849, "xla_temp_mb": 2197.265884399414, "xla_arg_mb": 140.625, "xla_out_mb": 46.875, "compile_s": 0.018467431887984276, "inner": 18, "median_ms": 2.527601784095168, "q1_ms": 2.5272539997887282, "q3_ms": 2.529592887084517, "tflops": 58.33830349695941 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 128, "seq": 3000, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00014616549015045166, "rel_err": 0.0006567650125361979, "xla_temp_mb": 384.00048828125, "xla_arg_mb": 187.5, "xla_out_mb": 140.6250228881836, "compile_s": 1.6993258642032743, "inner": 3, "median_ms": 13.03408636401097, "q1_ms": 13.031986192800105, "q3_ms": 13.149639319938919, "tflops": 39.59587082566961 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 128, "seq": 3000, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00014616549015045166, "rel_err": 0.0006567650125361979, "xla_temp_mb": 288.37548828125, "xla_arg_mb": 187.5, "xla_out_mb": 140.6250228881836, "compile_s": 1.0072060588281602, "inner": 3, "median_ms": 12.642177287489176, "q1_ms": 12.636112320857743, "q3_ms": 12.73723381261031, "tflops": 40.82334777180619 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 128, "seq": 3000, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00014616549015045166, "rel_err": 0.0006567650125361979, "xla_temp_mb": 288.37548828125, "xla_arg_mb": 187.5, "xla_out_mb": 140.6250228881836, "compile_s": 0.9832481809426099, "inner": 3, "median_ms": 12.639897332216302, "q1_ms": 12.630115534799794, "q3_ms": 12.75451535669466, "tflops": 40.8307113922979 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 128, "seq": 3000, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00014568865299224854, "rel_err": 0.0006590303382836282, "xla_temp_mb": 3296.265151977539, "xla_arg_mb": 187.5, "xla_out_mb": 140.6250228881836, "compile_s": 0.039638967951759696, "inner": 6, "median_ms": 7.500875663633148, "q1_ms": 7.489998824894428, "q3_ms": 7.50241382047534, "tflops": 68.80476668906975 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 128, "seq": 3000, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.000994861125946045, "rel_err": 0.0002707694948185235, "xla_temp_mb": 192.0, "xla_arg_mb": 140.625, "xla_out_mb": 46.875, "compile_s": 0.6427697730250657, "inner": 37, "median_ms": 1.1662199174532213, "q1_ms": 1.1638903650581032, "q3_ms": 1.1751127580331788, "tflops": 63.219637134140555 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 128, "seq": 3000, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.000994861125946045, "rel_err": 0.0002707694948185235, "xla_temp_mb": 192.0, "xla_arg_mb": 140.625, "xla_out_mb": 46.875, "compile_s": 0.6478092798497528, "inner": 38, "median_ms": 1.181195104053538, "q1_ms": 1.1730397236533463, "q3_ms": 1.1863187629435408, "tflops": 62.418138838355915 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 128, "seq": 3000, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.000994861125946045, "rel_err": 0.0002707694948185235, "xla_temp_mb": 192.0, "xla_arg_mb": 140.625, "xla_out_mb": 46.875, "compile_s": 0.6420365059748292, "inner": 38, "median_ms": 1.176551946340815, "q1_ms": 1.173105118428602, "q3_ms": 1.1846445666282976, "tflops": 62.66446647706535 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 128, "seq": 3000, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0010943412780761719, "rel_err": 0.00029784481739625335, "xla_temp_mb": 2197.265884399414, "xla_arg_mb": 140.625, "xla_out_mb": 46.875, "compile_s": 0.4276209711097181, "inner": 18, "median_ms": 2.532670610687799, "q1_ms": 2.5321348285716443, "q3_ms": 2.533513166579521, "tflops": 29.11077330343311 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 128, "seq": 3000, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0016607046127319336, "rel_err": 0.0007745111361145973, "xla_temp_mb": 384.00048828125, "xla_arg_mb": 187.5, "xla_out_mb": 140.6250228881836, "compile_s": 1.733971324050799, "inner": 6, "median_ms": 7.146249835689862, "q1_ms": 7.12258490966633, "q3_ms": 7.191819682096441, "tflops": 36.10956878547045 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 128, "seq": 3000, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0016607046127319336, "rel_err": 0.0007745111361145973, "xla_temp_mb": 288.37548828125, "xla_arg_mb": 187.5, "xla_out_mb": 140.6250228881836, "compile_s": 1.0315270361024886, "inner": 7, "median_ms": 6.779676702405725, "q1_ms": 6.7435989289411475, "q3_ms": 6.822551667158093, "tflops": 38.06199193959106 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 128, "seq": 3000, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0016607046127319336, "rel_err": 0.0007745111361145973, "xla_temp_mb": 288.37548828125, "xla_arg_mb": 187.5, "xla_out_mb": 140.6250228881836, "compile_s": 1.0158285561483353, "inner": 7, "median_ms": 6.797708138557417, "q1_ms": 6.772985931352846, "q3_ms": 6.849016495315092, "tflops": 37.96102962060415 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 128, "seq": 3000, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0015103816986083984, "rel_err": 0.0006280419765971601, "xla_temp_mb": 3296.265151977539, "xla_arg_mb": 187.5, "xla_out_mb": 140.6250228881836, "compile_s": 0.94312704494223, "inner": 6, "median_ms": 7.564132179444035, "q1_ms": 7.562299588850388, "q3_ms": 7.573660307874283, "tflops": 34.1146867715057 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 128, "seq": 4096, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00011526048183441162, "rel_err": 0.0007247517351061106, "xla_temp_mb": 0.0, "xla_arg_mb": 192.0, "xla_out_mb": 64.0, "compile_s": 0.5449543870054185, "inner": 16, "median_ms": 2.8027808730257675, "q1_ms": 2.7924168171011843, "q3_ms": 2.8370907821226865, "tflops": 98.07327771837298 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 128, "seq": 4096, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00011526048183441162, "rel_err": 0.0007247517351061106, "xla_temp_mb": 0.0, "xla_arg_mb": 192.0, "xla_out_mb": 64.0, "compile_s": 0.5439134398475289, "inner": 16, "median_ms": 2.7876633685082197, "q1_ms": 2.782008123176638, "q3_ms": 2.805431460728869, "tflops": 98.60512931699397 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 128, "seq": 4096, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00011526048183441162, "rel_err": 0.0007247517351061106, "xla_temp_mb": 0.0, "xla_arg_mb": 192.0, "xla_out_mb": 64.0, "compile_s": 0.550090943928808, "inner": 16, "median_ms": 2.7914265083381906, "q1_ms": 2.78130062361015, "q3_ms": 2.821318026690278, "tflops": 98.47219911501163 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 128, "seq": 4096, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 8.158385753631592e-05, "rel_err": 0.0005129949422553182, "xla_temp_mb": 4096.000259399414, "xla_arg_mb": 192.0, "xla_out_mb": 64.0, "compile_s": 0.021007294999435544, "inner": 10, "median_ms": 4.499579407274723, "q1_ms": 4.494767799042165, "q3_ms": 4.503576306160539, "tflops": 61.08968907173622 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 128, "seq": 4096, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0001957714557647705, "rel_err": 0.0008156931144185364, "xla_temp_mb": 65.000244140625, "xla_arg_mb": 256.0, "xla_out_mb": 192.0000228881836, "compile_s": 1.528247238136828, "inner": 2, "median_ms": 23.631118470802903, "q1_ms": 23.596198996528983, "q3_ms": 23.749095969833434, "tflops": 40.71210914086337 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 128, "seq": 4096, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0001957714557647705, "rel_err": 0.0008156931144185364, "xla_temp_mb": 1.00048828125, "xla_arg_mb": 256.0, "xla_out_mb": 192.0000228881836, "compile_s": 0.9375571189448237, "inner": 2, "median_ms": 22.748056100681424, "q1_ms": 22.65639853430912, "q3_ms": 22.858297976199538, "tflops": 42.29252249273206 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 128, "seq": 4096, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0001957714557647705, "rel_err": 0.0008156931144185364, "xla_temp_mb": 1.00048828125, "xla_arg_mb": 256.0, "xla_out_mb": 192.0000228881836, "compile_s": 0.9298848900943995, "inner": 2, "median_ms": 22.732580895535648, "q1_ms": 22.676675987895578, "q3_ms": 22.855000745039433, "tflops": 42.321313128723425 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 128, "seq": 4096, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00019532442092895508, "rel_err": 0.000810497731436044, "xla_temp_mb": 6144.000503540039, "xla_arg_mb": 256.0, "xla_out_mb": 192.0000228881836, "compile_s": 0.042555700056254864, "inner": 4, "median_ms": 11.638532218057662, "q1_ms": 11.63372749579139, "q3_ms": 11.646867264062166, "tflops": 82.66271521861705 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 128, "seq": 4096, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0009086132049560547, "rel_err": 0.0003274537157267332, "xla_temp_mb": 0.0, "xla_arg_mb": 192.0, "xla_out_mb": 64.0, "compile_s": 0.6313751770649105, "inner": 27, "median_ms": 1.6528972234852888, "q1_ms": 1.6449639030421774, "q3_ms": 1.6581538711088124, "tflops": 83.15033234927765 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 128, "seq": 4096, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0009086132049560547, "rel_err": 0.0003274537157267332, "xla_temp_mb": 0.0, "xla_arg_mb": 192.0, "xla_out_mb": 64.0, "compile_s": 0.6293572941794991, "inner": 27, "median_ms": 1.6473375985191927, "q1_ms": 1.641638352777119, "q3_ms": 1.6646620006886899, "tflops": 83.43095768320056 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 128, "seq": 4096, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0009086132049560547, "rel_err": 0.0003274537157267332, "xla_temp_mb": 0.0, "xla_arg_mb": 192.0, "xla_out_mb": 64.0, "compile_s": 0.6386108801234514, "inner": 27, "median_ms": 1.6543005196446623, "q1_ms": 1.6471492385284767, "q3_ms": 1.6559194040450231, "tflops": 83.07979828327768 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 128, "seq": 4096, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0013452768325805664, "rel_err": 0.0004848222597502172, "xla_temp_mb": 4096.000259399414, "xla_arg_mb": 192.0, "xla_out_mb": 64.0, "compile_s": 0.49407247197814286, "inner": 9, "median_ms": 5.168681110565861, "q1_ms": 5.14389277021918, "q3_ms": 5.172295508802765, "tflops": 26.590720249900144 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 128, "seq": 4096, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0018148422241210938, "rel_err": 0.0005610198713839054, "xla_temp_mb": 65.000244140625, "xla_arg_mb": 256.0, "xla_out_mb": 192.0000228881836, "compile_s": 1.7054514531046152, "inner": 4, "median_ms": 11.776766972616315, "q1_ms": 11.749882076401263, "q3_ms": 11.890386784216389, "tflops": 40.84621341922786 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 128, "seq": 4096, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0018148422241210938, "rel_err": 0.0005610198713839054, "xla_temp_mb": 1.00048828125, "xla_arg_mb": 256.0, "xla_out_mb": 192.0000228881836, "compile_s": 0.9878709260374308, "inner": 3, "median_ms": 12.245591652269164, "q1_ms": 12.196756821746629, "q3_ms": 12.329029850661755, "tflops": 39.282408789359046 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 128, "seq": 4096, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0018148422241210938, "rel_err": 0.0005610198713839054, "xla_temp_mb": 1.00048828125, "xla_arg_mb": 256.0, "xla_out_mb": 192.0000228881836, "compile_s": 0.9795288189779967, "inner": 3, "median_ms": 12.228675031413635, "q1_ms": 12.21265197576334, "q3_ms": 12.404587857114773, "tflops": 39.33675037698603 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 128, "seq": 4096, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0018155574798583984, "rel_err": 0.0004941796069033444, "xla_temp_mb": 6144.000503540039, "xla_arg_mb": 256.0, "xla_out_mb": 192.0000228881836, "compile_s": 0.05770303099416196, "inner": 3, "median_ms": 13.349995715543628, "q1_ms": 13.291968808819849, "q3_ms": 13.370943837799132, "tflops": 36.032695994944866 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 128, "seq": 8192, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 5.854666233062744e-05, "rel_err": 0.0005574174574576318, "xla_temp_mb": 0.0, "xla_arg_mb": 384.0, "xla_out_mb": 128.0, "compile_s": 0.5531918709166348, "inner": 4, "median_ms": 10.807484504766762, "q1_ms": 10.797157156048343, "q3_ms": 10.927760624326766, "tflops": 101.7361280778194 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 128, "seq": 8192, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 5.854666233062744e-05, "rel_err": 0.0005574174574576318, "xla_temp_mb": 0.0, "xla_arg_mb": 384.0, "xla_out_mb": 128.0, "compile_s": 0.5532487500458956, "inner": 4, "median_ms": 10.863259492907673, "q1_ms": 10.820389521541074, "q3_ms": 10.890014236792922, "tflops": 101.21378657058143 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 128, "seq": 8192, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 5.854666233062744e-05, "rel_err": 0.0005574174574576318, "xla_temp_mb": 0.0, "xla_arg_mb": 384.0, "xla_out_mb": 128.0, "compile_s": 0.557009199867025, "inner": 4, "median_ms": 10.916371771600097, "q1_ms": 10.819679504493251, "q3_ms": 10.963445500237867, "tflops": 100.72134320640089 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 128, "seq": 8192, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 6.587058305740356e-05, "rel_err": 0.0006271477905102074, "xla_temp_mb": 16384.00001525879, "xla_arg_mb": 384.0, "xla_out_mb": 128.0, "compile_s": 0.03300835518166423, "inner": 2, "median_ms": 17.207315540872514, "q1_ms": 17.18812552280724, "q3_ms": 17.227022966835648, "tflops": 63.897917438913204 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 128, "seq": 8192, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0001177862286567688, "rel_err": 0.0006473481771536171, "xla_temp_mb": 130.000244140625, "xla_arg_mb": 512.0, "xla_out_mb": 384.0000228881836, "compile_s": 1.6244349819608033, "inner": 1, "median_ms": 93.11620891094208, "q1_ms": 92.93037944007665, "q3_ms": 93.24985952116549, "tflops": 41.32782833648834 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 128, "seq": 8192, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0001177862286567688, "rel_err": 0.0006473481771536171, "xla_temp_mb": 2.00048828125, "xla_arg_mb": 512.0, "xla_out_mb": 384.0000228881836, "compile_s": 1.0339836480561644, "inner": 1, "median_ms": 89.34917999431491, "q1_ms": 89.0511239413172, "q3_ms": 89.69058294314891, "tflops": 43.070240795280476 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 128, "seq": 8192, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0001177862286567688, "rel_err": 0.0006473481771536171, "xla_temp_mb": 2.00048828125, "xla_arg_mb": 512.0, "xla_out_mb": 384.0000228881836, "compile_s": 1.0083862540777773, "inner": 1, "median_ms": 89.36970890499651, "q1_ms": 89.26827902905643, "q3_ms": 89.53973837196827, "tflops": 43.060347229136475 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 128, "seq": 8192, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00011768937110900879, "rel_err": 0.0006468158680945635, "xla_temp_mb": 24576.00050354004, "xla_arg_mb": 512.0, "xla_out_mb": 384.0000228881836, "compile_s": 0.07774913893081248, "inner": 1, "median_ms": 44.77674514055252, "q1_ms": 44.6724594803527, "q3_ms": 44.78795395698398, "tflops": 85.94395785438088 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 128, "seq": 8192, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0012829303741455078, "rel_err": 0.00042788597056642175, "xla_temp_mb": 0.0, "xla_arg_mb": 384.0, "xla_out_mb": 128.0, "compile_s": 0.6421466800384223, "inner": 8, "median_ms": 5.8045136102009565, "q1_ms": 5.7808206911431625, "q3_ms": 5.833356757648289, "tflops": 94.71177962643576 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 128, "seq": 8192, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0012829303741455078, "rel_err": 0.00042788597056642175, "xla_temp_mb": 0.0, "xla_arg_mb": 384.0, "xla_out_mb": 128.0, "compile_s": 0.632449665106833, "inner": 8, "median_ms": 5.815433745738119, "q1_ms": 5.811393741169013, "q3_ms": 5.8191212301608175, "tflops": 94.53393124646848 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 128, "seq": 8192, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0012829303741455078, "rel_err": 0.00042788597056642175, "xla_temp_mb": 0.0, "xla_arg_mb": 384.0, "xla_out_mb": 128.0, "compile_s": 0.6316442978568375, "inner": 8, "median_ms": 5.819944868562743, "q1_ms": 5.80175629875157, "q3_ms": 5.840738682309166, "tflops": 94.46065663913484 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 128, "seq": 8192, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0011343955993652344, "rel_err": 0.00037834630347788334, "xla_temp_mb": 16384.00001525879, "xla_arg_mb": 384.0, "xla_out_mb": 128.0, "compile_s": 0.48432000912725925, "inner": 2, "median_ms": 19.775293534621596, "q1_ms": 19.615144236013293, "q3_ms": 19.778891233727336, "tflops": 27.800134188932013 }, { "impl": "probjax_mha", "dtype": "float32", "head_dim": 128, "seq": 8192, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.001294255256652832, "rel_err": 0.0005421916139312088, "xla_temp_mb": 130.000244140625, "xla_arg_mb": 512.0, "xla_out_mb": 384.0000228881836, "compile_s": 1.7186019390355796, "inner": 1, "median_ms": 45.56924197822809, "q1_ms": 45.36722297780216, "q3_ms": 45.58686702512205, "tflops": 42.22465121380143 }, { "impl": "probjax_flex", "dtype": "float32", "head_dim": 128, "seq": 8192, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.001294255256652832, "rel_err": 0.0005421916139312088, "xla_temp_mb": 2.00048828125, "xla_arg_mb": 512.0, "xla_out_mb": 384.0000228881836, "compile_s": 1.0417713881470263, "inner": 1, "median_ms": 46.954998979344964, "q1_ms": 46.74274940043688, "q3_ms": 47.01076354831457, "tflops": 40.97849835870324 }, { "impl": "probjax_default", "dtype": "float32", "head_dim": 128, "seq": 8192, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.001294255256652832, "rel_err": 0.0005421916139312088, "xla_temp_mb": 2.00048828125, "xla_arg_mb": 512.0, "xla_out_mb": 384.0000228881836, "compile_s": 1.019721585093066, "inner": 1, "median_ms": 46.7894789762795, "q1_ms": 46.780164004303515, "q3_ms": 46.87690897844732, "tflops": 41.12346174197557 }, { "impl": "jax_xla", "dtype": "float32", "head_dim": 128, "seq": 8192, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.001294255256652832, "rel_err": 0.0004761589225381613, "xla_temp_mb": 24576.00050354004, "xla_arg_mb": 512.0, "xla_out_mb": 384.0000228881836, "compile_s": 1.0114301559515297, "inner": 1, "median_ms": 50.909418147057295, "q1_ms": 50.790804089047015, "q3_ms": 51.053018076345325, "tflops": 37.79546925973305 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 64, "seq": 128, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0021501779556274414, "rel_err": 0.002294705016538501, "peak_mb": 3.5166015625, "compile_s": 0.0001264710444957018, "inner": 200, "median_ms": 0.018255264731124043, "q1_ms": 0.018203314975835383, "q3_ms": 0.018353061750531197, "tflops": 7.352275082111928 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 64, "seq": 128, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0021501779556274414, "rel_err": 0.002294705016538501, "peak_mb": 3.5009765625, "compile_s": 9.628990665078163e-05, "inner": 200, "median_ms": 0.023398185148835182, "q1_ms": 0.02336325531359762, "q3_ms": 0.023423630045726895, "tflops": 5.736245232108597 }, { "impl": "torch_math", "dtype": "bfloat16", "head_dim": 64, "seq": 128, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.001876235008239746, "rel_err": 0.002002348657697439, "peak_mb": 43.50439453125, "compile_s": 0.0003122908528894186, "inner": 200, "median_ms": 0.15614709467627108, "q1_ms": 0.15610851522069424, "q3_ms": 0.1562904694583267, "tflops": 0.8595595600306511 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 64, "seq": 128, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0032709836959838867, "rel_err": 0.0035336874425411224, "peak_mb": 40.533203125, "compile_s": 0.0005487820599228144, "inner": 189, "median_ms": 0.1811436981100727, "q1_ms": 0.18080607950983066, "q3_ms": 0.1817318227723556, "tflops": 2.593311569219191 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 64, "seq": 128, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0038875937461853027, "rel_err": 0.003933617379516363, "peak_mb": 40.0341796875, "compile_s": 0.0003298309165984392, "inner": 194, "median_ms": 0.18432455679843412, "q1_ms": 0.1835250205686796, "q3_ms": 0.18495695903109982, "tflops": 2.5485592161965838 }, { "impl": "torch_math", "dtype": "bfloat16", "head_dim": 64, "seq": 128, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0027950406074523926, "rel_err": 0.0032178577966988087, "peak_mb": 79.5009765625, "compile_s": 0.0006930120289325714, "inner": 80, "median_ms": 0.5457642750116065, "q1_ms": 0.5450822689454071, "q3_ms": 0.5461065884446725, "tflops": 0.8607416599227017 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 64, "seq": 128, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.005635261535644531, "rel_err": 0.0019494958687573671, "peak_mb": 67.5166015625, "compile_s": 7.71600753068924e-05, "inner": 200, "median_ms": 0.018451964715495706, "q1_ms": 0.018431664793752134, "q3_ms": 0.018632264691405, "tflops": 3.6369495083438403 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 64, "seq": 128, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.005635261535644531, "rel_err": 0.0019494958687573671, "peak_mb": 67.5009765625, "compile_s": 7.827999070286751e-05, "inner": 200, "median_ms": 0.0234483799431473, "q1_ms": 0.023424680111929774, "q3_ms": 0.02356097975280136, "tflops": 2.861982966955988 }, { "impl": "torch_math", "dtype": "bfloat16", "head_dim": 64, "seq": 128, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0038852691650390625, "rel_err": 0.0013440931215882301, "peak_mb": 75.56689453125, "compile_s": 0.0003415709361433983, "inner": 194, "median_ms": 0.19711774255427503, "q1_ms": 0.19700604625353496, "q3_ms": 0.19817973712708842, "tflops": 0.34045065213509146 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 64, "seq": 128, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.009996175765991211, "rel_err": 0.005375280044972897, "peak_mb": 72.533203125, "compile_s": 0.0004657318349927664, "inner": 186, "median_ms": 0.18098197316610684, "q1_ms": 0.18081213134811613, "q3_ms": 0.1813932666204549, "tflops": 1.29781447229788 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 64, "seq": 128, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.012074887752532959, "rel_err": 0.005375280044972897, "peak_mb": 72.0341796875, "compile_s": 0.00036331103183329105, "inner": 200, "median_ms": 0.18479963997378945, "q1_ms": 0.1835015625692904, "q3_ms": 0.21260496228933334, "tflops": 1.2710036882826918 }, { "impl": "torch_math", "dtype": "bfloat16", "head_dim": 64, "seq": 128, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.009835958480834961, "rel_err": 0.003037187270820141, "peak_mb": 79.5009765625, "compile_s": 0.0008276330772787333, "inner": 66, "median_ms": 0.6673186805776574, "q1_ms": 0.6663665602294107, "q3_ms": 0.668714371523961, "tflops": 0.3519772948610965 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 64, "seq": 512, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.001060187816619873, "rel_err": 0.002345811342820525, "peak_mb": 78.0634765625, "compile_s": 8.422997780144215e-05, "inner": 200, "median_ms": 0.018421615241095424, "q1_ms": 0.018401240231469274, "q3_ms": 0.018478889833204448, "tflops": 116.57412338139258 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 64, "seq": 512, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0010578334331512451, "rel_err": 0.0023406019899994135, "peak_mb": 78.0009765625, "compile_s": 8.393009193241596e-05, "inner": 200, "median_ms": 0.03055760520510376, "q1_ms": 0.030533130047842864, "q3_ms": 0.030560606974177063, "tflops": 70.27656891258368 }, { "impl": "torch_math", "dtype": "bfloat16", "head_dim": 64, "seq": 512, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0009667873382568359, "rel_err": 0.002139149932190776, "peak_mb": 164.01611328125, "compile_s": 0.00030711106956005096, "inner": 180, "median_ms": 0.2173220337782469, "q1_ms": 0.21626538816942936, "q3_ms": 0.21738631120469004, "tflops": 9.881573491030686 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 64, "seq": 512, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.001962810754776001, "rel_err": 0.005092823877930641, "peak_mb": 98.126953125, "compile_s": 0.0003705208655446768, "inner": 199, "median_ms": 0.18196324607235703, "q1_ms": 0.18186977626896234, "q3_ms": 0.18228206551872456, "tflops": 41.306104008560126 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 64, "seq": 512, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0030225813388824463, "rel_err": 0.007842566817998886, "peak_mb": 96.130859375, "compile_s": 0.0003563119098544121, "inner": 139, "median_ms": 0.18386229564496082, "q1_ms": 0.18307592103378378, "q3_ms": 0.18416617899838328, "tflops": 40.87946765613007 }, { "impl": "torch_math", "dtype": "bfloat16", "head_dim": 64, "seq": 512, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0016031265258789062, "rel_err": 0.0033890039194375277, "peak_mb": 222.0009765625, "compile_s": 0.000748553080484271, "inner": 75, "median_ms": 0.544713626926144, "q1_ms": 0.5442396939421694, "q3_ms": 0.5449965596199036, "tflops": 13.798429847283952 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 64, "seq": 512, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.007361173629760742, "rel_err": 0.0025465681683272123, "peak_mb": 78.0634765625, "compile_s": 8.883094415068626e-05, "inner": 200, "median_ms": 0.018537763971835375, "q1_ms": 0.01849333755671978, "q3_ms": 0.018609915277920663, "tflops": 57.92186294050068 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 64, "seq": 512, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.007361173629760742, "rel_err": 0.0025465681683272123, "peak_mb": 78.0009765625, "compile_s": 8.430005982518196e-05, "inner": 200, "median_ms": 0.03019075491465628, "q1_ms": 0.030186105286702514, "q3_ms": 0.03020017989911139, "tflops": 35.5652525760045 }, { "impl": "torch_math", "dtype": "bfloat16", "head_dim": 64, "seq": 512, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.007361173629760742, "rel_err": 0.0025465681683272123, "peak_mb": 165.01611328125, "compile_s": 0.0003682610113173723, "inner": 146, "median_ms": 0.2580620619557696, "q1_ms": 0.2573787749821499, "q3_ms": 0.258170486880747, "tflops": 4.160789136777622 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 64, "seq": 512, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.01617288589477539, "rel_err": 0.0037439705338329077, "peak_mb": 98.126953125, "compile_s": 0.0003773909993469715, "inner": 181, "median_ms": 0.2162085977274904, "q1_ms": 0.21542553038681078, "q3_ms": 0.21721111373842092, "tflops": 17.381808232884012 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 64, "seq": 512, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.01617288589477539, "rel_err": 0.005760819651186466, "peak_mb": 96.130859375, "compile_s": 0.00036759208887815475, "inner": 174, "median_ms": 0.21624644206644134, "q1_ms": 0.21615701435861745, "q3_ms": 0.21628299179411997, "tflops": 17.37876631905616 }, { "impl": "torch_math", "dtype": "bfloat16", "head_dim": 64, "seq": 512, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.01617288589477539, "rel_err": 0.0037439705338329077, "peak_mb": 222.0009765625, "compile_s": 0.000866543035954237, "inner": 61, "median_ms": 0.6588881980383493, "q1_ms": 0.6582545818852596, "q3_ms": 0.6589945901918117, "tflops": 5.703693578347061 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 64, "seq": 600, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0009658932685852051, "rel_err": 0.002841358305886388, "peak_mb": 80.48046875, "compile_s": 8.392008021473885e-05, "inner": 200, "median_ms": 0.02268158015795052, "q1_ms": 0.022675730288028717, "q3_ms": 0.02268967975396663, "tflops": 130.02268710834292 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 64, "seq": 600, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0010502040386199951, "rel_err": 0.00308937462978065, "peak_mb": 80.4072265625, "compile_s": 9.455997496843338e-05, "inner": 200, "median_ms": 0.042531900107860565, "q1_ms": 0.042515869718044996, "q3_ms": 0.04253454739227891, "tflops": 69.3390135996994 }, { "impl": "torch_math", "dtype": "bfloat16", "head_dim": 64, "seq": 600, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0009658932685852051, "rel_err": 0.002841358305886388, "peak_mb": 195.81787109375, "compile_s": 0.00040515209548175335, "inner": 132, "median_ms": 0.32185089322879457, "q1_ms": 0.3207694848463165, "q3_ms": 0.3222987350344545, "tflops": 9.16300082443318 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 64, "seq": 600, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.002047628164291382, "rel_err": 0.004772087559103966, "peak_mb": 104.3095703125, "compile_s": 0.0006928930524736643, "inner": 124, "median_ms": 0.2180284993242352, "q1_ms": 0.2175279754033733, "q3_ms": 0.2184930565436521, "tflops": 47.34206781219934 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 64, "seq": 600, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00208398699760437, "rel_err": 0.004597815684974194, "peak_mb": 102.8994140625, "compile_s": 0.00035999203100800514, "inner": 150, "median_ms": 0.21633181255310774, "q1_ms": 0.215789246528099, "q3_ms": 0.21641322023545703, "tflops": 47.71337085462662 }, { "impl": "torch_math", "dtype": "bfloat16", "head_dim": 64, "seq": 600, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0015577971935272217, "rel_err": 0.003416633466258645, "peak_mb": 275.1572265625, "compile_s": 0.0010801439639180899, "inner": 50, "median_ms": 0.7431053835898638, "q1_ms": 0.7407562807202339, "q3_ms": 0.7436846802011132, "tflops": 13.890250599633516 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 64, "seq": 600, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.004250764846801758, "rel_err": 0.0016897449968382716, "peak_mb": 80.48046875, "compile_s": 7.864017970860004e-05, "inner": 200, "median_ms": 0.023153580259531736, "q1_ms": 0.0230793550144881, "q3_ms": 0.02332390460651368, "tflops": 63.686046972928146 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 64, "seq": 600, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.004250764846801758, "rel_err": 0.0016897449968382716, "peak_mb": 80.4072265625, "compile_s": 9.079999290406704e-05, "inner": 200, "median_ms": 0.03908613580279052, "q1_ms": 0.03905613499227911, "q3_ms": 0.03910611500032246, "tflops": 37.72590893712049 }, { "impl": "torch_math", "dtype": "bfloat16", "head_dim": 64, "seq": 600, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.003871440887451172, "rel_err": 0.0015389578184112906, "peak_mb": 198.16162109375, "compile_s": 0.0005051109474152327, "inner": 105, "median_ms": 0.386899725223581, "q1_ms": 0.38647429485406193, "q3_ms": 0.387002781610049, "tflops": 3.8112200755580368 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 64, "seq": 600, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.014807701110839844, "rel_err": 0.0035223369486629963, "peak_mb": 104.3095703125, "compile_s": 0.0007111618760973215, "inner": 161, "median_ms": 0.21815467018472112, "q1_ms": 0.21805588202261778, "q3_ms": 0.21887787325620095, "tflops": 23.65734364352589 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 64, "seq": 600, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.014807701110839844, "rel_err": 0.004309217445552349, "peak_mb": 102.8994140625, "compile_s": 0.0003500410821288824, "inner": 159, "median_ms": 0.21671716990710804, "q1_ms": 0.21660339333257583, "q3_ms": 0.21727031401399546, "tflops": 23.814264472963323 }, { "impl": "torch_math", "dtype": "bfloat16", "head_dim": 64, "seq": 600, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.014807701110839844, "rel_err": 0.0035223369486629963, "peak_mb": 275.1572265625, "compile_s": 0.0010221328120678663, "inner": 54, "median_ms": 0.8007507400449227, "q1_ms": 0.7986307415800791, "q3_ms": 0.8012144524535094, "tflops": 6.445151708146366 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 64, "seq": 1024, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00095328688621521, "rel_err": 0.003387594362720847, "peak_mb": 92.1259765625, "compile_s": 8.84511973708868e-05, "inner": 200, "median_ms": 0.03380881971679628, "q1_ms": 0.03373166488017887, "q3_ms": 0.03384831768926233, "tflops": 254.07377908943997 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 64, "seq": 1024, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0011154413223266602, "rel_err": 0.0039638252928853035, "peak_mb": 92.0009765625, "compile_s": 0.0001327809877693653, "inner": 200, "median_ms": 0.08246143581345677, "q1_ms": 0.07953952706884593, "q3_ms": 0.08526449440978467, "tflops": 104.16911259503222 }, { "impl": "torch_math", "dtype": "bfloat16", "head_dim": 64, "seq": 1024, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00095328688621521, "rel_err": 0.003387594362720847, "peak_mb": 408.03173828125, "compile_s": 0.0008773228619247675, "inner": 58, "median_ms": 0.8041388274909094, "q1_ms": 0.8020607136382625, "q3_ms": 0.8044292656544211, "tflops": 10.682153750494168 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 64, "seq": 1024, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0015018582344055176, "rel_err": 0.004941033665090799, "peak_mb": 132.251953125, "compile_s": 0.0006766521837562323, "inner": 150, "median_ms": 0.21570982101062935, "q1_ms": 0.21545337978750467, "q3_ms": 0.21684178694461784, "tflops": 139.37599563683528 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 64, "seq": 1024, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.002331942319869995, "rel_err": 0.005864798091351986, "peak_mb": 128.259765625, "compile_s": 0.0009197229519486427, "inner": 92, "median_ms": 0.4366847823131019, "q1_ms": 0.43509509819357295, "q3_ms": 0.4375154891496767, "tflops": 68.84776454253364 }, { "impl": "torch_math", "dtype": "bfloat16", "head_dim": 64, "seq": 1024, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0010400712490081787, "rel_err": 0.002775006927549839, "peak_mb": 636.0009765625, "compile_s": 0.0018440170679241419, "inner": 27, "median_ms": 1.7563134811266705, "q1_ms": 1.7490336731835097, "q3_ms": 1.7574468489598345, "tflops": 17.118112111007385 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 64, "seq": 1024, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.006038188934326172, "rel_err": 0.002593584591522813, "peak_mb": 92.1259765625, "compile_s": 8.431007154285908e-05, "inner": 200, "median_ms": 0.03343511954881251, "q1_ms": 0.033340115332975984, "q3_ms": 0.033462466672062874, "tflops": 128.45676504101334 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 64, "seq": 1024, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.006038188934326172, "rel_err": 0.002593584591522813, "peak_mb": 92.0009765625, "compile_s": 0.00012141000479459763, "inner": 200, "median_ms": 0.0681469903793186, "q1_ms": 0.06809973565395921, "q3_ms": 0.06859663757495582, "tflops": 63.02504735856165 }, { "impl": "torch_math", "dtype": "bfloat16", "head_dim": 64, "seq": 1024, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.006038188934326172, "rel_err": 0.002593584591522813, "peak_mb": 412.03173828125, "compile_s": 0.001048052916303277, "inner": 48, "median_ms": 0.9600320821239924, "q1_ms": 0.9582769780536182, "q3_ms": 0.960729164944496, "tflops": 4.473774758128641 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 64, "seq": 1024, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.012666940689086914, "rel_err": 0.0043540229089558125, "peak_mb": 132.251953125, "compile_s": 0.0005305020604282618, "inner": 151, "median_ms": 0.240490310794569, "q1_ms": 0.24036056938772368, "q3_ms": 0.2412149098251533, "tflops": 62.50723983986584 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 64, "seq": 1024, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.012666940689086914, "rel_err": 0.004499658010900021, "peak_mb": 128.259765625, "compile_s": 0.0006344220601022243, "inner": 112, "median_ms": 0.32479425135534257, "q1_ms": 0.3246415321232884, "q3_ms": 0.3249288482558248, "tflops": 46.282794332938344 }, { "impl": "torch_math", "dtype": "bfloat16", "head_dim": 64, "seq": 1024, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.009896278381347656, "rel_err": 0.0034243951085954905, "peak_mb": 636.0009765625, "compile_s": 0.0020024869590997696, "inner": 25, "median_ms": 1.8908610008656979, "q1_ms": 1.8890599580481648, "q3_ms": 1.8966828007251024, "tflops": 7.9500214606561155 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 64, "seq": 1500, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0006342977285385132, "rel_err": 0.002317151054739952, "peak_mb": 105.48095703125, "compile_s": 0.00014872010797262192, "inner": 200, "median_ms": 0.09493192424997687, "q1_ms": 0.09198124462272972, "q3_ms": 0.09625698497984558, "tflops": 194.1601852656483 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 64, "seq": 1500, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0006347447633743286, "rel_err": 0.0023187841288745403, "peak_mb": 105.2978515625, "compile_s": 0.00023049116134643555, "inner": 200, "median_ms": 0.18806560547091067, "q1_ms": 0.18543261976446956, "q3_ms": 0.19136351475026459, "tflops": 98.00835168050438 }, { "impl": "torch_math", "dtype": "bfloat16", "head_dim": 64, "seq": 1500, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0004876554012298584, "rel_err": 0.0017814523307606578, "peak_mb": 765.46484375, "compile_s": 0.0016973060555756092, "inner": 29, "median_ms": 1.687340036933792, "q1_ms": 1.6830689474490694, "q3_ms": 1.6907436195118675, "tflops": 10.923702156380015 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 64, "seq": 1500, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0012814700603485107, "rel_err": 0.003724043257534504, "peak_mb": 165.38818359375, "compile_s": 0.0007788820657879114, "inner": 104, "median_ms": 0.40101842247307873, "q1_ms": 0.39778407349681055, "q3_ms": 0.40385707800921333, "tflops": 160.87041488556758 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 64, "seq": 1500, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0013823211193084717, "rel_err": 0.004368611145764589, "peak_mb": 159.53662109375, "compile_s": 0.0009363829158246517, "inner": 62, "median_ms": 0.8111778203578245, "q1_ms": 0.8059314268100406, "q3_ms": 0.8223768860672511, "tflops": 79.52880167697712 }, { "impl": "torch_math", "dtype": "bfloat16", "head_dim": 64, "seq": 1500, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0010690689086914062, "rel_err": 0.003709746291860938, "peak_mb": 1251.6494140625, "compile_s": 0.0036995229311287403, "inner": 13, "median_ms": 3.760279227907841, "q1_ms": 3.749705340641622, "q3_ms": 3.7661149936656537, "tflops": 17.156172744089922 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 64, "seq": 1500, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.007173061370849609, "rel_err": 0.0024814915377646685, "peak_mb": 105.76220703125, "compile_s": 0.00012219115160405636, "inner": 200, "median_ms": 0.06593963014893234, "q1_ms": 0.06558940222021192, "q3_ms": 0.06658498256001621, "tflops": 139.7642052159618 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 64, "seq": 1500, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.007173061370849609, "rel_err": 0.0024814915377646685, "peak_mb": 105.5791015625, "compile_s": 0.00018487009219825268, "inner": 200, "median_ms": 0.12499098433181642, "q1_ms": 0.1248902577208355, "q3_ms": 0.1256324880523607, "tflops": 73.73331804103626 }, { "impl": "torch_math", "dtype": "bfloat16", "head_dim": 64, "seq": 1500, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.007173061370849609, "rel_err": 0.0024814915377646685, "peak_mb": 774.32958984375, "compile_s": 0.002054997021332383, "inner": 24, "median_ms": 2.002118630722786, "q1_ms": 1.997721546407168, "q3_ms": 2.0072898963311063, "tflops": 4.603123840205676 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 64, "seq": 1500, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.014818668365478516, "rel_err": 0.005191786680370569, "peak_mb": 165.38818359375, "compile_s": 0.0005801618099212646, "inner": 120, "median_ms": 0.3123844255848477, "q1_ms": 0.3120120003586635, "q3_ms": 0.3126282128505409, "tflops": 103.25738851932248 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 64, "seq": 1500, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.014818668365478516, "rel_err": 0.005191786680370569, "peak_mb": 159.53662109375, "compile_s": 0.0008567830082029104, "inner": 73, "median_ms": 0.6057344383419785, "q1_ms": 0.6051339457857692, "q3_ms": 0.6064642268577463, "tflops": 53.25105848082767 }, { "impl": "torch_math", "dtype": "bfloat16", "head_dim": 64, "seq": 1500, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.01280665397644043, "rel_err": 0.003944694995880127, "peak_mb": 1251.6494140625, "compile_s": 0.004035813966766, "inner": 12, "median_ms": 4.020092232773702, "q1_ms": 4.0073239263923215, "q3_ms": 4.028484002143765, "tflops": 8.023696505526358 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 64, "seq": 2048, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.000682830810546875, "rel_err": 0.0027175149880349636, "peak_mb": 120.2509765625, "compile_s": 0.0001811198890209198, "inner": 200, "median_ms": 0.14028718462213874, "q1_ms": 0.13607482251245528, "q3_ms": 0.14224757032934576, "tflops": 244.92428485572222 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 64, "seq": 2048, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.000682830810546875, "rel_err": 0.0027175149880349636, "peak_mb": 120.0009765625, "compile_s": 0.00033985101617872715, "inner": 160, "median_ms": 0.3143643436487764, "q1_ms": 0.31329655976151116, "q3_ms": 0.320857331098523, "tflops": 109.29909534011409 }, { "impl": "torch_math", "dtype": "bfloat16", "head_dim": 64, "seq": 2048, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.000682830810546875, "rel_err": 0.0027175149880349636, "peak_mb": 1328.06298828125, "compile_s": 0.002957721007987857, "inner": 16, "median_ms": 2.978512813569978, "q1_ms": 2.9714746851823293, "q3_ms": 2.9853375599486753, "tflops": 11.53587058999998 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 64, "seq": 2048, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0011211037635803223, "rel_err": 0.0036814571358263493, "peak_mb": 200.501953125, "compile_s": 0.00090932403691113, "inner": 80, "median_ms": 0.6202055374160409, "q1_ms": 0.6155351991765201, "q3_ms": 0.6257402375922538, "tflops": 193.90198415356753 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 64, "seq": 2048, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0010370314121246338, "rel_err": 0.004226936958730221, "peak_mb": 192.517578125, "compile_s": 0.0013717750553041697, "inner": 36, "median_ms": 1.3997256949854393, "q1_ms": 1.394726374807457, "q3_ms": 1.4348651263086747, "tflops": 85.91617966208086 }, { "impl": "torch_math", "dtype": "bfloat16", "head_dim": 64, "seq": 2048, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0008320212364196777, "rel_err": 0.0034037057776004076, "peak_mb": 2232.0009765625, "compile_s": 0.0064754721242934465, "inner": 7, "median_ms": 6.581751563187156, "q1_ms": 6.557532702572644, "q3_ms": 6.607694434933364, "tflops": 18.271592771843483 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 64, "seq": 2048, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.008076190948486328, "rel_err": 0.002793925581499934, "peak_mb": 120.2509765625, "compile_s": 0.00015183980576694012, "inner": 200, "median_ms": 0.10159745579585433, "q1_ms": 0.09867861750535667, "q3_ms": 0.10196662973612547, "tflops": 169.09743506294598 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 64, "seq": 2048, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.008076190948486328, "rel_err": 0.002793925581499934, "peak_mb": 120.0009765625, "compile_s": 0.0002556710969656706, "inner": 200, "median_ms": 0.20848752465099096, "q1_ms": 0.2061626425711438, "q3_ms": 0.21116628486197442, "tflops": 82.4023845683773 }, { "impl": "torch_math", "dtype": "bfloat16", "head_dim": 64, "seq": 2048, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.007548809051513672, "rel_err": 0.0026114799547940493, "peak_mb": 1344.06298828125, "compile_s": 0.003618963062763214, "inner": 13, "median_ms": 3.6034301424828863, "q1_ms": 3.5990101482289343, "q3_ms": 3.6044944650851765, "tflops": 4.7676431912628905 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 64, "seq": 2048, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.016491413116455078, "rel_err": 0.008016865700483322, "peak_mb": 200.501953125, "compile_s": 0.0008074629586189985, "inner": 94, "median_ms": 0.46349001027881465, "q1_ms": 0.4588506912733329, "q3_ms": 0.46387820707713356, "tflops": 129.73212110403153 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 64, "seq": 2048, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.016491413116455078, "rel_err": 0.008016865700483322, "peak_mb": 192.517578125, "compile_s": 0.0010533039458096027, "inner": 52, "median_ms": 0.9340872907509598, "q1_ms": 0.9269087951487074, "q3_ms": 0.9369308558794168, "tflops": 64.37250858606463 }, { "impl": "torch_math", "dtype": "bfloat16", "head_dim": 64, "seq": 2048, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.011342048645019531, "rel_err": 0.004965658765286207, "peak_mb": 2232.0009765625, "compile_s": 0.007114334963262081, "inner": 7, "median_ms": 7.109567589525666, "q1_ms": 7.097413197958044, "q3_ms": 7.12786049449018, "tflops": 8.457552641118038 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 64, "seq": 3000, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.000499039888381958, "rel_err": 0.0027869928162544966, "peak_mb": 147.5234375, "compile_s": 0.00032907119020819664, "inner": 166, "median_ms": 0.30521021774375295, "q1_ms": 0.3034518673982336, "q3_ms": 0.3111600568686354, "tflops": 241.56465188167527 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 64, "seq": 3000, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0005548298358917236, "rel_err": 0.0030985635239630938, "peak_mb": 147.1572265625, "compile_s": 0.0006928718648850918, "inner": 75, "median_ms": 0.7032124387721221, "q1_ms": 0.6771798199042678, "q3_ms": 0.7049405202269554, "tflops": 104.84456180658054 }, { "impl": "torch_math", "dtype": "bfloat16", "head_dim": 64, "seq": 3000, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.000477522611618042, "rel_err": 0.002666825195774436, "peak_mb": 2703.45361328125, "compile_s": 0.006385911954566836, "inner": 7, "median_ms": 6.525874137878418, "q1_ms": 6.5126112229855995, "q3_ms": 6.540420564955899, "tflops": 11.297796807336098 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 64, "seq": 3000, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0008320212364196777, "rel_err": 0.0034725868608802557, "peak_mb": 268.1806640625, "compile_s": 0.0013875749427825212, "inner": 39, "median_ms": 1.3016904346071756, "q1_ms": 1.2985364330980258, "q3_ms": 1.3266652553843765, "tflops": 198.2406823768923 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 64, "seq": 3000, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0011704117059707642, "rel_err": 0.005179732106626034, "peak_mb": 255.8916015625, "compile_s": 0.0027531099040061235, "inner": 18, "median_ms": 2.9140317734951773, "q1_ms": 2.903073696264376, "q3_ms": 2.9655763905288444, "tflops": 88.55359860763957 }, { "impl": "torch_math", "dtype": "bfloat16", "head_dim": 64, "seq": 3000, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0008320212364196777, "rel_err": 0.0033195260912179947, "peak_mb": 4637.6884765625, "compile_s": 0.01394927897490561, "inner": 3, "median_ms": 14.382906723767519, "q1_ms": 14.264579319084685, "q3_ms": 14.402950143752, "tflops": 17.941296912784665 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 64, "seq": 3000, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.006508588790893555, "rel_err": 0.002251619938760996, "peak_mb": 147.8046875, "compile_s": 0.00027146097272634506, "inner": 200, "median_ms": 0.21507400088012218, "q1_ms": 0.20862687495537102, "q3_ms": 0.21520309732295573, "tflops": 171.40147042016125 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 64, "seq": 3000, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.006508588790893555, "rel_err": 0.002251619938760996, "peak_mb": 147.4384765625, "compile_s": 0.00045557180419564247, "inner": 118, "median_ms": 0.41468279688792714, "q1_ms": 0.4077290028917714, "q3_ms": 0.4197537457033739, "tflops": 88.89686352231999 }, { "impl": "torch_math", "dtype": "bfloat16", "head_dim": 64, "seq": 3000, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.006508588790893555, "rel_err": 0.002251619938760996, "peak_mb": 2738.0673828125, "compile_s": 0.008091457886621356, "inner": 6, "median_ms": 8.131348333942393, "q1_ms": 8.115015671743702, "q3_ms": 8.137856668327004, "tflops": 4.533565466150299 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 64, "seq": 3000, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.01353597640991211, "rel_err": 0.0039185951463878155, "peak_mb": 268.4619140625, "compile_s": 0.001048634061589837, "inner": 57, "median_ms": 0.8831806154104701, "q1_ms": 0.877350587152729, "q3_ms": 0.8889814295495551, "tflops": 146.09016292781104 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 64, "seq": 3000, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.01353597640991211, "rel_err": 0.005356759298592806, "peak_mb": 256.1728515625, "compile_s": 0.0018611859995871782, "inner": 27, "median_ms": 1.8845780373171523, "q1_ms": 1.8776320783352412, "q3_ms": 1.8862048896995407, "tflops": 68.46307101385729 }, { "impl": "torch_math", "dtype": "bfloat16", "head_dim": 64, "seq": 3000, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.013211488723754883, "rel_err": 0.003183129010722041, "peak_mb": 4637.4072265625, "compile_s": 0.01549034402705729, "inner": 3, "median_ms": 15.754011382038396, "q1_ms": 15.748154682417711, "q3_ms": 15.806739839414757, "tflops": 8.189914103217165 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 64, "seq": 4096, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00045752525329589844, "rel_err": 0.0028358343988656998, "peak_mb": 176.5009765625, "compile_s": 0.0005635421257466078, "inner": 96, "median_ms": 0.5526864066875229, "q1_ms": 0.5474628912149153, "q3_ms": 0.5639459692853658, "tflops": 248.67438715514683 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 64, "seq": 4096, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00045752525329589844, "rel_err": 0.0028358343988656998, "peak_mb": 176.0009765625, "compile_s": 0.0011932139750570059, "inner": 43, "median_ms": 1.2451889762265045, "q1_ms": 1.2258464649189698, "q3_ms": 1.2498109758463363, "tflops": 110.37597994844386 }, { "impl": "torch_math", "dtype": "bfloat16", "head_dim": 64, "seq": 4096, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00045752525329589844, "rel_err": 0.0028358343988656998, "peak_mb": 4896.12548828125, "compile_s": 0.011528261005878448, "inner": 4, "median_ms": 11.722980765625834, "q1_ms": 11.67626425740309, "q3_ms": 11.76120963646099, "tflops": 11.723891407806365 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 64, "seq": 4096, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0007363855838775635, "rel_err": 0.003625117940828204, "peak_mb": 337.001953125, "compile_s": 0.002091257832944393, "inner": 24, "median_ms": 2.1785405018211654, "q1_ms": 2.1346642897697166, "q3_ms": 2.2031691041775048, "tflops": 220.80669914094983 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 64, "seq": 4096, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0007723420858383179, "rel_err": 0.004020660184323788, "peak_mb": 321.033203125, "compile_s": 0.004902817076072097, "inner": 10, "median_ms": 5.214192206040025, "q1_ms": 5.182590952608734, "q3_ms": 5.251033743843436, "tflops": 92.25519853195598 }, { "impl": "torch_math", "dtype": "bfloat16", "head_dim": 64, "seq": 4096, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0007155686616897583, "rel_err": 0.0036366344429552555, "peak_mb": 8496.0009765625, "compile_s": 0.02514912700280547, "inner": 1, "median_ms": 25.71429987438023, "q1_ms": 25.609944481402636, "q3_ms": 25.797009468078613, "tflops": 18.706958365655055 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 64, "seq": 4096, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.004129528999328613, "rel_err": 0.0016415518475696445, "peak_mb": 176.5009765625, "compile_s": 0.0003668610006570816, "inner": 150, "median_ms": 0.34344305905203026, "q1_ms": 0.33072891645133495, "q3_ms": 0.34487553328896564, "tflops": 200.08986912031105 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 64, "seq": 4096, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.004129528999328613, "rel_err": 0.0016415518475696445, "peak_mb": 176.0009765625, "compile_s": 0.0007257531397044659, "inner": 71, "median_ms": 0.7224498467493645, "q1_ms": 0.7054351981807971, "q3_ms": 0.7253191335467806, "tflops": 95.12006548994461 }, { "impl": "torch_math", "dtype": "bfloat16", "head_dim": 64, "seq": 4096, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0038406848907470703, "rel_err": 0.001526731881313026, "peak_mb": 4960.12548828125, "compile_s": 0.014684651046991348, "inner": 3, "median_ms": 14.720757647107044, "q1_ms": 14.698856005755564, "q3_ms": 14.731356369641919, "tflops": 4.668202437902706 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 64, "seq": 4096, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.015132904052734375, "rel_err": 0.003882892895489931, "peak_mb": 337.001953125, "compile_s": 0.001445394940674305, "inner": 36, "median_ms": 1.4000376363077924, "q1_ms": 1.3880827367150537, "q3_ms": 1.4044393319636583, "tflops": 171.79407348669523 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 64, "seq": 4096, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.015132904052734375, "rel_err": 0.004317702725529671, "peak_mb": 321.033203125, "compile_s": 0.0030076501425355673, "inner": 16, "median_ms": 3.102978938841261, "q1_ms": 3.0908635671949014, "q3_ms": 3.1239258678397164, "tflops": 77.51202096969959 }, { "impl": "torch_math", "dtype": "bfloat16", "head_dim": 64, "seq": 4096, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.015132904052734375, "rel_err": 0.0036865491420030594, "peak_mb": 8496.0009765625, "compile_s": 0.02796170790679753, "inner": 1, "median_ms": 28.33575801923871, "q1_ms": 28.31681352108717, "q3_ms": 28.395294095389545, "tflops": 8.488150146281562 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 64, "seq": 8192, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00033595412969589233, "rel_err": 0.0027797212824225426, "peak_mb": 289.0009765625, "compile_s": 0.002046847017481923, "inner": 25, "median_ms": 2.135543841868639, "q1_ms": 2.12305077817291, "q3_ms": 2.163923545740545, "tflops": 257.4312936637975 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 64, "seq": 8192, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0002998635172843933, "rel_err": 0.002481103641912341, "peak_mb": 288.0009765625, "compile_s": 0.004426114959642291, "inner": 11, "median_ms": 4.819575908847831, "q1_ms": 4.795282084325498, "q3_ms": 4.868545548313043, "tflops": 114.06725908782806 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 64, "seq": 8192, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0006060153245925903, "rel_err": 0.0036204783245921135, "peak_mb": 610.001953125, "compile_s": 0.00764953694306314, "inner": 6, "median_ms": 8.174963334264854, "q1_ms": 8.144224896871796, "q3_ms": 8.20597360143438, "tflops": 235.37051726495991 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 64, "seq": 8192, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0006060153245925903, "rel_err": 0.003806143067777157, "peak_mb": 578.064453125, "compile_s": 0.018140593077987432, "inner": 2, "median_ms": 19.50739254243672, "q1_ms": 19.4796952418983, "q3_ms": 20.111822173930705, "tflops": 98.63672679073746 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 64, "seq": 8192, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.007561683654785156, "rel_err": 0.003005886683240533, "peak_mb": 289.0009765625, "compile_s": 0.001156164100393653, "inner": 45, "median_ms": 1.2069239798519344, "q1_ms": 1.198076389523016, "q3_ms": 1.2186125671077104, "tflops": 227.75080413740892 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 64, "seq": 8192, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.007561683654785156, "rel_err": 0.003005886683240533, "peak_mb": 288.0009765625, "compile_s": 0.0024661088827997446, "inner": 20, "median_ms": 2.5547098484821618, "q1_ms": 2.5486140977591276, "q3_ms": 2.633699675789103, "tflops": 107.5965269039512 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 64, "seq": 8192, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.01398777961730957, "rel_err": 0.0038633455988019705, "peak_mb": 610.001953125, "compile_s": 0.004469125997275114, "inner": 11, "median_ms": 4.638597094030542, "q1_ms": 4.614819634959779, "q3_ms": 4.642388868060979, "tflops": 207.40595805186464 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 64, "seq": 8192, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.01398777961730957, "rel_err": 0.0048573496751487255, "peak_mb": 578.064453125, "compile_s": 0.01052279700525105, "inner": 4, "median_ms": 11.108876031357795, "q1_ms": 11.068948486354202, "q3_ms": 11.385190999135375, "tflops": 86.60396169588091 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 128, "seq": 128, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0027510523796081543, "rel_err": 0.0036955196410417557, "peak_mb": 71.0166015625, "compile_s": 7.912004366517067e-05, "inner": 200, "median_ms": 0.018523664912208915, "q1_ms": 0.018477565026842058, "q3_ms": 0.01859471492934972, "tflops": 14.491487363446886 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 128, "seq": 128, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0027510523796081543, "rel_err": 0.0036955196410417557, "peak_mb": 71.0009765625, "compile_s": 9.134015999734402e-05, "inner": 200, "median_ms": 0.023920584935694933, "q1_ms": 0.02383428276516497, "q3_ms": 0.02403033256996423, "tflops": 11.221943640660452 }, { "impl": "torch_math", "dtype": "bfloat16", "head_dim": 128, "seq": 128, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0019525885581970215, "rel_err": 0.002622934291139245, "peak_mb": 84.0, "compile_s": 0.0002440200187265873, "inner": 200, "median_ms": 0.15167142963036895, "q1_ms": 0.1516053796513006, "q3_ms": 0.15184697753284127, "tflops": 1.7698485248948401 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 128, "seq": 128, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.005645871162414551, "rel_err": 0.0045579178258776665, "peak_mb": 81.033203125, "compile_s": 0.00048444210551679134, "inner": 189, "median_ms": 0.2297949250393286, "q1_ms": 0.1872738517781414, "q3_ms": 0.24278383275799492, "tflops": 4.0885328335219056 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 128, "seq": 128, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.007427692413330078, "rel_err": 0.006798913236707449, "peak_mb": 80.03369140625, "compile_s": 0.0006694919429719448, "inner": 153, "median_ms": 0.24193999345347383, "q1_ms": 0.24138479737041435, "q3_ms": 0.24212058177035228, "tflops": 3.8832938803921837 }, { "impl": "torch_math", "dtype": "bfloat16", "head_dim": 128, "seq": 128, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.004670619964599609, "rel_err": 0.004275236278772354, "peak_mb": 88.0, "compile_s": 0.0012094841804355383, "inner": 72, "median_ms": 0.6213806369083209, "q1_ms": 0.6210193209375979, "q3_ms": 0.6215141798343717, "tflops": 1.5119944848532805 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 128, "seq": 128, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.005361318588256836, "rel_err": 0.001854726462624967, "peak_mb": 71.0166015625, "compile_s": 7.502082735300064e-05, "inner": 200, "median_ms": 0.018900964641943574, "q1_ms": 0.018846640596166253, "q3_ms": 0.019007193041034043, "tflops": 7.101104654846784 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 128, "seq": 128, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.005361318588256836, "rel_err": 0.001854726462624967, "peak_mb": 71.0009765625, "compile_s": 8.236989378929138e-05, "inner": 200, "median_ms": 0.023924284614622593, "q1_ms": 0.02387577958870679, "q3_ms": 0.023959309910424054, "tflops": 5.610104133185479 }, { "impl": "torch_math", "dtype": "bfloat16", "head_dim": 128, "seq": 128, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0039037466049194336, "rel_err": 0.0013504852540791035, "peak_mb": 84.0625, "compile_s": 0.0006702521350234747, "inner": 173, "median_ms": 0.1970823057469605, "q1_ms": 0.19701395043834097, "q3_ms": 0.19727430018364867, "tflops": 0.6810237351917624 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 128, "seq": 128, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.01621723175048828, "rel_err": 0.00388345867395401, "peak_mb": 81.033203125, "compile_s": 0.00047885114327073097, "inner": 187, "median_ms": 0.18786541734109907, "q1_ms": 0.18480896045837333, "q3_ms": 0.2259147966051484, "tflops": 2.5005243362437133 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 128, "seq": 128, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.01621723175048828, "rel_err": 0.005187384784221649, "peak_mb": 80.03369140625, "compile_s": 0.0006441220175474882, "inner": 158, "median_ms": 0.24159818363104818, "q1_ms": 0.24142998780560077, "q3_ms": 0.24168910105017166, "tflops": 1.9443939558642862 }, { "impl": "torch_math", "dtype": "bfloat16", "head_dim": 128, "seq": 128, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.01621723175048828, "rel_err": 0.00388345867395401, "peak_mb": 88.0625, "compile_s": 0.0013636040966957808, "inner": 65, "median_ms": 0.6669469392643524, "q1_ms": 0.6664104759693146, "q3_ms": 0.667005478261182, "tflops": 0.7043469582724993 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 128, "seq": 512, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0010970830917358398, "rel_err": 0.002655942225828767, "peak_mb": 92.0634765625, "compile_s": 8.169002830982208e-05, "inner": 200, "median_ms": 0.01927806413732469, "q1_ms": 0.019253442296758294, "q3_ms": 0.019304364686831832, "tflops": 222.79038317361017 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 128, "seq": 512, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0011283457279205322, "rel_err": 0.0027316263876855373, "peak_mb": 92.0009765625, "compile_s": 9.083002805709839e-05, "inner": 200, "median_ms": 0.0337223696988076, "q1_ms": 0.03371659491676837, "q3_ms": 0.033745214459486306, "tflops": 127.36255886999149 }, { "impl": "torch_math", "dtype": "bfloat16", "head_dim": 128, "seq": 512, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0009626150131225586, "rel_err": 0.002330406801775098, "peak_mb": 192.01611328125, "compile_s": 0.0003526520449668169, "inner": 154, "median_ms": 0.2825550127218102, "q1_ms": 0.2810980228090344, "q3_ms": 0.2825985807780321, "tflops": 15.200463989745648 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 128, "seq": 512, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0023974180221557617, "rel_err": 0.005364567507058382, "peak_mb": 132.126953125, "compile_s": 0.00037115090526640415, "inner": 160, "median_ms": 0.18366351287113503, "q1_ms": 0.18284010657225735, "q3_ms": 0.1843249840021599, "tflops": 81.8474246789958 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 128, "seq": 512, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0022614002227783203, "rel_err": 0.004969405476003885, "peak_mb": 128.12890625, "compile_s": 0.0003955711144953966, "inner": 134, "median_ms": 0.19014386613088757, "q1_ms": 0.18825927669349224, "q3_ms": 0.1915368953469529, "tflops": 79.0579567034379 }, { "impl": "torch_math", "dtype": "bfloat16", "head_dim": 128, "seq": 512, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0017199218273162842, "rel_err": 0.003803657367825508, "peak_mb": 252.0009765625, "compile_s": 0.0008568721823394299, "inner": 61, "median_ms": 0.684137803456578, "q1_ms": 0.6827522874580787, "q3_ms": 0.6858145237946119, "tflops": 21.972745052311815 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 128, "seq": 512, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.006621599197387695, "rel_err": 0.002405368722975254, "peak_mb": 92.0634765625, "compile_s": 7.808092050254345e-05, "inner": 200, "median_ms": 0.02087372587993741, "q1_ms": 0.02086939464788884, "q3_ms": 0.020888375001959503, "tflops": 102.87974750420739 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 128, "seq": 512, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.006621599197387695, "rel_err": 0.002405368722975254, "peak_mb": 92.0009765625, "compile_s": 9.877001866698265e-05, "inner": 200, "median_ms": 0.03311381442472339, "q1_ms": 0.033080289722420275, "q3_ms": 0.033125043264590204, "tflops": 64.85159397392313 }, { "impl": "torch_math", "dtype": "bfloat16", "head_dim": 128, "seq": 512, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.006621599197387695, "rel_err": 0.002405368722975254, "peak_mb": 193.01611328125, "compile_s": 0.00041562202386558056, "inner": 128, "median_ms": 0.3212770552636357, "q1_ms": 0.3197275191269, "q3_ms": 0.32142420695890905, "tflops": 6.684211065859662 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 128, "seq": 512, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.013897895812988281, "rel_err": 0.004398775286972523, "peak_mb": 132.126953125, "compile_s": 0.0005422618705779314, "inner": 154, "median_ms": 0.18406992887037915, "q1_ms": 0.18392863015148353, "q3_ms": 0.18446966884778693, "tflops": 40.83335509567591 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 128, "seq": 512, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.013897895812988281, "rel_err": 0.006165427155792713, "peak_mb": 128.12890625, "compile_s": 0.0005542309954762459, "inner": 132, "median_ms": 0.18724035577509893, "q1_ms": 0.18676648544373387, "q3_ms": 0.18767792481054185, "tflops": 40.141948763587955 }, { "impl": "torch_math", "dtype": "bfloat16", "head_dim": 128, "seq": 512, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.013897895812988281, "rel_err": 0.004398775286972523, "peak_mb": 252.0009765625, "compile_s": 0.0009070828091353178, "inner": 58, "median_ms": 0.71267438008739, "q1_ms": 0.7102383292633399, "q3_ms": 0.7139444732556826, "tflops": 10.54646129846613 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 128, "seq": 600, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.001120448112487793, "rel_err": 0.0022516855970025063, "peak_mb": 96.88671875, "compile_s": 9.516021236777306e-05, "inner": 200, "median_ms": 0.035866874968633056, "q1_ms": 0.035859100171364844, "q3_ms": 0.03587730054277927, "tflops": 164.44811557065498 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 128, "seq": 600, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0011000633239746094, "rel_err": 0.002210719743743539, "peak_mb": 96.8134765625, "compile_s": 0.0001154008787125349, "inner": 200, "median_ms": 0.059124710969626904, "q1_ms": 0.059099579812027514, "q3_ms": 0.05967728269752115, "tflops": 99.7593037373155 }, { "impl": "torch_math", "dtype": "bfloat16", "head_dim": 128, "seq": 600, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0009724795818328857, "rel_err": 0.00195432361215353, "peak_mb": 228.63037109375, "compile_s": 0.0004891620483249426, "inner": 109, "median_ms": 0.4272888256493238, "q1_ms": 0.4245613844757643, "q3_ms": 0.42742102859763925, "tflops": 13.803871400186555 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 128, "seq": 600, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.002751857042312622, "rel_err": 0.0050126416608691216, "peak_mb": 145.0908203125, "compile_s": 0.0005493620410561562, "inner": 153, "median_ms": 0.21736010501246436, "q1_ms": 0.21731216396552092, "q3_ms": 0.21747941854318567, "tflops": 94.97529456390441 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 128, "seq": 600, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0027889013290405273, "rel_err": 0.0050126416608691216, "peak_mb": 139.77685546875, "compile_s": 0.0006802030839025974, "inner": 98, "median_ms": 0.35554328600742985, "q1_ms": 0.35417149061983333, "q3_ms": 0.3577356383546578, "tflops": 58.06280363727245 }, { "impl": "torch_math", "dtype": "bfloat16", "head_dim": 128, "seq": 600, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0021083950996398926, "rel_err": 0.003840544493868947, "peak_mb": 310.3134765625, "compile_s": 0.0011064040008932352, "inner": 46, "median_ms": 1.0451577583570844, "q1_ms": 1.0413810226094464, "q3_ms": 1.0461848451877418, "tflops": 19.751888970762355 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 128, "seq": 600, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.006784915924072266, "rel_err": 0.0023472141474485397, "peak_mb": 96.88671875, "compile_s": 8.958997204899788e-05, "inner": 200, "median_ms": 0.03623557509854436, "q1_ms": 0.036159674637019634, "q3_ms": 0.036262349458411336, "tflops": 81.38742084207932 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 128, "seq": 600, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.006784915924072266, "rel_err": 0.0023472141474485397, "peak_mb": 96.8134765625, "compile_s": 9.890994988381863e-05, "inner": 200, "median_ms": 0.044454854214563966, "q1_ms": 0.04444130521733314, "q3_ms": 0.044462779769673944, "tflops": 66.33966193581244 }, { "impl": "torch_math", "dtype": "bfloat16", "head_dim": 128, "seq": 600, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.006784915924072266, "rel_err": 0.0023472141474485397, "peak_mb": 230.00390625, "compile_s": 0.0005701419431716204, "inner": 92, "median_ms": 0.48739473973198433, "q1_ms": 0.48626039642840624, "q3_ms": 0.48781615180854243, "tflops": 6.050783399143177 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 128, "seq": 600, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.012672662734985352, "rel_err": 0.00421340623870492, "peak_mb": 145.0908203125, "compile_s": 0.0005179820582270622, "inner": 141, "median_ms": 0.19288131247228343, "q1_ms": 0.19183034797245307, "q3_ms": 0.1937003900356432, "tflops": 53.51436003673624 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 128, "seq": 600, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.012672662734985352, "rel_err": 0.004915619734674692, "peak_mb": 139.77685546875, "compile_s": 0.0005522121209651232, "inner": 138, "median_ms": 0.26010713823463605, "q1_ms": 0.2596173907859602, "q3_ms": 0.2609087353356291, "tflops": 39.683339988496805 }, { "impl": "torch_math", "dtype": "bfloat16", "head_dim": 128, "seq": 600, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.012472391128540039, "rel_err": 0.0035177457612007856, "peak_mb": 310.3134765625, "compile_s": 0.0011875440832227468, "inner": 42, "median_ms": 1.0889175935604034, "q1_ms": 1.0851404540950345, "q3_ms": 1.0931201201553145, "tflops": 9.479064403992872 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 128, "seq": 1024, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0007398128509521484, "rel_err": 0.002683986909687519, "peak_mb": 120.1259765625, "compile_s": 0.00011593103408813477, "inner": 200, "median_ms": 0.06363057065755129, "q1_ms": 0.059806156787090003, "q3_ms": 0.064861502032727, "tflops": 269.99395112231633 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 128, "seq": 1024, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0007505267858505249, "rel_err": 0.002722856355831027, "peak_mb": 120.0009765625, "compile_s": 0.00017019989900290966, "inner": 200, "median_ms": 0.126865494530648, "q1_ms": 0.1184797374298796, "q3_ms": 0.13074428017716855, "tflops": 135.41798144214627 }, { "impl": "torch_math", "dtype": "bfloat16", "head_dim": 128, "seq": 1024, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0007398128509521484, "rel_err": 0.002683986909687519, "peak_mb": 464.03173828125, "compile_s": 0.0010388740338385105, "inner": 49, "median_ms": 1.024195615544307, "q1_ms": 1.0190375187263196, "q3_ms": 1.0251881630749118, "tflops": 16.7740116470522 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 128, "seq": 1024, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0013456344604492188, "rel_err": 0.0038550400640815496, "peak_mb": 200.251953125, "compile_s": 0.0007377930451184511, "inner": 123, "median_ms": 0.3560585541484075, "q1_ms": 0.353202285298063, "q3_ms": 0.35824681200632236, "tflops": 168.87543198565487 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 128, "seq": 1024, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0019114017486572266, "rel_err": 0.0041186693124473095, "peak_mb": 192.255859375, "compile_s": 0.0008016328793019056, "inner": 80, "median_ms": 0.5978059489279985, "q1_ms": 0.5967223376501352, "q3_ms": 0.5987292126519606, "tflops": 100.58371324645712 }, { "impl": "torch_math", "dtype": "bfloat16", "head_dim": 128, "seq": 1024, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00145760178565979, "rel_err": 0.003629300044849515, "peak_mb": 696.0009765625, "compile_s": 0.002345428103581071, "inner": 21, "median_ms": 2.3910445161163807, "q1_ms": 2.3856002205450624, "q3_ms": 2.3930354786681987, "tflops": 25.147813743620524 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 128, "seq": 1024, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.005032062530517578, "rel_err": 0.0017408216372132301, "peak_mb": 120.1259765625, "compile_s": 0.00011528981849551201, "inner": 200, "median_ms": 0.055277489591389894, "q1_ms": 0.055271139717660844, "q3_ms": 0.05552287038881332, "tflops": 155.3966118124507 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 128, "seq": 1024, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.005032062530517578, "rel_err": 0.0017408216372132301, "peak_mb": 120.0009765625, "compile_s": 0.00015218998305499554, "inner": 200, "median_ms": 0.093564975541085, "q1_ms": 0.09290412184782326, "q3_ms": 0.09367537451907992, "tflops": 91.80715906058354 }, { "impl": "torch_math", "dtype": "bfloat16", "head_dim": 128, "seq": 1024, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.005032062530517578, "rel_err": 0.0017408216372132301, "peak_mb": 468.03173828125, "compile_s": 0.0012181440833956003, "inner": 41, "median_ms": 1.1762992149536928, "q1_ms": 1.1709535823816932, "q3_ms": 1.1768861838457423, "tflops": 7.302508139766257 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 128, "seq": 1024, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.01667952537536621, "rel_err": 0.004910856951028109, "peak_mb": 200.251953125, "compile_s": 0.0006899519357830286, "inner": 118, "median_ms": 0.3158418733973877, "q1_ms": 0.3157652079885415, "q3_ms": 0.31594800372150234, "tflops": 95.18931340105414 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 128, "seq": 1024, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.013173818588256836, "rel_err": 0.004045982379466295, "peak_mb": 192.255859375, "compile_s": 0.0007752231322228909, "inner": 85, "median_ms": 0.49347324735101533, "q1_ms": 0.4934097708695952, "q3_ms": 0.49399848315207395, "tflops": 60.9248246655496 }, { "impl": "torch_math", "dtype": "bfloat16", "head_dim": 128, "seq": 1024, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.011829137802124023, "rel_err": 0.0024739813525229692, "peak_mb": 696.0009765625, "compile_s": 0.0025311079807579517, "inner": 20, "median_ms": 2.5039906962774694, "q1_ms": 2.467106096446514, "q3_ms": 2.5135692791081965, "tflops": 12.006742324041166 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 128, "seq": 1500, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0007283985614776611, "rel_err": 0.0028605633415281773, "peak_mb": 146.77783203125, "compile_s": 0.00020290003158152103, "inner": 200, "median_ms": 0.1522071799263358, "q1_ms": 0.14731288712937385, "q3_ms": 0.15404161531478167, "tflops": 242.19619611795704 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 128, "seq": 1500, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0007283985614776611, "rel_err": 0.0028605633415281773, "peak_mb": 146.5947265625, "compile_s": 0.00031171110458672047, "inner": 181, "median_ms": 0.2709151984022632, "q1_ms": 0.27016743127441867, "q3_ms": 0.2797224915210796, "tflops": 136.07210011622604 }, { "impl": "torch_math", "dtype": "bfloat16", "head_dim": 128, "seq": 1500, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0007283985614776611, "rel_err": 0.0028605633415281773, "peak_mb": 849.46484375, "compile_s": 0.0021618669852614403, "inner": 23, "median_ms": 2.20878682423221, "q1_ms": 2.1963198466793354, "q3_ms": 2.212448151903632, "tflops": 16.689704771674464 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 128, "seq": 1500, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0016910731792449951, "rel_err": 0.004783402197062969, "peak_mb": 267.81005859375, "compile_s": 0.0009645731188356876, "inner": 71, "median_ms": 0.7026197756050337, "q1_ms": 0.6973212322508786, "q3_ms": 0.7117471270735415, "tflops": 183.6327477246082 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 128, "seq": 1500, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.002030491828918457, "rel_err": 0.006870935205370188, "peak_mb": 255.24951171875, "compile_s": 0.0014179949648678303, "inner": 39, "median_ms": 1.2904024587418788, "q1_ms": 1.279879973317759, "q3_ms": 1.3121926398852315, "tflops": 99.98741022688091 }, { "impl": "torch_math", "dtype": "bfloat16", "head_dim": 128, "seq": 1500, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0011043697595596313, "rel_err": 0.002777058631181717, "peak_mb": 1340.6650390625, "compile_s": 0.00515755801461637, "inner": 9, "median_ms": 5.385788778464, "q1_ms": 5.378464223920471, "q3_ms": 5.393456564181381, "tflops": 23.956379521589223 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 128, "seq": 1500, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0072650909423828125, "rel_err": 0.0025133287999778986, "peak_mb": 147.62158203125, "compile_s": 0.00020018103532493114, "inner": 200, "median_ms": 0.12301922542974354, "q1_ms": 0.11934801819734275, "q3_ms": 0.1238459552405402, "tflops": 149.83023942486568 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 128, "seq": 1500, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0072650909423828125, "rel_err": 0.0025133287999778986, "peak_mb": 147.4384765625, "compile_s": 0.00025853095576167107, "inner": 200, "median_ms": 0.1841490902006626, "q1_ms": 0.18212106020655483, "q3_ms": 0.18429106450639665, "tflops": 100.09281055863546 }, { "impl": "torch_math", "dtype": "bfloat16", "head_dim": 128, "seq": 1500, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0072650909423828125, "rel_err": 0.0025133287999778986, "peak_mb": 858.04833984375, "compile_s": 0.002511268947273493, "inner": 20, "median_ms": 2.5269907899200916, "q1_ms": 2.518431277712807, "q3_ms": 2.5281910493504256, "tflops": 7.294051119427648 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 128, "seq": 1500, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.01283717155456543, "rel_err": 0.0036486154422163963, "peak_mb": 267.81005859375, "compile_s": 0.0008529829792678356, "inner": 80, "median_ms": 0.567060973844491, "q1_ms": 0.5664798503858037, "q3_ms": 0.5682517323293723, "tflops": 113.76554369917119 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 128, "seq": 1500, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.015437126159667969, "rel_err": 0.007170556113123894, "peak_mb": 255.24951171875, "compile_s": 0.001232924172654748, "inner": 50, "median_ms": 0.9429510962218046, "q1_ms": 0.9394264686852694, "q3_ms": 0.9446282917633653, "tflops": 68.41500079748063 }, { "impl": "torch_math", "dtype": "bfloat16", "head_dim": 128, "seq": 1500, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.011974096298217773, "rel_err": 0.0034811771474778652, "peak_mb": 1341.2275390625, "compile_s": 0.005446049151942134, "inner": 9, "median_ms": 5.600168343840374, "q1_ms": 5.585964445749092, "q3_ms": 5.61415232045369, "tflops": 11.519653703081403 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 128, "seq": 2048, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0006326138973236084, "rel_err": 0.0026098722591996193, "peak_mb": 176.2509765625, "compile_s": 0.00027498090639710426, "inner": 200, "median_ms": 0.24601425509899852, "q1_ms": 0.24410134996287525, "q3_ms": 0.2480948151787743, "tflops": 279.3312798412702 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 128, "seq": 2048, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0005683600902557373, "rel_err": 0.0023447908461093903, "peak_mb": 176.0009765625, "compile_s": 0.0005055719520896673, "inner": 108, "median_ms": 0.4744059988297522, "q1_ms": 0.4730408135974021, "q3_ms": 0.4829067318631267, "tflops": 144.85372635572645 }, { "impl": "torch_math", "dtype": "bfloat16", "head_dim": 128, "seq": 2048, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.000486910343170166, "rel_err": 0.0020087668672204018, "peak_mb": 1440.06298828125, "compile_s": 0.0038699228316545486, "inner": 12, "median_ms": 3.937010420486331, "q1_ms": 3.9223390825403235, "q3_ms": 3.96118675901865, "tflops": 17.454735801159302 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 128, "seq": 2048, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.001190096139907837, "rel_err": 0.004034040495753288, "peak_mb": 336.501953125, "compile_s": 0.0012560940813273191, "inner": 44, "median_ms": 1.1574594748460434, "q1_ms": 1.1529599353459412, "q3_ms": 1.1782964998432859, "tflops": 207.7983495776316 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 128, "seq": 2048, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0010787546634674072, "rel_err": 0.0036132358945906162, "peak_mb": 320.509765625, "compile_s": 0.00203651818446815, "inner": 24, "median_ms": 2.1350886672735214, "q1_ms": 2.1167102725788327, "q3_ms": 2.1411139362802105, "tflops": 112.65020149403836 }, { "impl": "torch_math", "dtype": "bfloat16", "head_dim": 128, "seq": 2048, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0007630288600921631, "rel_err": 0.0032062882091850042, "peak_mb": 2352.0009765625, "compile_s": 0.00884337117895484, "inner": 5, "median_ms": 9.20452419668436, "q1_ms": 9.15345789398998, "q3_ms": 9.213525103405118, "tflops": 26.130429279836008 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 128, "seq": 2048, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.007552385330200195, "rel_err": 0.0026127169840037823, "peak_mb": 176.2509765625, "compile_s": 0.0002413308247923851, "inner": 200, "median_ms": 0.1889779104385525, "q1_ms": 0.1865995500702411, "q3_ms": 0.1899885101011023, "tflops": 181.81880775516515 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 128, "seq": 2048, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.007552385330200195, "rel_err": 0.0026127169840037823, "peak_mb": 176.0009765625, "compile_s": 0.0003567810636013746, "inner": 163, "median_ms": 0.3048628411524318, "q1_ms": 0.3035750462919291, "q3_ms": 0.30497397564671525, "tflops": 112.70556371552047 }, { "impl": "torch_math", "dtype": "bfloat16", "head_dim": 128, "seq": 2048, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.007552385330200195, "rel_err": 0.0026127169840037823, "peak_mb": 1456.06298828125, "compile_s": 0.004507825011387467, "inner": 11, "median_ms": 4.526233914392916, "q1_ms": 4.524630822495303, "q3_ms": 4.540353092703629, "tflops": 7.591242303836726 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 128, "seq": 2048, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.013039588928222656, "rel_err": 0.00369092496111989, "peak_mb": 336.501953125, "compile_s": 0.0012056438717991114, "inner": 57, "median_ms": 0.851451194632733, "q1_ms": 0.8365007203170343, "q3_ms": 0.8645573847420644, "tflops": 141.2401380678934 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 128, "seq": 2048, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.013039588928222656, "rel_err": 0.004484264645725489, "peak_mb": 320.509765625, "compile_s": 0.0015355050563812256, "inner": 34, "median_ms": 1.4523027040173904, "q1_ms": 1.4470198794322857, "q3_ms": 1.4554144855698241, "tflops": 82.80579796163484 }, { "impl": "torch_math", "dtype": "bfloat16", "head_dim": 128, "seq": 2048, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.013872146606445312, "rel_err": 0.003926584962755442, "peak_mb": 2352.0009765625, "compile_s": 0.009425553027540445, "inner": 5, "median_ms": 9.652693616226315, "q1_ms": 9.645280591212213, "q3_ms": 9.66697761323303, "tflops": 12.45860368818117 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 128, "seq": 3000, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0005313754081726074, "rel_err": 0.001971340971067548, "peak_mb": 228.9921875, "compile_s": 0.0005581420846283436, "inner": 100, "median_ms": 0.5207774112932384, "q1_ms": 0.5142651440110058, "q3_ms": 0.5224056204315275, "tflops": 283.1459214673402 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 128, "seq": 3000, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0005130171775817871, "rel_err": 0.0019032340496778488, "peak_mb": 228.0634765625, "compile_s": 0.001031004125252366, "inner": 50, "median_ms": 1.038115220144391, "q1_ms": 1.0348128015175462, "q3_ms": 1.0464538400992751, "tflops": 142.04203650871278 }, { "impl": "torch_math", "dtype": "bfloat16", "head_dim": 128, "seq": 3000, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00048758089542388916, "rel_err": 0.0018088683718815446, "peak_mb": 2864.14111328125, "compile_s": 0.007804816821590066, "inner": 6, "median_ms": 8.05233969974021, "q1_ms": 8.042508057163406, "q3_ms": 8.061971418404331, "tflops": 18.312193163529514 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 128, "seq": 3000, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.00100022554397583, "rel_err": 0.00427584582939744, "peak_mb": 464.3056640625, "compile_s": 0.0022862779442220926, "inner": 22, "median_ms": 2.3310122249478646, "q1_ms": 2.3076423623768445, "q3_ms": 2.3847399014894934, "tflops": 221.40424424910213 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 128, "seq": 3000, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0010921657085418701, "rel_err": 0.004668879322707653, "peak_mb": 440.8720703125, "compile_s": 0.004519656067714095, "inner": 11, "median_ms": 4.785899453881112, "q1_ms": 4.745763356120071, "q3_ms": 4.827367279424586, "tflops": 107.83678281863472 }, { "impl": "torch_math", "dtype": "bfloat16", "head_dim": 128, "seq": 3000, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.000692903995513916, "rel_err": 0.003122720867395401, "peak_mb": 4810.0947265625, "compile_s": 0.01800558203831315, "inner": 2, "median_ms": 18.80805555265397, "q1_ms": 18.672622682061046, "q3_ms": 18.896743305958807, "tflops": 27.440157147301424 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 128, "seq": 3000, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0070171356201171875, "rel_err": 0.0024275495670735836, "peak_mb": 228.4296875, "compile_s": 0.0004134711343795061, "inner": 135, "median_ms": 0.37025691837900215, "q1_ms": 0.3671048735095947, "q3_ms": 0.3744570400427889, "tflops": 199.12659653406016 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 128, "seq": 3000, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0070171356201171875, "rel_err": 0.0024275495670735836, "peak_mb": 228.0634765625, "compile_s": 0.000635222066193819, "inner": 84, "median_ms": 0.6059271067247859, "q1_ms": 0.6002206604794732, "q3_ms": 0.6067078932011056, "tflops": 121.67800248865166 }, { "impl": "torch_math", "dtype": "bfloat16", "head_dim": 128, "seq": 3000, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0070171356201171875, "rel_err": 0.0024275495670735836, "peak_mb": 2898.4736328125, "compile_s": 0.009527093032374978, "inner": 5, "median_ms": 9.625205583870411, "q1_ms": 9.59375542588532, "q3_ms": 9.639732609502971, "tflops": 7.659888337715181 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 128, "seq": 3000, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.011933326721191406, "rel_err": 0.004095529671758413, "peak_mb": 465.4306640625, "compile_s": 0.001643936149775982, "inner": 31, "median_ms": 1.5861745473117599, "q1_ms": 1.5746535462958198, "q3_ms": 1.6003084864707724, "tflops": 162.68575260984886 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 128, "seq": 3000, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.012387275695800781, "rel_err": 0.004737239796668291, "peak_mb": 442.5595703125, "compile_s": 0.0029563000425696373, "inner": 17, "median_ms": 2.968290998764774, "q1_ms": 2.9459117350223307, "q3_ms": 3.0067610581788946, "tflops": 86.93487266153628 }, { "impl": "torch_math", "dtype": "bfloat16", "head_dim": 128, "seq": 3000, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.009859561920166016, "rel_err": 0.002931643510237336, "peak_mb": 4810.6572265625, "compile_s": 0.019537438172847033, "inner": 2, "median_ms": 19.87271453253925, "q1_ms": 19.85700154909864, "q3_ms": 19.874256569892168, "tflops": 12.985040346525208 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 128, "seq": 4096, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.000495031476020813, "rel_err": 0.003562052734196186, "peak_mb": 288.5009765625, "compile_s": 0.00097263278439641, "inner": 56, "median_ms": 0.952361712864201, "q1_ms": 0.9361327247461304, "q3_ms": 0.9531819835371736, "tflops": 288.627633000189 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 128, "seq": 4096, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.000495031476020813, "rel_err": 0.003562052734196186, "peak_mb": 288.0009765625, "compile_s": 0.001778926933184266, "inner": 29, "median_ms": 1.847762309400172, "q1_ms": 1.834807259126984, "q3_ms": 1.901773691723316, "tflops": 148.76259005046595 }, { "impl": "torch_math", "dtype": "bfloat16", "head_dim": 128, "seq": 4096, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.000481531023979187, "rel_err": 0.003464908804744482, "peak_mb": 5120.12548828125, "compile_s": 0.01399792917072773, "inner": 3, "median_ms": 14.445083681493998, "q1_ms": 14.373588336942095, "q3_ms": 14.503622311167419, "tflops": 19.02916680892294 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 128, "seq": 4096, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0008475780487060547, "rel_err": 0.004556367173790932, "peak_mb": 609.001953125, "compile_s": 0.003921613795682788, "inner": 13, "median_ms": 4.01835861759117, "q1_ms": 3.98119345593911, "q3_ms": 4.146808310626792, "tflops": 239.41931665639146 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 128, "seq": 4096, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0008551478385925293, "rel_err": 0.004597060848027468, "peak_mb": 577.017578125, "compile_s": 0.007407565834000707, "inner": 6, "median_ms": 7.850967347621918, "q1_ms": 7.819253912506004, "q3_ms": 7.992347003892064, "tflops": 122.5419278549687 }, { "impl": "torch_math", "dtype": "bfloat16", "head_dim": 128, "seq": 4096, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0006635189056396484, "rel_err": 0.0030714820604771376, "peak_mb": 8736.0009765625, "compile_s": 0.03236389299854636, "inner": 1, "median_ms": 33.789418172091246, "q1_ms": 33.70981244370341, "q3_ms": 34.01364339515567, "tflops": 28.47260255870978 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 128, "seq": 4096, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.006810665130615234, "rel_err": 0.002356122015044093, "peak_mb": 288.5009765625, "compile_s": 0.0006564720533788204, "inner": 89, "median_ms": 0.6089940316598402, "q1_ms": 0.5998846177493087, "q3_ms": 0.610441623485825, "tflops": 225.68193828994353 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 128, "seq": 4096, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.006810665130615234, "rel_err": 0.002356122015044093, "peak_mb": 288.0009765625, "compile_s": 0.0010580038651823997, "inner": 50, "median_ms": 1.0706921387463808, "q1_ms": 1.0654928116127849, "q3_ms": 1.0733142402023077, "tflops": 128.3645863253655 }, { "impl": "torch_math", "dtype": "bfloat16", "head_dim": 128, "seq": 4096, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.006810665130615234, "rel_err": 0.002356122015044093, "peak_mb": 5184.12548828125, "compile_s": 0.01716218004003167, "inner": 2, "median_ms": 17.286985530517995, "q1_ms": 17.214060004334897, "q3_ms": 17.306322522927076, "tflops": 7.9504291381148455 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 128, "seq": 4096, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.015851497650146484, "rel_err": 0.003916752990335226, "peak_mb": 609.001953125, "compile_s": 0.0025770689826458693, "inner": 20, "median_ms": 2.597698546014726, "q1_ms": 2.5397711026016623, "q3_ms": 2.6076948503032327, "tflops": 185.17789059472838 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 128, "seq": 4096, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.015851497650146484, "rel_err": 0.003916752990335226, "peak_mb": 577.017578125, "compile_s": 0.004574715858325362, "inner": 11, "median_ms": 4.7099381812255485, "q1_ms": 4.66930540278554, "q3_ms": 4.7753006854856554, "tflops": 102.13219763042241 }, { "impl": "torch_math", "dtype": "bfloat16", "head_dim": 128, "seq": 4096, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.015851497650146484, "rel_err": 0.003916752990335226, "peak_mb": 8736.0009765625, "compile_s": 0.035069752018898726, "inner": 1, "median_ms": 36.10978601500392, "q1_ms": 35.895834560506046, "q3_ms": 36.275601014494896, "tflops": 13.32149509144489 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 128, "seq": 8192, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0003153011202812195, "rel_err": 0.002931902650743723, "peak_mb": 513.0009765625, "compile_s": 0.003563513047993183, "inner": 14, "median_ms": 3.753031005284616, "q1_ms": 3.7450112868100405, "q3_ms": 3.880444248872144, "tflops": 292.9663054282753 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 128, "seq": 8192, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0003512650728225708, "rel_err": 0.003266321960836649, "peak_mb": 512.0009765625, "compile_s": 0.0067433028016239405, "inner": 7, "median_ms": 7.277723867446184, "q1_ms": 7.228206643568618, "q3_ms": 7.480728290309863, "tflops": 151.07905270962527 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 128, "seq": 8192, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0006157159805297852, "rel_err": 0.005359946750104427, "peak_mb": 1154.001953125, "compile_s": 0.013908217893913388, "inner": 3, "median_ms": 14.997265379255017, "q1_ms": 14.906550330730777, "q3_ms": 15.17653469151507, "tflops": 256.59949330090217 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 128, "seq": 8192, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.0006157159805297852, "rel_err": 0.005359946750104427, "peak_mb": 1090.033203125, "compile_s": 0.028337948955595493, "inner": 1, "median_ms": 30.677677132189274, "q1_ms": 30.391750391572714, "q3_ms": 31.038912944495678, "tflops": 125.44270156550056 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 128, "seq": 8192, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.004915475845336914, "rel_err": 0.0017004889668896794, "peak_mb": 513.0009765625, "compile_s": 0.0020226668566465378, "inner": 25, "median_ms": 2.100210525095463, "q1_ms": 2.0495289750397205, "q3_ms": 2.114935778081417, "tflops": 261.76224112723725 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 128, "seq": 8192, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.004915475845336914, "rel_err": 0.0017004889668896794, "peak_mb": 512.0009765625, "compile_s": 0.003631601808592677, "inner": 14, "median_ms": 3.7831060056175505, "q1_ms": 3.7777417455799878, "q3_ms": 3.823466890025884, "tflops": 145.31863846047804 }, { "impl": "torch_flash", "dtype": "bfloat16", "head_dim": 128, "seq": 8192, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.01760244369506836, "rel_err": 0.005058928392827511, "peak_mb": 1154.001953125, "compile_s": 0.00797761813737452, "inner": 6, "median_ms": 8.399719333586594, "q1_ms": 8.341048242679488, "q3_ms": 8.435636817011982, "tflops": 229.07257637933597 }, { "impl": "torch_efficient", "dtype": "bfloat16", "head_dim": 128, "seq": 8192, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 0.01760244369506836, "rel_err": 0.005591186229139566, "peak_mb": 1090.033203125, "compile_s": 0.015571523923426867, "inner": 3, "median_ms": 16.697434708476067, "q1_ms": 16.619109509823225, "q3_ms": 17.007144206824403, "tflops": 115.23598577877667 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 64, "seq": 128, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 8.642673492431641e-07, "rel_err": 9.249482673112652e-07, "peak_mb": 70.0009765625, "compile_s": 0.0002504209987819195, "inner": 200, "median_ms": 0.02362438477575779, "q1_ms": 0.02359755744691938, "q3_ms": 0.02364963002037257, "tflops": 5.681321620604816 }, { "impl": "torch_math", "dtype": "float32", "head_dim": 64, "seq": 128, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 4.470348358154297e-07, "rel_err": 4.784215548170323e-07, "peak_mb": 74.50439453125, "compile_s": 0.00034217117354273796, "inner": 200, "median_ms": 0.12004217016510665, "q1_ms": 0.11972441745456308, "q3_ms": 0.12016419204883277, "tflops": 1.1180881503174775 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 64, "seq": 128, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 1.3113021850585938e-06, "rel_err": 1.412345000062487e-06, "peak_mb": 78.048828125, "compile_s": 0.0006458619609475136, "inner": 168, "median_ms": 0.21242061934788667, "q1_ms": 0.2115813923800098, "q3_ms": 0.21318598217996104, "tflops": 2.2114710400625404 }, { "impl": "torch_math", "dtype": "float32", "head_dim": 64, "seq": 128, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 4.172325134277344e-07, "rel_err": 5.33427794380259e-07, "peak_mb": 83.0009765625, "compile_s": 0.0009195129387080669, "inner": 83, "median_ms": 0.5134731076702654, "q1_ms": 0.512921494005405, "q3_ms": 0.5140231083517512, "tflops": 0.9148717644267844 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 64, "seq": 128, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 1.0728836059570312e-06, "rel_err": 3.539919930517499e-07, "peak_mb": 70.0009765625, "compile_s": 7.76299275457859e-05, "inner": 200, "median_ms": 0.023842635564506054, "q1_ms": 0.023810809943825006, "q3_ms": 0.023938632803037763, "tflops": 2.8146579608801017 }, { "impl": "torch_math", "dtype": "float32", "head_dim": 64, "seq": 128, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 5.960464477539062e-07, "rel_err": 1.966622136251317e-07, "peak_mb": 74.56689453125, "compile_s": 0.0002820009831339121, "inner": 200, "median_ms": 0.16564138000831008, "q1_ms": 0.1652208500308916, "q3_ms": 0.16578647715505213, "tflops": 0.4051455258138588 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 64, "seq": 128, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 1.9073486328125e-06, "rel_err": 1.1423012438172009e-06, "peak_mb": 78.048828125, "compile_s": 0.0006659519858658314, "inner": 139, "median_ms": 0.21468175628738437, "q1_ms": 0.21453711569389639, "q3_ms": 0.2147643098259251, "tflops": 1.094089353757549 }, { "impl": "torch_math", "dtype": "float32", "head_dim": 64, "seq": 128, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 1.6689300537109375e-06, "rel_err": 6.106315595388878e-07, "peak_mb": 83.0009765625, "compile_s": 0.0009549139067530632, "inner": 74, "median_ms": 0.5657015701856565, "q1_ms": 0.564699128861665, "q3_ms": 0.5662910265822869, "tflops": 0.41520306178912475 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 64, "seq": 512, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 4.3213367462158203e-07, "rel_err": 1.052677021107229e-06, "peak_mb": 88.0009765625, "compile_s": 0.00027654110454022884, "inner": 200, "median_ms": 0.08671590592712164, "q1_ms": 0.08331324032042176, "q3_ms": 0.08792492968495935, "tflops": 24.764587592555426 }, { "impl": "torch_math", "dtype": "float32", "head_dim": 64, "seq": 512, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 3.5762786865234375e-07, "rel_err": 8.711809869055287e-07, "peak_mb": 160.01611328125, "compile_s": 0.00029198103584349155, "inner": 200, "median_ms": 0.1847137906588614, "q1_ms": 0.18444031768012792, "q3_ms": 0.18473402014933527, "tflops": 11.626006051524758 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 64, "seq": 512, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 1.1324882507324219e-06, "rel_err": 1.943797542480752e-06, "peak_mb": 120.189453125, "compile_s": 0.00046231108717620373, "inner": 126, "median_ms": 0.3104060001316524, "q1_ms": 0.30845571919122616, "q3_ms": 0.31622142128143754, "tflops": 24.21407049094463 }, { "impl": "torch_math", "dtype": "float32", "head_dim": 64, "seq": 512, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 7.748603820800781e-07, "rel_err": 1.32996672164154e-06, "peak_mb": 236.0009765625, "compile_s": 0.0009548128582537174, "inner": 72, "median_ms": 0.5102496943436563, "q1_ms": 0.5093980273361214, "q3_ms": 0.5104296320092141, "tflops": 14.730420912193232 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 64, "seq": 512, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 9.5367431640625e-07, "rel_err": 3.336561462674581e-07, "peak_mb": 88.0009765625, "compile_s": 0.00012456090189516544, "inner": 200, "median_ms": 0.07598516414873302, "q1_ms": 0.0759210146497935, "q3_ms": 0.07602011493872851, "tflops": 14.13094037539042 }, { "impl": "torch_math", "dtype": "float32", "head_dim": 64, "seq": 512, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 3.8743019104003906e-07, "rel_err": 1.35547807644798e-07, "peak_mb": 161.01611328125, "compile_s": 0.0003584008663892746, "inner": 161, "median_ms": 0.2254874341374969, "q1_ms": 0.2246417137013274, "q3_ms": 0.22552808043865127, "tflops": 4.761869893580224 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 64, "seq": 512, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 3.337860107421875e-06, "rel_err": 8.472996455566317e-07, "peak_mb": 120.189453125, "compile_s": 0.0004957718774676323, "inner": 101, "median_ms": 0.283485336118553, "q1_ms": 0.28139479356900887, "q3_ms": 0.2843596036592037, "tflops": 13.25675759972421 }, { "impl": "torch_math", "dtype": "float32", "head_dim": 64, "seq": 512, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 3.337860107421875e-06, "rel_err": 1.1120807812403655e-06, "peak_mb": 236.0009765625, "compile_s": 0.000997783849015832, "inner": 67, "median_ms": 0.5612946567195121, "q1_ms": 0.5607912768679324, "q3_ms": 0.5623374096076212, "tflops": 6.695407374736477 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 64, "seq": 600, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 6.556510925292969e-07, "rel_err": 1.4932564909031498e-06, "peak_mb": 92.1259765625, "compile_s": 0.00018202094361186028, "inner": 200, "median_ms": 0.14346285024657845, "q1_ms": 0.13988448772579432, "q3_ms": 0.14353267499245703, "tflops": 20.556680666326965 }, { "impl": "torch_math", "dtype": "float32", "head_dim": 64, "seq": 600, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 3.2782554626464844e-07, "rel_err": 7.466282454515749e-07, "peak_mb": 191.13037109375, "compile_s": 0.00045193196274340153, "inner": 146, "median_ms": 0.27608109497758626, "q1_ms": 0.27508373700771227, "q3_ms": 0.27764720159697615, "tflops": 10.682078757473146 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 64, "seq": 600, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 1.1920928955078125e-06, "rel_err": 1.6196685237446218e-06, "peak_mb": 130.091796875, "compile_s": 0.0008710429538041353, "inner": 78, "median_ms": 0.573317383010036, "q1_ms": 0.5722364174942367, "q3_ms": 0.5739781533917174, "tflops": 18.003849710273503 }, { "impl": "torch_math", "dtype": "float32", "head_dim": 64, "seq": 600, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 7.152557373046875e-07, "rel_err": 1.127220912167104e-06, "peak_mb": 291.5634765625, "compile_s": 0.0008931539487093687, "inner": 62, "median_ms": 0.6753675181478743, "q1_ms": 0.6706425732362175, "q3_ms": 0.6761293569129081, "tflops": 15.283411953697152 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 64, "seq": 600, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 9.834766387939453e-07, "rel_err": 4.1586460497455846e-07, "peak_mb": 92.1259765625, "compile_s": 0.00016660080291330814, "inner": 200, "median_ms": 0.11586465407162905, "q1_ms": 0.11577532975934446, "q3_ms": 0.11592113005463034, "tflops": 12.726573188475646 }, { "impl": "torch_math", "dtype": "float32", "head_dim": 64, "seq": 600, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 4.172325134277344e-07, "rel_err": 1.76427406017865e-07, "peak_mb": 192.50390625, "compile_s": 0.0004569909069687128, "inner": 118, "median_ms": 0.33564022860600284, "q1_ms": 0.3351817160087117, "q3_ms": 0.33592379659856275, "tflops": 4.393275520411286 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 64, "seq": 600, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 4.291534423828125e-06, "rel_err": 1.113206849367998e-06, "peak_mb": 130.091796875, "compile_s": 0.0006817721296101809, "inner": 106, "median_ms": 0.3887775841713795, "q1_ms": 0.3886285758542143, "q3_ms": 0.38894442050665057, "tflops": 13.274839420075631 }, { "impl": "torch_math", "dtype": "float32", "head_dim": 64, "seq": 600, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 2.86102294921875e-06, "rel_err": 9.239368523594749e-07, "peak_mb": 291.5634765625, "compile_s": 0.0009545630309730768, "inner": 57, "median_ms": 0.7257674363229359, "q1_ms": 0.7248962984273308, "q3_ms": 0.7274864801091322, "tflops": 7.111038249591002 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 64, "seq": 1024, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 4.917383193969727e-07, "rel_err": 1.6426118918388966e-06, "peak_mb": 112.0009765625, "compile_s": 0.00036197109147906303, "inner": 148, "median_ms": 0.3467041816612756, "q1_ms": 0.3381184093037469, "q3_ms": 0.3498712195888967, "tflops": 24.775976311679532 }, { "impl": "torch_math", "dtype": "float32", "head_dim": 64, "seq": 1024, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 3.0547380447387695e-07, "rel_err": 1.0204104228250799e-06, "peak_mb": 400.03173828125, "compile_s": 0.0007442520000040531, "inner": 69, "median_ms": 0.6803780274492675, "q1_ms": 0.6787150430128626, "q3_ms": 0.6806785274512959, "tflops": 12.625238096244239 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 64, "seq": 1024, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 5.662441253662109e-07, "rel_err": 1.656587301113177e-06, "peak_mb": 176.376953125, "compile_s": 0.0011464538984000683, "inner": 46, "median_ms": 1.1192656314486393, "q1_ms": 1.1108146433758996, "q3_ms": 1.120340957245587, "tflops": 26.861158090852722 }, { "impl": "torch_math", "dtype": "float32", "head_dim": 64, "seq": 1024, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 5.21540641784668e-07, "rel_err": 1.420848320776713e-06, "peak_mb": 664.0009765625, "compile_s": 0.0016887059900909662, "inner": 30, "median_ms": 1.6110586312909922, "q1_ms": 1.605118267859022, "q3_ms": 1.6139626347770295, "tflops": 18.661500263282253 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 64, "seq": 1024, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 8.344650268554688e-07, "rel_err": 2.8424688025552314e-07, "peak_mb": 112.0009765625, "compile_s": 0.0002734509762376547, "inner": 200, "median_ms": 0.2352216700091958, "q1_ms": 0.23139487951993942, "q3_ms": 0.23611372045706958, "tflops": 18.25923307079697 }, { "impl": "torch_math", "dtype": "float32", "head_dim": 64, "seq": 1024, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 4.172325134277344e-07, "rel_err": 1.4212344012776157e-07, "peak_mb": 404.03173828125, "compile_s": 0.0009740130044519901, "inner": 54, "median_ms": 0.8359917965545146, "q1_ms": 0.8335306839500036, "q3_ms": 0.8402919996081403, "tflops": 5.137571102612999 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 64, "seq": 1024, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 2.6226043701171875e-06, "rel_err": 1.0977464626193978e-06, "peak_mb": 176.376953125, "compile_s": 0.0010022141505032778, "inner": 59, "median_ms": 0.7922188655110235, "q1_ms": 0.7856414748071614, "q3_ms": 0.7994617631469491, "tflops": 18.97504110344975 }, { "impl": "torch_math", "dtype": "float32", "head_dim": 64, "seq": 1024, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 5.7220458984375e-06, "rel_err": 1.7685914599496755e-06, "peak_mb": 664.0009765625, "compile_s": 0.00188536592759192, "inner": 27, "median_ms": 1.7427072204925396, "q1_ms": 1.7358903014273555, "q3_ms": 1.743542185674111, "tflops": 8.62588124914718 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 64, "seq": 1500, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 3.129243850708008e-07, "rel_err": 1.2418921642165515e-06, "peak_mb": 134.8759765625, "compile_s": 0.0007993530016392469, "inner": 75, "median_ms": 0.6883786680797735, "q1_ms": 0.684875319711864, "q3_ms": 0.7048857941602666, "tflops": 26.775960462888698 }, { "impl": "torch_math", "dtype": "float32", "head_dim": 64, "seq": 1500, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 4.470348358154297e-07, "rel_err": 1.774131760612363e-06, "peak_mb": 752.90234375, "compile_s": 0.0014752349816262722, "inner": 34, "median_ms": 1.4560915570815698, "q1_ms": 1.4518962364972514, "q3_ms": 1.4570797656607979, "tflops": 12.658544657001569 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 64, "seq": 1500, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 6.258487701416016e-07, "rel_err": 1.2473811921154265e-06, "peak_mb": 231.25537109375, "compile_s": 0.002198738045990467, "inner": 23, "median_ms": 2.3383616024385327, "q1_ms": 2.3004675850920053, "q3_ms": 2.401890764620317, "tflops": 27.588547439679314 }, { "impl": "torch_math", "dtype": "float32", "head_dim": 64, "seq": 1500, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 5.662441253662109e-07, "rel_err": 1.207150944537716e-06, "peak_mb": 1292.3837890625, "compile_s": 0.003490183036774397, "inner": 14, "median_ms": 3.4938871394842863, "q1_ms": 3.479736782277801, "q3_ms": 3.498352922698749, "tflops": 18.4642483928437 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 64, "seq": 1500, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 1.3709068298339844e-06, "rel_err": 4.6151853894116357e-07, "peak_mb": 135.1572265625, "compile_s": 0.0004839219618588686, "inner": 110, "median_ms": 0.4726121900603175, "q1_ms": 0.4720527773977003, "q3_ms": 0.4744641144167293, "tflops": 19.50013180748427 }, { "impl": "torch_math", "dtype": "float32", "head_dim": 64, "seq": 1500, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 4.76837158203125e-07, "rel_err": 1.6052818807565927e-07, "peak_mb": 762.04833984375, "compile_s": 0.0018467758782207966, "inner": 27, "median_ms": 1.7767914105206728, "q1_ms": 1.773673208878824, "q3_ms": 1.7777728477354955, "tflops": 5.186877843640258 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 64, "seq": 1500, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 3.814697265625e-06, "rel_err": 1.5930278323139646e-06, "peak_mb": 231.81787109375, "compile_s": 0.0015991448890417814, "inner": 33, "median_ms": 1.5325150311444744, "q1_ms": 1.5294197276516845, "q3_ms": 1.5384794572706926, "tflops": 21.04775440663142 }, { "impl": "torch_math", "dtype": "float32", "head_dim": 64, "seq": 1500, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 4.5299530029296875e-06, "rel_err": 1.4902518614690052e-06, "peak_mb": 1292.9462890625, "compile_s": 0.003793272888287902, "inner": 13, "median_ms": 3.7606962275906253, "q1_ms": 3.7441972750597277, "q3_ms": 3.7633573117021184, "tflops": 8.577135202612611 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 64, "seq": 2048, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 3.650784492492676e-07, "rel_err": 1.8326874169360963e-06, "peak_mb": 160.0009765625, "compile_s": 0.0013113049790263176, "inner": 43, "median_ms": 1.223551228555829, "q1_ms": 1.2163784311607828, "q3_ms": 1.28655819774627, "tflops": 28.081977743224677 }, { "impl": "torch_math", "dtype": "float32", "head_dim": 64, "seq": 2048, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 3.129243850708008e-07, "rel_err": 1.5708749288023682e-06, "peak_mb": 1312.06298828125, "compile_s": 0.002593049081042409, "inner": 19, "median_ms": 2.5731552588312248, "q1_ms": 2.560485999933199, "q3_ms": 2.5740616011286255, "tflops": 13.353153973152338 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 64, "seq": 2048, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 8.046627044677734e-07, "rel_err": 2.781386911010486e-06, "peak_mb": 288.751953125, "compile_s": 0.003842714009806514, "inner": 13, "median_ms": 4.083483462007, "q1_ms": 4.011702070084329, "q3_ms": 4.246142889874486, "tflops": 29.45012154619909 }, { "impl": "torch_math", "dtype": "float32", "head_dim": 64, "seq": 2048, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 5.513429641723633e-07, "rel_err": 1.648229385864397e-06, "peak_mb": 2288.0009765625, "compile_s": 0.006075580837205052, "inner": 8, "median_ms": 6.172191497171298, "q1_ms": 6.1595890001626685, "q3_ms": 6.179542760946788, "tflops": 19.484017037241067 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 64, "seq": 2048, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 1.0728836059570312e-06, "rel_err": 4.1042079601538717e-07, "peak_mb": 160.0009765625, "compile_s": 0.0007774629630148411, "inner": 66, "median_ms": 0.8227604399011893, "q1_ms": 0.7746961906157208, "q3_ms": 0.8282217595754473, "tflops": 20.880767172110563 }, { "impl": "torch_math", "dtype": "float32", "head_dim": 64, "seq": 2048, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 3.5762786865234375e-07, "rel_err": 1.3680693200512906e-07, "peak_mb": 1328.06298828125, "compile_s": 0.0032765408977866173, "inner": 15, "median_ms": 3.1879911199212074, "q1_ms": 3.182321429873506, "q3_ms": 3.199441839630405, "tflops": 5.388932571563942 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 64, "seq": 2048, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 4.291534423828125e-06, "rel_err": 1.452910396437801e-06, "peak_mb": 288.751953125, "compile_s": 0.002470399020239711, "inner": 20, "median_ms": 2.562994894105941, "q1_ms": 2.5127337779849768, "q3_ms": 2.580012724502012, "tflops": 23.46065623551514 }, { "impl": "torch_math", "dtype": "float32", "head_dim": 64, "seq": 2048, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 6.9141387939453125e-06, "rel_err": 1.8645683894646936e-06, "peak_mb": 2288.0009765625, "compile_s": 0.0066794229205697775, "inner": 7, "median_ms": 6.671917425202472, "q1_ms": 6.658732508575278, "q3_ms": 6.684344710915216, "tflops": 9.012333083869851 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 64, "seq": 3000, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 3.203749656677246e-07, "rel_err": 1.7645199932303512e-06, "peak_mb": 204.6259765625, "compile_s": 0.002492329105734825, "inner": 20, "median_ms": 2.6443487033247948, "q1_ms": 2.6249506801832467, "q3_ms": 2.7629886229988188, "tflops": 27.881345568116732 }, { "impl": "torch_math", "dtype": "float32", "head_dim": 64, "seq": 3000, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 3.501772880554199e-07, "rel_err": 1.928661276906496e-06, "peak_mb": 2676.64111328125, "compile_s": 0.0055732091423124075, "inner": 8, "median_ms": 5.728391115553677, "q1_ms": 5.706727446522564, "q3_ms": 5.739556305343285, "tflops": 12.870629555969805 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 64, "seq": 3000, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 8.940696716308594e-07, "rel_err": 3.1199194836517563e-06, "peak_mb": 397.1640625, "compile_s": 0.00815481785684824, "inner": 6, "median_ms": 8.799173985607922, "q1_ms": 8.68706857242311, "q3_ms": 8.971464645583183, "tflops": 29.326389093120294 }, { "impl": "torch_math", "dtype": "float32", "head_dim": 64, "seq": 3000, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 9.5367431640625e-07, "rel_err": 3.3279141007369617e-06, "peak_mb": 4718.0322265625, "compile_s": 0.013125655939802527, "inner": 3, "median_ms": 13.503773603588343, "q1_ms": 13.488893668788176, "q3_ms": 13.516920308272045, "tflops": 19.109325109792216 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 64, "seq": 3000, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 1.0728836059570312e-06, "rel_err": 3.909314330030611e-07, "peak_mb": 205.7509765625, "compile_s": 0.0014793460723012686, "inner": 33, "median_ms": 1.4972803676783135, "q1_ms": 1.4769408964749537, "q3_ms": 1.5098863161366545, "tflops": 24.620639391112437 }, { "impl": "torch_math", "dtype": "float32", "head_dim": 64, "seq": 3000, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 4.76837158203125e-07, "rel_err": 1.7374730987285147e-07, "peak_mb": 2712.0986328125, "compile_s": 0.00727512501180172, "inner": 6, "median_ms": 7.3390121494109435, "q1_ms": 7.317342174549898, "q3_ms": 7.358233172756931, "tflops": 5.02301934504344 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 64, "seq": 3000, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 3.337860107421875e-06, "rel_err": 1.4926331459719222e-06, "peak_mb": 397.7265625, "compile_s": 0.004741176031529903, "inner": 10, "median_ms": 5.075963702984154, "q1_ms": 4.940533253829926, "q3_ms": 5.135991401039064, "tflops": 25.41862147756236 }, { "impl": "torch_math", "dtype": "float32", "head_dim": 64, "seq": 3000, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 1.0013580322265625e-05, "rel_err": 3.607196958910208e-06, "peak_mb": 4717.4697265625, "compile_s": 0.01472575101070106, "inner": 3, "median_ms": 14.942205355813106, "q1_ms": 14.864451834000647, "q3_ms": 14.952602136569718, "tflops": 8.634869949087173 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 64, "seq": 4096, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 3.0547380447387695e-07, "rel_err": 1.3578446669271216e-06, "peak_mb": 256.0009765625, "compile_s": 0.004493615124374628, "inner": 11, "median_ms": 4.80164035054093, "q1_ms": 4.771486136384985, "q3_ms": 4.918959415094419, "tflops": 28.623333577350657 }, { "impl": "torch_math", "dtype": "float32", "head_dim": 64, "seq": 4096, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 4.3213367462158203e-07, "rel_err": 1.9208534922654508e-06, "peak_mb": 4864.12548828125, "compile_s": 0.009986025048419833, "inner": 5, "median_ms": 10.261655785143375, "q1_ms": 10.239298571832478, "q3_ms": 10.274386801756918, "tflops": 13.393448031162908 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 64, "seq": 4096, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 7.450580596923828e-07, "rel_err": 3.5952757571067195e-06, "peak_mb": 513.501953125, "compile_s": 0.014772261027246714, "inner": 3, "median_ms": 16.048466010640066, "q1_ms": 15.948287366578976, "q3_ms": 16.386293650915228, "tflops": 29.973976131617498 }, { "impl": "torch_math", "dtype": "float32", "head_dim": 64, "seq": 4096, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 6.407499313354492e-07, "rel_err": 2.838011823769193e-06, "peak_mb": 8608.0009765625, "compile_s": 0.02360473293811083, "inner": 2, "median_ms": 24.2412545485422, "q1_ms": 24.225266999565065, "q3_ms": 24.28581955609843, "tflops": 19.843706363824644 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 64, "seq": 4096, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 1.2516975402832031e-06, "rel_err": 4.4143843069832656e-07, "peak_mb": 256.0009765625, "compile_s": 0.002543109003454447, "inner": 19, "median_ms": 2.632124472017351, "q1_ms": 2.6117693694112334, "q3_ms": 2.6394502308807875, "tflops": 26.107988990099326 }, { "impl": "torch_math", "dtype": "float32", "head_dim": 64, "seq": 4096, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 5.364418029785156e-07, "rel_err": 1.8918790090083348e-07, "peak_mb": 4928.12548828125, "compile_s": 0.013117486145347357, "inner": 3, "median_ms": 13.247382982323566, "q1_ms": 13.181196215252083, "q3_ms": 13.269819513273736, "tflops": 5.187400170108672 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 64, "seq": 4096, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 6.198883056640625e-06, "rel_err": 1.7577269773028092e-06, "peak_mb": 513.501953125, "compile_s": 0.008394569158554077, "inner": 6, "median_ms": 9.02271484180043, "q1_ms": 8.908431006905932, "q3_ms": 9.114861759978037, "tflops": 26.656962210722597 }, { "impl": "torch_math", "dtype": "float32", "head_dim": 64, "seq": 4096, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 7.152557373046875e-06, "rel_err": 2.6365905796410516e-06, "peak_mb": 8608.0009765625, "compile_s": 0.026418392080813646, "inner": 1, "median_ms": 26.856482960283756, "q1_ms": 26.774498401209712, "q3_ms": 26.880873483605683, "tflops": 8.955683770346479 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 64, "seq": 8192, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 2.8312206268310547e-07, "rel_err": 2.5178296709782444e-06, "peak_mb": 448.0009765625, "compile_s": 0.017671552021056414, "inner": 2, "median_ms": 19.266702001914382, "q1_ms": 18.917513254564255, "q3_ms": 19.646203552838415, "tflops": 28.533986451514902 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 64, "seq": 8192, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 5.140900611877441e-07, "rel_err": 3.7055394841445377e-06, "peak_mb": 963.001953125, "compile_s": 0.06370388204231858, "inner": 1, "median_ms": 65.28051709756255, "q1_ms": 64.56094444729388, "q3_ms": 65.68088906351477, "tflops": 29.475032278502646 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 64, "seq": 8192, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 1.0728836059570312e-06, "rel_err": 4.066610586050956e-07, "peak_mb": 448.0009765625, "compile_s": 0.00926714320667088, "inner": 5, "median_ms": 9.78560820221901, "q1_ms": 9.700239822268486, "q3_ms": 9.871178423054516, "tflops": 28.090017632390797 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 64, "seq": 8192, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 1.0013580322265625e-05, "rel_err": 3.3173118936247192e-06, "peak_mb": 963.001953125, "compile_s": 0.030912618152797222, "inner": 1, "median_ms": 33.85834791697562, "q1_ms": 32.77042938861996, "q3_ms": 34.478735411539674, "tflops": 28.41463726059841 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 128, "seq": 128, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 1.2516975402832031e-06, "rel_err": 1.4724993206982617e-06, "peak_mb": 76.0009765625, "compile_s": 8.078105747699738e-05, "inner": 200, "median_ms": 0.023915680358186364, "q1_ms": 0.02386606065556407, "q3_ms": 0.02404293220024556, "tflops": 11.224245013297907 }, { "impl": "torch_math", "dtype": "float32", "head_dim": 128, "seq": 128, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 6.556510925292969e-07, "rel_err": 7.71309146330168e-07, "peak_mb": 80.50439453125, "compile_s": 0.00023110094480216503, "inner": 200, "median_ms": 0.11793151032179594, "q1_ms": 0.11782603512983769, "q3_ms": 0.11809981253463775, "tflops": 2.276197898827283 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 128, "seq": 128, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 1.0728836059570312e-06, "rel_err": 1.3375477010413306e-06, "peak_mb": 92.048828125, "compile_s": 0.0003423320595175028, "inner": 200, "median_ms": 0.18094983068294823, "q1_ms": 0.18037427798844874, "q3_ms": 0.18110623059328645, "tflops": 5.192180022794218 }, { "impl": "torch_math", "dtype": "float32", "head_dim": 128, "seq": 128, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 7.748603820800781e-07, "rel_err": 1.1700042250595288e-06, "peak_mb": 94.0009765625, "compile_s": 0.0006389529444277287, "inner": 90, "median_ms": 0.4545270206613673, "q1_ms": 0.4544341380500959, "q3_ms": 0.4549371337311135, "tflops": 2.0670368389384848 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 128, "seq": 128, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 1.7881393432617188e-06, "rel_err": 5.69636540603824e-07, "peak_mb": 76.0009765625, "compile_s": 8.159107528626919e-05, "inner": 200, "median_ms": 0.023726829094812274, "q1_ms": 0.02366200729738921, "q3_ms": 0.02378920733463019, "tflops": 5.656791620307405 }, { "impl": "torch_math", "dtype": "float32", "head_dim": 128, "seq": 128, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 8.344650268554688e-07, "rel_err": 2.6583037993077596e-07, "peak_mb": 80.56689453125, "compile_s": 0.0002543199807405472, "inner": 200, "median_ms": 0.16387411975301802, "q1_ms": 0.16366726718842983, "q3_ms": 0.16408714465796947, "tflops": 0.8190294367547818 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 128, "seq": 128, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 2.0265579223632812e-06, "rel_err": 1.1511345974213327e-06, "peak_mb": 92.048828125, "compile_s": 0.0003652020823210478, "inner": 164, "median_ms": 0.24139395029079624, "q1_ms": 0.24123425615338107, "q3_ms": 0.24203748780120982, "tflops": 1.946039026388603 }, { "impl": "torch_math", "dtype": "float32", "head_dim": 128, "seq": 128, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 2.1457672119140625e-06, "rel_err": 5.560196996157174e-07, "peak_mb": 94.0009765625, "compile_s": 0.0007665730081498623, "inner": 77, "median_ms": 0.5684124952932069, "q1_ms": 0.5682615923102606, "q3_ms": 0.5697562803768299, "tflops": 0.8264456743824401 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 128, "seq": 512, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 1.043081283569336e-06, "rel_err": 1.5308652336898376e-06, "peak_mb": 112.0009765625, "compile_s": 0.00015983101911842823, "inner": 200, "median_ms": 0.11063883546739817, "q1_ms": 0.1091509050456807, "q3_ms": 0.11332932277582586, "tflops": 38.81970808763251 }, { "impl": "torch_math", "dtype": "float32", "head_dim": 128, "seq": 512, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 6.556510925292969e-07, "rel_err": 9.62258127401583e-07, "peak_mb": 184.01611328125, "compile_s": 0.00033552106469869614, "inner": 174, "median_ms": 0.24457935688096558, "q1_ms": 0.24175966337815613, "q3_ms": 0.24518093974704203, "tflops": 17.56062879047605 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 128, "seq": 512, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 1.1920928955078125e-06, "rel_err": 1.6162218798854155e-06, "peak_mb": 176.189453125, "compile_s": 0.0007928130216896534, "inner": 88, "median_ms": 0.5086307487958533, "q1_ms": 0.5073242090557787, "q3_ms": 0.5089380230691114, "tflops": 29.554614170669172 }, { "impl": "torch_math", "dtype": "float32", "head_dim": 128, "seq": 512, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 1.0132789611816406e-06, "rel_err": 1.4112234794083633e-06, "peak_mb": 280.0009765625, "compile_s": 0.0007927222177386284, "inner": 70, "median_ms": 0.6212083024105856, "q1_ms": 0.6174437951163522, "q3_ms": 0.6218113786807018, "tflops": 24.19862303460393 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 128, "seq": 512, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 1.3709068298339844e-06, "rel_err": 3.929849867745361e-07, "peak_mb": 112.0009765625, "compile_s": 0.00014639017172157764, "inner": 200, "median_ms": 0.09296397445723414, "q1_ms": 0.0927384226815775, "q3_ms": 0.09322032448835671, "tflops": 23.100170367478196 }, { "impl": "torch_math", "dtype": "float32", "head_dim": 128, "seq": 512, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 1.1920928955078125e-06, "rel_err": 3.417260927562893e-07, "peak_mb": 185.01611328125, "compile_s": 0.000417881878092885, "inner": 138, "median_ms": 0.28233526045105595, "q1_ms": 0.28083140901087417, "q3_ms": 0.28238982633701054, "tflops": 7.606147544480282 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 128, "seq": 512, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 2.86102294921875e-06, "rel_err": 1.1809877378254896e-06, "peak_mb": 176.189453125, "compile_s": 0.0007221931591629982, "inner": 93, "median_ms": 0.44471853893370394, "q1_ms": 0.4436302572108245, "q3_ms": 0.4452176827697023, "tflops": 16.901010661758065 }, { "impl": "torch_math", "dtype": "float32", "head_dim": 128, "seq": 512, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 5.7220458984375e-06, "rel_err": 1.1296403954474954e-06, "peak_mb": 280.0009765625, "compile_s": 0.0009199630003422499, "inner": 62, "median_ms": 0.647451450687743, "q1_ms": 0.6440295837819576, "q3_ms": 0.6481156451627612, "tflops": 11.608890149239866 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 128, "seq": 600, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 1.0132789611816406e-06, "rel_err": 2.266329829581082e-06, "peak_mb": 120.2509765625, "compile_s": 0.0002457709051668644, "inner": 200, "median_ms": 0.19300852087326348, "q1_ms": 0.1910482649691403, "q3_ms": 0.19716463750228286, "tflops": 30.559479826660098 }, { "impl": "torch_math", "dtype": "float32", "head_dim": 128, "seq": 600, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 1.9371509552001953e-06, "rel_err": 4.332689513830701e-06, "peak_mb": 219.25537109375, "compile_s": 0.0004335809499025345, "inner": 122, "median_ms": 0.37018818109006174, "q1_ms": 0.3681385777050965, "q3_ms": 0.37046592514656607, "tflops": 15.933085661006121 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 128, "seq": 600, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 8.344650268554688e-07, "rel_err": 2.030101541095064e-06, "peak_mb": 196.6494140625, "compile_s": 0.0010569228325039148, "inner": 56, "median_ms": 0.8356834462444697, "q1_ms": 0.8294014551211149, "q3_ms": 0.8410132868448272, "tflops": 24.702942355472818 }, { "impl": "torch_math", "dtype": "float32", "head_dim": 128, "seq": 600, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 1.6093254089355469e-06, "rel_err": 3.709342990987352e-06, "peak_mb": 343.1259765625, "compile_s": 0.0010863838251680136, "inner": 50, "median_ms": 0.9533789195120335, "q1_ms": 0.952206018846482, "q3_ms": 0.9548228303901851, "tflops": 21.65334221000618 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 128, "seq": 600, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 1.430511474609375e-06, "rel_err": 4.378110531888524e-07, "peak_mb": 120.2509765625, "compile_s": 0.00018307100981473923, "inner": 200, "median_ms": 0.12817055103369057, "q1_ms": 0.12766657047905028, "q3_ms": 0.1289932255167514, "tflops": 23.0093416640208 }, { "impl": "torch_math", "dtype": "float32", "head_dim": 128, "seq": 600, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 1.1920928955078125e-06, "rel_err": 3.648425490609952e-07, "peak_mb": 220.57421875, "compile_s": 0.0005673819687217474, "inner": 100, "median_ms": 0.4257543897256255, "q1_ms": 0.42375633027404547, "q3_ms": 0.426843436434865, "tflops": 6.926810553616465 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 128, "seq": 600, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 3.814697265625e-06, "rel_err": 1.427154188604618e-06, "peak_mb": 196.029296875, "compile_s": 0.001164924120530486, "inner": 78, "median_ms": 0.588014100988706, "q1_ms": 0.5856958913425796, "q3_ms": 0.5887311537009783, "tflops": 17.5538647502575 }, { "impl": "torch_math", "dtype": "float32", "head_dim": 128, "seq": 600, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 4.0531158447265625e-06, "rel_err": 1.3379570873439661e-06, "peak_mb": 343.1259765625, "compile_s": 0.0011359942145645618, "inner": 45, "median_ms": 0.9994523754964272, "q1_ms": 0.993645020450155, "q3_ms": 1.0001786001440554, "tflops": 10.327575633479395 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 128, "seq": 1024, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 1.1324882507324219e-06, "rel_err": 3.026040076292702e-06, "peak_mb": 160.0009765625, "compile_s": 0.0004651418421417475, "inner": 113, "median_ms": 0.43520497248829465, "q1_ms": 0.4344467388927541, "q3_ms": 0.43684479665110065, "tflops": 39.47535131727401 }, { "impl": "torch_math", "dtype": "float32", "head_dim": 128, "seq": 1024, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 9.5367431640625e-07, "rel_err": 2.5482443106739083e-06, "peak_mb": 448.03173828125, "compile_s": 0.0009045230690389872, "inner": 57, "median_ms": 0.8726218069919891, "q1_ms": 0.8643151142454722, "q3_ms": 0.8755334032078583, "tflops": 19.687645949647592 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 128, "seq": 1024, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 1.1324882507324219e-06, "rel_err": 2.0152297111053485e-06, "peak_mb": 288.376953125, "compile_s": 0.0019387570209801197, "inner": 27, "median_ms": 1.8505145864630188, "q1_ms": 1.8414349459042705, "q3_ms": 1.8537079622210175, "tflops": 32.49341701160465 }, { "impl": "torch_math", "dtype": "float32", "head_dim": 128, "seq": 1024, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 1.1324882507324219e-06, "rel_err": 1.5304817679862026e-06, "peak_mb": 752.0009765625, "compile_s": 0.0021943680476397276, "inner": 23, "median_ms": 2.183836743073619, "q1_ms": 2.1772180258741844, "q3_ms": 2.1960930848170235, "tflops": 27.53389983693162 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 128, "seq": 1024, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 2.384185791015625e-06, "rel_err": 6.987323217799712e-07, "peak_mb": 160.0009765625, "compile_s": 0.0003359310794621706, "inner": 161, "median_ms": 0.29376835377181165, "q1_ms": 0.2872070560798697, "q3_ms": 0.2966114060999536, "tflops": 29.24050355223879 }, { "impl": "torch_math", "dtype": "float32", "head_dim": 128, "seq": 1024, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 1.2516975402832031e-06, "rel_err": 3.668344845664251e-07, "peak_mb": 452.03173828125, "compile_s": 0.0011320130433887243, "inner": 46, "median_ms": 1.0269992154739473, "q1_ms": 1.0213150956627466, "q3_ms": 1.0286602166319347, "tflops": 8.364110179028572 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 128, "seq": 1024, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 3.337860107421875e-06, "rel_err": 1.2475003359213588e-06, "peak_mb": 288.376953125, "compile_s": 0.0015868251211941242, "inner": 38, "median_ms": 1.2713744194786016, "q1_ms": 1.2636369057489854, "q3_ms": 1.2757228687405586, "tflops": 23.64745633652889 }, { "impl": "torch_math", "dtype": "float32", "head_dim": 128, "seq": 1024, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 5.0067901611328125e-06, "rel_err": 1.368881839880487e-06, "peak_mb": 752.0009765625, "compile_s": 0.0023932689800858498, "inner": 21, "median_ms": 2.3014999127813747, "q1_ms": 2.2897825055267838, "q3_ms": 2.3107516213453243, "tflops": 13.063120665369294 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 128, "seq": 1500, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 3.8743019104003906e-07, "rel_err": 1.7500067315268097e-06, "peak_mb": 204.6259765625, "compile_s": 0.0009566931985318661, "inner": 54, "median_ms": 0.9459149635707339, "q1_ms": 0.9428555171729789, "q3_ms": 0.9465463506920193, "tflops": 38.97179072085096 }, { "impl": "torch_math", "dtype": "float32", "head_dim": 128, "seq": 1500, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 4.917383193969727e-07, "rel_err": 2.2211622763279593e-06, "peak_mb": 822.65234375, "compile_s": 0.001905607059597969, "inner": 26, "median_ms": 1.9495956545982223, "q1_ms": 1.9356859624027631, "q3_ms": 1.9510737351643352, "tflops": 18.90853619469983 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 128, "seq": 1500, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 1.0728836059570312e-06, "rel_err": 2.3929278540890664e-06, "peak_mb": 396.44287109375, "compile_s": 0.0038050829898566008, "inner": 13, "median_ms": 3.951994612669716, "q1_ms": 3.9439075715983143, "q3_ms": 3.97830311423884, "tflops": 32.647817784559074 }, { "impl": "torch_math", "dtype": "float32", "head_dim": 128, "seq": 1500, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 6.258487701416016e-07, "rel_err": 1.3958745057607302e-06, "peak_mb": 1422.1337890625, "compile_s": 0.0048929378390312195, "inner": 10, "median_ms": 5.077809700742364, "q1_ms": 4.999782459344715, "q3_ms": 5.105415813159198, "tflops": 25.409380737749387 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 128, "seq": 1500, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 1.1324882507324219e-06, "rel_err": 4.1614308088355756e-07, "peak_mb": 205.7509765625, "compile_s": 0.0006167721003293991, "inner": 86, "median_ms": 0.5923821753280801, "q1_ms": 0.5725708324462175, "q3_ms": 0.5936216059422423, "tflops": 31.11504830440209 }, { "impl": "torch_math", "dtype": "float32", "head_dim": 128, "seq": 1500, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 1.1920928955078125e-06, "rel_err": 4.3804536176139663e-07, "peak_mb": 832.36083984375, "compile_s": 0.002259668195620179, "inner": 22, "median_ms": 2.261708776297217, "q1_ms": 2.2437905479984526, "q3_ms": 2.2648894968865947, "tflops": 8.149590342120069 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 128, "seq": 1500, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 4.0531158447265625e-06, "rel_err": 1.161187242360029e-06, "peak_mb": 396.44287109375, "compile_s": 0.002469097962602973, "inner": 20, "median_ms": 2.430452441330999, "q1_ms": 2.397373819258064, "q3_ms": 2.469054877292365, "tflops": 26.543206072640128 }, { "impl": "torch_math", "dtype": "float32", "head_dim": 128, "seq": 1500, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 8.344650268554688e-06, "rel_err": 3.1262734410120174e-06, "peak_mb": 1421.5712890625, "compile_s": 0.005177047802135348, "inner": 9, "median_ms": 5.312195202956597, "q1_ms": 5.234366489781275, "q3_ms": 5.322719109244645, "tflops": 12.144132046219742 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 128, "seq": 2048, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 4.917383193969727e-07, "rel_err": 2.690290102691506e-06, "peak_mb": 256.0009765625, "compile_s": 0.0016171259339898825, "inner": 31, "median_ms": 1.6876803861270027, "q1_ms": 1.6739008414949623, "q3_ms": 1.6979959177514239, "tflops": 40.71830027823092 }, { "impl": "torch_math", "dtype": "float32", "head_dim": 128, "seq": 2048, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 3.948807716369629e-07, "rel_err": 2.160384610760957e-06, "peak_mb": 1408.06298828125, "compile_s": 0.0033906218595802784, "inner": 14, "median_ms": 3.5070329405633465, "q1_ms": 3.465040287535105, "q3_ms": 3.5205080057494342, "tflops": 19.594762267890577 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 128, "seq": 2048, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 5.364418029785156e-07, "rel_err": 2.2752760742150713e-06, "peak_mb": 512.751953125, "compile_s": 0.0066629331558942795, "inner": 7, "median_ms": 6.991048576310277, "q1_ms": 6.965533647287105, "q3_ms": 7.0492681282173315, "tflops": 34.40373299522119 }, { "impl": "torch_math", "dtype": "float32", "head_dim": 128, "seq": 2048, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 7.450580596923828e-07, "rel_err": 2.6057530249090632e-06, "peak_mb": 2464.0009765625, "compile_s": 0.00838470901362598, "inner": 6, "median_ms": 8.705052003885308, "q1_ms": 8.695133593088638, "q3_ms": 8.719807005642602, "tflops": 27.62972219679446 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 128, "seq": 2048, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 1.4901161193847656e-06, "rel_err": 4.2636250441319135e-07, "peak_mb": 256.0009765625, "compile_s": 0.0009689729195088148, "inner": 53, "median_ms": 0.9800873582585241, "q1_ms": 0.9694677153777964, "q3_ms": 0.9857473020740556, "tflops": 35.057832425318054 }, { "impl": "torch_math", "dtype": "float32", "head_dim": 128, "seq": 2048, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 1.1920928955078125e-06, "rel_err": 3.410899864775274e-07, "peak_mb": 1424.06298828125, "compile_s": 0.004049513954669237, "inner": 12, "median_ms": 4.110767661283414, "q1_ms": 4.0819592056019856, "q3_ms": 4.11451599211432, "tflops": 8.358472479875601 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 128, "seq": 2048, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 4.76837158203125e-06, "rel_err": 2.3267602955456823e-06, "peak_mb": 512.751953125, "compile_s": 0.004017294151708484, "inner": 12, "median_ms": 4.129275252732138, "q1_ms": 4.109684746557226, "q3_ms": 4.187713745826235, "tflops": 29.123533048185752 }, { "impl": "torch_math", "dtype": "float32", "head_dim": 128, "seq": 2048, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 8.106231689453125e-06, "rel_err": 2.5594363251002505e-06, "peak_mb": 2464.0009765625, "compile_s": 0.008907550945878029, "inner": 5, "median_ms": 9.144032001495361, "q1_ms": 9.13836790714413, "q3_ms": 9.159589908085763, "tflops": 13.151647355163842 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 128, "seq": 3000, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 2.980232238769531e-07, "rel_err": 1.6828230400278699e-06, "peak_mb": 345.2509765625, "compile_s": 0.003641453105956316, "inner": 14, "median_ms": 3.6766542720475366, "q1_ms": 3.6451612499409487, "q3_ms": 3.7093125616333316, "tflops": 40.10602822274106 }, { "impl": "torch_math", "dtype": "float32", "head_dim": 128, "seq": 3000, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 4.023313522338867e-07, "rel_err": 2.271811126774992e-06, "peak_mb": 2817.26611328125, "compile_s": 0.006907572969794273, "inner": 7, "median_ms": 7.205962269966092, "q1_ms": 7.193533005192876, "q3_ms": 7.214637356810272, "tflops": 20.463054686614928 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 128, "seq": 3000, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 6.705522537231445e-07, "rel_err": 2.276585064464598e-06, "peak_mb": 722.1337890625, "compile_s": 0.01450857101008296, "inner": 3, "median_ms": 15.2967000224938, "q1_ms": 15.27787318142752, "q3_ms": 15.377470292150974, "tflops": 33.73904170449056 }, { "impl": "torch_math", "dtype": "float32", "head_dim": 128, "seq": 3000, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 7.748603820800781e-07, "rel_err": 3.1148110792855732e-06, "peak_mb": 4974.1572265625, "compile_s": 0.017096309922635555, "inner": 2, "median_ms": 17.938082572072744, "q1_ms": 17.864822526462376, "q3_ms": 17.944912484381348, "tflops": 28.770968018816806 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 128, "seq": 3000, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 1.1920928955078125e-06, "rel_err": 3.2444970088363334e-07, "peak_mb": 345.2509765625, "compile_s": 0.0018942670430988073, "inner": 26, "median_ms": 1.9723553067216506, "q1_ms": 1.951237003175685, "q3_ms": 1.9853480625897646, "tflops": 37.3806888387402 }, { "impl": "torch_math", "dtype": "float32", "head_dim": 128, "seq": 3000, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 1.430511474609375e-06, "rel_err": 3.893396467447019e-07, "peak_mb": 2851.5986328125, "compile_s": 0.008582910057157278, "inner": 5, "median_ms": 8.817460807040334, "q1_ms": 8.738191495649517, "q3_ms": 8.838045783340931, "tflops": 8.361590894867557 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 128, "seq": 3000, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 5.245208740234375e-06, "rel_err": 1.8902719602920115e-06, "peak_mb": 722.1337890625, "compile_s": 0.008051138138398528, "inner": 6, "median_ms": 8.352162355246643, "q1_ms": 8.33516069299852, "q3_ms": 8.367331756744534, "tflops": 30.895951135085394 }, { "impl": "torch_math", "dtype": "float32", "head_dim": 128, "seq": 3000, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 8.106231689453125e-06, "rel_err": 3.6137550978310173e-06, "peak_mb": 4974.1572265625, "compile_s": 0.01854202407412231, "inner": 2, "median_ms": 19.152006483636796, "q1_ms": 19.080254016444087, "q3_ms": 19.22895445022732, "tflops": 13.473679649203993 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 128, "seq": 4096, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 4.3213367462158203e-07, "rel_err": 2.717233428484178e-06, "peak_mb": 448.0009765625, "compile_s": 0.006133242044597864, "inner": 8, "median_ms": 6.603000481845811, "q1_ms": 6.555150946951471, "q3_ms": 6.658499492914416, "tflops": 41.6292422967021 }, { "impl": "torch_math", "dtype": "float32", "head_dim": 128, "seq": 4096, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 9.313225746154785e-07, "rel_err": 5.856106781720882e-06, "peak_mb": 5056.12548828125, "compile_s": 0.012404982931911945, "inner": 4, "median_ms": 12.972047785297036, "q1_ms": 12.899474997539073, "q3_ms": 13.077626877930015, "tflops": 21.190016525806826 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 128, "seq": 4096, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 8.642673492431641e-07, "rel_err": 3.2987438771669986e-06, "peak_mb": 961.501953125, "compile_s": 0.02545953798107803, "inner": 1, "median_ms": 28.199238004162908, "q1_ms": 27.85786206368357, "q3_ms": 28.277418576180935, "tflops": 34.116974159442684 }, { "impl": "torch_math", "dtype": "float32", "head_dim": 128, "seq": 4096, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 1.6689300537109375e-06, "rel_err": 6.3699881138745695e-06, "peak_mb": 8960.0009765625, "compile_s": 0.03075951710343361, "inner": 1, "median_ms": 32.2564619127661, "q1_ms": 32.12913160677999, "q3_ms": 32.28197793941945, "tflops": 29.825734666927055 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 128, "seq": 4096, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 1.9073486328125e-06, "rel_err": 6.873864322187728e-07, "peak_mb": 448.0009765625, "compile_s": 0.0033186920918524265, "inner": 15, "median_ms": 3.5456183987359204, "q1_ms": 3.5354570019990206, "q3_ms": 3.5929498340313635, "tflops": 38.7630415955083 }, { "impl": "torch_math", "dtype": "float32", "head_dim": 128, "seq": 4096, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 1.1920928955078125e-06, "rel_err": 4.2961653434758773e-07, "peak_mb": 5120.12548828125, "compile_s": 0.015518823871389031, "inner": 3, "median_ms": 15.868368713806074, "q1_ms": 15.804218361154199, "q3_ms": 15.904895340402922, "tflops": 8.661189814200812 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 128, "seq": 4096, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 1.049041748046875e-05, "rel_err": 2.534422037570039e-06, "peak_mb": 961.501953125, "compile_s": 0.013820677995681763, "inner": 3, "median_ms": 14.587153991063436, "q1_ms": 14.490672193157177, "q3_ms": 14.610684321572384, "tflops": 32.976709332519455 }, { "impl": "torch_math", "dtype": "float32", "head_dim": 128, "seq": 4096, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 9.059906005859375e-06, "rel_err": 3.395472504053032e-06, "peak_mb": 8960.0009765625, "compile_s": 0.03335608611814678, "inner": 1, "median_ms": 34.38678989186883, "q1_ms": 34.31901440490037, "q3_ms": 34.55825534183532, "tflops": 13.988986429516842 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 128, "seq": 8192, "causal": false, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 4.3213367462158203e-07, "rel_err": 4.1143052840197925e-06, "peak_mb": 832.0009765625, "compile_s": 0.024457076098769903, "inner": 1, "median_ms": 26.495682075619698, "q1_ms": 26.008145534433424, "q3_ms": 26.889564003795385, "tflops": 41.49776649032666 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 128, "seq": 8192, "causal": false, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 9.5367431640625e-07, "rel_err": 5.241354301688261e-06, "peak_mb": 1859.001953125, "compile_s": 0.10603432892821729, "inner": 1, "median_ms": 109.30961091071367, "q1_ms": 108.35058242082596, "q3_ms": 119.50461042579263, "tflops": 35.20541940597852 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 128, "seq": 8192, "causal": true, "mode": "fwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 1.430511474609375e-06, "rel_err": 4.771075623466459e-07, "peak_mb": 832.0009765625, "compile_s": 0.012361462926492095, "inner": 3, "median_ms": 13.414806686341763, "q1_ms": 13.363354839384556, "q3_ms": 13.429690307627123, "tflops": 40.98126993121205 }, { "impl": "torch_efficient", "dtype": "float32", "head_dim": 128, "seq": 8192, "causal": true, "mode": "fwd_bwd", "batch": 4, "heads": 8, "tag": "h100_nvl", "max_abs_err": 1.4066696166992188e-05, "rel_err": 4.82915083921398e-06, "peak_mb": 1859.001953125, "compile_s": 0.05496867117471993, "inner": 1, "median_ms": 54.8763619735837, "q1_ms": 54.627465549856424, "q3_ms": 54.950681631453335, "tflops": 35.06328188326774 } ] }