{
  "device": {
    "jax": "0.11.2",
    "backend": "gpu",
    "devices": [
      "cuda:0"
    ],
    "python": "3.14.8",
    "device_kind": "NVIDIA GeForce RTX 3090",
    "gpu": "NVIDIA GeForce RTX 3090",
    "driver": "615.78.08",
    "gpu_memory_gb": 24.0,
    "torch": "2.14.1+cu130",
    "torch_cuda": "13.0"
  },
  "config": {
    "output": "docs/benchmarks/gpu_nn.json",
    "suites": [
      "transformer",
      "diffusion",
      "moe",
      "ssm",
      "unet"
    ],
    "dtypes": [
      "float32",
      "bfloat16"
    ],
    "tokens": [
      256,
      1024,
      4096
    ],
    "moe_tokens": [
      4096,
      32768
    ],
    "ssm_tokens": [
      1024,
      4096
    ],
    "batch": 8,
    "model_dim": 256,
    "ssm_rerun": "ssm suite re-run after the SSDCell rework (Mamba-2 decay, real heads) with Mamba2Cell and GatedDeltaNetCell added; includes the mamba tile fix of 9386bfb",
    "attention_rerun": "transformer and diffusion suites re-run with the single-pass attention backward defaults"
  },
  "results": [
    {
      "suite": "moe",
      "impl": "pallas",
      "dtype": "float32",
      "tokens": 4096,
      "experts": 8,
      "top_k": 2,
      "d_model": 512,
      "d_hidden": 1024,
      "op": "forward",
      "params": 12587008,
      "finite": true,
      "compile_s": 0.6929407709976658,
      "inner": 34,
      "median_ms": 1.2867375878747278,
      "q1_ms": 1.2792476762329257,
      "q3_ms": 1.2877615588556237
    },
    {
      "suite": "moe",
      "impl": "pallas",
      "dtype": "float32",
      "tokens": 4096,
      "experts": 8,
      "top_k": 2,
      "d_model": 512,
      "d_hidden": 1024,
      "op": "train_step",
      "params": 12587008,
      "finite": true,
      "compile_s": 1.3518774940021103,
      "inner": 10,
      "median_ms": 4.636383200704586,
      "q1_ms": 4.629410550114699,
      "q3_ms": 4.6569676997023635
    },
    {
      "suite": "moe",
      "impl": "blocked",
      "dtype": "float32",
      "tokens": 4096,
      "experts": 8,
      "top_k": 2,
      "d_model": 512,
      "d_hidden": 1024,
      "op": "forward",
      "params": 12587008,
      "finite": true,
      "compile_s": 1.0746354180009803,
      "inner": 19,
      "median_ms": 2.3680073163153508,
      "q1_ms": 2.362134894371449,
      "q3_ms": 2.369730526544644
    },
    {
      "suite": "moe",
      "impl": "blocked",
      "dtype": "float32",
      "tokens": 4096,
      "experts": 8,
      "top_k": 2,
      "d_model": 512,
      "d_hidden": 1024,
      "op": "train_step",
      "params": 12587008,
      "finite": true,
      "compile_s": 2.0710619279998355,
      "inner": 8,
      "median_ms": 5.801188750410802,
      "q1_ms": 5.774773062512395,
      "q3_ms": 5.803139187264605
    },
    {
      "suite": "moe",
      "impl": "ragged_dot",
      "dtype": "float32",
      "tokens": 4096,
      "experts": 8,
      "top_k": 2,
      "d_model": 512,
      "d_hidden": 1024,
      "op": "forward",
      "params": 12587008,
      "finite": true,
      "compile_s": 1.683056295994902,
      "inner": 7,
      "median_ms": 6.6610715709560155,
      "q1_ms": 6.657143427998692,
      "q3_ms": 6.684189642588276
    },
    {
      "suite": "moe",
      "impl": "ragged_dot",
      "dtype": "float32",
      "tokens": 4096,
      "experts": 8,
      "top_k": 2,
      "d_model": 512,
      "d_hidden": 1024,
      "op": "train_step",
      "params": 12587008,
      "finite": true,
      "compile_s": 1.6449657390039647,
      "inner": 2,
      "median_ms": 19.414862494159024,
      "q1_ms": 19.33579774777172,
      "q3_ms": 19.42831274936907
    },
    {
      "suite": "moe",
      "impl": "pallas",
      "dtype": "float32",
      "tokens": 32768,
      "experts": 8,
      "top_k": 2,
      "d_model": 512,
      "d_hidden": 1024,
      "op": "forward",
      "params": 12587008,
      "finite": true,
      "compile_s": 0.7317303099989658,
      "inner": 5,
      "median_ms": 9.293828398222104,
      "q1_ms": 9.274444698530715,
      "q3_ms": 9.323246800340712
    },
    {
      "suite": "moe",
      "impl": "pallas",
      "dtype": "float32",
      "tokens": 32768,
      "experts": 8,
      "top_k": 2,
      "d_model": 512,
      "d_hidden": 1024,
      "op": "train_step",
      "params": 12587008,
      "finite": true,
      "compile_s": 0.9519799150002655,
      "inner": 1,
      "median_ms": 31.90415899734944,
      "q1_ms": 31.881849499768578,
      "q3_ms": 32.014468997658696
    },
    {
      "suite": "moe",
      "impl": "blocked",
      "dtype": "float32",
      "tokens": 32768,
      "experts": 8,
      "top_k": 2,
      "d_model": 512,
      "d_hidden": 1024,
      "op": "forward",
      "params": 12587008,
      "finite": true,
      "compile_s": 1.5316107539983932,
      "inner": 3,
      "median_ms": 15.649532666429877,
      "q1_ms": 15.626277000895547,
      "q3_ms": 15.704024333293395
    },
    {
      "suite": "moe",
      "impl": "blocked",
      "dtype": "float32",
      "tokens": 32768,
      "experts": 8,
      "top_k": 2,
      "d_model": 512,
      "d_hidden": 1024,
      "op": "train_step",
      "params": 12587008,
      "finite": true,
      "compile_s": 2.072070813999744,
      "inner": 1,
      "median_ms": 36.4371899922844,
      "q1_ms": 36.39235400623875,
      "q3_ms": 36.688644999230746
    },
    {
      "suite": "moe",
      "impl": "ragged_dot",
      "dtype": "float32",
      "tokens": 32768,
      "experts": 8,
      "top_k": 2,
      "d_model": 512,
      "d_hidden": 1024,
      "op": "forward",
      "params": 12587008,
      "finite": true,
      "compile_s": 3.564945031001116,
      "inner": 1,
      "median_ms": 48.87978000624571,
      "q1_ms": 48.793999005283695,
      "q3_ms": 48.911796999163926
    },
    {
      "suite": "moe",
      "impl": "ragged_dot",
      "dtype": "float32",
      "tokens": 32768,
      "experts": 8,
      "top_k": 2,
      "d_model": 512,
      "d_hidden": 1024,
      "op": "train_step",
      "params": 12587008,
      "finite": true,
      "compile_s": 5.213303296986851,
      "inner": 1,
      "median_ms": 146.39188699948136,
      "q1_ms": 146.05545149970567,
      "q3_ms": 146.5268894971814
    },
    {
      "suite": "moe",
      "impl": "pallas",
      "dtype": "bfloat16",
      "tokens": 4096,
      "experts": 8,
      "top_k": 2,
      "d_model": 512,
      "d_hidden": 1024,
      "op": "forward",
      "params": 12587008,
      "finite": true,
      "compile_s": 0.5255591810127953,
      "inner": 50,
      "median_ms": 0.8565408797585405,
      "q1_ms": 0.8471217699116096,
      "q3_ms": 0.8688591799000278
    },
    {
      "suite": "moe",
      "impl": "pallas",
      "dtype": "bfloat16",
      "tokens": 4096,
      "experts": 8,
      "top_k": 2,
      "d_model": 512,
      "d_hidden": 1024,
      "op": "train_step",
      "params": 12587008,
      "finite": true,
      "compile_s": 1.9471824780048337,
      "inner": 15,
      "median_ms": 2.7281380675655478,
      "q1_ms": 2.7258618002330577,
      "q3_ms": 2.7314480001223274
    },
    {
      "suite": "moe",
      "impl": "blocked",
      "dtype": "bfloat16",
      "tokens": 4096,
      "experts": 8,
      "top_k": 2,
      "d_model": 512,
      "d_hidden": 1024,
      "op": "forward",
      "params": 12587008,
      "finite": true,
      "compile_s": 0.317491265013814,
      "inner": 20,
      "median_ms": 2.3470161504519638,
      "q1_ms": 2.330037549836561,
      "q3_ms": 2.3481748499762034
    },
    {
      "suite": "moe",
      "impl": "blocked",
      "dtype": "bfloat16",
      "tokens": 4096,
      "experts": 8,
      "top_k": 2,
      "d_model": 512,
      "d_hidden": 1024,
      "op": "train_step",
      "params": 12587008,
      "finite": true,
      "compile_s": 0.7045912940084236,
      "inner": 8,
      "median_ms": 5.689258374331985,
      "q1_ms": 5.681336250745517,
      "q3_ms": 5.692750874914054
    },
    {
      "suite": "moe",
      "impl": "ragged_dot",
      "dtype": "bfloat16",
      "tokens": 4096,
      "experts": 8,
      "top_k": 2,
      "d_model": 512,
      "d_hidden": 1024,
      "op": "forward",
      "params": 12587008,
      "finite": true,
      "compile_s": 0.9072506819939008,
      "inner": 7,
      "median_ms": 6.691526284807229,
      "q1_ms": 6.654023999414806,
      "q3_ms": 6.700823928570442
    },
    {
      "suite": "moe",
      "impl": "ragged_dot",
      "dtype": "bfloat16",
      "tokens": 4096,
      "experts": 8,
      "top_k": 2,
      "d_model": 512,
      "d_hidden": 1024,
      "op": "train_step",
      "params": 12587008,
      "finite": true,
      "compile_s": 0.38875615000142716,
      "inner": 2,
      "median_ms": 19.247461001214106,
      "q1_ms": 19.143010496918578,
      "q3_ms": 19.459742248727707
    },
    {
      "suite": "moe",
      "impl": "pallas",
      "dtype": "bfloat16",
      "tokens": 32768,
      "experts": 8,
      "top_k": 2,
      "d_model": 512,
      "d_hidden": 1024,
      "op": "forward",
      "params": 12587008,
      "finite": true,
      "compile_s": 0.5342372029990656,
      "inner": 9,
      "median_ms": 5.064477890199568,
      "q1_ms": 5.008436833223742,
      "q3_ms": 5.073005055439555
    },
    {
      "suite": "moe",
      "impl": "pallas",
      "dtype": "bfloat16",
      "tokens": 32768,
      "experts": 8,
      "top_k": 2,
      "d_model": 512,
      "d_hidden": 1024,
      "op": "train_step",
      "params": 12587008,
      "finite": true,
      "compile_s": 0.7440002350049326,
      "inner": 3,
      "median_ms": 15.890558667403337,
      "q1_ms": 15.851311164927512,
      "q3_ms": 15.929411834804341
    },
    {
      "suite": "moe",
      "impl": "blocked",
      "dtype": "bfloat16",
      "tokens": 32768,
      "experts": 8,
      "top_k": 2,
      "d_model": 512,
      "d_hidden": 1024,
      "op": "forward",
      "params": 12587008,
      "finite": true,
      "compile_s": 0.4564471269986825,
      "inner": 3,
      "median_ms": 15.603857003346397,
      "q1_ms": 15.560789999047604,
      "q3_ms": 15.621281165901262
    },
    {
      "suite": "moe",
      "impl": "blocked",
      "dtype": "bfloat16",
      "tokens": 32768,
      "experts": 8,
      "top_k": 2,
      "d_model": 512,
      "d_hidden": 1024,
      "op": "train_step",
      "params": 12587008,
      "finite": true,
      "compile_s": 0.6418208740069531,
      "inner": 1,
      "median_ms": 36.39911899517756,
      "q1_ms": 36.17615850089351,
      "q3_ms": 36.46825099713169
    },
    {
      "suite": "moe",
      "impl": "ragged_dot",
      "dtype": "bfloat16",
      "tokens": 32768,
      "experts": 8,
      "top_k": 2,
      "d_model": 512,
      "d_hidden": 1024,
      "op": "forward",
      "params": 12587008,
      "finite": true,
      "compile_s": 2.306932612991659,
      "inner": 1,
      "median_ms": 48.80831499758642,
      "q1_ms": 48.715516000811476,
      "q3_ms": 49.11557849845849
    },
    {
      "suite": "moe",
      "impl": "ragged_dot",
      "dtype": "bfloat16",
      "tokens": 32768,
      "experts": 8,
      "top_k": 2,
      "d_model": 512,
      "d_hidden": 1024,
      "op": "train_step",
      "params": 12587008,
      "finite": true,
      "compile_s": 0.6062190339871449,
      "inner": 1,
      "median_ms": 144.37941899814177,
      "q1_ms": 144.17565800249577,
      "q3_ms": 144.4914255043841
    },
    {
      "suite": "unet",
      "dtype": "float32",
      "resolution": 32,
      "batch": 8,
      "op": "forward",
      "params": 2803968,
      "finite": true,
      "compile_s": 1.5739056510064984,
      "inner": 26,
      "median_ms": 1.6434738846818129,
      "q1_ms": 1.643215327091569,
      "q3_ms": 1.645635250147629
    },
    {
      "suite": "unet",
      "dtype": "float32",
      "resolution": 32,
      "batch": 8,
      "op": "train_step",
      "params": 2803968,
      "finite": true,
      "compile_s": 3.872681888999068,
      "inner": 7,
      "median_ms": 5.414999426907992,
      "q1_ms": 5.407580286763342,
      "q3_ms": 5.420358714348238
    },
    {
      "suite": "unet",
      "dtype": "float32",
      "resolution": 64,
      "batch": 8,
      "op": "forward",
      "params": 2803968,
      "finite": true,
      "compile_s": 1.4467379430134315,
      "inner": 10,
      "median_ms": 4.701456300972495,
      "q1_ms": 4.6875168998667505,
      "q3_ms": 4.70342004991835
    },
    {
      "suite": "unet",
      "dtype": "float32",
      "resolution": 64,
      "batch": 8,
      "op": "train_step",
      "params": 2803968,
      "finite": true,
      "compile_s": 4.302772442009882,
      "inner": 3,
      "median_ms": 15.298468667121293,
      "q1_ms": 15.286026832958061,
      "q3_ms": 15.315878668237321
    },
    {
      "suite": "unet",
      "dtype": "bfloat16",
      "resolution": 32,
      "batch": 8,
      "op": "forward",
      "params": 2803968,
      "finite": true,
      "compile_s": 1.3643167490081396,
      "inner": 33,
      "median_ms": 1.1525202121860771,
      "q1_ms": 1.146889757585472,
      "q3_ms": 1.1537830758959788
    },
    {
      "suite": "unet",
      "dtype": "bfloat16",
      "resolution": 32,
      "batch": 8,
      "op": "train_step",
      "params": 2803968,
      "finite": true,
      "compile_s": 11.503543283004547,
      "inner": 9,
      "median_ms": 3.969393555760487,
      "q1_ms": 3.967053055627427,
      "q3_ms": 3.994162000405291
    },
    {
      "suite": "unet",
      "dtype": "bfloat16",
      "resolution": 64,
      "batch": 8,
      "op": "forward",
      "params": 2803968,
      "finite": true,
      "compile_s": 1.2621258269937243,
      "inner": 15,
      "median_ms": 2.9126043334448086,
      "q1_ms": 2.9011519669438712,
      "q3_ms": 2.914901299906584
    },
    {
      "suite": "unet",
      "dtype": "bfloat16",
      "resolution": 64,
      "batch": 8,
      "op": "train_step",
      "params": 2803968,
      "finite": true,
      "compile_s": 13.316741700997227,
      "inner": 4,
      "median_ms": 10.122343250259291,
      "q1_ms": 9.585231126038707,
      "q3_ms": 10.277462748490507
    },
    {
      "suite": "ssm",
      "cell": "mamba",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 128,
      "op": "forward",
      "params": 930704,
      "finite": true,
      "compile_s": 2.535260616001324,
      "inner": 9,
      "median_ms": 5.090220333335714,
      "q1_ms": 5.083352500352905,
      "q3_ms": 5.099769777088012
    },
    {
      "suite": "ssm",
      "cell": "mamba",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 128,
      "op": "train_step",
      "params": 930704,
      "finite": true,
      "compile_s": 6.378781646999414,
      "inner": 1,
      "median_ms": 37.17908399994485,
      "q1_ms": 36.2261460031732,
      "q3_ms": 37.60828750091605
    },
    {
      "suite": "ssm",
      "cell": "ssd",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 128,
      "op": "forward",
      "params": 863644,
      "finite": true,
      "compile_s": 4.928017325990368,
      "inner": 22,
      "median_ms": 1.8967950911081228,
      "q1_ms": 1.8944996134616137,
      "q3_ms": 1.9502686138201335
    },
    {
      "suite": "ssm",
      "cell": "ssd",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 128,
      "op": "train_step",
      "params": 863644,
      "finite": true,
      "compile_s": 10.403776341991033,
      "inner": 5,
      "median_ms": 7.167033400037326,
      "q1_ms": 7.12944340048125,
      "q3_ms": 7.183451599848922
    },
    {
      "suite": "ssm",
      "cell": "mamba2",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 128,
      "op": "forward",
      "params": 1136064,
      "finite": true,
      "compile_s": 3.1587284540000837,
      "inner": 14,
      "median_ms": 3.119161071871141,
      "q1_ms": 3.1146600714626924,
      "q3_ms": 3.1197063212208116
    },
    {
      "suite": "ssm",
      "cell": "mamba2",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 128,
      "op": "train_step",
      "params": 1136064,
      "finite": true,
      "compile_s": 11.34031775299809,
      "inner": 3,
      "median_ms": 12.05367100192234,
      "q1_ms": 11.944654666876886,
      "q3_ms": 12.06799150046815
    },
    {
      "suite": "ssm",
      "cell": "gdn",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 128,
      "op": "forward",
      "params": 1006144,
      "finite": true,
      "compile_s": 4.350665006990312,
      "inner": 15,
      "median_ms": 2.6240658674699566,
      "q1_ms": 2.5553386997974785,
      "q3_ms": 2.6303016003415296
    },
    {
      "suite": "ssm",
      "cell": "gdn",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 128,
      "op": "train_step",
      "params": 1006144,
      "finite": true,
      "compile_s": 8.332819101007772,
      "inner": 3,
      "median_ms": 11.888078001599448,
      "q1_ms": 11.72045650067351,
      "q3_ms": 12.118783832799334
    },
    {
      "suite": "ssm",
      "cell": "lru",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 128,
      "op": "forward",
      "params": 995728,
      "finite": true,
      "compile_s": 1.4752462699980242,
      "inner": 25,
      "median_ms": 1.6149592801230028,
      "q1_ms": 1.59560669970233,
      "q3_ms": 1.615880580211524
    },
    {
      "suite": "ssm",
      "cell": "lru",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 128,
      "op": "train_step",
      "params": 995728,
      "finite": true,
      "compile_s": 5.073489720001817,
      "inner": 7,
      "median_ms": 5.433400142854745,
      "q1_ms": 5.430114570897006,
      "q3_ms": 5.46140378511544
    },
    {
      "suite": "ssm",
      "cell": "mamba",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 128,
      "op": "forward",
      "params": 930704,
      "finite": true,
      "compile_s": 2.584006473000045,
      "inner": 2,
      "median_ms": 18.733760000031907,
      "q1_ms": 18.69672550310497,
      "q3_ms": 18.76008749968605
    },
    {
      "suite": "ssm",
      "cell": "mamba",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 128,
      "op": "train_step",
      "params": 930704,
      "finite": true,
      "compile_s": 6.294068172996049,
      "inner": 1,
      "median_ms": 146.9874169997638,
      "q1_ms": 146.92009500140557,
      "q3_ms": 147.36518050631275
    },
    {
      "suite": "ssm",
      "cell": "ssd",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 128,
      "op": "forward",
      "params": 863644,
      "finite": true,
      "compile_s": 4.985784574993886,
      "inner": 6,
      "median_ms": 7.178327333046279,
      "q1_ms": 7.176713501394261,
      "q3_ms": 7.184441500915758
    },
    {
      "suite": "ssm",
      "cell": "ssd",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 128,
      "op": "train_step",
      "params": 863644,
      "finite": true,
      "compile_s": 10.849746813997626,
      "inner": 1,
      "median_ms": 25.15106400824152,
      "q1_ms": 24.954316490038764,
      "q3_ms": 25.31041899783304
    },
    {
      "suite": "ssm",
      "cell": "mamba2",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 128,
      "op": "forward",
      "params": 1136064,
      "finite": true,
      "compile_s": 2.5223453880025772,
      "inner": 4,
      "median_ms": 11.678336752083851,
      "q1_ms": 11.67504899785854,
      "q3_ms": 11.683306374834501
    },
    {
      "suite": "ssm",
      "cell": "mamba2",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 128,
      "op": "train_step",
      "params": 1136064,
      "finite": true,
      "compile_s": 10.097311682999134,
      "inner": 1,
      "median_ms": 44.587811004021205,
      "q1_ms": 44.463332000304945,
      "q3_ms": 44.85170800035121
    },
    {
      "suite": "ssm",
      "cell": "gdn",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 128,
      "op": "forward",
      "params": 1006144,
      "finite": true,
      "compile_s": 4.27845235800487,
      "inner": 4,
      "median_ms": 10.15257949984516,
      "q1_ms": 10.064585625514155,
      "q3_ms": 10.210838376224274
    },
    {
      "suite": "ssm",
      "cell": "gdn",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 128,
      "op": "train_step",
      "params": 1006144,
      "finite": true,
      "compile_s": 8.486224993990618,
      "inner": 1,
      "median_ms": 41.05898599664215,
      "q1_ms": 40.60319449490635,
      "q3_ms": 41.55566850386094
    },
    {
      "suite": "ssm",
      "cell": "lru",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 128,
      "op": "forward",
      "params": 995728,
      "finite": true,
      "compile_s": 1.4345465529913781,
      "inner": 7,
      "median_ms": 6.198856000472525,
      "q1_ms": 6.188342643976544,
      "q3_ms": 6.2006215718741124
    },
    {
      "suite": "ssm",
      "cell": "lru",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 128,
      "op": "train_step",
      "params": 995728,
      "finite": true,
      "compile_s": 5.5283588339952985,
      "inner": 2,
      "median_ms": 18.328359998122323,
      "q1_ms": 18.293581255420577,
      "q3_ms": 18.404120000923285
    },
    {
      "suite": "transformer",
      "attention": "probjax_default",
      "dtype": "float32",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "compile_s": 3.0159296169877052,
      "inner": 43,
      "median_ms": 0.8751056745897459,
      "q1_ms": 0.8726496163817924,
      "q3_ms": 0.9127092093559589
    },
    {
      "suite": "transformer",
      "attention": "probjax_default",
      "dtype": "float32",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "compile_s": 4.341343791005784,
      "inner": 10,
      "median_ms": 3.51241300086258,
      "q1_ms": 3.4612381001352333,
      "q3_ms": 3.6860066509689204
    },
    {
      "suite": "transformer",
      "attention": "probjax_flex",
      "dtype": "float32",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "compile_s": 0.44885639699350577,
      "inner": 42,
      "median_ms": 0.9144670714574333,
      "q1_ms": 0.9142493928222739,
      "q3_ms": 0.9146390714574539
    },
    {
      "suite": "transformer",
      "attention": "probjax_flex",
      "dtype": "float32",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "compile_s": 1.3350546970032156,
      "inner": 11,
      "median_ms": 3.452447454699061,
      "q1_ms": 3.4410080453116363,
      "q3_ms": 3.457577590862374
    },
    {
      "suite": "transformer",
      "attention": "jax_xla",
      "dtype": "float32",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "compile_s": 1.7058044050063472,
      "inner": 35,
      "median_ms": 1.1813117428183821,
      "q1_ms": 1.180469428696337,
      "q3_ms": 1.18182672844601
    },
    {
      "suite": "transformer",
      "attention": "jax_xla",
      "dtype": "float32",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "compile_s": 4.152636011000141,
      "inner": 9,
      "median_ms": 4.310298444276365,
      "q1_ms": 4.289107222575694,
      "q3_ms": 4.315235611607527
    },
    {
      "suite": "transformer",
      "attention": "probjax_default",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "compile_s": 3.1275674009957584,
      "inner": 11,
      "median_ms": 4.180793909736994,
      "q1_ms": 4.155784726952499,
      "q3_ms": 4.197865272792775
    },
    {
      "suite": "transformer",
      "attention": "probjax_default",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "compile_s": 4.8805038330028765,
      "inner": 3,
      "median_ms": 15.344121668022126,
      "q1_ms": 15.328625665764168,
      "q3_ms": 15.393496832984965
    },
    {
      "suite": "transformer",
      "attention": "probjax_flex",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "compile_s": 0.45220315401093103,
      "inner": 11,
      "median_ms": 4.17148809091569,
      "q1_ms": 4.138330818519039,
      "q3_ms": 4.1822679545889105
    },
    {
      "suite": "transformer",
      "attention": "probjax_flex",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "compile_s": 1.424535525002284,
      "inner": 3,
      "median_ms": 15.16980500309728,
      "q1_ms": 15.165695501006365,
      "q3_ms": 15.202256334305275
    },
    {
      "suite": "transformer",
      "attention": "jax_xla",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "compile_s": 1.7812377130030654,
      "inner": 5,
      "median_ms": 8.193036000011489,
      "q1_ms": 8.190126599220093,
      "q3_ms": 8.1991614002618
    },
    {
      "suite": "transformer",
      "attention": "jax_xla",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "compile_s": 4.201021461994969,
      "inner": 2,
      "median_ms": 24.403920004260726,
      "q1_ms": 24.266869251732714,
      "q3_ms": 24.490942250849912
    },
    {
      "suite": "transformer",
      "attention": "probjax_default",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "compile_s": 2.494445395001094,
      "inner": 1,
      "median_ms": 28.234352997969836,
      "q1_ms": 28.22843250760343,
      "q3_ms": 28.248689995962195
    },
    {
      "suite": "transformer",
      "attention": "probjax_default",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "compile_s": 4.856268552000984,
      "inner": 1,
      "median_ms": 126.47872600064147,
      "q1_ms": 126.02909900306258,
      "q3_ms": 126.97880700579844
    },
    {
      "suite": "transformer",
      "attention": "probjax_flex",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "compile_s": 0.47537216599448584,
      "inner": 1,
      "median_ms": 28.766800998710096,
      "q1_ms": 28.630943495954853,
      "q3_ms": 28.79478449904127
    },
    {
      "suite": "transformer",
      "attention": "probjax_flex",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "compile_s": 1.584059853994404,
      "inner": 1,
      "median_ms": 126.88366099609993,
      "q1_ms": 126.34663549397374,
      "q3_ms": 126.97694400412729
    },
    {
      "suite": "transformer",
      "attention": "jax_xla",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "compile_s": 2.3306790240021655,
      "inner": 1,
      "median_ms": 92.29319299629424,
      "q1_ms": 92.27712349820649,
      "q3_ms": 92.36131700163241
    },
    {
      "suite": "transformer",
      "attention": "jax_xla",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 24.72GiB with allocator GPU_0_bfc on device 0. [executable_name='jit_train_step'] [tf-allocator-allocation-error='']"
    },
    {
      "suite": "transformer",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "compile_s": 2.009445103001781,
      "inner": 64,
      "median_ms": 0.5326239217993134,
      "q1_ms": 0.5224169843813797,
      "q3_ms": 0.5329380468310774
    },
    {
      "suite": "transformer",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "compile_s": 3.7366857799934223,
      "inner": 14,
      "median_ms": 2.6996751423991685,
      "q1_ms": 2.674502142846385,
      "q3_ms": 2.704945785808377
    },
    {
      "suite": "transformer",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "compile_s": 0.4181942879949929,
      "inner": 61,
      "median_ms": 0.5074882293959744,
      "q1_ms": 0.4972154261462284,
      "q3_ms": 0.532521942599875
    },
    {
      "suite": "transformer",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "compile_s": 1.3166419720073463,
      "inner": 15,
      "median_ms": 2.4225523996089278,
      "q1_ms": 2.421006500177706,
      "q3_ms": 2.4294898665781757
    },
    {
      "suite": "transformer",
      "attention": "jax_xla",
      "dtype": "bfloat16",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "compile_s": 1.2389919249981176,
      "inner": 50,
      "median_ms": 0.7476612398750149,
      "q1_ms": 0.7464539498323575,
      "q3_ms": 0.7479478800087236
    },
    {
      "suite": "transformer",
      "attention": "jax_xla",
      "dtype": "bfloat16",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "compile_s": 3.4133094779972453,
      "inner": 11,
      "median_ms": 3.2769840909168124,
      "q1_ms": 3.2132604544379073,
      "q3_ms": 3.2886434088471685
    },
    {
      "suite": "transformer",
      "attention": "jax_cudnn",
      "dtype": "bfloat16",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "compile_s": 0.4345828739897115,
      "inner": 57,
      "median_ms": 0.5820050001810223,
      "q1_ms": 0.5819007456649846,
      "q3_ms": 0.5823227718856447
    },
    {
      "suite": "transformer",
      "attention": "jax_cudnn",
      "dtype": "bfloat16",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "compile_s": 1.3215085629926762,
      "inner": 13,
      "median_ms": 2.6641109994111156,
      "q1_ms": 2.65823100030852,
      "q3_ms": 2.6713886153839455
    },
    {
      "suite": "transformer",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "compile_s": 1.8504378710058518,
      "inner": 19,
      "median_ms": 2.2013078943066495,
      "q1_ms": 2.19515655235706,
      "q3_ms": 2.2452947370383254
    },
    {
      "suite": "transformer",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "compile_s": 3.7282698479975807,
      "inner": 6,
      "median_ms": 7.147405834984966,
      "q1_ms": 7.131010333978338,
      "q3_ms": 7.186735000383729
    },
    {
      "suite": "transformer",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "compile_s": 0.43046123199746944,
      "inner": 18,
      "median_ms": 2.1941402778288146,
      "q1_ms": 2.18770441642846,
      "q3_ms": 2.245436111277539
    },
    {
      "suite": "transformer",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "compile_s": 1.4229581670078915,
      "inner": 6,
      "median_ms": 7.187043665908277,
      "q1_ms": 7.178570333053358,
      "q3_ms": 7.214774583796194
    },
    {
      "suite": "transformer",
      "attention": "jax_xla",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "compile_s": 1.2383655610028654,
      "inner": 8,
      "median_ms": 5.584799500866211,
      "q1_ms": 5.569576936977683,
      "q3_ms": 5.600903625236242
    },
    {
      "suite": "transformer",
      "attention": "jax_xla",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "compile_s": 3.461723376996815,
      "inner": 2,
      "median_ms": 17.697276001854334,
      "q1_ms": 17.694977253995603,
      "q3_ms": 17.764257499948144
    },
    {
      "suite": "transformer",
      "attention": "jax_cudnn",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "compile_s": 0.43187422401388176,
      "inner": 15,
      "median_ms": 2.945366933514985,
      "q1_ms": 2.9443727335698595,
      "q3_ms": 2.946316366433166
    },
    {
      "suite": "transformer",
      "attention": "jax_cudnn",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "compile_s": 1.342600832998869,
      "inner": 4,
      "median_ms": 9.719177000079071,
      "q1_ms": 9.706190750875976,
      "q3_ms": 9.726266249344917
    },
    {
      "suite": "transformer",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "compile_s": 1.8561287310003536,
      "inner": 3,
      "median_ms": 14.80953766925571,
      "q1_ms": 14.793742998638965,
      "q3_ms": 14.888114666973706
    },
    {
      "suite": "transformer",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "compile_s": 3.856487154000206,
      "inner": 1,
      "median_ms": 50.052847000188194,
      "q1_ms": 49.96063299040543,
      "q3_ms": 50.634581995836925
    },
    {
      "suite": "transformer",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "compile_s": 0.32275911400211044,
      "inner": 3,
      "median_ms": 14.80052966993147,
      "q1_ms": 14.766539832635317,
      "q3_ms": 14.872460332602106
    },
    {
      "suite": "transformer",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "compile_s": 1.4444008529972052,
      "inner": 1,
      "median_ms": 50.2538479922805,
      "q1_ms": 50.13236750528449,
      "q3_ms": 51.304934997460805
    },
    {
      "suite": "transformer",
      "attention": "jax_xla",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "compile_s": 1.7992907450097846,
      "inner": 1,
      "median_ms": 66.63321699306834,
      "q1_ms": 66.47365999378962,
      "q3_ms": 66.68258399440674
    },
    {
      "suite": "transformer",
      "attention": "jax_xla",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "compile_s": 4.830112277995795,
      "inner": 1,
      "median_ms": 225.48831001040526,
      "q1_ms": 224.68751050473657,
      "q3_ms": 225.64970100211212
    },
    {
      "suite": "transformer",
      "attention": "jax_cudnn",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "compile_s": 0.25789609500498045,
      "inner": 2,
      "median_ms": 22.510267503093928,
      "q1_ms": 22.390895497665042,
      "q3_ms": 23.318206000112696
    },
    {
      "suite": "transformer",
      "attention": "jax_cudnn",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "compile_s": 1.3923948200099403,
      "inner": 1,
      "median_ms": 87.48136999201961,
      "q1_ms": 87.03713599970797,
      "q3_ms": 87.82588350004517
    },
    {
      "suite": "diffusion",
      "attention": "probjax_default",
      "dtype": "float32",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "compile_s": 2.7067594140098663,
      "inner": 38,
      "median_ms": 0.9928530263622587,
      "q1_ms": 0.9525644868278361,
      "q3_ms": 0.994759263113864
    },
    {
      "suite": "diffusion",
      "attention": "probjax_default",
      "dtype": "float32",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "compile_s": 4.946941300993785,
      "inner": 8,
      "median_ms": 4.126619875023607,
      "q1_ms": 4.019329563561769,
      "q3_ms": 4.245928625095985
    },
    {
      "suite": "diffusion",
      "attention": "probjax_flex",
      "dtype": "float32",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "compile_s": 0.5914399769972079,
      "inner": 38,
      "median_ms": 0.9505658687741839,
      "q1_ms": 0.947662460499774,
      "q3_ms": 0.9921435527135865
    },
    {
      "suite": "diffusion",
      "attention": "probjax_flex",
      "dtype": "float32",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "compile_s": 2.1303029999980936,
      "inner": 8,
      "median_ms": 4.029504125355743,
      "q1_ms": 4.006763750112441,
      "q3_ms": 4.221511750074569
    },
    {
      "suite": "diffusion",
      "attention": "jax_xla",
      "dtype": "float32",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "compile_s": 0.6209033380000619,
      "inner": 31,
      "median_ms": 1.2620448063488208,
      "q1_ms": 1.2429168550184957,
      "q3_ms": 1.2628723547564848
    },
    {
      "suite": "diffusion",
      "attention": "jax_xla",
      "dtype": "float32",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "compile_s": 3.6158850789943244,
      "inner": 8,
      "median_ms": 4.859829499764601,
      "q1_ms": 4.844910125029855,
      "q3_ms": 4.862718749791384
    },
    {
      "suite": "diffusion",
      "attention": "probjax_default",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "compile_s": 2.3125920010061236,
      "inner": 10,
      "median_ms": 4.254563098947983,
      "q1_ms": 4.21912554957089,
      "q3_ms": 4.255894199741306
    },
    {
      "suite": "diffusion",
      "attention": "probjax_default",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "compile_s": 4.389988146998803,
      "inner": 2,
      "median_ms": 16.281446994980797,
      "q1_ms": 16.193066752748564,
      "q3_ms": 16.34936425034539
    },
    {
      "suite": "diffusion",
      "attention": "probjax_flex",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "compile_s": 0.5967186910129385,
      "inner": 10,
      "median_ms": 4.251874999317806,
      "q1_ms": 4.215903499425622,
      "q3_ms": 4.255094600375742
    },
    {
      "suite": "diffusion",
      "attention": "probjax_flex",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "compile_s": 2.2454585709929233,
      "inner": 2,
      "median_ms": 16.615508000541013,
      "q1_ms": 16.45614849985577,
      "q3_ms": 16.667266503645806
    },
    {
      "suite": "diffusion",
      "attention": "jax_xla",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "compile_s": 0.6557802060124232,
      "inner": 5,
      "median_ms": 8.276676799869165,
      "q1_ms": 8.273848200042266,
      "q3_ms": 8.282884399523027
    },
    {
      "suite": "diffusion",
      "attention": "jax_xla",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "compile_s": 3.953129337998689,
      "inner": 1,
      "median_ms": 26.149704004637897,
      "q1_ms": 26.091633502801415,
      "q3_ms": 26.201381500868592
    },
    {
      "suite": "diffusion",
      "attention": "probjax_default",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "compile_s": 1.865923051009304,
      "inner": 1,
      "median_ms": 28.142533003119752,
      "q1_ms": 28.088970997487195,
      "q3_ms": 28.277612502279226
    },
    {
      "suite": "diffusion",
      "attention": "probjax_default",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "compile_s": 4.053781788999913,
      "inner": 1,
      "median_ms": 128.0265609966591,
      "q1_ms": 127.85816250834614,
      "q3_ms": 128.54248649819056
    },
    {
      "suite": "diffusion",
      "attention": "probjax_flex",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "compile_s": 0.6272168380091898,
      "inner": 1,
      "median_ms": 28.828368987888098,
      "q1_ms": 28.749515498930123,
      "q3_ms": 28.872573508124333
    },
    {
      "suite": "diffusion",
      "attention": "probjax_flex",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "compile_s": 2.334638886997709,
      "inner": 1,
      "median_ms": 128.33284999942407,
      "q1_ms": 127.57923150638817,
      "q3_ms": 128.85183449543547
    },
    {
      "suite": "diffusion",
      "attention": "jax_xla",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "compile_s": 1.3372413820034126,
      "inner": 1,
      "median_ms": 92.36951099592261,
      "q1_ms": 92.26851799758151,
      "q3_ms": 92.5147120069596
    },
    {
      "suite": "diffusion",
      "attention": "jax_xla",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 24.78GiB with allocator GPU_0_bfc on device 0. [executable_name='jit_train_step'] [tf-allocator-allocation-error='']"
    },
    {
      "suite": "diffusion",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "compile_s": 1.130481519998284,
      "inner": 51,
      "median_ms": 0.6281467646533879,
      "q1_ms": 0.6041946961012596,
      "q3_ms": 0.6308485391463026
    },
    {
      "suite": "diffusion",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "compile_s": 3.0500928959954763,
      "inner": 11,
      "median_ms": 3.0447854554619305,
      "q1_ms": 3.0375536364524374,
      "q3_ms": 3.049834135956851
    },
    {
      "suite": "diffusion",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "compile_s": 0.5784909789945232,
      "inner": 49,
      "median_ms": 0.631548673370663,
      "q1_ms": 0.6309131734139686,
      "q3_ms": 0.6321200408096121
    },
    {
      "suite": "diffusion",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "compile_s": 2.1234387180011254,
      "inner": 10,
      "median_ms": 3.201515699038282,
      "q1_ms": 3.082250599982217,
      "q3_ms": 3.3444639993831515
    },
    {
      "suite": "diffusion",
      "attention": "jax_xla",
      "dtype": "bfloat16",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "compile_s": 0.7624866170081077,
      "inner": 42,
      "median_ms": 0.8472110237649065,
      "q1_ms": 0.8454546190304904,
      "q3_ms": 0.8476394763246866
    },
    {
      "suite": "diffusion",
      "attention": "jax_xla",
      "dtype": "bfloat16",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "compile_s": 2.504101037993678,
      "inner": 9,
      "median_ms": 3.970304222699876,
      "q1_ms": 3.9158188333709,
      "q3_ms": 3.9801454444790982
    },
    {
      "suite": "diffusion",
      "attention": "jax_cudnn",
      "dtype": "bfloat16",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "compile_s": 0.5799755360058043,
      "inner": 48,
      "median_ms": 0.6286365833148011,
      "q1_ms": 0.6207534374880197,
      "q3_ms": 0.6779553646083514
    },
    {
      "suite": "diffusion",
      "attention": "jax_cudnn",
      "dtype": "bfloat16",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "compile_s": 2.1030871600087266,
      "inner": 9,
      "median_ms": 3.542160000910776,
      "q1_ms": 3.3675302224259616,
      "q3_ms": 3.644647110680429
    },
    {
      "suite": "diffusion",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "compile_s": 0.8106683380028699,
      "inner": 18,
      "median_ms": 2.353405166811879,
      "q1_ms": 2.3015287224552594,
      "q3_ms": 2.3624669997338463
    },
    {
      "suite": "diffusion",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "compile_s": 2.7929619759961497,
      "inner": 4,
      "median_ms": 8.518011498381384,
      "q1_ms": 8.484390626108507,
      "q3_ms": 8.580703999541583
    },
    {
      "suite": "diffusion",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "compile_s": 0.5825420199980726,
      "inner": 18,
      "median_ms": 2.3134183890457884,
      "q1_ms": 2.283062972209235,
      "q3_ms": 2.352701027929369
    },
    {
      "suite": "diffusion",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "compile_s": 2.2882482830027584,
      "inner": 5,
      "median_ms": 8.187151400488801,
      "q1_ms": 8.136278099846095,
      "q3_ms": 8.27863539889222
    },
    {
      "suite": "diffusion",
      "attention": "jax_xla",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "compile_s": 0.7766323190007824,
      "inner": 8,
      "median_ms": 5.685483374691103,
      "q1_ms": 5.6816850637915195,
      "q3_ms": 5.6947516868604
    },
    {
      "suite": "diffusion",
      "attention": "jax_xla",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "compile_s": 2.5654922199901193,
      "inner": 2,
      "median_ms": 19.166833495546598,
      "q1_ms": 18.872243250370957,
      "q3_ms": 19.363352254003985
    },
    {
      "suite": "diffusion",
      "attention": "jax_cudnn",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "compile_s": 0.5977813710051123,
      "inner": 14,
      "median_ms": 3.0201673570055783,
      "q1_ms": 2.8669236432636223,
      "q3_ms": 3.0625253925141545
    },
    {
      "suite": "diffusion",
      "attention": "jax_cudnn",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "compile_s": 2.1635887589945924,
      "inner": 3,
      "median_ms": 11.440148669256208,
      "q1_ms": 11.395668332018733,
      "q3_ms": 11.475181165830389
    },
    {
      "suite": "diffusion",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "compile_s": 0.6835816740058362,
      "inner": 3,
      "median_ms": 15.046338332467712,
      "q1_ms": 15.031824502026819,
      "q3_ms": 15.165985333927285
    },
    {
      "suite": "diffusion",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "compile_s": 2.7474685000051977,
      "inner": 1,
      "median_ms": 52.42963599448558,
      "q1_ms": 51.99871899822028,
      "q3_ms": 54.04746299609542
    },
    {
      "suite": "diffusion",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "compile_s": 0.47644102400226984,
      "inner": 3,
      "median_ms": 15.042611999281993,
      "q1_ms": 15.03286583465524,
      "q3_ms": 15.078039999934845
    },
    {
      "suite": "diffusion",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "compile_s": 2.318484596995404,
      "inner": 1,
      "median_ms": 52.58262201095931,
      "q1_ms": 52.139709987386595,
      "q3_ms": 52.59103499702178
    },
    {
      "suite": "diffusion",
      "attention": "jax_xla",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "compile_s": 0.8066992509993725,
      "inner": 1,
      "median_ms": 66.7369730072096,
      "q1_ms": 66.64561949583003,
      "q3_ms": 66.93883800471667
    },
    {
      "suite": "diffusion",
      "attention": "jax_xla",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "compile_s": 2.786761048992048,
      "inner": 1,
      "median_ms": 228.29814099532086,
      "q1_ms": 227.9836055022315,
      "q3_ms": 228.497158997925
    },
    {
      "suite": "diffusion",
      "attention": "jax_cudnn",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "compile_s": 0.44737984499079175,
      "inner": 2,
      "median_ms": 22.671435493975878,
      "q1_ms": 22.42518774801283,
      "q3_ms": 22.884593498019967
    },
    {
      "suite": "diffusion",
      "attention": "jax_cudnn",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "compile_s": 2.2972284559946274,
      "inner": 1,
      "median_ms": 88.53064700087998,
      "q1_ms": 87.85052049643127,
      "q3_ms": 88.67890850524418
    }
  ]
}