{
  "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/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
  },
  "results": [
    {
      "suite": "transformer",
      "attention": "probjax_default",
      "dtype": "float32",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "compile_s": 4.268841214012355,
      "inner": 86,
      "median_ms": 0.2957669189632979,
      "q1_ms": 0.29521964918172294,
      "q3_ms": 0.29649872056663384
    },
    {
      "suite": "transformer",
      "attention": "probjax_default",
      "dtype": "float32",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "compile_s": 7.274503361899406,
      "inner": 19,
      "median_ms": 1.6445793750646867,
      "q1_ms": 1.6423303964792897,
      "q3_ms": 1.6472006559764085
    },
    {
      "suite": "transformer",
      "attention": "probjax_flex",
      "dtype": "float32",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "compile_s": 0.08127954602241516,
      "inner": 90,
      "median_ms": 0.2957075316872862,
      "q1_ms": 0.2947006930804087,
      "q3_ms": 0.2957473612493939
    },
    {
      "suite": "transformer",
      "attention": "probjax_flex",
      "dtype": "float32",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "compile_s": 2.2935487960930914,
      "inner": 19,
      "median_ms": 1.757221725328188,
      "q1_ms": 1.714807395872317,
      "q3_ms": 1.789396794169749
    },
    {
      "suite": "transformer",
      "attention": "jax_xla",
      "dtype": "float32",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "compile_s": 2.697267451090738,
      "inner": 79,
      "median_ms": 0.33820782422642165,
      "q1_ms": 0.337109336617721,
      "q3_ms": 0.3403299104988198
    },
    {
      "suite": "transformer",
      "attention": "jax_xla",
      "dtype": "float32",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "compile_s": 7.367556503973901,
      "inner": 17,
      "median_ms": 2.1157734068658423,
      "q1_ms": 2.0586635218933225,
      "q3_ms": 2.3694424958461346
    },
    {
      "suite": "transformer",
      "attention": "probjax_default",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "compile_s": 4.100193199003115,
      "inner": 36,
      "median_ms": 1.1500505562354293,
      "q1_ms": 1.1338998188471627,
      "q3_ms": 1.162577792355377
    },
    {
      "suite": "transformer",
      "attention": "probjax_default",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "compile_s": 8.335730341030285,
      "inner": 9,
      "median_ms": 4.270396997324294,
      "q1_ms": 4.205449665379193,
      "q3_ms": 4.345601618600389
    },
    {
      "suite": "transformer",
      "attention": "probjax_flex",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "compile_s": 0.07775661791674793,
      "inner": 38,
      "median_ms": 1.1514511792675446,
      "q1_ms": 1.1306342907789113,
      "q3_ms": 1.1580682874314094
    },
    {
      "suite": "transformer",
      "attention": "probjax_flex",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "compile_s": 2.0748407740611583,
      "inner": 10,
      "median_ms": 4.144280101172626,
      "q1_ms": 4.12594445515424,
      "q3_ms": 4.226923105306923
    },
    {
      "suite": "transformer",
      "attention": "jax_xla",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "compile_s": 2.8164549078792334,
      "inner": 16,
      "median_ms": 2.8479186876211315,
      "q1_ms": 2.8445849384297617,
      "q3_ms": 2.849187185347546
    },
    {
      "suite": "transformer",
      "attention": "jax_xla",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "compile_s": 7.845192312961444,
      "inner": 5,
      "median_ms": 8.542026579380035,
      "q1_ms": 8.502196590416133,
      "q3_ms": 8.542774710804224
    },
    {
      "suite": "transformer",
      "attention": "probjax_default",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "compile_s": 4.32047687494196,
      "inner": 5,
      "median_ms": 9.204390179365873,
      "q1_ms": 9.113331907428801,
      "q3_ms": 9.335810388438404
    },
    {
      "suite": "transformer",
      "attention": "probjax_default",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "compile_s": 8.290095594944432,
      "inner": 1,
      "median_ms": 39.1898259986192,
      "q1_ms": 38.617014535702765,
      "q3_ms": 39.73091649822891
    },
    {
      "suite": "transformer",
      "attention": "probjax_flex",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "compile_s": 0.09325919416733086,
      "inner": 5,
      "median_ms": 9.35259060934186,
      "q1_ms": 9.23522119410336,
      "q3_ms": 9.390745498239994
    },
    {
      "suite": "transformer",
      "attention": "probjax_flex",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "compile_s": 2.222939470084384,
      "inner": 1,
      "median_ms": 38.989725057035685,
      "q1_ms": 38.721299497410655,
      "q3_ms": 39.31024100165814
    },
    {
      "suite": "transformer",
      "attention": "jax_xla",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "compile_s": 3.386729684891179,
      "inner": 1,
      "median_ms": 35.07479699328542,
      "q1_ms": 35.07113200612366,
      "q3_ms": 35.08421150036156
    },
    {
      "suite": "transformer",
      "attention": "jax_xla",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "compile_s": 7.814029133878648,
      "inner": 1,
      "median_ms": 99.7496978379786,
      "q1_ms": 99.70814303960651,
      "q3_ms": 99.77171861100942
    },
    {
      "suite": "transformer",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "compile_s": 3.0825555899646133,
      "inner": 90,
      "median_ms": 0.2701087001090249,
      "q1_ms": 0.2680722559388313,
      "q3_ms": 0.27159458905872375
    },
    {
      "suite": "transformer",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "compile_s": 6.463239670963958,
      "inner": 22,
      "median_ms": 1.621447630565275,
      "q1_ms": 1.6192223834381863,
      "q3_ms": 1.6254508412781763
    },
    {
      "suite": "transformer",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "compile_s": 0.08440196514129639,
      "inner": 107,
      "median_ms": 0.25312017832195094,
      "q1_ms": 0.2528214391982444,
      "q3_ms": 0.2545875194700104
    },
    {
      "suite": "transformer",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "compile_s": 1.7991376938298345,
      "inner": 23,
      "median_ms": 1.5978147883130156,
      "q1_ms": 1.5968678458391325,
      "q3_ms": 1.6031546072791452
    },
    {
      "suite": "transformer",
      "attention": "jax_xla",
      "dtype": "bfloat16",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "compile_s": 2.040558691835031,
      "inner": 98,
      "median_ms": 0.27276815285868183,
      "q1_ms": 0.27253743432158106,
      "q3_ms": 0.2740147293127161
    },
    {
      "suite": "transformer",
      "attention": "jax_xla",
      "dtype": "bfloat16",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "compile_s": 6.308257213095203,
      "inner": 21,
      "median_ms": 1.6837674747442914,
      "q1_ms": 1.6833620223527155,
      "q3_ms": 1.6896729801027548
    },
    {
      "suite": "transformer",
      "attention": "jax_cudnn",
      "dtype": "bfloat16",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "compile_s": 2.2498388790991157,
      "inner": 97,
      "median_ms": 0.3013408565221681,
      "q1_ms": 0.3000913252669972,
      "q3_ms": 0.3020590106409389
    },
    {
      "suite": "transformer",
      "attention": "jax_cudnn",
      "dtype": "bfloat16",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "compile_s": 1.838921841001138,
      "inner": 21,
      "median_ms": 1.7255966196812333,
      "q1_ms": 1.7211706643658027,
      "q3_ms": 1.7494328563944215
    },
    {
      "suite": "transformer",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "compile_s": 2.950744451954961,
      "inner": 62,
      "median_ms": 0.5546515326826803,
      "q1_ms": 0.5498656368393812,
      "q3_ms": 0.5560608959997133
    },
    {
      "suite": "transformer",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "compile_s": 6.595839271089062,
      "inner": 14,
      "median_ms": 2.1525240070851788,
      "q1_ms": 2.145828995188432,
      "q3_ms": 2.1712097888147195
    },
    {
      "suite": "transformer",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "compile_s": 0.0844168751500547,
      "inner": 61,
      "median_ms": 0.5537948197089747,
      "q1_ms": 0.5481771888119764,
      "q3_ms": 0.5578598175503191
    },
    {
      "suite": "transformer",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "compile_s": 1.9740024460479617,
      "inner": 14,
      "median_ms": 2.164505422115326,
      "q1_ms": 2.124528566907559,
      "q3_ms": 2.170150101717029
    },
    {
      "suite": "transformer",
      "attention": "jax_xla",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "compile_s": 1.9747671380173415,
      "inner": 22,
      "median_ms": 1.9395115510137244,
      "q1_ms": 1.936414226127619,
      "q3_ms": 1.9401108412156727
    },
    {
      "suite": "transformer",
      "attention": "jax_xla",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "compile_s": 6.543627845821902,
      "inner": 6,
      "median_ms": 6.292081981276472,
      "q1_ms": 6.28595968009904,
      "q3_ms": 6.311103003099561
    },
    {
      "suite": "transformer",
      "attention": "jax_cudnn",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "compile_s": 0.4757821799721569,
      "inner": 58,
      "median_ms": 0.567121069942569,
      "q1_ms": 0.5553646619303216,
      "q3_ms": 0.5691627939713412
    },
    {
      "suite": "transformer",
      "attention": "jax_cudnn",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "compile_s": 1.8622224708087742,
      "inner": 14,
      "median_ms": 2.160407569525497,
      "q1_ms": 2.1479864621401896,
      "q3_ms": 2.1826893872847517
    },
    {
      "suite": "transformer",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "compile_s": 2.982668901095167,
      "inner": 11,
      "median_ms": 4.132539986378767,
      "q1_ms": 4.119778094305234,
      "q3_ms": 4.140067178840664
    },
    {
      "suite": "transformer",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "compile_s": 6.803192052058876,
      "inner": 3,
      "median_ms": 14.816536024833718,
      "q1_ms": 14.576375174025696,
      "q3_ms": 14.989681192673743
    },
    {
      "suite": "transformer",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "compile_s": 0.06829849001951516,
      "inner": 11,
      "median_ms": 4.155409183691849,
      "q1_ms": 4.129947173747826,
      "q3_ms": 4.224006499333138
    },
    {
      "suite": "transformer",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "compile_s": 0.4152241989504546,
      "inner": 3,
      "median_ms": 14.847022403652469,
      "q1_ms": 14.619117137044668,
      "q3_ms": 14.90962435491383
    },
    {
      "suite": "transformer",
      "attention": "jax_xla",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "compile_s": 2.097621134016663,
      "inner": 1,
      "median_ms": 25.289945071563125,
      "q1_ms": 25.27925046160817,
      "q3_ms": 25.302415480837226
    },
    {
      "suite": "transformer",
      "attention": "jax_xla",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "compile_s": 7.144707886967808,
      "inner": 1,
      "median_ms": 73.96565191447735,
      "q1_ms": 73.91790195833892,
      "q3_ms": 73.96934204734862
    },
    {
      "suite": "transformer",
      "attention": "jax_cudnn",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 3159040,
      "finite": true,
      "compile_s": 0.06267417687922716,
      "inner": 12,
      "median_ms": 4.172951662136863,
      "q1_ms": 4.145365751658877,
      "q3_ms": 4.20907591857637
    },
    {
      "suite": "transformer",
      "attention": "jax_cudnn",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 3159040,
      "finite": true,
      "compile_s": 0.3546207568142563,
      "inner": 3,
      "median_ms": 14.525334971646467,
      "q1_ms": 14.242937827172378,
      "q3_ms": 14.765394191878537
    },
    {
      "suite": "diffusion",
      "attention": "probjax_default",
      "dtype": "float32",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "compile_s": 3.418572586029768,
      "inner": 68,
      "median_ms": 0.3913836485213216,
      "q1_ms": 0.39030416176983096,
      "q3_ms": 0.39179850650458214
    },
    {
      "suite": "diffusion",
      "attention": "probjax_default",
      "dtype": "float32",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "compile_s": 6.827028300147504,
      "inner": 14,
      "median_ms": 2.6422665022047505,
      "q1_ms": 2.637996113792594,
      "q3_ms": 2.6481111833293522
    },
    {
      "suite": "diffusion",
      "attention": "probjax_flex",
      "dtype": "float32",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "compile_s": 0.1197075319942087,
      "inner": 70,
      "median_ms": 0.3637227861742888,
      "q1_ms": 0.36360593678961906,
      "q3_ms": 0.36465037181707366
    },
    {
      "suite": "diffusion",
      "attention": "probjax_flex",
      "dtype": "float32",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "compile_s": 3.4138689159881324,
      "inner": 14,
      "median_ms": 2.677595858196063,
      "q1_ms": 2.6714716056760937,
      "q3_ms": 2.6802816095628907
    },
    {
      "suite": "diffusion",
      "attention": "jax_xla",
      "dtype": "float32",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "compile_s": 1.1858957069925964,
      "inner": 64,
      "median_ms": 0.40171558066504076,
      "q1_ms": 0.40003251524467487,
      "q3_ms": 0.40244001502287574
    },
    {
      "suite": "diffusion",
      "attention": "jax_xla",
      "dtype": "float32",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "compile_s": 4.746893126051873,
      "inner": 14,
      "median_ms": 2.766631793097726,
      "q1_ms": 2.764718174668295,
      "q3_ms": 2.7696867556577285
    },
    {
      "suite": "diffusion",
      "attention": "probjax_default",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "compile_s": 2.6757436329498887,
      "inner": 33,
      "median_ms": 1.1461985819606166,
      "q1_ms": 1.1247809567121845,
      "q3_ms": 1.1513458642487726
    },
    {
      "suite": "diffusion",
      "attention": "probjax_default",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "compile_s": 6.113771890988573,
      "inner": 8,
      "median_ms": 4.498762486036867,
      "q1_ms": 4.483217926463112,
      "q3_ms": 4.5343568053795025
    },
    {
      "suite": "diffusion",
      "attention": "probjax_flex",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "compile_s": 0.12287991098128259,
      "inner": 34,
      "median_ms": 1.154267235987765,
      "q1_ms": 1.137650290997151,
      "q3_ms": 1.1559957809581913
    },
    {
      "suite": "diffusion",
      "attention": "probjax_flex",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "compile_s": 3.1662080511450768,
      "inner": 8,
      "median_ms": 4.530255013378337,
      "q1_ms": 4.499166127061471,
      "q3_ms": 4.560208195471205
    },
    {
      "suite": "diffusion",
      "attention": "jax_xla",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "compile_s": 0.8675949468743056,
      "inner": 15,
      "median_ms": 2.8832403322060904,
      "q1_ms": 2.8745153375590844,
      "q3_ms": 2.8889250087862215
    },
    {
      "suite": "diffusion",
      "attention": "jax_xla",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "compile_s": 5.168496812926605,
      "inner": 4,
      "median_ms": 9.319333010353148,
      "q1_ms": 9.302693011704832,
      "q3_ms": 9.321527992142364
    },
    {
      "suite": "diffusion",
      "attention": "probjax_default",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "compile_s": 2.939137122128159,
      "inner": 5,
      "median_ms": 9.214658429846168,
      "q1_ms": 9.129624092020094,
      "q3_ms": 9.37567159999162
    },
    {
      "suite": "diffusion",
      "attention": "probjax_default",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "compile_s": 6.543305556057021,
      "inner": 1,
      "median_ms": 40.8210689201951,
      "q1_ms": 40.423189057037234,
      "q3_ms": 41.233430383726954
    },
    {
      "suite": "diffusion",
      "attention": "probjax_flex",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "compile_s": 0.1387539750430733,
      "inner": 5,
      "median_ms": 9.270586213096976,
      "q1_ms": 9.180301218293607,
      "q3_ms": 9.322654502466321
    },
    {
      "suite": "diffusion",
      "attention": "probjax_flex",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "compile_s": 3.621909182984382,
      "inner": 1,
      "median_ms": 40.74035887606442,
      "q1_ms": 40.49571952782571,
      "q3_ms": 40.97451490815729
    },
    {
      "suite": "diffusion",
      "attention": "jax_xla",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "compile_s": 1.6285917200148106,
      "inner": 1,
      "median_ms": 35.23293696343899,
      "q1_ms": 35.15530750155449,
      "q3_ms": 35.274536930955946
    },
    {
      "suite": "diffusion",
      "attention": "jax_xla",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "compile_s": 5.6231800420209765,
      "inner": 1,
      "median_ms": 101.64565895684063,
      "q1_ms": 101.57248354516923,
      "q3_ms": 101.91625449806452
    },
    {
      "suite": "diffusion",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "compile_s": 2.428649205947295,
      "inner": 62,
      "median_ms": 0.41014092012999515,
      "q1_ms": 0.40985511112657763,
      "q3_ms": 0.41085415549816623
    },
    {
      "suite": "diffusion",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "compile_s": 5.53545910702087,
      "inner": 9,
      "median_ms": 4.832555332945453,
      "q1_ms": 4.726432845927775,
      "q3_ms": 4.845186447103818
    },
    {
      "suite": "diffusion",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "compile_s": 0.12781559792347252,
      "inner": 75,
      "median_ms": 0.40491558611392975,
      "q1_ms": 0.40354224697997176,
      "q3_ms": 0.4053183800230424
    },
    {
      "suite": "diffusion",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "compile_s": 3.426769602112472,
      "inner": 9,
      "median_ms": 4.435327757770817,
      "q1_ms": 3.8198936704753175,
      "q3_ms": 4.492692882195115
    },
    {
      "suite": "diffusion",
      "attention": "jax_xla",
      "dtype": "bfloat16",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "compile_s": 1.0435392381623387,
      "inner": 74,
      "median_ms": 0.40443254181662125,
      "q1_ms": 0.40408902065325025,
      "q3_ms": 0.40492503368925004
    },
    {
      "suite": "diffusion",
      "attention": "jax_xla",
      "dtype": "bfloat16",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "compile_s": 3.6044282619841397,
      "inner": 13,
      "median_ms": 2.802333239322672,
      "q1_ms": 2.792348610595442,
      "q3_ms": 2.8069609274657874
    },
    {
      "suite": "diffusion",
      "attention": "jax_cudnn",
      "dtype": "bfloat16",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "compile_s": 0.9401905660051852,
      "inner": 67,
      "median_ms": 0.4432459271499025,
      "q1_ms": 0.4430350225502208,
      "q3_ms": 0.4435704774292769
    },
    {
      "suite": "diffusion",
      "attention": "jax_cudnn",
      "dtype": "bfloat16",
      "tokens": 256,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "compile_s": 2.9974947459995747,
      "inner": 13,
      "median_ms": 2.904601914521593,
      "q1_ms": 2.90072384254577,
      "q3_ms": 2.913116584890164
    },
    {
      "suite": "diffusion",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "compile_s": 1.302193981129676,
      "inner": 51,
      "median_ms": 0.6230708038178729,
      "q1_ms": 0.6192354211474167,
      "q3_ms": 0.6257169994105602
    },
    {
      "suite": "diffusion",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "compile_s": 4.477760627167299,
      "inner": 12,
      "median_ms": 2.858195589700093,
      "q1_ms": 2.853687289947023,
      "q3_ms": 2.864071837393567
    },
    {
      "suite": "diffusion",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "compile_s": 0.12839585007168353,
      "inner": 51,
      "median_ms": 0.6250267044878474,
      "q1_ms": 0.6193524710468802,
      "q3_ms": 0.6258261307854863
    },
    {
      "suite": "diffusion",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "compile_s": 3.3391988170333207,
      "inner": 12,
      "median_ms": 2.8792674226375916,
      "q1_ms": 2.8759073369049775,
      "q3_ms": 2.8836576966568828
    },
    {
      "suite": "diffusion",
      "attention": "jax_xla",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "compile_s": 1.0328834208194166,
      "inner": 19,
      "median_ms": 2.034759320514767,
      "q1_ms": 2.0301440534623048,
      "q3_ms": 2.0364164190955067
    },
    {
      "suite": "diffusion",
      "attention": "jax_xla",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "compile_s": 4.094373821048066,
      "inner": 5,
      "median_ms": 7.180062402039766,
      "q1_ms": 7.168528181500733,
      "q3_ms": 7.190397311933339
    },
    {
      "suite": "diffusion",
      "attention": "jax_cudnn",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "compile_s": 0.7955109460745007,
      "inner": 42,
      "median_ms": 0.6376823820617228,
      "q1_ms": 0.6333020144319605,
      "q3_ms": 0.644382652050505
    },
    {
      "suite": "diffusion",
      "attention": "jax_cudnn",
      "dtype": "bfloat16",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "compile_s": 3.727262402884662,
      "inner": 9,
      "median_ms": 5.485824542120099,
      "q1_ms": 5.3910993867450285,
      "q3_ms": 5.502325780172315
    },
    {
      "suite": "diffusion",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "compile_s": 1.2993347530718893,
      "inner": 10,
      "median_ms": 4.3945379089564085,
      "q1_ms": 4.303190764039755,
      "q3_ms": 4.399334394838661
    },
    {
      "suite": "diffusion",
      "attention": "probjax_default",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "compile_s": 4.8415347568225116,
      "inner": 2,
      "median_ms": 16.70770300552249,
      "q1_ms": 16.3392917602323,
      "q3_ms": 16.79277518996969
    },
    {
      "suite": "diffusion",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "compile_s": 0.11156829097308218,
      "inner": 11,
      "median_ms": 4.3859670827673245,
      "q1_ms": 4.27460055586628,
      "q3_ms": 4.409803141077811
    },
    {
      "suite": "diffusion",
      "attention": "probjax_flex",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "compile_s": 0.6613791729323566,
      "inner": 2,
      "median_ms": 16.486541950143874,
      "q1_ms": 16.017088724765927,
      "q3_ms": 16.666505252942443
    },
    {
      "suite": "diffusion",
      "attention": "jax_xla",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "compile_s": 1.0320312699768692,
      "inner": 1,
      "median_ms": 25.354817043989897,
      "q1_ms": 25.272877071984112,
      "q3_ms": 25.389702059328556
    },
    {
      "suite": "diffusion",
      "attention": "jax_xla",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "compile_s": 4.101622796850279,
      "inner": 1,
      "median_ms": 77.03261310234666,
      "q1_ms": 77.00025907251984,
      "q3_ms": 77.15891895350069
    },
    {
      "suite": "diffusion",
      "attention": "jax_cudnn",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "forward",
      "params": 5817896,
      "finite": true,
      "compile_s": 0.10625343001447618,
      "inner": 11,
      "median_ms": 4.222215915268118,
      "q1_ms": 4.154286230914295,
      "q3_ms": 4.2534714501181785
    },
    {
      "suite": "diffusion",
      "attention": "jax_cudnn",
      "dtype": "bfloat16",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 256,
      "op": "train_step",
      "params": 5817896,
      "finite": true,
      "compile_s": 0.6275697171222419,
      "inner": 3,
      "median_ms": 15.92022933376332,
      "q1_ms": 15.629103484873971,
      "q3_ms": 16.016122846243285
    },
    {
      "suite": "moe",
      "impl": "pallas",
      "dtype": "float32",
      "tokens": 4096,
      "experts": 8,
      "top_k": 2,
      "d_model": 512,
      "d_hidden": 1024,
      "op": "forward",
      "params": 12587008,
      "error": "ValueError: Only f16 WGMMA supports transposes"
    },
    {
      "suite": "moe",
      "impl": "pallas",
      "dtype": "float32",
      "tokens": 4096,
      "experts": 8,
      "top_k": 2,
      "d_model": 512,
      "d_hidden": 1024,
      "op": "train_step",
      "params": 12587008,
      "error": "ValueError: Only f16 WGMMA supports transposes"
    },
    {
      "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": 2.191271889023483,
      "inner": 51,
      "median_ms": 0.7775892165726891,
      "q1_ms": 0.777192461742636,
      "q3_ms": 0.7783088435911957
    },
    {
      "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": 3.189351776149124,
      "inner": 22,
      "median_ms": 1.9306630552323027,
      "q1_ms": 1.9303387274372983,
      "q3_ms": 1.9341434976509349
    },
    {
      "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.833937434013933,
      "inner": 47,
      "median_ms": 0.9435864036625371,
      "q1_ms": 0.9382885231814804,
      "q3_ms": 0.9455673589787268
    },
    {
      "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": 2.4434416140429676,
      "inner": 13,
      "median_ms": 3.427038613993388,
      "q1_ms": 3.406894303715,
      "q3_ms": 3.4663725393609357
    },
    {
      "suite": "moe",
      "impl": "pallas",
      "dtype": "float32",
      "tokens": 32768,
      "experts": 8,
      "top_k": 2,
      "d_model": 512,
      "d_hidden": 1024,
      "op": "forward",
      "params": 12587008,
      "error": "ValueError: Only f16 WGMMA supports transposes"
    },
    {
      "suite": "moe",
      "impl": "pallas",
      "dtype": "float32",
      "tokens": 32768,
      "experts": 8,
      "top_k": 2,
      "d_model": 512,
      "d_hidden": 1024,
      "op": "train_step",
      "params": 12587008,
      "error": "ValueError: Only f16 WGMMA supports transposes"
    },
    {
      "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": 2.2637228248640895,
      "inner": 10,
      "median_ms": 4.693677695468068,
      "q1_ms": 4.690285201650113,
      "q3_ms": 4.698079254012555
    },
    {
      "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.7547336439602077,
      "inner": 4,
      "median_ms": 10.80352975986898,
      "q1_ms": 10.787169710965827,
      "q3_ms": 10.813060769578442
    },
    {
      "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": 1.6705325599759817,
      "inner": 6,
      "median_ms": 7.538145485644539,
      "q1_ms": 7.508992159273475,
      "q3_ms": 7.578951500666639
    },
    {
      "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": 2.726175168994814,
      "inner": 1,
      "median_ms": 27.02720183879137,
      "q1_ms": 26.901651523075998,
      "q3_ms": 27.109607006423175
    },
    {
      "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": 2.1855060250964016,
      "inner": 122,
      "median_ms": 0.2412569430488788,
      "q1_ms": 0.2408778031555111,
      "q3_ms": 0.24213383191830065
    },
    {
      "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": 3.9647967249620706,
      "inner": 58,
      "median_ms": 0.620962618368453,
      "q1_ms": 0.6205921908358818,
      "q3_ms": 0.6223114303730686
    },
    {
      "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.42725331406109035,
      "inner": 52,
      "median_ms": 0.7750594612354269,
      "q1_ms": 0.7744855180275269,
      "q3_ms": 0.7755010021635547
    },
    {
      "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": 2.224657795159146,
      "inner": 21,
      "median_ms": 1.98499071190045,
      "q1_ms": 1.9845147838904746,
      "q3_ms": 1.9858809564972209
    },
    {
      "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": 1.051568079041317,
      "inner": 47,
      "median_ms": 0.9361782973196279,
      "q1_ms": 0.9278113720938563,
      "q3_ms": 0.9376878820438969
    },
    {
      "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.4999170599039644,
      "inner": 13,
      "median_ms": 3.423840165711366,
      "q1_ms": 3.420460494593359,
      "q3_ms": 3.448720546797491
    },
    {
      "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": 1.8604267439804971,
      "inner": 30,
      "median_ms": 1.4808197040110826,
      "q1_ms": 1.478803033630053,
      "q3_ms": 1.4812130325784287
    },
    {
      "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": 3.9678224420640618,
      "inner": 14,
      "median_ms": 3.1937307877732173,
      "q1_ms": 3.1802201410755515,
      "q3_ms": 3.201687284412661
    },
    {
      "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.38860378693789244,
      "inner": 10,
      "median_ms": 4.636616609059274,
      "q1_ms": 4.633530543651432,
      "q3_ms": 4.6377601101994514
    },
    {
      "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.7764446220826358,
      "inner": 4,
      "median_ms": 10.689024231396616,
      "q1_ms": 10.669896728359163,
      "q3_ms": 10.694879398215562
    },
    {
      "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": 1.160050816135481,
      "inner": 6,
      "median_ms": 7.460380322299898,
      "q1_ms": 7.440103586607924,
      "q3_ms": 7.5300479074940085
    },
    {
      "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.46036737598478794,
      "inner": 1,
      "median_ms": 26.322080055251718,
      "q1_ms": 26.299004908651114,
      "q3_ms": 26.395480148494244
    },
    {
      "suite": "ssm",
      "cell": "mamba",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 128,
      "op": "forward",
      "params": 930704,
      "finite": true,
      "compile_s": 4.2124762260355055,
      "inner": 15,
      "median_ms": 2.9267892707139254,
      "q1_ms": 2.9219652991741896,
      "q3_ms": 2.9304123638818664
    },
    {
      "suite": "ssm",
      "cell": "mamba",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 128,
      "op": "train_step",
      "params": 930704,
      "finite": true,
      "compile_s": 9.524919492891058,
      "inner": 2,
      "median_ms": 16.218837001360953,
      "q1_ms": 16.203036764636636,
      "q3_ms": 16.264109523035586
    },
    {
      "suite": "ssm",
      "cell": "ssd",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 128,
      "op": "forward",
      "params": 863644,
      "finite": true,
      "compile_s": 9.163385071093217,
      "inner": 38,
      "median_ms": 0.7733183447271585,
      "q1_ms": 0.7723362375001765,
      "q3_ms": 0.7741412356201756
    },
    {
      "suite": "ssm",
      "cell": "ssd",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 128,
      "op": "train_step",
      "params": 863644,
      "finite": true,
      "compile_s": 22.71358431293629,
      "inner": 7,
      "median_ms": 7.063654584011862,
      "q1_ms": 7.035073142365686,
      "q3_ms": 7.134411284433944
    },
    {
      "suite": "ssm",
      "cell": "mamba2",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 128,
      "op": "forward",
      "params": 1136064,
      "finite": true,
      "compile_s": 6.745558636961505,
      "inner": 30,
      "median_ms": 1.1950268332536023,
      "q1_ms": 1.194640800046424,
      "q3_ms": 1.195814130672564
    },
    {
      "suite": "ssm",
      "cell": "mamba2",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 128,
      "op": "train_step",
      "params": 1136064,
      "finite": true,
      "compile_s": 23.262761533958837,
      "inner": 7,
      "median_ms": 5.001291732436845,
      "q1_ms": 4.976624506525695,
      "q3_ms": 5.005426072914686
    },
    {
      "suite": "ssm",
      "cell": "gdn",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 128,
      "op": "forward",
      "params": 1006144,
      "finite": true,
      "compile_s": 4.821296344976872,
      "inner": 33,
      "median_ms": 0.9931562694184709,
      "q1_ms": 0.9885105008090085,
      "q3_ms": 0.9945687118240378
    },
    {
      "suite": "ssm",
      "cell": "gdn",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 128,
      "op": "train_step",
      "params": 1006144,
      "finite": true,
      "compile_s": 10.852999893017113,
      "inner": 6,
      "median_ms": 5.120315472595394,
      "q1_ms": 5.0543943264832105,
      "q3_ms": 5.188637995161116
    },
    {
      "suite": "ssm",
      "cell": "lru",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 128,
      "op": "forward",
      "params": 995728,
      "finite": true,
      "compile_s": 1.970310227945447,
      "inner": 56,
      "median_ms": 0.4938238416798413,
      "q1_ms": 0.4926940101930606,
      "q3_ms": 0.49507027766334694
    },
    {
      "suite": "ssm",
      "cell": "lru",
      "dtype": "float32",
      "tokens": 1024,
      "batch": 8,
      "model_dim": 128,
      "op": "train_step",
      "params": 995728,
      "finite": true,
      "compile_s": 7.133343748981133,
      "inner": 12,
      "median_ms": 2.958629668379823,
      "q1_ms": 2.9494930737807104,
      "q3_ms": 2.9948289641955244
    },
    {
      "suite": "ssm",
      "cell": "mamba",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 128,
      "op": "forward",
      "params": 930704,
      "finite": true,
      "compile_s": 3.8943453792016953,
      "inner": 5,
      "median_ms": 8.967100037261844,
      "q1_ms": 8.961290796287358,
      "q3_ms": 8.967679808847606
    },
    {
      "suite": "ssm",
      "cell": "mamba",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 128,
      "op": "train_step",
      "params": 930704,
      "finite": true,
      "compile_s": 10.067664484959096,
      "inner": 1,
      "median_ms": 58.21725097484887,
      "q1_ms": 58.18973702844232,
      "q3_ms": 58.23011649772525
    },
    {
      "suite": "ssm",
      "cell": "ssd",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 128,
      "op": "forward",
      "params": 863644,
      "finite": true,
      "compile_s": 9.016833167988807,
      "inner": 16,
      "median_ms": 2.818231252604164,
      "q1_ms": 2.810241530823987,
      "q3_ms": 2.8223999615875073
    },
    {
      "suite": "ssm",
      "cell": "ssd",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 128,
      "op": "train_step",
      "params": 863644,
      "finite": true,
      "compile_s": 22.606969696935266,
      "inner": 4,
      "median_ms": 10.663783468771726,
      "q1_ms": 10.629782162141055,
      "q3_ms": 10.67106836126186
    },
    {
      "suite": "ssm",
      "cell": "mamba2",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 128,
      "op": "forward",
      "params": 1136064,
      "finite": true,
      "compile_s": 5.308368764817715,
      "inner": 10,
      "median_ms": 4.642555303871632,
      "q1_ms": 4.6252352418377995,
      "q3_ms": 4.644395294599235
    },
    {
      "suite": "ssm",
      "cell": "mamba2",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 128,
      "op": "train_step",
      "params": 1136064,
      "finite": true,
      "compile_s": 20.822938313009217,
      "inner": 2,
      "median_ms": 18.185483990237117,
      "q1_ms": 18.139544292353094,
      "q3_ms": 18.202516774181277
    },
    {
      "suite": "ssm",
      "cell": "gdn",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 128,
      "op": "forward",
      "params": 1006144,
      "finite": true,
      "compile_s": 4.692825552076101,
      "inner": 10,
      "median_ms": 4.2630873853340745,
      "q1_ms": 4.258935351390392,
      "q3_ms": 4.283675947226584
    },
    {
      "suite": "ssm",
      "cell": "gdn",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 128,
      "op": "train_step",
      "params": 1006144,
      "finite": true,
      "compile_s": 11.525399594800547,
      "inner": 2,
      "median_ms": 17.64018298126757,
      "q1_ms": 17.60484278202057,
      "q3_ms": 17.697940522339195
    },
    {
      "suite": "ssm",
      "cell": "lru",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 128,
      "op": "forward",
      "params": 995728,
      "finite": true,
      "compile_s": 2.0093409458640963,
      "inner": 22,
      "median_ms": 1.8619627226144075,
      "q1_ms": 1.8452103831805289,
      "q3_ms": 1.8684297683648765
    },
    {
      "suite": "ssm",
      "cell": "lru",
      "dtype": "float32",
      "tokens": 4096,
      "batch": 8,
      "model_dim": 128,
      "op": "train_step",
      "params": 995728,
      "finite": true,
      "compile_s": 8.09308628900908,
      "inner": 5,
      "median_ms": 6.4851398114115,
      "q1_ms": 6.4694717060774565,
      "q3_ms": 6.499150767922401
    },
    {
      "suite": "unet",
      "dtype": "float32",
      "resolution": 32,
      "batch": 8,
      "op": "forward",
      "params": 2803968,
      "finite": true,
      "compile_s": 3.3344795419834554,
      "inner": 56,
      "median_ms": 0.6517535158699113,
      "q1_ms": 0.6515102057684479,
      "q3_ms": 0.652405125687697
    },
    {
      "suite": "unet",
      "dtype": "float32",
      "resolution": 32,
      "batch": 8,
      "op": "train_step",
      "params": 2803968,
      "finite": true,
      "compile_s": 6.824007142102346,
      "inner": 10,
      "median_ms": 3.247366799041629,
      "q1_ms": 3.244873706717044,
      "q3_ms": 3.2514633028768003
    },
    {
      "suite": "unet",
      "dtype": "float32",
      "resolution": 64,
      "batch": 8,
      "op": "forward",
      "params": 2803968,
      "finite": true,
      "compile_s": 3.39439590414986,
      "inner": 39,
      "median_ms": 0.841735593544749,
      "q1_ms": 0.8380693079044039,
      "q3_ms": 0.8428934125755078
    },
    {
      "suite": "unet",
      "dtype": "float32",
      "resolution": 64,
      "batch": 8,
      "op": "train_step",
      "params": 2803968,
      "finite": true,
      "compile_s": 7.78964430000633,
      "inner": 9,
      "median_ms": 3.7151865350703397,
      "q1_ms": 3.703416557982564,
      "q3_ms": 3.847746333728234
    },
    {
      "suite": "unet",
      "dtype": "bfloat16",
      "resolution": 32,
      "batch": 8,
      "op": "forward",
      "params": 2803968,
      "finite": true,
      "compile_s": 2.233096922049299,
      "inner": 52,
      "median_ms": 0.6702881342229935,
      "q1_ms": 0.6695515006028402,
      "q3_ms": 0.6707146732458988
    },
    {
      "suite": "unet",
      "dtype": "bfloat16",
      "resolution": 32,
      "batch": 8,
      "op": "train_step",
      "params": 2803968,
      "finite": true,
      "compile_s": 6.678437325870618,
      "inner": 10,
      "median_ms": 3.2692787935957313,
      "q1_ms": 3.2530897529795766,
      "q3_ms": 3.2773737562820315
    },
    {
      "suite": "unet",
      "dtype": "bfloat16",
      "resolution": 64,
      "batch": 8,
      "op": "forward",
      "params": 2803968,
      "finite": true,
      "compile_s": 2.399156136903912,
      "inner": 40,
      "median_ms": 0.7724813476670533,
      "q1_ms": 0.7722923386609182,
      "q3_ms": 0.773055077297613
    },
    {
      "suite": "unet",
      "dtype": "bfloat16",
      "resolution": 64,
      "batch": 8,
      "op": "train_step",
      "params": 2803968,
      "finite": true,
      "compile_s": 7.448797419201583,
      "inner": 10,
      "median_ms": 3.445769310928881,
      "q1_ms": 3.4052222035825253,
      "q3_ms": 3.4688113490119576
    }
  ]
}