{
 "device": {
  "device_kind": "NVIDIA H100 (NVL, 94 GB)",
  "cc": "9.0",
  "jax": "0.11.2"
 },
 "description": "MoELayer (dropless, bf16) forward and train step (value_and_grad w.r.t. params and input) by dropless_impl / dropless_dispatch; errors are max relative differences to triton/sorted.",
 "script": "docs/scripts/benchmark_gpu_moe_dispatch.py",
 "results": [
  {
   "label": "v10_decoder",
   "host": "uc3n088.localdomain",
   "gpu": "NVIDIA H100",
   "jax": "0.11.2",
   "d_model": 512,
   "width": 128,
   "experts": 64,
   "top_k": 8,
   "shared": 1024,
   "tokens": 16384,
   "impl": "triton",
   "dispatch": "sorted",
   "fwd_ms": 2.83,
   "train_ms": 4.835,
   "out_err": 0.0,
   "gx_err": 0.0,
   "gw_err": 0.0,
   "status": "ok"
  },
  {
   "label": "v10_decoder",
   "host": "uc3n088.localdomain",
   "gpu": "NVIDIA H100",
   "jax": "0.11.2",
   "d_model": 512,
   "width": 128,
   "experts": 64,
   "top_k": 8,
   "shared": 1024,
   "tokens": 16384,
   "impl": "triton",
   "dispatch": "gather",
   "fwd_ms": 1.332,
   "train_ms": 3.68,
   "out_err": 0.0040650395676493645,
   "gx_err": 0.011363634839653969,
   "gw_err": 0.011363636702299118,
   "status": "ok"
  },
  {
   "label": "v10_decoder",
   "host": "uc3n088.localdomain",
   "gpu": "NVIDIA H100",
   "jax": "0.11.2",
   "d_model": 512,
   "width": 128,
   "experts": 64,
   "top_k": 8,
   "shared": 1024,
   "tokens": 16384,
   "impl": "triton",
   "dispatch": "fused",
   "fwd_ms": 1.274,
   "train_ms": 3.313,
   "out_err": 0.0040650395676493645,
   "gx_err": 0.011363634839653969,
   "gw_err": 0.011363636702299118,
   "status": "ok"
  },
  {
   "label": "v10_decoder",
   "host": "uc3n088.localdomain",
   "gpu": "NVIDIA H100",
   "jax": "0.11.2",
   "d_model": 512,
   "width": 128,
   "experts": 64,
   "top_k": 8,
   "shared": 1024,
   "tokens": 16384,
   "impl": "triton",
   "dispatch": "fused_gsum",
   "fwd_ms": 1.07,
   "train_ms": 2.888,
   "out_err": 0.0040650395676493645,
   "gx_err": 0.011363634839653969,
   "gw_err": 0.011363636702299118,
   "status": "ok"
  },
  {
   "label": "v10_decoder",
   "host": "uc3n088.localdomain",
   "gpu": "NVIDIA H100",
   "jax": "0.11.2",
   "d_model": 512,
   "width": 128,
   "experts": 64,
   "top_k": 8,
   "shared": 1024,
   "tokens": 16384,
   "impl": "mosaic",
   "dispatch": "sorted",
   "fwd_ms": 1.051,
   "train_ms": 2.668,
   "out_err": 0.005081299692392349,
   "gx_err": 0.008522726595401764,
   "gw_err": 0.009892881847918034,
   "status": "ok"
  },
  {
   "label": "v10_decoder",
   "host": "uc3n088.localdomain",
   "gpu": "NVIDIA H100",
   "jax": "0.11.2",
   "d_model": 512,
   "width": 128,
   "experts": 64,
   "top_k": 8,
   "shared": 1024,
   "tokens": 16384,
   "impl": "pallas",
   "dispatch": "sorted",
   "fwd_ms": 1.05,
   "train_ms": 2.628,
   "out_err": 0.005081299692392349,
   "gx_err": 0.008522726595401764,
   "gw_err": 0.009892881847918034,
   "status": "ok"
  },
  {
   "label": "v10_decoder",
   "host": "uc3n088.localdomain",
   "gpu": "NVIDIA H100",
   "jax": "0.11.2",
   "d_model": 512,
   "width": 128,
   "experts": 64,
   "top_k": 8,
   "shared": 1024,
   "tokens": 16384,
   "impl": "ragged_dot",
   "dispatch": "sorted",
   "fwd_ms": 31.312,
   "train_ms": 147.746,
   "out_err": 0.0040650395676493645,
   "gx_err": 0.008522726595401764,
   "gw_err": 0.010842466726899147,
   "status": "ok"
  },
  {
   "label": "v10_decoder",
   "host": "uc3n088.localdomain",
   "gpu": "NVIDIA H100",
   "jax": "0.11.2",
   "d_model": 512,
   "width": 128,
   "experts": 64,
   "top_k": 8,
   "shared": 1024,
   "tokens": 131072,
   "impl": "triton",
   "dispatch": "sorted",
   "fwd_ms": 10.26,
   "train_ms": 28.456,
   "out_err": 0.0,
   "gx_err": 0.0,
   "gw_err": 0.0,
   "status": "ok"
  },
  {
   "label": "v10_decoder",
   "host": "uc3n088.localdomain",
   "gpu": "NVIDIA H100",
   "jax": "0.11.2",
   "d_model": 512,
   "width": 128,
   "experts": 64,
   "top_k": 8,
   "shared": 1024,
   "tokens": 131072,
   "impl": "triton",
   "dispatch": "gather",
   "fwd_ms": 9.971,
   "train_ms": 30.569,
   "out_err": 0.008130079135298729,
   "gx_err": 0.010695185512304306,
   "gw_err": 0.008947395719587803,
   "status": "ok"
  },
  {
   "label": "v10_decoder",
   "host": "uc3n088.localdomain",
   "gpu": "NVIDIA H100",
   "jax": "0.11.2",
   "d_model": 512,
   "width": 128,
   "experts": 64,
   "top_k": 8,
   "shared": 1024,
   "tokens": 131072,
   "impl": "triton",
   "dispatch": "fused",
   "fwd_ms": 8.373,
   "train_ms": 26.505,
   "out_err": 0.008130079135298729,
   "gx_err": 0.010695185512304306,
   "gw_err": 0.008947395719587803,
   "status": "ok"
  },
  {
   "label": "v10_decoder",
   "host": "uc3n088.localdomain",
   "gpu": "NVIDIA H100",
   "jax": "0.11.2",
   "d_model": 512,
   "width": 128,
   "experts": 64,
   "top_k": 8,
   "shared": 1024,
   "tokens": 131072,
   "impl": "triton",
   "dispatch": "fused_gsum",
   "fwd_ms": 6.516,
   "train_ms": 22.161,
   "out_err": 0.008130079135298729,
   "gx_err": 0.010695185512304306,
   "gw_err": 0.008947395719587803,
   "status": "ok"
  },
  {
   "label": "v10_decoder",
   "host": "uc3n088.localdomain",
   "gpu": "NVIDIA H100",
   "jax": "0.11.2",
   "d_model": 512,
   "width": 128,
   "experts": 64,
   "top_k": 8,
   "shared": 1024,
   "tokens": 131072,
   "impl": "mosaic",
   "dispatch": "sorted",
   "fwd_ms": 7.126,
   "train_ms": 19.163,
   "out_err": 0.008130079135298729,
   "gx_err": 0.010695185512304306,
   "gw_err": 0.009543013758957386,
   "status": "ok"
  },
  {
   "label": "v10_decoder",
   "host": "uc3n088.localdomain",
   "gpu": "NVIDIA H100",
   "jax": "0.11.2",
   "d_model": 512,
   "width": 128,
   "experts": 64,
   "top_k": 8,
   "shared": 1024,
   "tokens": 131072,
   "impl": "pallas",
   "dispatch": "sorted",
   "fwd_ms": 7.139,
   "train_ms": 19.126,
   "out_err": 0.008130079135298729,
   "gx_err": 0.010695185512304306,
   "gw_err": 0.009543013758957386,
   "status": "ok"
  },
  {
   "label": "v10_decoder",
   "host": "uc3n088.localdomain",
   "gpu": "NVIDIA H100",
   "jax": "0.11.2",
   "d_model": 512,
   "width": 128,
   "experts": 64,
   "top_k": 8,
   "shared": 1024,
   "tokens": 131072,
   "impl": "ragged_dot",
   "dispatch": "sorted",
   "status": "error",
   "error": "JaxRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 128.00GiB with allocator GPU_0_bfc on device 0. [tf-allocator-allocation-error='']"
  },
  {
   "label": "v10_decoder",
   "host": "uc3n088.localdomain",
   "gpu": "NVIDIA H100",
   "jax": "0.11.2",
   "d_model": 512,
   "width": 128,
   "experts": 64,
   "top_k": 8,
   "shared": 1024,
   "tokens": 1048576,
   "impl": "triton",
   "dispatch": "sorted",
   "fwd_ms": 82.037,
   "train_ms": 267.54,
   "out_err": 0.0,
   "gx_err": 0.0,
   "gw_err": 0.0,
   "status": "ok"
  },
  {
   "label": "v10_decoder",
   "host": "uc3n088.localdomain",
   "gpu": "NVIDIA H100",
   "jax": "0.11.2",
   "d_model": 512,
   "width": 128,
   "experts": 64,
   "top_k": 8,
   "shared": 1024,
   "tokens": 1048576,
   "impl": "triton",
   "dispatch": "gather",
   "fwd_ms": 78.916,
   "train_ms": 247.879,
   "out_err": 0.006329112686216831,
   "gx_err": 0.010416665114462376,
   "gw_err": 0.008285319432616234,
   "status": "ok"
  },
  {
   "label": "v10_decoder",
   "host": "uc3n088.localdomain",
   "gpu": "NVIDIA H100",
   "jax": "0.11.2",
   "d_model": 512,
   "width": 128,
   "experts": 64,
   "top_k": 8,
   "shared": 1024,
   "tokens": 1048576,
   "impl": "triton",
   "dispatch": "fused",
   "fwd_ms": 71.17,
   "train_ms": 227.472,
   "out_err": 0.006329112686216831,
   "gx_err": 0.010416665114462376,
   "gw_err": 0.008285319432616234,
   "status": "ok"
  },
  {
   "label": "v10_decoder",
   "host": "uc3n088.localdomain",
   "gpu": "NVIDIA H100",
   "jax": "0.11.2",
   "d_model": 512,
   "width": 128,
   "experts": 64,
   "top_k": 8,
   "shared": 1024,
   "tokens": 1048576,
   "impl": "triton",
   "dispatch": "fused_gsum",
   "fwd_ms": 52.496,
   "train_ms": 188.12,
   "out_err": 0.006329112686216831,
   "gx_err": 0.010416665114462376,
   "gw_err": 0.008285319432616234,
   "status": "ok"
  },
  {
   "label": "v10_decoder",
   "host": "uc3n088.localdomain",
   "gpu": "NVIDIA H100",
   "jax": "0.11.2",
   "d_model": 512,
   "width": 128,
   "experts": 64,
   "top_k": 8,
   "shared": 1024,
   "tokens": 1048576,
   "impl": "mosaic",
   "dispatch": "sorted",
   "fwd_ms": 56.257,
   "train_ms": 150.781,
   "out_err": 0.006329112686216831,
   "gx_err": 0.010416665114462376,
   "gw_err": 0.009663057513535023,
   "status": "ok"
  },
  {
   "label": "v10_decoder",
   "host": "uc3n088.localdomain",
   "gpu": "NVIDIA H100",
   "jax": "0.11.2",
   "d_model": 512,
   "width": 128,
   "experts": 64,
   "top_k": 8,
   "shared": 1024,
   "tokens": 1048576,
   "impl": "pallas",
   "dispatch": "sorted",
   "fwd_ms": 56.465,
   "train_ms": 151.718,
   "out_err": 0.006329112686216831,
   "gx_err": 0.010416665114462376,
   "gw_err": 0.009663057513535023,
   "status": "ok"
  },
  {
   "label": "v10_decoder",
   "host": "uc3n088.localdomain",
   "gpu": "NVIDIA H100",
   "jax": "0.11.2",
   "d_model": 512,
   "width": 128,
   "experts": 64,
   "top_k": 8,
   "shared": 1024,
   "tokens": 1048576,
   "impl": "ragged_dot",
   "dispatch": "sorted",
   "status": "error",
   "error": "JaxRuntimeError: NOT_FOUND: Failed to get configs for: 3 out of 16 instructions. See logs for all failures. Example failure: \nAll configs failed during profiling or were excluded from selection.\nFailures (3):\nEXECUTION FAILED: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 256.00GiB with allocator GPU_0_"
  }
 ]
}