{ "device": { "device_kind": "NVIDIA H100 (NVL, 94 GB)", "cc": "9.0", "jax": "0.11.2", "cudnn": "9.27", "cuda": "13.4" }, "description": "Attention fwd and fwd+bwd on H100, bf16, by implementation; labels pr123_* ran on perf/gpu-cpu-kernels 849e1d2 (generic dispatch), h100br_* on the device-defaults branch. Errors are max |difference| to flex_attention triton_split. cuDNN rows labelled main_* ran in the same session (cuDNN does not depend on the probjax revision). Labels h100_torch_*: probjax backends (flex:default/cudnn/triton) and PyTorch 2.14.1+cu130 (torch:flex = compiled torch.nn.attention.flex_attention with a create_block_mask from the KV lengths; torch:flash / torch:efficient = SDPA pinned to one backend, dense boolean mask for the latter), on the same inputs; torch rows ran without JAX preallocation. flex:mosaic rows: flex_attention(implementation='mosaic') with the Mosaic KV-length kernels.", "script": "docs/scripts/benchmark_gpu_attention_device.py", "results": [ { "label": "pr123_isab", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "isab_ind_rows", "B": 56, "Lq": 256, "Lkv": 32768, "H": 8, "D": 32, "masked": true, "impl": "flex:triton_split", "fwd_ms": 3.093, "fwd_bwd_ms": 13.414, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_isab", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "isab_ind_rows", "B": 56, "Lq": 256, "Lkv": 32768, "H": 8, "D": 32, "masked": true, "impl": "flex:triton_onepass", "fwd_ms": 3.054, "fwd_bwd_ms": 12.876, "out_err": 0.0, "grad_err": [ 0.000244140625, 6.103515625e-05, 3.0517578125e-05 ], "status": "ok" }, { "label": "pr123_isab", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "isab_ind_rows", "B": 56, "Lq": 256, "Lkv": 32768, "H": 8, "D": 32, "masked": true, "impl": "flex:auto", "fwd_ms": 3.102, "fwd_bwd_ms": 13.461, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_mla", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "mla_32k84", "B": 2, "Lq": 131072, "Lkv": 32768, "H": 1, "D": 128, "masked": true, "impl": "flex:triton_split", "fwd_ms": 14.264, "fwd_bwd_ms": 51.156, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_mla", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "mla_32k84", "B": 2, "Lq": 131072, "Lkv": 32768, "H": 1, "D": 128, "masked": true, "impl": "flex:triton_onepass", "fwd_ms": 14.216, "fwd_bwd_ms": 81.046, "out_err": 0.0, "grad_err": [ 0.000244140625, 0.00048828125, 0.0009765625 ], "status": "ok" }, { "label": "pr123_mla", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "mla_32k84", "B": 2, "Lq": 131072, "Lkv": 32768, "H": 1, "D": 128, "masked": true, "impl": "flex:auto", "fwd_ms": 13.995, "fwd_bwd_ms": 89.688, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d32", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d32_t512", "B": 4, "Lq": 512, "Lkv": 512, "H": 8, "D": 32, "masked": false, "impl": "flex:triton_split", "fwd_ms": 0.059, "fwd_bwd_ms": 0.111, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d32", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d32_t512", "B": 4, "Lq": 512, "Lkv": 512, "H": 8, "D": 32, "masked": false, "impl": "flex:triton_fused", "fwd_ms": 0.06, "fwd_bwd_ms": 0.108, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_mla", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "mla_128k36", "B": 1, "Lq": 65536, "Lkv": 131072, "H": 1, "D": 128, "masked": true, "impl": "flex:triton_split", "fwd_ms": 15.201, "fwd_bwd_ms": 54.484, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d32", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d32_t512", "B": 4, "Lq": 512, "Lkv": 512, "H": 8, "D": 32, "masked": false, "impl": "flex:auto", "fwd_ms": 0.06, "fwd_bwd_ms": 0.126, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d32", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d32_t512", "B": 4, "Lq": 512, "Lkv": 512, "H": 8, "D": 32, "masked": false, "impl": "cudnn", "fwd_ms": 0.107, "fwd_bwd_ms": 0.177, "out_err": 0.0009765625, "grad_err": [ 0.0009765625, 0.001953125, 0.0009765625 ], "status": "ok" }, { "label": "pr123_mla", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "mla_128k36", "B": 1, "Lq": 65536, "Lkv": 131072, "H": 1, "D": 128, "masked": true, "impl": "flex:triton_onepass", "fwd_ms": 15.311, "fwd_bwd_ms": 79.417, "out_err": 0.0, "grad_err": [ 0.0001220703125, 0.0001220703125, 0.0001220703125 ], "status": "ok" }, { "label": "pr123_generic_d32", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d32_t512", "B": 4, "Lq": 512, "Lkv": 512, "H": 8, "D": 32, "masked": false, "impl": "flex:triton_onepass", "fwd_ms": 0.06, "fwd_bwd_ms": 0.114, "out_err": 0.0, "grad_err": [ 0.0009765625, 0.0009765625, 0.00048828125 ], "status": "ok" }, { "label": "pr123_generic_d32", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d32_t512", "B": 4, "Lq": 512, "Lkv": 512, "H": 8, "D": 32, "masked": false, "impl": "flex:triton_split:64,128,128,64", "fwd_ms": 0.061, "fwd_bwd_ms": 0.115, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_mla", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "mla_128k36", "B": 1, "Lq": 65536, "Lkv": 131072, "H": 1, "D": 128, "masked": true, "impl": "flex:auto", "fwd_ms": 15.289, "fwd_bwd_ms": 75.325, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d32", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d32_t2048", "B": 4, "Lq": 2048, "Lkv": 2048, "H": 8, "D": 32, "masked": false, "impl": "flex:triton_split", "fwd_ms": 0.158, "fwd_bwd_ms": 0.414, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d32", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d32_t2048", "B": 4, "Lq": 2048, "Lkv": 2048, "H": 8, "D": 32, "masked": false, "impl": "flex:triton_fused", "fwd_ms": 0.159, "fwd_bwd_ms": 0.415, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d32", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d32_t2048", "B": 4, "Lq": 2048, "Lkv": 2048, "H": 8, "D": 32, "masked": false, "impl": "flex:auto", "fwd_ms": 0.163, "fwd_bwd_ms": 0.42, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d32", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d32_t2048", "B": 4, "Lq": 2048, "Lkv": 2048, "H": 8, "D": 32, "masked": false, "impl": "cudnn", "fwd_ms": 0.177, "fwd_bwd_ms": 0.419, "out_err": 0.0009765625, "grad_err": [ 0.0009765625, 0.0009765625, 0.0009765625 ], "status": "ok" }, { "label": "pr123_generic_d32", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d32_t2048", "B": 4, "Lq": 2048, "Lkv": 2048, "H": 8, "D": 32, "masked": false, "impl": "flex:triton_onepass", "fwd_ms": 0.158, "fwd_bwd_ms": 0.437, "out_err": 0.0, "grad_err": [ 0.0009765625, 0.0009765625, 0.00048828125 ], "status": "ok" }, { "label": "pr123_generic_d32", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d32_t2048", "B": 4, "Lq": 2048, "Lkv": 2048, "H": 8, "D": 32, "masked": false, "impl": "flex:triton_split:64,128,128,64", "fwd_ms": 0.158, "fwd_bwd_ms": 0.394, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d32", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d32_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 32, "masked": false, "impl": "flex:triton_split", "fwd_ms": 1.394, "fwd_bwd_ms": 5.368, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d32", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d32_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 32, "masked": false, "impl": "flex:triton_fused", "fwd_ms": 1.394, "fwd_bwd_ms": 5.504, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d32", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d32_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 32, "masked": false, "impl": "flex:auto", "fwd_ms": 1.391, "fwd_bwd_ms": 5.526, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d32", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d32_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 32, "masked": false, "impl": "cudnn", "fwd_ms": 1.373, "fwd_bwd_ms": 4.771, "out_err": 0.00048828125, "grad_err": [ 0.00048828125, 0.00048828125, 0.0009765625 ], "status": "ok" }, { "label": "pr123_generic_d32", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d32_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 32, "masked": false, "impl": "flex:triton_onepass", "fwd_ms": 1.4, "fwd_bwd_ms": 5.905, "out_err": 0.0, "grad_err": [ 0.00048828125, 0.00048828125, 0.00048828125 ], "status": "ok" }, { "label": "pr123_generic_d32", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d32_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 32, "masked": false, "impl": "flex:triton_split:64,128,128,64", "fwd_ms": 1.42, "fwd_bwd_ms": 5.201, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d32", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d32_t16384", "B": 4, "Lq": 16384, "Lkv": 16384, "H": 8, "D": 32, "masked": false, "impl": "flex:triton_split", "fwd_ms": 5.26, "fwd_bwd_ms": 22.081, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d32", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d32_t16384", "B": 4, "Lq": 16384, "Lkv": 16384, "H": 8, "D": 32, "masked": false, "impl": "flex:triton_fused", "fwd_ms": 5.342, "fwd_bwd_ms": 22.068, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d32", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d32_t16384", "B": 4, "Lq": 16384, "Lkv": 16384, "H": 8, "D": 32, "masked": false, "impl": "flex:auto", "fwd_ms": 5.28, "fwd_bwd_ms": 22.023, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d32", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d32_t16384", "B": 4, "Lq": 16384, "Lkv": 16384, "H": 8, "D": 32, "masked": false, "impl": "cudnn", "fwd_ms": 5.18, "fwd_bwd_ms": 19.142, "out_err": 0.000244140625, "grad_err": [ 0.00048828125, 0.00048828125, 0.00048828125 ], "status": "ok" }, { "label": "pr123_generic_d32", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d32_t16384", "B": 4, "Lq": 16384, "Lkv": 16384, "H": 8, "D": 32, "masked": false, "impl": "flex:triton_onepass", "fwd_ms": 5.512, "fwd_bwd_ms": 24.235, "out_err": 0.0, "grad_err": [ 0.00048828125, 0.00048828125, 0.00048828125 ], "status": "ok" }, { "label": "pr123_generic_d32", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d32_t16384", "B": 4, "Lq": 16384, "Lkv": 16384, "H": 8, "D": 32, "masked": false, "impl": "flex:triton_split:64,128,128,64", "fwd_ms": 5.517, "fwd_bwd_ms": 22.554, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_mla", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "mla_1k756", "B": 4, "Lq": 262144, "Lkv": 1024, "H": 1, "D": 128, "masked": true, "impl": "flex:triton_split", "fwd_ms": 1.839, "fwd_bwd_ms": 11.316, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d32", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d32_t512", "B": 4, "Lq": 512, "Lkv": 512, "H": 8, "D": 32, "masked": true, "impl": "flex:triton_split", "fwd_ms": 0.062, "fwd_bwd_ms": 0.112, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d32", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d32_t512", "B": 4, "Lq": 512, "Lkv": 512, "H": 8, "D": 32, "masked": true, "impl": "flex:triton_fused", "fwd_ms": 0.064, "fwd_bwd_ms": 0.112, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d32", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d32_t512", "B": 4, "Lq": 512, "Lkv": 512, "H": 8, "D": 32, "masked": true, "impl": "flex:auto", "fwd_ms": 0.076, "fwd_bwd_ms": 0.11, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d32", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d32_t512", "B": 4, "Lq": 512, "Lkv": 512, "H": 8, "D": 32, "masked": true, "impl": "cudnn", "fwd_ms": 0.094, "fwd_bwd_ms": 0.165, "out_err": 0.00048828125, "grad_err": [ 0.001953125, 0.001953125, 0.001953125 ], "status": "ok" }, { "label": "pr123_generic_d32", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d32_t512", "B": 4, "Lq": 512, "Lkv": 512, "H": 8, "D": 32, "masked": true, "impl": "flex:triton_onepass", "fwd_ms": 0.077, "fwd_bwd_ms": 0.111, "out_err": 0.0, "grad_err": [ 0.0009765625, 0.001953125, 0.0009765625 ], "status": "ok" }, { "label": "pr123_generic_d32", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d32_t512", "B": 4, "Lq": 512, "Lkv": 512, "H": 8, "D": 32, "masked": true, "impl": "flex:triton_split:64,128,128,64", "fwd_ms": 0.063, "fwd_bwd_ms": 0.117, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d32", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d32_t2048", "B": 4, "Lq": 2048, "Lkv": 2048, "H": 8, "D": 32, "masked": true, "impl": "flex:triton_split", "fwd_ms": 0.161, "fwd_bwd_ms": 0.427, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d32", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d32_t2048", "B": 4, "Lq": 2048, "Lkv": 2048, "H": 8, "D": 32, "masked": true, "impl": "flex:triton_fused", "fwd_ms": 0.157, "fwd_bwd_ms": 0.463, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d32", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d32_t2048", "B": 4, "Lq": 2048, "Lkv": 2048, "H": 8, "D": 32, "masked": true, "impl": "flex:auto", "fwd_ms": 0.157, "fwd_bwd_ms": 0.463, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d32", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d32_t2048", "B": 4, "Lq": 2048, "Lkv": 2048, "H": 8, "D": 32, "masked": true, "impl": "cudnn", "fwd_ms": 0.184, "fwd_bwd_ms": 0.432, "out_err": 0.0009765625, "grad_err": [ 0.0009765625, 0.0009765625, 0.0009765625 ], "status": "ok" }, { "label": "pr123_generic_d32", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d32_t2048", "B": 4, "Lq": 2048, "Lkv": 2048, "H": 8, "D": 32, "masked": true, "impl": "flex:triton_onepass", "fwd_ms": 0.158, "fwd_bwd_ms": 0.443, "out_err": 0.0, "grad_err": [ 0.0009765625, 0.0009765625, 0.0009765625 ], "status": "ok" }, { "label": "pr123_generic_d64", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d64_t512", "B": 4, "Lq": 512, "Lkv": 512, "H": 8, "D": 64, "masked": false, "impl": "flex:triton_split", "fwd_ms": 0.063, "fwd_bwd_ms": 0.119, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d32", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d32_t2048", "B": 4, "Lq": 2048, "Lkv": 2048, "H": 8, "D": 32, "masked": true, "impl": "flex:triton_split:64,128,128,64", "fwd_ms": 0.16, "fwd_bwd_ms": 0.418, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d64", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d64_t512", "B": 4, "Lq": 512, "Lkv": 512, "H": 8, "D": 64, "masked": false, "impl": "flex:triton_fused", "fwd_ms": 0.079, "fwd_bwd_ms": 0.116, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d64", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d64_t512", "B": 4, "Lq": 512, "Lkv": 512, "H": 8, "D": 64, "masked": false, "impl": "flex:auto", "fwd_ms": 0.08, "fwd_bwd_ms": 0.119, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d64", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d64_t512", "B": 4, "Lq": 512, "Lkv": 512, "H": 8, "D": 64, "masked": false, "impl": "cudnn", "fwd_ms": 0.087, "fwd_bwd_ms": 0.165, "out_err": 0.00048828125, "grad_err": [ 0.0009765625, 0.001953125, 0.0009765625 ], "status": "ok" }, { "label": "pr123_generic_d32", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d32_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 32, "masked": true, "impl": "flex:triton_split", "fwd_ms": 1.45, "fwd_bwd_ms": 5.739, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d64", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d64_t512", "B": 4, "Lq": 512, "Lkv": 512, "H": 8, "D": 64, "masked": false, "impl": "flex:triton_onepass", "fwd_ms": 0.084, "fwd_bwd_ms": 0.129, "out_err": 0.0, "grad_err": [ 0.0009765625, 0.0009765625, 0.0009765625 ], "status": "ok" }, { "label": "pr123_generic_d32", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d32_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 32, "masked": true, "impl": "flex:triton_fused", "fwd_ms": 1.458, "fwd_bwd_ms": 6.648, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d64", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d64_t512", "B": 4, "Lq": 512, "Lkv": 512, "H": 8, "D": 64, "masked": false, "impl": "flex:triton_split:64,128,128,64", "fwd_ms": 0.081, "fwd_bwd_ms": 0.132, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d32", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d32_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 32, "masked": true, "impl": "flex:auto", "fwd_ms": 1.47, "fwd_bwd_ms": 6.558, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d32", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d32_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 32, "masked": true, "impl": "cudnn", "fwd_ms": 1.118, "fwd_bwd_ms": 3.716, "out_err": 0.00048828125, "grad_err": [ 0.0009765625, 0.0009765625, 0.0009765625 ], "status": "ok" }, { "label": "pr123_generic_d64", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d64_t2048", "B": 4, "Lq": 2048, "Lkv": 2048, "H": 8, "D": 64, "masked": false, "impl": "flex:triton_split", "fwd_ms": 0.187, "fwd_bwd_ms": 0.537, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d32", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d32_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 32, "masked": true, "impl": "flex:triton_onepass", "fwd_ms": 1.459, "fwd_bwd_ms": 5.81, "out_err": 0.0, "grad_err": [ 0.00048828125, 0.0009765625, 0.0009765625 ], "status": "ok" }, { "label": "pr123_mla", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "mla_1k756", "B": 4, "Lq": 262144, "Lkv": 1024, "H": 1, "D": 128, "masked": true, "impl": "flex:triton_onepass", "fwd_ms": 1.833, "fwd_bwd_ms": 19.989, "out_err": 0.0, "grad_err": [ 0.001953125, 0.015625, 0.0625 ], "status": "ok" }, { "label": "pr123_generic_d64", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d64_t2048", "B": 4, "Lq": 2048, "Lkv": 2048, "H": 8, "D": 64, "masked": false, "impl": "flex:triton_fused", "fwd_ms": 0.189, "fwd_bwd_ms": 0.574, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d32", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d32_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 32, "masked": true, "impl": "flex:triton_split:64,128,128,64", "fwd_ms": 1.462, "fwd_bwd_ms": 5.459, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d64", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d64_t2048", "B": 4, "Lq": 2048, "Lkv": 2048, "H": 8, "D": 64, "masked": false, "impl": "flex:auto", "fwd_ms": 0.189, "fwd_bwd_ms": 0.574, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d64", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d64_t2048", "B": 4, "Lq": 2048, "Lkv": 2048, "H": 8, "D": 64, "masked": false, "impl": "cudnn", "fwd_ms": 0.179, "fwd_bwd_ms": 0.461, "out_err": 0.00048828125, "grad_err": [ 0.0009765625, 0.0009765625, 0.0009765625 ], "status": "ok" }, { "label": "pr123_generic_d64", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d64_t2048", "B": 4, "Lq": 2048, "Lkv": 2048, "H": 8, "D": 64, "masked": false, "impl": "flex:triton_onepass", "fwd_ms": 0.192, "fwd_bwd_ms": 0.609, "out_err": 0.0, "grad_err": [ 0.0009765625, 0.0009765625, 0.0009765625 ], "status": "ok" }, { "label": "pr123_mla", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "mla_1k756", "B": 4, "Lq": 262144, "Lkv": 1024, "H": 1, "D": 128, "masked": true, "impl": "flex:auto", "fwd_ms": 1.869, "fwd_bwd_ms": 11.395, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d32", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d32_t16384", "B": 4, "Lq": 16384, "Lkv": 16384, "H": 8, "D": 32, "masked": true, "impl": "flex:triton_split", "fwd_ms": 5.493, "fwd_bwd_ms": 22.95, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d64", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d64_t2048", "B": 4, "Lq": 2048, "Lkv": 2048, "H": 8, "D": 64, "masked": false, "impl": "flex:triton_split:64,128,128,64", "fwd_ms": 0.186, "fwd_bwd_ms": 0.594, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d32", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d32_t16384", "B": 4, "Lq": 16384, "Lkv": 16384, "H": 8, "D": 32, "masked": true, "impl": "flex:triton_fused", "fwd_ms": 5.739, "fwd_bwd_ms": 26.626, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_isab", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "isab_rows_ind", "B": 56, "Lq": 32768, "Lkv": 256, "H": 8, "D": 32, "masked": false, "impl": "flex:triton_split", "fwd_ms": 2.829, "fwd_bwd_ms": 13.016, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d32", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d32_t16384", "B": 4, "Lq": 16384, "Lkv": 16384, "H": 8, "D": 32, "masked": true, "impl": "flex:auto", "fwd_ms": 5.776, "fwd_bwd_ms": 26.786, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d32", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d32_t16384", "B": 4, "Lq": 16384, "Lkv": 16384, "H": 8, "D": 32, "masked": true, "impl": "cudnn", "fwd_ms": 4.51, "fwd_bwd_ms": 16.265, "out_err": 0.000244140625, "grad_err": [ 0.00048828125, 0.0009765625, 0.00048828125 ], "status": "ok" }, { "label": "pr123_generic_d64", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d64_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 64, "masked": false, "impl": "flex:triton_split", "fwd_ms": 1.883, "fwd_bwd_ms": 7.489, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d32", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d32_t16384", "B": 4, "Lq": 16384, "Lkv": 16384, "H": 8, "D": 32, "masked": true, "impl": "flex:triton_onepass", "fwd_ms": 5.554, "fwd_bwd_ms": 24.51, "out_err": 0.0, "grad_err": [ 0.00048828125, 0.0009765625, 0.00048828125 ], "status": "ok" }, { "label": "pr123_generic_d64", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d64_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 64, "masked": false, "impl": "flex:triton_fused", "fwd_ms": 1.893, "fwd_bwd_ms": 8.491, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d32", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d32_t16384", "B": 4, "Lq": 16384, "Lkv": 16384, "H": 8, "D": 32, "masked": true, "impl": "flex:triton_split:64,128,128,64", "fwd_ms": 5.51, "fwd_bwd_ms": 22.953, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d64", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d64_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 64, "masked": false, "impl": "flex:auto", "fwd_ms": 1.839, "fwd_bwd_ms": 8.316, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d64", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d64_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 64, "masked": false, "impl": "cudnn", "fwd_ms": 1.435, "fwd_bwd_ms": 5.06, "out_err": 0.000244140625, "grad_err": [ 0.00048828125, 0.00048828125, 0.00048828125 ], "status": "ok" }, { "label": "pr123_generic_d32", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "cross_d32_256x16384", "B": 4, "Lq": 256, "Lkv": 16384, "H": 8, "D": 32, "masked": true, "impl": "flex:triton_split", "fwd_ms": 0.357, "fwd_bwd_ms": 0.831, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d64", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d64_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 64, "masked": false, "impl": "flex:triton_onepass", "fwd_ms": 1.858, "fwd_bwd_ms": 8.843, "out_err": 0.0, "grad_err": [ 0.00048828125, 0.00048828125, 0.00048828125 ], "status": "ok" }, { "label": "pr123_generic_d64", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d64_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 64, "masked": false, "impl": "flex:triton_split:64,128,128,64", "fwd_ms": 1.843, "fwd_bwd_ms": 8.648, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d64", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d64_t16384", "B": 4, "Lq": 16384, "Lkv": 16384, "H": 8, "D": 64, "masked": false, "impl": "flex:triton_split", "fwd_ms": 7.645, "fwd_bwd_ms": 34.164, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d64", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d64_t16384", "B": 4, "Lq": 16384, "Lkv": 16384, "H": 8, "D": 64, "masked": false, "impl": "flex:triton_fused", "fwd_ms": 7.104, "fwd_bwd_ms": 32.84, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d64", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d64_t16384", "B": 4, "Lq": 16384, "Lkv": 16384, "H": 8, "D": 64, "masked": false, "impl": "flex:auto", "fwd_ms": 7.134, "fwd_bwd_ms": 33.123, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d64", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d64_t16384", "B": 4, "Lq": 16384, "Lkv": 16384, "H": 8, "D": 64, "masked": false, "impl": "cudnn", "fwd_ms": 5.389, "fwd_bwd_ms": 20.604, "out_err": 0.000244140625, "grad_err": [ 0.00048828125, 0.00048828125, 0.00048828125 ], "status": "ok" }, { "label": "pr123_generic_d64", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d64_t16384", "B": 4, "Lq": 16384, "Lkv": 16384, "H": 8, "D": 64, "masked": false, "impl": "flex:triton_onepass", "fwd_ms": 7.814, "fwd_bwd_ms": 35.434, "out_err": 0.0, "grad_err": [ 0.00048828125, 0.00048828125, 0.00048828125 ], "status": "ok" }, { "label": "pr123_generic_d64", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d64_t16384", "B": 4, "Lq": 16384, "Lkv": 16384, "H": 8, "D": 64, "masked": false, "impl": "flex:triton_split:64,128,128,64", "fwd_ms": 7.78, "fwd_bwd_ms": 33.437, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d64", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d64_t512", "B": 4, "Lq": 512, "Lkv": 512, "H": 8, "D": 64, "masked": true, "impl": "flex:triton_split", "fwd_ms": 0.083, "fwd_bwd_ms": 0.122, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d128", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d128_t512", "B": 4, "Lq": 512, "Lkv": 512, "H": 8, "D": 128, "masked": false, "impl": "flex:triton_split", "fwd_ms": 0.089, "fwd_bwd_ms": 0.147, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d64", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d64_t512", "B": 4, "Lq": 512, "Lkv": 512, "H": 8, "D": 64, "masked": true, "impl": "flex:triton_fused", "fwd_ms": 0.082, "fwd_bwd_ms": 0.121, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d128", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d128_t512", "B": 4, "Lq": 512, "Lkv": 512, "H": 8, "D": 128, "masked": false, "impl": "flex:triton_fused", "fwd_ms": 0.088, "fwd_bwd_ms": 0.147, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d64", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d64_t512", "B": 4, "Lq": 512, "Lkv": 512, "H": 8, "D": 64, "masked": true, "impl": "flex:auto", "fwd_ms": 0.08, "fwd_bwd_ms": 0.118, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d64", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d64_t512", "B": 4, "Lq": 512, "Lkv": 512, "H": 8, "D": 64, "masked": true, "impl": "cudnn", "fwd_ms": 0.098, "fwd_bwd_ms": 0.17, "out_err": 0.00048828125, "grad_err": [ 0.001953125, 0.001953125, 0.001953125 ], "status": "ok" }, { "label": "pr123_generic_d128", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d128_t512", "B": 4, "Lq": 512, "Lkv": 512, "H": 8, "D": 128, "masked": false, "impl": "flex:auto", "fwd_ms": 0.089, "fwd_bwd_ms": 0.146, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d64", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d64_t512", "B": 4, "Lq": 512, "Lkv": 512, "H": 8, "D": 64, "masked": true, "impl": "flex:triton_onepass", "fwd_ms": 0.066, "fwd_bwd_ms": 0.126, "out_err": 0.0, "grad_err": [ 0.001953125, 0.0009765625, 0.001953125 ], "status": "ok" }, { "label": "pr123_generic_d128", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d128_t512", "B": 4, "Lq": 512, "Lkv": 512, "H": 8, "D": 128, "masked": false, "impl": "cudnn", "fwd_ms": 0.088, "fwd_bwd_ms": 0.19, "out_err": 0.001953125, "grad_err": [ 0.001953125, 0.001953125, 0.0009765625 ], "status": "ok" }, { "label": "pr123_generic_d64", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d64_t512", "B": 4, "Lq": 512, "Lkv": 512, "H": 8, "D": 64, "masked": true, "impl": "flex:triton_split:64,128,128,64", "fwd_ms": 0.085, "fwd_bwd_ms": 0.141, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d128", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d128_t512", "B": 4, "Lq": 512, "Lkv": 512, "H": 8, "D": 128, "masked": false, "impl": "flex:triton_onepass", "fwd_ms": 0.089, "fwd_bwd_ms": 0.188, "out_err": 0.0, "grad_err": [ 0.0009765625, 0.0009765625, 0.0009765625 ], "status": "ok" }, { "label": "pr123_generic_d128", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d128_t512", "B": 4, "Lq": 512, "Lkv": 512, "H": 8, "D": 128, "masked": false, "impl": "flex:triton_split:64,128,128,64", "fwd_ms": 0.092, "fwd_bwd_ms": 0.218, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d64", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d64_t2048", "B": 4, "Lq": 2048, "Lkv": 2048, "H": 8, "D": 64, "masked": true, "impl": "flex:triton_split", "fwd_ms": 0.195, "fwd_bwd_ms": 0.548, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d64", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d64_t2048", "B": 4, "Lq": 2048, "Lkv": 2048, "H": 8, "D": 64, "masked": true, "impl": "flex:triton_fused", "fwd_ms": 0.197, "fwd_bwd_ms": 0.585, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d128", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d128_t2048", "B": 4, "Lq": 2048, "Lkv": 2048, "H": 8, "D": 128, "masked": false, "impl": "flex:triton_split", "fwd_ms": 0.296, "fwd_bwd_ms": 0.89, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d64", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d64_t2048", "B": 4, "Lq": 2048, "Lkv": 2048, "H": 8, "D": 64, "masked": true, "impl": "flex:auto", "fwd_ms": 0.195, "fwd_bwd_ms": 0.59, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d64", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d64_t2048", "B": 4, "Lq": 2048, "Lkv": 2048, "H": 8, "D": 64, "masked": true, "impl": "cudnn", "fwd_ms": 0.182, "fwd_bwd_ms": 0.472, "out_err": 0.00048828125, "grad_err": [ 0.0009765625, 0.0009765625, 0.0009765625 ], "status": "ok" }, { "label": "pr123_generic_d128", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d128_t2048", "B": 4, "Lq": 2048, "Lkv": 2048, "H": 8, "D": 128, "masked": false, "impl": "flex:triton_fused", "fwd_ms": 0.295, "fwd_bwd_ms": 0.896, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d64", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d64_t2048", "B": 4, "Lq": 2048, "Lkv": 2048, "H": 8, "D": 64, "masked": true, "impl": "flex:triton_onepass", "fwd_ms": 0.195, "fwd_bwd_ms": 0.603, "out_err": 0.0, "grad_err": [ 0.0009765625, 0.0009765625, 0.0009765625 ], "status": "ok" }, { "label": "pr123_generic_d128", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d128_t2048", "B": 4, "Lq": 2048, "Lkv": 2048, "H": 8, "D": 128, "masked": false, "impl": "flex:auto", "fwd_ms": 0.298, "fwd_bwd_ms": 0.896, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d128", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d128_t2048", "B": 4, "Lq": 2048, "Lkv": 2048, "H": 8, "D": 128, "masked": false, "impl": "cudnn", "fwd_ms": 0.221, "fwd_bwd_ms": 0.697, "out_err": 0.0009765625, "grad_err": [ 0.0009765625, 0.0009765625, 0.0009765625 ], "status": "ok" }, { "label": "pr123_generic_d64", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d64_t2048", "B": 4, "Lq": 2048, "Lkv": 2048, "H": 8, "D": 64, "masked": true, "impl": "flex:triton_split:64,128,128,64", "fwd_ms": 0.195, "fwd_bwd_ms": 0.642, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d128", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d128_t2048", "B": 4, "Lq": 2048, "Lkv": 2048, "H": 8, "D": 128, "masked": false, "impl": "flex:triton_onepass", "fwd_ms": 0.291, "fwd_bwd_ms": 1.45, "out_err": 0.0, "grad_err": [ 0.0009765625, 0.0009765625, 0.0009765625 ], "status": "ok" }, { "label": "pr123_generic_d128", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d128_t2048", "B": 4, "Lq": 2048, "Lkv": 2048, "H": 8, "D": 128, "masked": false, "impl": "flex:triton_split:64,128,128,64", "fwd_ms": 0.292, "fwd_bwd_ms": 1.948, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_isab", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "isab_rows_ind", "B": 56, "Lq": 32768, "Lkv": 256, "H": 8, "D": 32, "masked": false, "impl": "flex:triton_onepass", "fwd_ms": 2.81, "fwd_bwd_ms": 19.257, "out_err": 0.0, "grad_err": [ 0.00390625, 0.03125, 0.03125 ], "status": "ok" }, { "label": "pr123_generic_d64", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d64_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 64, "masked": true, "impl": "flex:triton_split", "fwd_ms": 1.895, "fwd_bwd_ms": 8.308, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d64", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d64_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 64, "masked": true, "impl": "flex:triton_fused", "fwd_ms": 1.903, "fwd_bwd_ms": 8.433, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d64", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d64_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 64, "masked": true, "impl": "flex:auto", "fwd_ms": 1.927, "fwd_bwd_ms": 8.35, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d128", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d128_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 128, "masked": false, "impl": "flex:triton_split", "fwd_ms": 3.223, "fwd_bwd_ms": 13.014, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d64", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d64_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 64, "masked": true, "impl": "cudnn", "fwd_ms": 1.155, "fwd_bwd_ms": 4.048, "out_err": 0.00048828125, "grad_err": [ 0.00048828125, 0.0009765625, 0.0009765625 ], "status": "ok" }, { "label": "pr123_generic_d64", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d64_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 64, "masked": true, "impl": "flex:triton_onepass", "fwd_ms": 1.905, "fwd_bwd_ms": 8.607, "out_err": 0.0, "grad_err": [ 0.00048828125, 0.00048828125, 0.0009765625 ], "status": "ok" }, { "label": "pr123_isab", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "isab_rows_ind", "B": 56, "Lq": 32768, "Lkv": 256, "H": 8, "D": 32, "masked": false, "impl": "flex:auto", "fwd_ms": 2.835, "fwd_bwd_ms": 13.032, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d128", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d128_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 128, "masked": false, "impl": "flex:triton_fused", "fwd_ms": 3.192, "fwd_bwd_ms": 13.525, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d64", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d64_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 64, "masked": true, "impl": "flex:triton_split:64,128,128,64", "fwd_ms": 1.897, "fwd_bwd_ms": 9.392, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d128", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d128_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 128, "masked": false, "impl": "flex:auto", "fwd_ms": 3.221, "fwd_bwd_ms": 13.274, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d128", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d128_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 128, "masked": false, "impl": "cudnn", "fwd_ms": 2.199, "fwd_bwd_ms": 8.315, "out_err": 0.00048828125, "grad_err": [ 0.00048828125, 0.00048828125, 0.00048828125 ], "status": "ok" }, { "label": "pr123_generic_d64", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d64_t16384", "B": 4, "Lq": 16384, "Lkv": 16384, "H": 8, "D": 64, "masked": true, "impl": "flex:triton_split", "fwd_ms": 7.736, "fwd_bwd_ms": 33.468, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d128", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d128_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 128, "masked": false, "impl": "flex:triton_onepass", "fwd_ms": 3.206, "fwd_bwd_ms": 21.258, "out_err": 0.0, "grad_err": [ 0.00048828125, 0.00048828125, 0.00048828125 ], "status": "ok" }, { "label": "pr123_generic_d64", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d64_t16384", "B": 4, "Lq": 16384, "Lkv": 16384, "H": 8, "D": 64, "masked": true, "impl": "flex:triton_fused", "fwd_ms": 7.871, "fwd_bwd_ms": 33.268, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d64", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d64_t16384", "B": 4, "Lq": 16384, "Lkv": 16384, "H": 8, "D": 64, "masked": true, "impl": "flex:auto", "fwd_ms": 7.781, "fwd_bwd_ms": 33.584, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d128", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d128_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 128, "masked": false, "impl": "flex:triton_split:64,128,128,64", "fwd_ms": 3.226, "fwd_bwd_ms": 30.473, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d64", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d64_t16384", "B": 4, "Lq": 16384, "Lkv": 16384, "H": 8, "D": 64, "masked": true, "impl": "cudnn", "fwd_ms": 4.712, "fwd_bwd_ms": 16.887, "out_err": 0.000244140625, "grad_err": [ 0.00048828125, 0.00048828125, 0.00048828125 ], "status": "ok" }, { "label": "pr123_generic_d64", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d64_t16384", "B": 4, "Lq": 16384, "Lkv": 16384, "H": 8, "D": 64, "masked": true, "impl": "flex:triton_onepass", "fwd_ms": 7.73, "fwd_bwd_ms": 34.193, "out_err": 0.0, "grad_err": [ 0.00048828125, 0.00048828125, 0.00048828125 ], "status": "ok" }, { "label": "pr123_generic_d128", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d128_t16384", "B": 4, "Lq": 16384, "Lkv": 16384, "H": 8, "D": 128, "masked": false, "impl": "flex:triton_split", "fwd_ms": 14.061, "fwd_bwd_ms": 54.519, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d64", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d64_t16384", "B": 4, "Lq": 16384, "Lkv": 16384, "H": 8, "D": 64, "masked": true, "impl": "flex:triton_split:64,128,128,64", "fwd_ms": 7.888, "fwd_bwd_ms": 36.101, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d64", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "cross_d64_256x16384", "B": 4, "Lq": 256, "Lkv": 16384, "H": 8, "D": 64, "masked": true, "impl": "flex:triton_split", "fwd_ms": 0.399, "fwd_bwd_ms": 0.981, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_isab", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "isab_1k_ind", "B": 1008, "Lq": 1024, "Lkv": 256, "H": 8, "D": 32, "masked": false, "impl": "flex:triton_split", "fwd_ms": 1.78, "fwd_bwd_ms": 10.118, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d128", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d128_t16384", "B": 4, "Lq": 16384, "Lkv": 16384, "H": 8, "D": 128, "masked": false, "impl": "flex:triton_fused", "fwd_ms": 13.779, "fwd_bwd_ms": 57.446, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d128", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d128_t16384", "B": 4, "Lq": 16384, "Lkv": 16384, "H": 8, "D": 128, "masked": false, "impl": "flex:auto", "fwd_ms": 14.155, "fwd_bwd_ms": 54.175, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d128", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d128_t16384", "B": 4, "Lq": 16384, "Lkv": 16384, "H": 8, "D": 128, "masked": false, "impl": "cudnn", "fwd_ms": 8.835, "fwd_bwd_ms": 34.161, "out_err": 0.00048828125, "grad_err": [ 0.00048828125, 0.00048828125, 0.00048828125 ], "status": "ok" }, { "label": "pr123_generic_d128", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d128_t16384", "B": 4, "Lq": 16384, "Lkv": 16384, "H": 8, "D": 128, "masked": false, "impl": "flex:triton_onepass", "fwd_ms": 14.174, "fwd_bwd_ms": 83.593, "out_err": 0.0, "grad_err": [ 0.00048828125, 0.000244140625, 0.00048828125 ], "status": "ok" }, { "label": "pr123_isab", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "isab_1k_ind", "B": 1008, "Lq": 1024, "Lkv": 256, "H": 8, "D": 32, "masked": false, "impl": "flex:triton_onepass", "fwd_ms": 1.766, "fwd_bwd_ms": 11.557, "out_err": 0.0, "grad_err": [ 0.00390625, 0.0078125, 0.00390625 ], "status": "ok" }, { "label": "pr123_generic_d128", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d128_t16384", "B": 4, "Lq": 16384, "Lkv": 16384, "H": 8, "D": 128, "masked": false, "impl": "flex:triton_split:64,128,128,64", "fwd_ms": 13.947, "fwd_bwd_ms": 120.909, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d128", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d128_t512", "B": 4, "Lq": 512, "Lkv": 512, "H": 8, "D": 128, "masked": true, "impl": "flex:triton_split", "fwd_ms": 0.085, "fwd_bwd_ms": 0.138, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d128", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d128_t512", "B": 4, "Lq": 512, "Lkv": 512, "H": 8, "D": 128, "masked": true, "impl": "flex:triton_fused", "fwd_ms": 0.095, "fwd_bwd_ms": 0.139, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d128", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d128_t512", "B": 4, "Lq": 512, "Lkv": 512, "H": 8, "D": 128, "masked": true, "impl": "flex:auto", "fwd_ms": 0.086, "fwd_bwd_ms": 0.141, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d128", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d128_t512", "B": 4, "Lq": 512, "Lkv": 512, "H": 8, "D": 128, "masked": true, "impl": "cudnn", "fwd_ms": 0.094, "fwd_bwd_ms": 0.188, "out_err": 0.0009765625, "grad_err": [ 0.001953125, 0.001953125, 0.001953125 ], "status": "ok" }, { "label": "pr123_generic_d128", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d128_t512", "B": 4, "Lq": 512, "Lkv": 512, "H": 8, "D": 128, "masked": true, "impl": "flex:triton_onepass", "fwd_ms": 0.086, "fwd_bwd_ms": 0.175, "out_err": 0.0, "grad_err": [ 0.0009765625, 0.001953125, 0.001953125 ], "status": "ok" }, { "label": "pr123_generic_d128", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d128_t512", "B": 4, "Lq": 512, "Lkv": 512, "H": 8, "D": 128, "masked": true, "impl": "flex:triton_split:64,128,128,64", "fwd_ms": 0.087, "fwd_bwd_ms": 0.226, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d128", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d128_t2048", "B": 4, "Lq": 2048, "Lkv": 2048, "H": 8, "D": 128, "masked": true, "impl": "flex:triton_split", "fwd_ms": 0.307, "fwd_bwd_ms": 0.909, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d128", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d128_t2048", "B": 4, "Lq": 2048, "Lkv": 2048, "H": 8, "D": 128, "masked": true, "impl": "flex:triton_fused", "fwd_ms": 0.304, "fwd_bwd_ms": 0.921, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d128", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d128_t2048", "B": 4, "Lq": 2048, "Lkv": 2048, "H": 8, "D": 128, "masked": true, "impl": "flex:auto", "fwd_ms": 0.304, "fwd_bwd_ms": 0.918, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d128", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d128_t2048", "B": 4, "Lq": 2048, "Lkv": 2048, "H": 8, "D": 128, "masked": true, "impl": "cudnn", "fwd_ms": 0.216, "fwd_bwd_ms": 0.689, "out_err": 0.00048828125, "grad_err": [ 0.0009765625, 0.0009765625, 0.0009765625 ], "status": "ok" }, { "label": "pr123_generic_d128", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d128_t2048", "B": 4, "Lq": 2048, "Lkv": 2048, "H": 8, "D": 128, "masked": true, "impl": "flex:triton_onepass", "fwd_ms": 0.306, "fwd_bwd_ms": 1.491, "out_err": 0.0, "grad_err": [ 0.0009765625, 0.0009765625, 0.0009765625 ], "status": "ok" }, { "label": "pr123_generic_d128", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d128_t2048", "B": 4, "Lq": 2048, "Lkv": 2048, "H": 8, "D": 128, "masked": true, "impl": "flex:triton_split:64,128,128,64", "fwd_ms": 0.304, "fwd_bwd_ms": 2.092, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d128", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d128_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 128, "masked": true, "impl": "flex:triton_split", "fwd_ms": 3.376, "fwd_bwd_ms": 13.743, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d128", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d128_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 128, "masked": true, "impl": "flex:triton_fused", "fwd_ms": 3.375, "fwd_bwd_ms": 14.333, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d128", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d128_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 128, "masked": true, "impl": "flex:auto", "fwd_ms": 3.369, "fwd_bwd_ms": 13.671, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d128", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d128_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 128, "masked": true, "impl": "cudnn", "fwd_ms": 1.721, "fwd_bwd_ms": 6.562, "out_err": 0.00048828125, "grad_err": [ 0.0009765625, 0.0009765625, 0.0009765625 ], "status": "ok" }, { "label": "pr123_isab", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "isab_1k_ind", "B": 1008, "Lq": 1024, "Lkv": 256, "H": 8, "D": 32, "masked": false, "impl": "flex:auto", "fwd_ms": 1.78, "fwd_bwd_ms": 19.447, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d128", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d128_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 128, "masked": true, "impl": "flex:triton_onepass", "fwd_ms": 3.364, "fwd_bwd_ms": 21.272, "out_err": 0.0, "grad_err": [ 0.0009765625, 0.0009765625, 0.0009765625 ], "status": "ok" }, { "label": "pr123_isab", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "isab_ind_1k", "B": 1008, "Lq": 256, "Lkv": 1024, "H": 8, "D": 32, "masked": true, "impl": "flex:triton_split", "fwd_ms": 1.684, "fwd_bwd_ms": 9.861, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d128", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d128_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 128, "masked": true, "impl": "flex:triton_split:64,128,128,64", "fwd_ms": 3.374, "fwd_bwd_ms": 32.987, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_isab", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "isab_ind_1k", "B": 1008, "Lq": 256, "Lkv": 1024, "H": 8, "D": 32, "masked": true, "impl": "flex:triton_onepass", "fwd_ms": 1.68, "fwd_bwd_ms": 11.218, "out_err": 0.0, "grad_err": [ 0.001953125, 0.001953125, 0.0009765625 ], "status": "ok" }, { "label": "pr123_generic_d128", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d128_t16384", "B": 4, "Lq": 16384, "Lkv": 16384, "H": 8, "D": 128, "masked": true, "impl": "flex:triton_split", "fwd_ms": 15.069, "fwd_bwd_ms": 54.992, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_isab", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "isab_ind_1k", "B": 1008, "Lq": 256, "Lkv": 1024, "H": 8, "D": 32, "masked": true, "impl": "flex:auto", "fwd_ms": 1.67, "fwd_bwd_ms": 18.336, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d128", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d128_t16384", "B": 4, "Lq": 16384, "Lkv": 16384, "H": 8, "D": 128, "masked": true, "impl": "flex:triton_fused", "fwd_ms": 14.789, "fwd_bwd_ms": 54.874, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d128", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d128_t16384", "B": 4, "Lq": 16384, "Lkv": 16384, "H": 8, "D": 128, "masked": true, "impl": "flex:auto", "fwd_ms": 15.027, "fwd_bwd_ms": 55.229, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d128", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d128_t16384", "B": 4, "Lq": 16384, "Lkv": 16384, "H": 8, "D": 128, "masked": true, "impl": "cudnn", "fwd_ms": 7.116, "fwd_bwd_ms": 27.939, "out_err": 0.000244140625, "grad_err": [ 0.00048828125, 0.00048828125, 0.00048828125 ], "status": "ok" }, { "label": "pr123_generic_d128", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d128_t16384", "B": 4, "Lq": 16384, "Lkv": 16384, "H": 8, "D": 128, "masked": true, "impl": "flex:triton_onepass", "fwd_ms": 14.738, "fwd_bwd_ms": 83.654, "out_err": 0.0, "grad_err": [ 0.00048828125, 0.00048828125, 0.00048828125 ], "status": "ok" }, { "label": "pr123_generic_d128", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d128_t16384", "B": 4, "Lq": 16384, "Lkv": 16384, "H": 8, "D": 128, "masked": true, "impl": "flex:triton_split:64,128,128,64", "fwd_ms": 14.941, "fwd_bwd_ms": 132.526, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d128", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "cross_d128_256x16384", "B": 4, "Lq": 256, "Lkv": 16384, "H": 8, "D": 128, "masked": true, "impl": "flex:triton_split", "fwd_ms": 0.532, "fwd_bwd_ms": 1.378, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d32", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "cross_d32_256x16384", "B": 4, "Lq": 256, "Lkv": 16384, "H": 8, "D": 32, "masked": true, "impl": "flex:triton_fused", "fwd_ms": 183.989, "fwd_bwd_ms": 25429.915, "out_err": 0.000244140625, "grad_err": [ 0.00048828125, 6.103515625e-05, 6.103515625e-05 ], "status": "ok" }, { "label": "pr123_generic_d32", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "cross_d32_256x16384", "B": 4, "Lq": 256, "Lkv": 16384, "H": 8, "D": 32, "masked": true, "impl": "flex:auto", "fwd_ms": 0.355, "fwd_bwd_ms": 0.831, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d32", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "cross_d32_256x16384", "B": 4, "Lq": 256, "Lkv": 16384, "H": 8, "D": 32, "masked": true, "impl": "cudnn", "fwd_ms": 0.286, "fwd_bwd_ms": 0.563, "out_err": 0.0001220703125, "grad_err": [ 0.000244140625, 0.0001220703125, 6.103515625e-05 ], "status": "ok" }, { "label": "pr123_generic_d32", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "cross_d32_256x16384", "B": 4, "Lq": 256, "Lkv": 16384, "H": 8, "D": 32, "masked": true, "impl": "flex:triton_onepass", "fwd_ms": 0.355, "fwd_bwd_ms": 0.653, "out_err": 0.0, "grad_err": [ 0.000244140625, 6.103515625e-05, 3.0517578125e-05 ], "status": "ok" }, { "label": "pr123_generic_d32", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "cross_d32_256x16384", "B": 4, "Lq": 256, "Lkv": 16384, "H": 8, "D": 32, "masked": true, "impl": "flex:triton_split:64,128,128,64", "fwd_ms": 0.357, "fwd_bwd_ms": 0.905, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d32", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "cross_d32_16384x256", "B": 4, "Lq": 16384, "Lkv": 256, "H": 8, "D": 32, "masked": true, "impl": "flex:triton_split", "fwd_ms": 0.18, "fwd_bwd_ms": 0.666, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d64", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "cross_d64_256x16384", "B": 4, "Lq": 256, "Lkv": 16384, "H": 8, "D": 64, "masked": true, "impl": "flex:triton_fused", "fwd_ms": 199.126, "fwd_bwd_ms": 28183.424, "out_err": 0.000244140625, "grad_err": [ 0.00048828125, 0.0001220703125, 6.103515625e-05 ], "status": "ok" }, { "label": "pr123_generic_d64", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "cross_d64_256x16384", "B": 4, "Lq": 256, "Lkv": 16384, "H": 8, "D": 64, "masked": true, "impl": "flex:auto", "fwd_ms": 0.404, "fwd_bwd_ms": 0.982, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d64", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "cross_d64_256x16384", "B": 4, "Lq": 256, "Lkv": 16384, "H": 8, "D": 64, "masked": true, "impl": "cudnn", "fwd_ms": 0.291, "fwd_bwd_ms": 0.632, "out_err": 0.000244140625, "grad_err": [ 0.000244140625, 6.103515625e-05, 6.103515625e-05 ], "status": "ok" }, { "label": "pr123_generic_d64", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "cross_d64_256x16384", "B": 4, "Lq": 256, "Lkv": 16384, "H": 8, "D": 64, "masked": true, "impl": "flex:triton_onepass", "fwd_ms": 0.398, "fwd_bwd_ms": 0.821, "out_err": 0.0, "grad_err": [ 0.000244140625, 6.103515625e-05, 6.103515625e-05 ], "status": "ok" }, { "label": "pr123_generic_d64", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "cross_d64_256x16384", "B": 4, "Lq": 256, "Lkv": 16384, "H": 8, "D": 64, "masked": true, "impl": "flex:triton_split:64,128,128,64", "fwd_ms": 0.398, "fwd_bwd_ms": 1.196, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d64", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "cross_d64_16384x256", "B": 4, "Lq": 16384, "Lkv": 256, "H": 8, "D": 64, "masked": true, "impl": "flex:triton_split", "fwd_ms": 0.22, "fwd_bwd_ms": 0.814, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d128", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "cross_d128_256x16384", "B": 4, "Lq": 256, "Lkv": 16384, "H": 8, "D": 128, "masked": true, "impl": "flex:triton_fused", "fwd_ms": 178.213, "fwd_bwd_ms": 25686.06, "out_err": 0.00048828125, "grad_err": [ 0.00048828125, 0.0001220703125, 6.103515625e-05 ], "status": "ok" }, { "label": "pr123_generic_d128", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "cross_d128_256x16384", "B": 4, "Lq": 256, "Lkv": 16384, "H": 8, "D": 128, "masked": true, "impl": "flex:auto", "fwd_ms": 0.523, "fwd_bwd_ms": 1.361, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d128", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "cross_d128_256x16384", "B": 4, "Lq": 256, "Lkv": 16384, "H": 8, "D": 128, "masked": true, "impl": "cudnn", "fwd_ms": 0.347, "fwd_bwd_ms": 0.889, "out_err": 0.0001220703125, "grad_err": [ 0.00048828125, 6.103515625e-05, 6.103515625e-05 ], "status": "ok" }, { "label": "pr123_generic_d128", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "cross_d128_256x16384", "B": 4, "Lq": 256, "Lkv": 16384, "H": 8, "D": 128, "masked": true, "impl": "flex:triton_onepass", "fwd_ms": 0.527, "fwd_bwd_ms": 1.646, "out_err": 0.0, "grad_err": [ 0.00048828125, 6.103515625e-05, 6.103515625e-05 ], "status": "ok" }, { "label": "pr123_generic_d128", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "cross_d128_256x16384", "B": 4, "Lq": 256, "Lkv": 16384, "H": 8, "D": 128, "masked": true, "impl": "flex:triton_split:64,128,128,64", "fwd_ms": 0.519, "fwd_bwd_ms": 3.053, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d128", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "cross_d128_16384x256", "B": 4, "Lq": 16384, "Lkv": 256, "H": 8, "D": 128, "masked": true, "impl": "flex:triton_split", "fwd_ms": 0.4, "fwd_bwd_ms": 1.373, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d32", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "cross_d32_16384x256", "B": 4, "Lq": 16384, "Lkv": 256, "H": 8, "D": 32, "masked": true, "impl": "flex:triton_fused", "fwd_ms": 353.182, "fwd_bwd_ms": 33921.567, "out_err": 0.001953125, "grad_err": [ 0.001953125, 0.03125, 0.03125 ], "status": "ok" }, { "label": "pr123_generic_d32", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "cross_d32_16384x256", "B": 4, "Lq": 16384, "Lkv": 256, "H": 8, "D": 32, "masked": true, "impl": "flex:auto", "fwd_ms": 0.185, "fwd_bwd_ms": 0.68, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d32", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "cross_d32_16384x256", "B": 4, "Lq": 16384, "Lkv": 256, "H": 8, "D": 32, "masked": true, "impl": "cudnn", "fwd_ms": 0.226, "fwd_bwd_ms": 0.851, "out_err": 0.001953125, "grad_err": [ 0.001953125, 0.03125, 0.03125 ], "status": "ok" }, { "label": "pr123_generic_d32", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "cross_d32_16384x256", "B": 4, "Lq": 16384, "Lkv": 256, "H": 8, "D": 32, "masked": true, "impl": "flex:triton_onepass", "fwd_ms": 0.185, "fwd_bwd_ms": 0.793, "out_err": 0.0, "grad_err": [ 0.001953125, 0.015625, 0.015625 ], "status": "ok" }, { "label": "pr123_generic_d32", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "cross_d32_16384x256", "B": 4, "Lq": 16384, "Lkv": 256, "H": 8, "D": 32, "masked": true, "impl": "flex:triton_split:64,128,128,64", "fwd_ms": 0.188, "fwd_bwd_ms": 0.842, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d32", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "cross_d32_1024x65536", "B": 4, "Lq": 1024, "Lkv": 65536, "H": 8, "D": 32, "masked": true, "impl": "flex:triton_split", "fwd_ms": 1.572, "fwd_bwd_ms": 5.979, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d64", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "cross_d64_16384x256", "B": 4, "Lq": 16384, "Lkv": 256, "H": 8, "D": 64, "masked": true, "impl": "flex:triton_fused", "fwd_ms": 394.404, "fwd_bwd_ms": 41932.816, "out_err": 0.00390625, "grad_err": [ 0.00390625, 0.03125, 0.03125 ], "status": "ok" }, { "label": "pr123_generic_d64", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "cross_d64_16384x256", "B": 4, "Lq": 16384, "Lkv": 256, "H": 8, "D": 64, "masked": true, "impl": "flex:auto", "fwd_ms": 0.216, "fwd_bwd_ms": 0.818, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d64", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "cross_d64_16384x256", "B": 4, "Lq": 16384, "Lkv": 256, "H": 8, "D": 64, "masked": true, "impl": "cudnn", "fwd_ms": 0.242, "fwd_bwd_ms": 1.037, "out_err": 0.001953125, "grad_err": [ 0.00390625, 0.015625, 0.03125 ], "status": "ok" }, { "label": "pr123_generic_d64", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "cross_d64_16384x256", "B": 4, "Lq": 16384, "Lkv": 256, "H": 8, "D": 64, "masked": true, "impl": "flex:triton_onepass", "fwd_ms": 0.214, "fwd_bwd_ms": 1.059, "out_err": 0.0, "grad_err": [ 0.00390625, 0.015625, 0.015625 ], "status": "ok" }, { "label": "pr123_generic_d64", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "cross_d64_16384x256", "B": 4, "Lq": 16384, "Lkv": 256, "H": 8, "D": 64, "masked": true, "impl": "flex:triton_split:64,128,128,64", "fwd_ms": 0.225, "fwd_bwd_ms": 1.311, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d64", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "cross_d64_1024x65536", "B": 4, "Lq": 1024, "Lkv": 65536, "H": 8, "D": 64, "masked": true, "impl": "flex:triton_split", "fwd_ms": 1.995, "fwd_bwd_ms": 8.275, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d128", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "cross_d128_16384x256", "B": 4, "Lq": 16384, "Lkv": 256, "H": 8, "D": 128, "masked": true, "impl": "flex:triton_fused", "fwd_ms": 342.877, "fwd_bwd_ms": 37171.267, "out_err": 0.00390625, "grad_err": [ 0.00390625, 0.03125, 0.03125 ], "status": "ok" }, { "label": "pr123_generic_d128", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "cross_d128_16384x256", "B": 4, "Lq": 16384, "Lkv": 256, "H": 8, "D": 128, "masked": true, "impl": "flex:auto", "fwd_ms": 0.399, "fwd_bwd_ms": 1.383, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d128", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "cross_d128_16384x256", "B": 4, "Lq": 16384, "Lkv": 256, "H": 8, "D": 128, "masked": true, "impl": "cudnn", "fwd_ms": 0.331, "fwd_bwd_ms": 1.619, "out_err": 0.001953125, "grad_err": [ 0.00390625, 0.015625, 0.03125 ], "status": "ok" }, { "label": "pr123_generic_d128", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "cross_d128_16384x256", "B": 4, "Lq": 16384, "Lkv": 256, "H": 8, "D": 128, "masked": true, "impl": "flex:triton_onepass", "fwd_ms": 0.398, "fwd_bwd_ms": 1.943, "out_err": 0.0, "grad_err": [ 0.001953125, 0.015625, 0.015625 ], "status": "ok" }, { "label": "h100br_smallq", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "sq_isab_256x131072", "B": 28, "Lq": 256, "Lkv": 131072, "H": 8, "D": 32, "masked": true, "impl": "flex:default", "fwd_ms": 5.503, "fwd_bwd_ms": 20.03, "out_err": 0.0001220703125, "grad_err": [ 0.0001220703125, 3.0517578125e-05, 1.52587890625e-05 ], "status": "ok" }, { "label": "pr123_generic_d128", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "cross_d128_16384x256", "B": 4, "Lq": 16384, "Lkv": 256, "H": 8, "D": 128, "masked": true, "impl": "flex:triton_split:64,128,128,64", "fwd_ms": 0.397, "fwd_bwd_ms": 3.475, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "h100br_smallq", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "sq_isab_256x131072", "B": 28, "Lq": 256, "Lkv": 131072, "H": 8, "D": 32, "masked": true, "impl": "flex:triton_generic", "fwd_ms": 6.142, "fwd_bwd_ms": 26.684, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "h100br_smallq", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "sq_isab_256x131072", "B": 28, "Lq": 256, "Lkv": 131072, "H": 8, "D": 32, "masked": true, "impl": "flex:triton_split", "fwd_ms": 6.12, "fwd_bwd_ms": 26.77, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "pr123_generic_d128", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "cross_d128_1024x65536", "B": 4, "Lq": 1024, "Lkv": 65536, "H": 8, "D": 128, "masked": true, "impl": "flex:triton_split", "fwd_ms": 3.694, "fwd_bwd_ms": 14.786, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "h100br_smallq", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "sq_isab_256x131072", "B": 28, "Lq": 256, "Lkv": 131072, "H": 8, "D": 32, "masked": true, "impl": "flex:triton_onepass", "fwd_ms": 6.253, "fwd_bwd_ms": 24.313, "out_err": 0.0, "grad_err": [ 0.0001220703125, 3.0517578125e-05, 1.52587890625e-05 ], "status": "ok" }, { "label": "h100br_smallq", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "sq_isab_256x131072", "B": 28, "Lq": 256, "Lkv": 131072, "H": 8, "D": 32, "masked": true, "impl": "flex:triton_fused", "fwd_ms": 6.202, "fwd_bwd_ms": 26.931, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "h100br_smallq", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "sq_isab_256x131072", "B": 28, "Lq": 256, "Lkv": 131072, "H": 8, "D": 32, "masked": true, "impl": "cudnn", "fwd_ms": 5.432, "fwd_bwd_ms": 20.095, "out_err": 0.0001220703125, "grad_err": [ 0.0001220703125, 3.0517578125e-05, 1.52587890625e-05 ], "status": "ok" }, { "label": "h100br_smallq", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "sq_isab_256x131072", "B": 28, "Lq": 256, "Lkv": 131072, "H": 8, "D": 32, "masked": true, "impl": "mosaic_fa3", "note": "unmasked (no mask support)", "status": "error", "error": "ValueError: Expected dtypes to match but src had type bfloat16 and dst had type float16" }, { "label": "h100br_smallq", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "sq_isab_256x32768", "B": 56, "Lq": 256, "Lkv": 32768, "H": 8, "D": 32, "masked": true, "impl": "flex:default", "fwd_ms": 2.281, "fwd_bwd_ms": 9.296, "out_err": 0.000244140625, "grad_err": [ 0.000244140625, 0.0001220703125, 6.103515625e-05 ], "status": "ok" }, { "label": "h100br_smallq", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "sq_isab_256x32768", "B": 56, "Lq": 256, "Lkv": 32768, "H": 8, "D": 32, "masked": true, "impl": "flex:triton_generic", "fwd_ms": 3.12, "fwd_bwd_ms": 13.591, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "h100br_smallq", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "sq_isab_256x32768", "B": 56, "Lq": 256, "Lkv": 32768, "H": 8, "D": 32, "masked": true, "impl": "flex:triton_split", "fwd_ms": 3.132, "fwd_bwd_ms": 13.857, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "h100br_smallq", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "sq_isab_256x32768", "B": 56, "Lq": 256, "Lkv": 32768, "H": 8, "D": 32, "masked": true, "impl": "flex:triton_onepass", "fwd_ms": 3.128, "fwd_bwd_ms": 13.127, "out_err": 0.0, "grad_err": [ 0.000244140625, 6.103515625e-05, 3.0517578125e-05 ], "status": "ok" }, { "label": "h100br_smallq", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "sq_isab_256x32768", "B": 56, "Lq": 256, "Lkv": 32768, "H": 8, "D": 32, "masked": true, "impl": "flex:triton_fused", "fwd_ms": 3.163, "fwd_bwd_ms": 13.523, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "h100br_smallq", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "sq_isab_256x32768", "B": 56, "Lq": 256, "Lkv": 32768, "H": 8, "D": 32, "masked": true, "impl": "cudnn", "fwd_ms": 2.31, "fwd_bwd_ms": 9.346, "out_err": 0.000244140625, "grad_err": [ 0.000244140625, 0.0001220703125, 6.103515625e-05 ], "status": "ok" }, { "label": "h100br_smallq", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "sq_isab_256x32768", "B": 56, "Lq": 256, "Lkv": 32768, "H": 8, "D": 32, "masked": true, "impl": "mosaic_fa3", "note": "unmasked (no mask support)", "status": "error", "error": "ValueError: Expected dtypes to match but src had type bfloat16 and dst had type float16" }, { "label": "h100br_smallq", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "sq_128x262144_d64", "B": 4, "Lq": 128, "Lkv": 262144, "H": 8, "D": 64, "masked": true, "impl": "flex:default", "fwd_ms": 2.753, "fwd_bwd_ms": 6.402, "out_err": 0.0001220703125, "grad_err": [ 6.103515625e-05, 7.62939453125e-06, 7.62939453125e-06 ], "status": "ok" }, { "label": "h100br_smallq", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "sq_128x262144_d64", "B": 4, "Lq": 128, "Lkv": 262144, "H": 8, "D": 64, "masked": true, "impl": "flex:triton_generic", "fwd_ms": 5.266, "fwd_bwd_ms": 12.916, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "h100br_smallq", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "sq_128x262144_d64", "B": 4, "Lq": 128, "Lkv": 262144, "H": 8, "D": 64, "masked": true, "impl": "flex:triton_split", "fwd_ms": 5.269, "fwd_bwd_ms": 12.935, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "h100br_smallq", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "sq_128x262144_d64", "B": 4, "Lq": 128, "Lkv": 262144, "H": 8, "D": 64, "masked": true, "impl": "flex:triton_onepass", "fwd_ms": 5.212, "fwd_bwd_ms": 9.288, "out_err": 0.0, "grad_err": [ 6.103515625e-05, 3.814697265625e-06, 3.814697265625e-06 ], "status": "ok" }, { "label": "h100br_smallq", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "sq_128x262144_d64", "B": 4, "Lq": 128, "Lkv": 262144, "H": 8, "D": 64, "masked": true, "impl": "flex:triton_fused", "fwd_ms": 5.278, "fwd_bwd_ms": 12.912, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "h100br_smallq", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "sq_128x262144_d64", "B": 4, "Lq": 128, "Lkv": 262144, "H": 8, "D": 64, "masked": true, "impl": "cudnn", "fwd_ms": 2.779, "fwd_bwd_ms": 6.417, "out_err": 0.0001220703125, "grad_err": [ 6.103515625e-05, 7.62939453125e-06, 7.62939453125e-06 ], "status": "ok" }, { "label": "h100br_smallq", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "sq_128x262144_d64", "B": 4, "Lq": 128, "Lkv": 262144, "H": 8, "D": 64, "masked": true, "impl": "mosaic_fa3", "note": "unmasked (no mask support)", "status": "error", "error": "ValueError: Expected dtypes to match but src had type bfloat16 and dst had type float16" }, { "label": "h100br_smallq", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "sq_64x131072_d64", "B": 8, "Lq": 64, "Lkv": 131072, "H": 8, "D": 64, "masked": true, "impl": "flex:default", "fwd_ms": 1.585, "fwd_bwd_ms": 4.849, "out_err": 0.0001220703125, "grad_err": [ 0.0001220703125, 7.62939453125e-06, 7.62939453125e-06 ], "status": "ok" }, { "label": "h100br_smallq", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "sq_64x131072_d64", "B": 8, "Lq": 64, "Lkv": 131072, "H": 8, "D": 64, "masked": true, "impl": "flex:triton_generic", "fwd_ms": 2.063, "fwd_bwd_ms": 6.909, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "h100br_smallq", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "sq_64x131072_d64", "B": 8, "Lq": 64, "Lkv": 131072, "H": 8, "D": 64, "masked": true, "impl": "flex:triton_split", "fwd_ms": 2.042, "fwd_bwd_ms": 6.947, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "h100br_smallq", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "sq_64x131072_d64", "B": 8, "Lq": 64, "Lkv": 131072, "H": 8, "D": 64, "masked": true, "impl": "flex:triton_onepass", "fwd_ms": 2.03, "fwd_bwd_ms": 4.511, "out_err": 0.0, "grad_err": [ 0.0001220703125, 0.0, 0.0 ], "status": "ok" }, { "label": "h100br_smallq", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "sq_64x131072_d64", "B": 8, "Lq": 64, "Lkv": 131072, "H": 8, "D": 64, "masked": true, "impl": "flex:triton_fused", "fwd_ms": 2.056, "fwd_bwd_ms": 6.945, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "h100br_smallq", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "sq_64x131072_d64", "B": 8, "Lq": 64, "Lkv": 131072, "H": 8, "D": 64, "masked": true, "impl": "cudnn", "fwd_ms": 1.576, "fwd_bwd_ms": 4.811, "out_err": 0.0001220703125, "grad_err": [ 0.0001220703125, 7.62939453125e-06, 7.62939453125e-06 ], "status": "ok" }, { "label": "h100br_smallq", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "sq_64x131072_d64", "B": 8, "Lq": 64, "Lkv": 131072, "H": 8, "D": 64, "masked": true, "impl": "mosaic_fa3", "note": "unmasked (no mask support)", "status": "error", "error": "ValueError: Expected dtypes to match but src had type bfloat16 and dst had type float16" }, { "label": "h100br_smallq", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "sq_16x65536_d128", "B": 16, "Lq": 16, "Lkv": 65536, "H": 8, "D": 128, "masked": true, "impl": "flex:default", "fwd_ms": 1.728, "fwd_bwd_ms": 7.385, "out_err": 0.000244140625, "grad_err": [ 0.000244140625, 1.52587890625e-05, 1.52587890625e-05 ], "status": "ok" }, { "label": "h100br_smallq", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "sq_16x65536_d128", "B": 16, "Lq": 16, "Lkv": 65536, "H": 8, "D": 128, "masked": true, "impl": "flex:triton_generic", "fwd_ms": 2.066, "fwd_bwd_ms": 8.188, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "h100br_smallq", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "sq_16x65536_d128", "B": 16, "Lq": 16, "Lkv": 65536, "H": 8, "D": 128, "masked": true, "impl": "flex:triton_split", "fwd_ms": 2.021, "fwd_bwd_ms": 8.212, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "h100br_smallq", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "sq_16x65536_d128", "B": 16, "Lq": 16, "Lkv": 65536, "H": 8, "D": 128, "masked": true, "impl": "flex:triton_onepass", "fwd_ms": 2.014, "fwd_bwd_ms": 6.477, "out_err": 0.0, "grad_err": [ 0.000244140625, 0.0, 0.0 ], "status": "ok" }, { "label": "h100br_smallq", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "sq_16x65536_d128", "B": 16, "Lq": 16, "Lkv": 65536, "H": 8, "D": 128, "masked": true, "impl": "flex:triton_fused", "fwd_ms": 2.018, "fwd_bwd_ms": 8.23, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "h100br_smallq", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "sq_16x65536_d128", "B": 16, "Lq": 16, "Lkv": 65536, "H": 8, "D": 128, "masked": true, "impl": "cudnn", "fwd_ms": 1.704, "fwd_bwd_ms": 7.377, "out_err": 0.000244140625, "grad_err": [ 0.000244140625, 1.52587890625e-05, 1.52587890625e-05 ], "status": "ok" }, { "label": "h100br_smallq", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "sq_16x65536_d128", "B": 16, "Lq": 16, "Lkv": 65536, "H": 8, "D": 128, "masked": true, "impl": "mosaic_fa3", "note": "unmasked (no mask support)", "status": "error", "error": "ValueError: Expected dtypes to match but src had type bfloat16 and dst had type float16" }, { "label": "h100br_smallq", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "sq_512x131072_d32", "B": 8, "Lq": 512, "Lkv": 131072, "H": 8, "D": 32, "masked": true, "impl": "flex:default", "fwd_ms": 3.043, "fwd_bwd_ms": 10.585, "out_err": 0.0001220703125, "grad_err": [ 0.0001220703125, 3.0517578125e-05, 1.52587890625e-05 ], "status": "ok" }, { "label": "h100br_smallq", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "sq_512x131072_d32", "B": 8, "Lq": 512, "Lkv": 131072, "H": 8, "D": 32, "masked": true, "impl": "flex:triton_generic", "fwd_ms": 3.089, "fwd_bwd_ms": 13.4, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "h100br_smallq", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "sq_512x131072_d32", "B": 8, "Lq": 512, "Lkv": 131072, "H": 8, "D": 32, "masked": true, "impl": "flex:triton_split", "fwd_ms": 3.083, "fwd_bwd_ms": 13.227, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "h100br_smallq", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "sq_512x131072_d32", "B": 8, "Lq": 512, "Lkv": 131072, "H": 8, "D": 32, "masked": true, "impl": "flex:triton_onepass", "fwd_ms": 3.082, "fwd_bwd_ms": 12.422, "out_err": 0.0, "grad_err": [ 0.0001220703125, 3.0517578125e-05, 3.0517578125e-05 ], "status": "ok" }, { "label": "h100br_smallq", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "sq_512x131072_d32", "B": 8, "Lq": 512, "Lkv": 131072, "H": 8, "D": 32, "masked": true, "impl": "flex:triton_fused", "fwd_ms": 3.087, "fwd_bwd_ms": 13.359, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "h100br_smallq", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "sq_512x131072_d32", "B": 8, "Lq": 512, "Lkv": 131072, "H": 8, "D": 32, "masked": true, "impl": "cudnn", "fwd_ms": 3.042, "fwd_bwd_ms": 10.598, "out_err": 0.0001220703125, "grad_err": [ 0.0001220703125, 3.0517578125e-05, 1.52587890625e-05 ], "status": "ok" }, { "label": "h100br_smallq", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "sq_512x131072_d32", "B": 8, "Lq": 512, "Lkv": 131072, "H": 8, "D": 32, "masked": true, "impl": "mosaic_fa3", "note": "unmasked (no mask support)", "status": "error", "error": "ValueError: Expected dtypes to match but src had type bfloat16 and dst had type float16" }, { "label": "h100br_mosaic_unmasked", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "mla_32k84", "B": 2, "Lq": 131072, "Lkv": 32768, "H": 1, "D": 128, "masked": false, "impl": "mosaic_fa3", "fwd_ms": 17.135, "fwd_bwd_ms": 138.562, "status": "ok" }, { "label": "h100br_mosaic_unmasked", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "mla_32k84", "B": 2, "Lq": 131072, "Lkv": 32768, "H": 1, "D": 128, "masked": false, "impl": "cudnn", "fwd_ms": 8.987, "fwd_bwd_ms": 32.833, "out_err": 0.000244140625, "grad_err": [ 0.000244140625, 0.00048828125, 0.0009765625 ], "status": "ok" }, { "label": "h100br_mosaic_unmasked", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "mla_32k84", "B": 2, "Lq": 131072, "Lkv": 32768, "H": 1, "D": 128, "masked": false, "impl": "flex:triton_split", "fwd_ms": 13.774, "fwd_bwd_ms": 50.579, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "h100br_mosaic_unmasked", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "mla_32k84", "B": 2, "Lq": 131072, "Lkv": 32768, "H": 1, "D": 128, "masked": false, "impl": "flex:default", "fwd_ms": 8.874, "fwd_bwd_ms": 33.987, "out_err": 0.000244140625, "grad_err": [ 0.000244140625, 0.00048828125, 0.0009765625 ], "status": "ok" }, { "label": "h100br_mosaic_unmasked", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "mla_128k36", "B": 1, "Lq": 65536, "Lkv": 131072, "H": 1, "D": 128, "masked": false, "impl": "mosaic_fa3", "fwd_ms": 17.038, "fwd_bwd_ms": 136.316, "status": "ok" }, { "label": "h100br_mosaic_unmasked", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "mla_128k36", "B": 1, "Lq": 65536, "Lkv": 131072, "H": 1, "D": 128, "masked": false, "impl": "cudnn", "fwd_ms": 8.053, "fwd_bwd_ms": 33.27, "out_err": 0.0001220703125, "grad_err": [ 0.0001220703125, 6.103515625e-05, 0.0001220703125 ], "status": "ok" }, { "label": "h100br_mosaic_unmasked", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "mla_128k36", "B": 1, "Lq": 65536, "Lkv": 131072, "H": 1, "D": 128, "masked": false, "impl": "flex:triton_split", "fwd_ms": 14.764, "fwd_bwd_ms": 52.314, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "h100br_mosaic_unmasked", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "mla_128k36", "B": 1, "Lq": 65536, "Lkv": 131072, "H": 1, "D": 128, "masked": false, "impl": "flex:default", "fwd_ms": 8.047, "fwd_bwd_ms": 32.753, "out_err": 0.0001220703125, "grad_err": [ 0.0001220703125, 6.103515625e-05, 0.0001220703125 ], "status": "ok" }, { "label": "h100br_mosaic_unmasked", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "sq_isab_256x131072", "B": 28, "Lq": 256, "Lkv": 131072, "H": 8, "D": 32, "masked": false, "impl": "mosaic_fa3", "fwd_ms": 10.985, "fwd_bwd_ms": 49.966, "status": "ok" }, { "label": "h100br_mosaic_unmasked", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "sq_isab_256x131072", "B": 28, "Lq": 256, "Lkv": 131072, "H": 8, "D": 32, "masked": false, "impl": "cudnn", "fwd_ms": 5.353, "fwd_bwd_ms": 22.61, "out_err": 0.000244140625, "grad_err": [ 0.0001220703125, 3.0517578125e-05, 1.52587890625e-05 ], "status": "ok" }, { "label": "h100br_mosaic_unmasked", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "sq_isab_256x131072", "B": 28, "Lq": 256, "Lkv": 131072, "H": 8, "D": 32, "masked": false, "impl": "flex:triton_split", "fwd_ms": 5.833, "fwd_bwd_ms": 23.679, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "h100br_mosaic_unmasked", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "sq_isab_256x131072", "B": 28, "Lq": 256, "Lkv": 131072, "H": 8, "D": 32, "masked": false, "impl": "flex:default", "fwd_ms": 5.405, "fwd_bwd_ms": 21.747, "out_err": 0.000244140625, "grad_err": [ 0.0001220703125, 3.0517578125e-05, 1.52587890625e-05 ], "status": "ok" }, { "label": "h100br_mosaic_unmasked", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "sq_128x262144_d64", "B": 4, "Lq": 128, "Lkv": 262144, "H": 8, "D": 64, "masked": false, "impl": "mosaic_fa3", "fwd_ms": 4.549, "fwd_bwd_ms": 16.35, "status": "ok" }, { "label": "h100br_mosaic_unmasked", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "sq_128x262144_d64", "B": 4, "Lq": 128, "Lkv": 262144, "H": 8, "D": 64, "masked": false, "impl": "cudnn", "fwd_ms": 2.351, "fwd_bwd_ms": 6.158, "out_err": 0.0001220703125, "grad_err": [ 6.103515625e-05, 7.62939453125e-06, 3.814697265625e-06 ], "status": "ok" }, { "label": "h100br_mosaic_unmasked", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "sq_128x262144_d64", "B": 4, "Lq": 128, "Lkv": 262144, "H": 8, "D": 64, "masked": false, "impl": "flex:triton_split", "fwd_ms": 4.882, "fwd_bwd_ms": 12.032, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "h100br_mosaic_unmasked", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "sq_128x262144_d64", "B": 4, "Lq": 128, "Lkv": 262144, "H": 8, "D": 64, "masked": false, "impl": "flex:default", "fwd_ms": 2.403, "fwd_bwd_ms": 6.151, "out_err": 0.0001220703125, "grad_err": [ 6.103515625e-05, 7.62939453125e-06, 3.814697265625e-06 ], "status": "ok" }, { "label": "h100br_mosaic_unmasked", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d64_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 64, "masked": false, "impl": "mosaic_fa3", "fwd_ms": 1.416, "fwd_bwd_ms": 9.411, "status": "ok" }, { "label": "h100br_mosaic_unmasked", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d64_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 64, "masked": false, "impl": "cudnn", "fwd_ms": 1.425, "fwd_bwd_ms": 5.014, "out_err": 0.000244140625, "grad_err": [ 0.00048828125, 0.00048828125, 0.00048828125 ], "status": "ok" }, { "label": "h100br_mosaic_unmasked", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d64_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 64, "masked": false, "impl": "flex:triton_split", "fwd_ms": 1.883, "fwd_bwd_ms": 8.027, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "h100br_mosaic_unmasked", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d64_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 64, "masked": false, "impl": "flex:default", "fwd_ms": 1.426, "fwd_bwd_ms": 5.038, "out_err": 0.000244140625, "grad_err": [ 0.00048828125, 0.00048828125, 0.00048828125 ], "status": "ok" }, { "label": "h100br_mosaic_unmasked", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d128_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 128, "masked": false, "impl": "mosaic_fa3", "fwd_ms": 4.098, "fwd_bwd_ms": 39.612, "status": "ok" }, { "label": "h100br_mosaic_unmasked", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d128_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 128, "masked": false, "impl": "cudnn", "fwd_ms": 2.133, "fwd_bwd_ms": 8.28, "out_err": 0.00048828125, "grad_err": [ 0.00048828125, 0.00048828125, 0.00048828125 ], "status": "ok" }, { "label": "h100br_mosaic_unmasked", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d128_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 128, "masked": false, "impl": "flex:triton_split", "fwd_ms": 3.189, "fwd_bwd_ms": 12.826, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "h100br_mosaic_unmasked", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d128_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 128, "masked": false, "impl": "flex:default", "fwd_ms": 2.15, "fwd_bwd_ms": 8.114, "out_err": 0.00048828125, "grad_err": [ 0.00048828125, 0.00048828125, 0.00048828125 ], "status": "ok" }, { "label": "h100br_mosaic_unmasked", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d128_t16384", "B": 4, "Lq": 16384, "Lkv": 16384, "H": 8, "D": 128, "masked": false, "impl": "mosaic_fa3", "fwd_ms": 17.572, "fwd_bwd_ms": 160.889, "status": "ok" }, { "label": "h100br_mosaic_unmasked", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d128_t16384", "B": 4, "Lq": 16384, "Lkv": 16384, "H": 8, "D": 128, "masked": false, "impl": "cudnn", "fwd_ms": 8.944, "fwd_bwd_ms": 33.601, "out_err": 0.00048828125, "grad_err": [ 0.00048828125, 0.00048828125, 0.00048828125 ], "status": "ok" }, { "label": "h100br_mosaic_unmasked", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d128_t16384", "B": 4, "Lq": 16384, "Lkv": 16384, "H": 8, "D": 128, "masked": false, "impl": "flex:triton_split", "fwd_ms": 13.427, "fwd_bwd_ms": 55.406, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "h100br_mosaic_unmasked", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d128_t16384", "B": 4, "Lq": 16384, "Lkv": 16384, "H": 8, "D": 128, "masked": false, "impl": "flex:default", "fwd_ms": 9.159, "fwd_bwd_ms": 33.316, "out_err": 0.00048828125, "grad_err": [ 0.00048828125, 0.00048828125, 0.00048828125 ], "status": "ok" }, { "label": "h100br_mosaic_unmasked", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "cross_d128_256x16384", "B": 4, "Lq": 256, "Lkv": 16384, "H": 8, "D": 128, "masked": false, "impl": "mosaic_fa3", "fwd_ms": 1.133, "fwd_bwd_ms": 3.769, "status": "ok" }, { "label": "h100br_mosaic_unmasked", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "cross_d128_256x16384", "B": 4, "Lq": 256, "Lkv": 16384, "H": 8, "D": 128, "masked": false, "impl": "cudnn", "fwd_ms": 0.298, "fwd_bwd_ms": 0.899, "out_err": 0.000244140625, "grad_err": [ 0.00048828125, 6.103515625e-05, 6.103515625e-05 ], "status": "ok" }, { "label": "h100br_mosaic_unmasked", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "cross_d128_256x16384", "B": 4, "Lq": 256, "Lkv": 16384, "H": 8, "D": 128, "masked": false, "impl": "flex:triton_split", "fwd_ms": 0.504, "fwd_bwd_ms": 1.308, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "h100br_mosaic_unmasked", "host": "uc3n105.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "cross_d128_256x16384", "B": 4, "Lq": 256, "Lkv": 16384, "H": 8, "D": 128, "masked": false, "impl": "flex:default", "fwd_ms": 0.304, "fwd_bwd_ms": 0.899, "out_err": 0.000244140625, "grad_err": [ 0.00048828125, 6.103515625e-05, 6.103515625e-05 ], "status": "ok" }, { "label": "h100br_mosaic_sweep", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d128_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 128, "masked": false, "impl": "cudnn", "fwd_ms": 2.147, "fwd_bwd_ms": 8.166, "out_err": 0.00048828125, "grad_err": [ 0.00048828125, 0.00048828125, 0.00048828125 ], "status": "ok" }, { "label": "h100br_mosaic_sweep", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d128_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 128, "masked": false, "impl": "mosaic:128,128,2,1", "fwd_ms": 4.093, "fwd_bwd_ms": 39.245, "status": "ok" }, { "label": "h100br_mosaic_sweep", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d128_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 128, "masked": false, "impl": "mosaic:64,128,2,1", "fwd_ms": 2.255, "fwd_bwd_ms": 17.16, "status": "ok" }, { "label": "h100br_mosaic_sweep", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d128_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 128, "masked": false, "impl": "mosaic:64,64,2,1", "fwd_ms": 2.528, "fwd_bwd_ms": 10.865, "status": "ok" }, { "label": "h100br_mosaic_sweep", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d128_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 128, "masked": false, "impl": "mosaic:128,64,2,1", "fwd_ms": 2.473, "fwd_bwd_ms": 15.94, "status": "ok" }, { "label": "h100br_mosaic_sweep", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d128_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 128, "masked": false, "impl": "mosaic:64,128,3,1", "fwd_ms": 2.267, "status": "error", "error": "ValueError: Mosaic GPU kernel exceeds available shared memory: smem_bytes=232568 > max_smem_bytes=232448" }, { "label": "h100br_mosaic_sweep", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d128_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 128, "masked": false, "impl": "mosaic:64,128,4,1", "status": "error", "error": "ValueError: Mosaic GPU kernel exceeds available shared memory: smem_bytes=297112 > max_smem_bytes=232448" }, { "label": "h100br_mosaic_sweep", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d128_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 128, "masked": false, "impl": "mosaic:128,128,2,2", "fwd_ms": 4.089, "status": "error", "error": "ValueError: Mosaic GPU kernel exceeds available shared memory: smem_bytes=264304 > max_smem_bytes=232448" }, { "label": "h100br_mosaic_sweep", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d128_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 128, "masked": false, "impl": "mosaic:64,128,2,2", "fwd_ms": 2.238, "fwd_bwd_ms": 19.593, "status": "ok" }, { "label": "h100br_mosaic_sweep", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d128_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 128, "masked": false, "impl": "mosaic:64,128,2,2,64,128,64,128", "fwd_ms": 2.226, "fwd_bwd_ms": 19.526, "status": "ok" }, { "label": "h100br_mosaic_sweep", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d128_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 128, "masked": false, "impl": "mosaic:64,128,2,2,128,64,64,128", "fwd_ms": 2.244, "fwd_bwd_ms": 9.516, "status": "ok" }, { "label": "h100br_mosaic_sweep", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d128_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 128, "masked": false, "impl": "mosaic:64,128,2,1,64,64,64,64", "fwd_ms": 2.218, "fwd_bwd_ms": 11.27, "status": "ok" }, { "label": "h100br_mosaic_sweep", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "mla_32k84", "B": 2, "Lq": 131072, "Lkv": 32768, "H": 1, "D": 128, "masked": false, "impl": "cudnn", "fwd_ms": 8.304, "fwd_bwd_ms": 33.183, "out_err": 0.000244140625, "grad_err": [ 0.000244140625, 0.00048828125, 0.0009765625 ], "status": "ok" }, { "label": "h100br_mosaic_sweep", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "mla_32k84", "B": 2, "Lq": 131072, "Lkv": 32768, "H": 1, "D": 128, "masked": false, "impl": "mosaic:128,128,2,1", "fwd_ms": 16.991, "fwd_bwd_ms": 136.506, "status": "ok" }, { "label": "h100br_mosaic_sweep", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "mla_32k84", "B": 2, "Lq": 131072, "Lkv": 32768, "H": 1, "D": 128, "masked": false, "impl": "mosaic:64,128,2,1", "fwd_ms": 8.459, "fwd_bwd_ms": 72.004, "status": "ok" }, { "label": "h100br_mosaic_sweep", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "mla_32k84", "B": 2, "Lq": 131072, "Lkv": 32768, "H": 1, "D": 128, "masked": false, "impl": "mosaic:64,64,2,1", "fwd_ms": 9.691, "fwd_bwd_ms": 45.689, "status": "ok" }, { "label": "h100br_mosaic_sweep", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "mla_32k84", "B": 2, "Lq": 131072, "Lkv": 32768, "H": 1, "D": 128, "masked": false, "impl": "mosaic:128,64,2,1", "fwd_ms": 9.36, "fwd_bwd_ms": 65.697, "status": "ok" }, { "label": "h100br_mosaic_sweep", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "mla_32k84", "B": 2, "Lq": 131072, "Lkv": 32768, "H": 1, "D": 128, "masked": false, "impl": "mosaic:64,128,3,1", "fwd_ms": 8.784, "status": "error", "error": "ValueError: Mosaic GPU kernel exceeds available shared memory: smem_bytes=232568 > max_smem_bytes=232448" }, { "label": "h100br_mosaic_sweep", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "mla_32k84", "B": 2, "Lq": 131072, "Lkv": 32768, "H": 1, "D": 128, "masked": false, "impl": "mosaic:64,128,4,1", "status": "error", "error": "ValueError: Mosaic GPU kernel exceeds available shared memory: smem_bytes=297112 > max_smem_bytes=232448" }, { "label": "h100br_mosaic_sweep", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "mla_32k84", "B": 2, "Lq": 131072, "Lkv": 32768, "H": 1, "D": 128, "masked": false, "impl": "mosaic:128,128,2,2", "fwd_ms": 17.151, "status": "error", "error": "ValueError: Mosaic GPU kernel exceeds available shared memory: smem_bytes=264304 > max_smem_bytes=232448" }, { "label": "h100br_mosaic_sweep", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "mla_32k84", "B": 2, "Lq": 131072, "Lkv": 32768, "H": 1, "D": 128, "masked": false, "impl": "mosaic:64,128,2,2", "fwd_ms": 8.231, "fwd_bwd_ms": 77.6, "status": "ok" }, { "label": "h100br_mosaic_sweep", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "mla_32k84", "B": 2, "Lq": 131072, "Lkv": 32768, "H": 1, "D": 128, "masked": false, "impl": "mosaic:64,128,2,2,64,128,64,128", "fwd_ms": 8.148, "fwd_bwd_ms": 77.963, "status": "ok" }, { "label": "h100br_mosaic_sweep", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "mla_32k84", "B": 2, "Lq": 131072, "Lkv": 32768, "H": 1, "D": 128, "masked": false, "impl": "mosaic:64,128,2,2,128,64,64,128", "fwd_ms": 8.31, "fwd_bwd_ms": 37.369, "status": "ok" }, { "label": "h100br_mosaic_sweep", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "mla_32k84", "B": 2, "Lq": 131072, "Lkv": 32768, "H": 1, "D": 128, "masked": false, "impl": "mosaic:64,128,2,1,64,64,64,64", "fwd_ms": 8.277, "fwd_bwd_ms": 42.954, "status": "ok" }, { "label": "h100br_mosaic_sweep", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d64_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 64, "masked": false, "impl": "cudnn", "fwd_ms": 1.436, "fwd_bwd_ms": 4.981, "out_err": 0.000244140625, "grad_err": [ 0.00048828125, 0.00048828125, 0.00048828125 ], "status": "ok" }, { "label": "h100br_mosaic_sweep", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d64_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 64, "masked": false, "impl": "mosaic:128,128,2,1", "fwd_ms": 1.445, "fwd_bwd_ms": 9.671, "status": "ok" }, { "label": "h100br_mosaic_sweep", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d64_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 64, "masked": false, "impl": "mosaic:64,128,2,1", "fwd_ms": 1.489, "fwd_bwd_ms": 9.531, "status": "ok" }, { "label": "h100br_mosaic_sweep", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d64_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 64, "masked": false, "impl": "mosaic:64,64,2,1", "fwd_ms": 1.715, "fwd_bwd_ms": 7.045, "status": "ok" }, { "label": "h100br_mosaic_sweep", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d64_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 64, "masked": false, "impl": "mosaic:128,64,2,1", "fwd_ms": 1.504, "fwd_bwd_ms": 6.392, "status": "ok" }, { "label": "h100br_mosaic_sweep", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d64_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 64, "masked": false, "impl": "mosaic:64,128,3,1", "fwd_ms": 1.519, "fwd_bwd_ms": 9.48, "status": "ok" }, { "label": "h100br_mosaic_sweep", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d64_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 64, "masked": false, "impl": "mosaic:64,128,4,1", "fwd_ms": 1.503, "fwd_bwd_ms": 9.746, "status": "ok" }, { "label": "h100br_mosaic_sweep", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d64_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 64, "masked": false, "impl": "mosaic:128,128,2,2", "fwd_ms": 1.425, "fwd_bwd_ms": 10.725, "status": "ok" }, { "label": "h100br_mosaic_sweep", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d64_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 64, "masked": false, "impl": "mosaic:64,128,2,2", "fwd_ms": 1.48, "fwd_bwd_ms": 5.899, "status": "ok" }, { "label": "h100br_mosaic_sweep", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d64_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 64, "masked": false, "impl": "mosaic:64,128,2,2,64,128,64,128", "fwd_ms": 1.477, "fwd_bwd_ms": 5.97, "status": "ok" }, { "label": "h100br_mosaic_sweep", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d64_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 64, "masked": false, "impl": "mosaic:64,128,2,2,128,64,64,128", "fwd_ms": 1.494, "fwd_bwd_ms": 5.402, "status": "ok" }, { "label": "h100br_mosaic_sweep", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d64_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 64, "masked": false, "impl": "mosaic:64,128,2,1,64,64,64,64", "fwd_ms": 1.464, "fwd_bwd_ms": 6.967, "status": "ok" }, { "label": "main_isab", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "isab_ind_rows", "B": 56, "Lq": 256, "Lkv": 32768, "H": 8, "D": 32, "masked": true, "impl": "cudnn", "fwd_ms": 2.199, "fwd_bwd_ms": 9.056, "out_err": 0.000244140625, "grad_err": [ 0.000244140625, 0.0001220703125, 6.103515625e-05 ], "status": "ok" }, { "label": "main_mla", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "mla_32k84", "B": 2, "Lq": 131072, "Lkv": 32768, "H": 1, "D": 128, "masked": true, "impl": "cudnn", "fwd_ms": 6.479, "fwd_bwd_ms": 24.884, "out_err": 0.000244140625, "grad_err": [ 0.000244140625, 0.00048828125, 0.0009765625 ], "status": "ok" }, { "label": "main_mla", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "mla_128k36", "B": 1, "Lq": 65536, "Lkv": 131072, "H": 1, "D": 128, "masked": true, "impl": "cudnn", "fwd_ms": 7.052, "fwd_bwd_ms": 25.811, "out_err": 0.0001220703125, "grad_err": [ 0.0001220703125, 0.0001220703125, 0.0001220703125 ], "status": "ok" }, { "label": "main_mla", "host": "uc3n089.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "mla_1k756", "B": 4, "Lq": 262144, "Lkv": 1024, "H": 1, "D": 128, "masked": true, "impl": "cudnn", "fwd_ms": 1.117, "fwd_bwd_ms": 11.207, "out_err": 0.0009765625, "grad_err": [ 0.001953125, 0.03125, 0.0625 ], "status": "ok" }, { "label": "main_isab", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "isab_rows_ind", "B": 56, "Lq": 32768, "Lkv": 256, "H": 8, "D": 32, "masked": false, "impl": "cudnn", "fwd_ms": 3.368, "fwd_bwd_ms": 15.11, "out_err": 0.00390625, "grad_err": [ 0.00390625, 0.03125, 0.03125 ], "status": "ok" }, { "label": "main_isab", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "isab_1k_ind", "B": 1008, "Lq": 1024, "Lkv": 256, "H": 8, "D": 32, "masked": false, "impl": "cudnn", "fwd_ms": 1.923, "fwd_bwd_ms": 8.791, "out_err": 0.00390625, "grad_err": [ 0.015625, 0.015625, 0.0078125 ], "status": "ok" }, { "label": "main_isab", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "isab_ind_1k", "B": 1008, "Lq": 256, "Lkv": 1024, "H": 8, "D": 32, "masked": true, "impl": "cudnn", "fwd_ms": 1.462, "fwd_bwd_ms": 6.436, "out_err": 0.0009765625, "grad_err": [ 0.00390625, 0.001953125, 0.001953125 ], "status": "ok" }, { "label": "h100_torch_B", "host": "uc3n106.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d64_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 64, "masked": true, "impl": "flex:default", "fwd_ms": 1.16, "fwd_bwd_ms": 4.086, "out_err": 0.00048828125, "grad_err": [ 0.00048828125, 0.0009765625, 0.0009765625 ], "status": "ok" }, { "label": "h100_torch_A", "host": "uc3n106.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "mla_32k84", "B": 2, "Lq": 131072, "Lkv": 32768, "H": 1, "D": 128, "masked": true, "impl": "flex:default", "fwd_ms": 6.796, "fwd_bwd_ms": 26.052, "out_err": 0.000244140625, "grad_err": [ 0.000244140625, 0.00048828125, 0.0009765625 ], "status": "ok" }, { "label": "h100_torch_B", "host": "uc3n106.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d64_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 64, "masked": true, "impl": "flex:cudnn", "fwd_ms": 1.189, "fwd_bwd_ms": 4.118, "out_err": 0.00048828125, "grad_err": [ 0.00048828125, 0.0009765625, 0.0009765625 ], "status": "ok" }, { "label": "h100_torch_B", "host": "uc3n106.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d64_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 64, "masked": true, "impl": "flex:triton", "fwd_ms": 1.941, "fwd_bwd_ms": 8.233, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "h100_torch_A", "host": "uc3n106.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "mla_32k84", "B": 2, "Lq": 131072, "Lkv": 32768, "H": 1, "D": 128, "masked": true, "impl": "flex:cudnn", "fwd_ms": 6.464, "fwd_bwd_ms": 26.308, "out_err": 0.000244140625, "grad_err": [ 0.000244140625, 0.00048828125, 0.0009765625 ], "status": "ok" }, { "label": "h100_torch_A", "host": "uc3n106.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "mla_32k84", "B": 2, "Lq": 131072, "Lkv": 32768, "H": 1, "D": 128, "masked": true, "impl": "flex:triton", "fwd_ms": 14.477, "fwd_bwd_ms": 89.072, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "h100_torch_A", "host": "uc3n106.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "mla_128k36", "B": 1, "Lq": 65536, "Lkv": 131072, "H": 1, "D": 128, "masked": true, "impl": "flex:default", "fwd_ms": 6.215, "fwd_bwd_ms": 25.206, "out_err": 0.0001220703125, "grad_err": [ 0.0001220703125, 0.0001220703125, 0.0001220703125 ], "status": "ok" }, { "label": "h100_torch_A", "host": "uc3n106.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "mla_128k36", "B": 1, "Lq": 65536, "Lkv": 131072, "H": 1, "D": 128, "masked": true, "impl": "flex:cudnn", "fwd_ms": 6.554, "fwd_bwd_ms": 25.364, "out_err": 0.0001220703125, "grad_err": [ 0.0001220703125, 0.0001220703125, 0.0001220703125 ], "status": "ok" }, { "label": "h100_torch_A", "host": "uc3n106.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "mla_128k36", "B": 1, "Lq": 65536, "Lkv": 131072, "H": 1, "D": 128, "masked": true, "impl": "flex:triton", "fwd_ms": 15.201, "fwd_bwd_ms": 72.78, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "h100_torch_B", "host": "uc3n106.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d128_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 128, "masked": false, "impl": "flex:default", "fwd_ms": 2.223, "fwd_bwd_ms": 8.335, "out_err": 0.00048828125, "grad_err": [ 0.00048828125, 0.00048828125, 0.00048828125 ], "status": "ok" }, { "label": "h100_torch_B", "host": "uc3n106.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d128_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 128, "masked": false, "impl": "flex:cudnn", "fwd_ms": 2.179, "fwd_bwd_ms": 8.473, "out_err": 0.00048828125, "grad_err": [ 0.00048828125, 0.00048828125, 0.00048828125 ], "status": "ok" }, { "label": "h100_torch_B", "host": "uc3n106.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d128_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 128, "masked": false, "impl": "flex:triton", "fwd_ms": 3.221, "fwd_bwd_ms": 13.114, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "h100_torch_A", "host": "uc3n106.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "mla_1k756", "B": 4, "Lq": 262144, "Lkv": 1024, "H": 1, "D": 128, "masked": true, "impl": "flex:default", "fwd_ms": 1.124, "fwd_bwd_ms": 11.295, "out_err": 0.0009765625, "grad_err": [ 0.001953125, 0.03125, 0.0625 ], "status": "ok" }, { "label": "h100_torch_B", "host": "uc3n106.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d128_t16384", "B": 4, "Lq": 16384, "Lkv": 16384, "H": 8, "D": 128, "masked": true, "impl": "flex:default", "fwd_ms": 7.009, "fwd_bwd_ms": 29.058, "out_err": 0.000244140625, "grad_err": [ 0.00048828125, 0.00048828125, 0.00048828125 ], "status": "ok" }, { "label": "h100_torch_B", "host": "uc3n106.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d128_t16384", "B": 4, "Lq": 16384, "Lkv": 16384, "H": 8, "D": 128, "masked": true, "impl": "flex:cudnn", "fwd_ms": 7.698, "fwd_bwd_ms": 28.097, "out_err": 0.000244140625, "grad_err": [ 0.00048828125, 0.00048828125, 0.00048828125 ], "status": "ok" }, { "label": "h100_torch_A", "host": "uc3n106.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "mla_1k756", "B": 4, "Lq": 262144, "Lkv": 1024, "H": 1, "D": 128, "masked": true, "impl": "flex:cudnn", "fwd_ms": 1.127, "fwd_bwd_ms": 11.254, "out_err": 0.0009765625, "grad_err": [ 0.001953125, 0.03125, 0.0625 ], "status": "ok" }, { "label": "h100_torch_B", "host": "uc3n106.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d128_t16384", "B": 4, "Lq": 16384, "Lkv": 16384, "H": 8, "D": 128, "masked": true, "impl": "flex:triton", "fwd_ms": 14.498, "fwd_bwd_ms": 54.993, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "h100_torch_B", "host": "uc3n106.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "cross_d128_256x16384", "B": 4, "Lq": 256, "Lkv": 16384, "H": 8, "D": 128, "masked": true, "impl": "flex:default", "fwd_ms": 0.352, "fwd_bwd_ms": 0.892, "out_err": 0.0001220703125, "grad_err": [ 0.00048828125, 6.103515625e-05, 6.103515625e-05 ], "status": "ok" }, { "label": "h100_torch_A", "host": "uc3n106.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "mla_1k756", "B": 4, "Lq": 262144, "Lkv": 1024, "H": 1, "D": 128, "masked": true, "impl": "flex:triton", "fwd_ms": 1.858, "fwd_bwd_ms": 11.296, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "h100_torch_B", "host": "uc3n106.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "cross_d128_256x16384", "B": 4, "Lq": 256, "Lkv": 16384, "H": 8, "D": 128, "masked": true, "impl": "flex:cudnn", "fwd_ms": 0.349, "fwd_bwd_ms": 0.886, "out_err": 0.0001220703125, "grad_err": [ 0.00048828125, 6.103515625e-05, 6.103515625e-05 ], "status": "ok" }, { "label": "h100_torch_B", "host": "uc3n106.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "cross_d128_256x16384", "B": 4, "Lq": 256, "Lkv": 16384, "H": 8, "D": 128, "masked": true, "impl": "flex:triton", "fwd_ms": 0.519, "fwd_bwd_ms": 1.369, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "h100_torch_A", "host": "uc3n106.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "isab_ind_rows", "B": 56, "Lq": 256, "Lkv": 32768, "H": 8, "D": 32, "masked": true, "impl": "flex:default", "fwd_ms": 2.225, "fwd_bwd_ms": 9.296, "out_err": 0.000244140625, "grad_err": [ 0.000244140625, 0.0001220703125, 6.103515625e-05 ], "status": "ok" }, { "label": "h100_torch_A", "host": "uc3n106.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "isab_ind_rows", "B": 56, "Lq": 256, "Lkv": 32768, "H": 8, "D": 32, "masked": true, "impl": "flex:cudnn", "fwd_ms": 2.22, "fwd_bwd_ms": 9.126, "out_err": 0.000244140625, "grad_err": [ 0.000244140625, 0.0001220703125, 6.103515625e-05 ], "status": "ok" }, { "label": "h100_torch_A", "host": "uc3n106.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "isab_ind_rows", "B": 56, "Lq": 256, "Lkv": 32768, "H": 8, "D": 32, "masked": true, "impl": "flex:triton", "fwd_ms": 3.129, "fwd_bwd_ms": 13.679, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "h100_torch_B", "host": "uc3n106.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "cross_d64_16384x256", "B": 4, "Lq": 16384, "Lkv": 256, "H": 8, "D": 64, "masked": true, "impl": "flex:default", "fwd_ms": 0.215, "fwd_bwd_ms": 0.812, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "h100_torch_B", "host": "uc3n106.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "cross_d64_16384x256", "B": 4, "Lq": 16384, "Lkv": 256, "H": 8, "D": 64, "masked": true, "impl": "flex:cudnn", "fwd_ms": 0.241, "fwd_bwd_ms": 1.028, "out_err": 0.001953125, "grad_err": [ 0.00390625, 0.015625, 0.03125 ], "status": "ok" }, { "label": "h100_torch_B", "host": "uc3n106.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "cross_d64_16384x256", "B": 4, "Lq": 16384, "Lkv": 256, "H": 8, "D": 64, "masked": true, "impl": "flex:triton", "fwd_ms": 0.222, "fwd_bwd_ms": 0.817, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "h100_torch_A", "host": "uc3n106.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "isab_1k_ind", "B": 1008, "Lq": 1024, "Lkv": 256, "H": 8, "D": 32, "masked": false, "impl": "flex:default", "fwd_ms": 1.943, "fwd_bwd_ms": 8.882, "out_err": 0.00390625, "grad_err": [ 0.015625, 0.015625, 0.0078125 ], "status": "ok" }, { "label": "h100_torch_A", "host": "uc3n106.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "isab_1k_ind", "B": 1008, "Lq": 1024, "Lkv": 256, "H": 8, "D": 32, "masked": false, "impl": "flex:cudnn", "fwd_ms": 1.966, "fwd_bwd_ms": 8.874, "out_err": 0.00390625, "grad_err": [ 0.015625, 0.015625, 0.0078125 ], "status": "ok" }, { "label": "h100_torch_A", "host": "uc3n106.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "isab_1k_ind", "B": 1008, "Lq": 1024, "Lkv": 256, "H": 8, "D": 32, "masked": false, "impl": "flex:triton", "fwd_ms": 1.759, "fwd_bwd_ms": 19.495, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "h100_torch_A", "host": "uc3n106.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "sq_isab_256x131072", "B": 28, "Lq": 256, "Lkv": 131072, "H": 8, "D": 32, "masked": true, "impl": "flex:default", "fwd_ms": 5.408, "fwd_bwd_ms": 19.878, "out_err": 0.0001220703125, "grad_err": [ 0.0001220703125, 3.0517578125e-05, 1.52587890625e-05 ], "status": "ok" }, { "label": "h100_torch_A", "host": "uc3n106.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "sq_isab_256x131072", "B": 28, "Lq": 256, "Lkv": 131072, "H": 8, "D": 32, "masked": true, "impl": "flex:cudnn", "fwd_ms": 5.419, "fwd_bwd_ms": 19.921, "out_err": 0.0001220703125, "grad_err": [ 0.0001220703125, 3.0517578125e-05, 1.52587890625e-05 ], "status": "ok" }, { "label": "h100_torch_A", "host": "uc3n106.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "sq_isab_256x131072", "B": 28, "Lq": 256, "Lkv": 131072, "H": 8, "D": 32, "masked": true, "impl": "flex:triton", "fwd_ms": 6.163, "fwd_bwd_ms": 26.437, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "h100_torch_A", "host": "uc3n106.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "sq_128x262144_d64", "B": 4, "Lq": 128, "Lkv": 262144, "H": 8, "D": 64, "masked": true, "impl": "flex:default", "fwd_ms": 2.718, "fwd_bwd_ms": 6.412, "out_err": 0.0001220703125, "grad_err": [ 6.103515625e-05, 7.62939453125e-06, 7.62939453125e-06 ], "status": "ok" }, { "label": "h100_torch_A", "host": "uc3n106.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "sq_128x262144_d64", "B": 4, "Lq": 128, "Lkv": 262144, "H": 8, "D": 64, "masked": true, "impl": "flex:cudnn", "fwd_ms": 2.718, "fwd_bwd_ms": 6.422, "out_err": 0.0001220703125, "grad_err": [ 6.103515625e-05, 7.62939453125e-06, 7.62939453125e-06 ], "status": "ok" }, { "label": "h100_torch_A", "host": "uc3n106.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "sq_128x262144_d64", "B": 4, "Lq": 128, "Lkv": 262144, "H": 8, "D": 64, "masked": true, "impl": "flex:triton", "fwd_ms": 5.264, "fwd_bwd_ms": 12.795, "out_err": 0.0, "grad_err": [ 0.0, 0.0, 0.0 ], "status": "ok" }, { "label": "h100_torch_A", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "mla_32k84", "B": 2, "Lq": 131072, "Lkv": 32768, "H": 1, "D": 128, "masked": true, "impl": "torch:flex", "fwd_ms": 9.159, "fwd_bwd_ms": 33.308, "status": "ok", "out_err": 0.000244140625, "grad_err": [ 0.00048828125, 0.0009765625, 0.0009765625 ] }, { "label": "h100_torch_A", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "mla_32k84", "B": 2, "Lq": 131072, "Lkv": 32768, "H": 1, "D": 128, "masked": true, "impl": "torch:flash", "status": "error", "error": "RuntimeError: SDPA flash takes no attention mask" }, { "label": "h100_torch_A", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "mla_32k84", "B": 2, "Lq": 131072, "Lkv": 32768, "H": 1, "D": 128, "masked": true, "impl": "torch:efficient", "fwd_ms": 55.678, "fwd_bwd_ms": 158.263, "status": "ok", "out_err": 0.000244140625, "grad_err": [ 0.00048828125, 0.0009765625, 0.0009765625 ] }, { "label": "h100_torch_B", "host": "uc3n106.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d64_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 64, "masked": true, "impl": "torch:flex", "fwd_ms": 1.244, "fwd_bwd_ms": 4.816, "status": "ok", "out_err": 0.00048828125, "grad_err": [ 0.0009765625, 0.0009765625, 0.0009765625 ] }, { "label": "h100_torch_B", "host": "uc3n106.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d64_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 64, "masked": true, "impl": "torch:flash", "status": "error", "error": "RuntimeError: SDPA flash takes no attention mask" }, { "label": "h100_torch_B", "host": "uc3n106.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d64_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 64, "masked": true, "impl": "torch:efficient", "fwd_ms": 7.032, "fwd_bwd_ms": 23.674, "status": "ok", "out_err": 0.0009765625, "grad_err": [ 0.001953125, 0.001953125, 0.0009765625 ] }, { "label": "h100_torch_A", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "mla_128k36", "B": 1, "Lq": 65536, "Lkv": 131072, "H": 1, "D": 128, "masked": true, "impl": "torch:flex", "fwd_ms": 8.622, "fwd_bwd_ms": 32.313, "status": "ok", "out_err": 0.000244140625, "grad_err": [ 0.000244140625, 0.000244140625, 0.0001220703125 ] }, { "label": "h100_torch_A", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "mla_128k36", "B": 1, "Lq": 65536, "Lkv": 131072, "H": 1, "D": 128, "masked": true, "impl": "torch:flash", "status": "error", "error": "RuntimeError: SDPA flash takes no attention mask" }, { "label": "h100_torch_A", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "mla_128k36", "B": 1, "Lq": 65536, "Lkv": 131072, "H": 1, "D": 128, "masked": true, "impl": "torch:efficient", "fwd_ms": 55.889, "fwd_bwd_ms": 159.263, "status": "ok", "out_err": 0.000244140625, "grad_err": [ 0.000244140625, 0.000244140625, 0.000244140625 ] }, { "label": "h100_torch_B", "host": "uc3n106.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d128_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 128, "masked": false, "impl": "torch:flex", "fwd_ms": 2.716, "fwd_bwd_ms": 10.295, "status": "ok", "out_err": 0.00048828125, "grad_err": [ 0.0009765625, 0.0009765625, 0.00048828125 ] }, { "label": "h100_torch_B", "host": "uc3n106.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d128_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 128, "masked": false, "impl": "torch:flash", "fwd_ms": 3.52, "fwd_bwd_ms": 14.994, "status": "ok", "out_err": 0.00048828125, "grad_err": [ 0.0009765625, 0.0009765625, 0.00048828125 ] }, { "label": "h100_torch_B", "host": "uc3n106.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d128_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 128, "masked": false, "impl": "torch:efficient", "fwd_ms": 7.39, "fwd_bwd_ms": 30.132, "status": "ok", "out_err": 0.00048828125, "grad_err": [ 0.0009765625, 0.0009765625, 0.00048828125 ] }, { "label": "h100_torch_B", "host": "uc3n106.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d128_t16384", "B": 4, "Lq": 16384, "Lkv": 16384, "H": 8, "D": 128, "masked": true, "impl": "torch:flex", "fwd_ms": 9.764, "fwd_bwd_ms": 36.167, "status": "ok", "out_err": 0.00048828125, "grad_err": [ 0.0009765625, 0.0009765625, 0.00048828125 ] }, { "label": "h100_torch_B", "host": "uc3n106.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d128_t16384", "B": 4, "Lq": 16384, "Lkv": 16384, "H": 8, "D": 128, "masked": true, "impl": "torch:flash", "status": "error", "error": "RuntimeError: SDPA flash takes no attention mask" }, { "label": "h100_torch_A", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "mla_1k756", "B": 4, "Lq": 262144, "Lkv": 1024, "H": 1, "D": 128, "masked": true, "impl": "torch:flex", "fwd_ms": 1.396, "fwd_bwd_ms": 11.637, "status": "ok", "out_err": 0.001953125, "grad_err": [ 0.0078125, 0.03125, 0.03125 ] }, { "label": "h100_torch_A", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "mla_1k756", "B": 4, "Lq": 262144, "Lkv": 1024, "H": 1, "D": 128, "masked": true, "impl": "torch:flash", "status": "error", "error": "RuntimeError: SDPA flash takes no attention mask" }, { "label": "h100_torch_A", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "mla_1k756", "B": 4, "Lq": 262144, "Lkv": 1024, "H": 1, "D": 128, "masked": true, "impl": "torch:efficient", "fwd_ms": 7.379, "fwd_bwd_ms": 54.009, "status": "ok", "out_err": 0.001953125, "grad_err": [ 0.0078125, 0.03125, 0.0625 ] }, { "label": "h100_torch_B", "host": "uc3n106.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d128_t16384", "B": 4, "Lq": 16384, "Lkv": 16384, "H": 8, "D": 128, "masked": true, "impl": "torch:efficient", "fwd_ms": 40.179, "fwd_bwd_ms": 146.065, "status": "ok", "out_err": 0.00048828125, "grad_err": [ 0.0009765625, 0.0009765625, 0.00048828125 ] }, { "label": "h100_torch_A", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "isab_ind_rows", "B": 56, "Lq": 256, "Lkv": 32768, "H": 8, "D": 32, "masked": true, "impl": "torch:flex", "fwd_ms": 1.938, "fwd_bwd_ms": 7.354, "status": "ok", "out_err": 0.00048828125, "grad_err": [ 0.0009765625, 0.0009765625, 6.103515625e-05 ] }, { "label": "h100_torch_B", "host": "uc3n106.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "cross_d128_256x16384", "B": 4, "Lq": 256, "Lkv": 16384, "H": 8, "D": 128, "masked": true, "impl": "torch:flex", "fwd_ms": 0.366, "fwd_bwd_ms": 1.204, "status": "ok", "out_err": 0.00048828125, "grad_err": [ 0.00048828125, 0.000244140625, 6.103515625e-05 ] }, { "label": "h100_torch_A", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "isab_ind_rows", "B": 56, "Lq": 256, "Lkv": 32768, "H": 8, "D": 32, "masked": true, "impl": "torch:flash", "status": "error", "error": "RuntimeError: SDPA flash takes no attention mask" }, { "label": "h100_torch_B", "host": "uc3n106.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "cross_d128_256x16384", "B": 4, "Lq": 256, "Lkv": 16384, "H": 8, "D": 128, "masked": true, "impl": "torch:flash", "status": "error", "error": "RuntimeError: SDPA flash takes no attention mask" }, { "label": "h100_torch_B", "host": "uc3n106.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "cross_d128_256x16384", "B": 4, "Lq": 256, "Lkv": 16384, "H": 8, "D": 128, "masked": true, "impl": "torch:efficient", "fwd_ms": 0.859, "fwd_bwd_ms": 12.371, "status": "ok", "out_err": 0.000244140625, "grad_err": [ 0.00048828125, 0.00048828125, 6.103515625e-05 ] }, { "label": "h100_torch_A", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "isab_ind_rows", "B": 56, "Lq": 256, "Lkv": 32768, "H": 8, "D": 32, "masked": true, "impl": "torch:efficient", "fwd_ms": 11.846, "fwd_bwd_ms": 378.029, "status": "ok", "out_err": 0.000244140625, "grad_err": [ 0.00048828125, 0.00048828125, 6.103515625e-05 ] }, { "label": "h100_torch_B", "host": "uc3n106.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "cross_d64_16384x256", "B": 4, "Lq": 16384, "Lkv": 256, "H": 8, "D": 64, "masked": true, "impl": "torch:flex", "fwd_ms": 0.213, "fwd_bwd_ms": 1.017, "status": "ok", "out_err": 0.0078125, "grad_err": [ 0.0078125, 0.03125, 0.015625 ] }, { "label": "h100_torch_B", "host": "uc3n106.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "cross_d64_16384x256", "B": 4, "Lq": 16384, "Lkv": 256, "H": 8, "D": 64, "masked": true, "impl": "torch:flash", "status": "error", "error": "RuntimeError: SDPA flash takes no attention mask" }, { "label": "h100_torch_B", "host": "uc3n106.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "cross_d64_16384x256", "B": 4, "Lq": 16384, "Lkv": 256, "H": 8, "D": 64, "masked": true, "impl": "torch:efficient", "fwd_ms": 0.484, "fwd_bwd_ms": 3.081, "status": "ok", "out_err": 0.0078125, "grad_err": [ 0.01171875, 0.03125, 0.03125 ] }, { "label": "h100_torch_A", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "isab_1k_ind", "B": 1008, "Lq": 1024, "Lkv": 256, "H": 8, "D": 32, "masked": false, "impl": "torch:flex", "fwd_ms": 1.421, "fwd_bwd_ms": 5.526, "status": "ok", "out_err": 0.0078125, "grad_err": [ 0.015625, 0.015625, 0.0078125 ] }, { "label": "h100_torch_A", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "isab_1k_ind", "B": 1008, "Lq": 1024, "Lkv": 256, "H": 8, "D": 32, "masked": false, "impl": "torch:flash", "fwd_ms": 1.752, "fwd_bwd_ms": 9.256, "status": "ok", "out_err": 0.0078125, "grad_err": [ 0.015625, 0.015625, 0.0078125 ] }, { "label": "h100_torch_A", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "isab_1k_ind", "B": 1008, "Lq": 1024, "Lkv": 256, "H": 8, "D": 32, "masked": false, "impl": "torch:efficient", "fwd_ms": 4.181, "fwd_bwd_ms": 19.784, "status": "ok", "out_err": 0.0078125, "grad_err": [ 0.015625, 0.015625, 0.0078125 ] }, { "label": "h100_torch_A", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "sq_isab_256x131072", "B": 28, "Lq": 256, "Lkv": 131072, "H": 8, "D": 32, "masked": true, "impl": "torch:flex", "fwd_ms": 4.328, "fwd_bwd_ms": 15.67, "status": "ok", "out_err": 0.000244140625, "grad_err": [ 0.000244140625, 0.0001220703125, 3.0517578125e-05 ] }, { "label": "h100_torch_A", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "sq_isab_256x131072", "B": 28, "Lq": 256, "Lkv": 131072, "H": 8, "D": 32, "masked": true, "impl": "torch:flash", "status": "error", "error": "RuntimeError: SDPA flash takes no attention mask" }, { "label": "h100_torch_A", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "sq_isab_256x131072", "B": 28, "Lq": 256, "Lkv": 131072, "H": 8, "D": 32, "masked": true, "impl": "torch:efficient", "fwd_ms": 23.621, "fwd_bwd_ms": 1320.625, "status": "ok", "out_err": 0.000244140625, "grad_err": [ 0.000244140625, 0.0001220703125, 1.52587890625e-05 ] }, { "label": "h100_torch_A", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "sq_128x262144_d64", "B": 4, "Lq": 128, "Lkv": 262144, "H": 8, "D": 64, "masked": true, "impl": "torch:flex", "fwd_ms": 3.406, "fwd_bwd_ms": 10.403, "status": "ok", "out_err": 0.0001220703125, "grad_err": [ 6.103515625e-05, 7.62939453125e-06, 3.814697265625e-06 ] }, { "label": "h100_torch_A", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "sq_128x262144_d64", "B": 4, "Lq": 128, "Lkv": 262144, "H": 8, "D": 64, "masked": true, "impl": "torch:flash", "status": "error", "error": "RuntimeError: SDPA flash takes no attention mask" }, { "label": "h100_torch_A", "host": "uc3n088.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "sq_128x262144_d64", "B": 4, "Lq": 128, "Lkv": 262144, "H": 8, "D": 64, "masked": true, "impl": "torch:efficient", "fwd_ms": 16.504, "fwd_bwd_ms": 409.298, "status": "ok", "out_err": 0.0001220703125, "grad_err": [ 6.103515625e-05, 3.0517578125e-05, 7.62939453125e-06 ] }, { "label": "h100_torch_B", "host": "uc3n106.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d64_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 64, "masked": true, "impl": "flex:mosaic", "fwd_ms": 1.126, "fwd_bwd_ms": 4.724, "out_err": 0.00048828125, "grad_err": [ 0.00048828125, 0.0009765625, 0.0009765625 ], "status": "ok" }, { "label": "h100_torch_A", "host": "uc3n106.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "mla_32k84", "B": 2, "Lq": 131072, "Lkv": 32768, "H": 1, "D": 128, "masked": true, "impl": "flex:mosaic", "fwd_ms": 6.206, "fwd_bwd_ms": 30.787, "out_err": 0.00048828125, "grad_err": [ 0.00048828125, 0.0009765625, 0.0009765625 ], "status": "ok" }, { "label": "h100_torch_A", "host": "uc3n106.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "mla_128k36", "B": 1, "Lq": 65536, "Lkv": 131072, "H": 1, "D": 128, "masked": true, "impl": "flex:mosaic", "fwd_ms": 5.867, "fwd_bwd_ms": 31.298, "out_err": 0.000244140625, "grad_err": [ 0.000244140625, 0.000244140625, 0.000244140625 ], "status": "ok" }, { "label": "h100_torch_B", "host": "uc3n106.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "self_d128_t8192", "B": 4, "Lq": 8192, "Lkv": 8192, "H": 8, "D": 128, "masked": false, "impl": "flex:mosaic", "fwd_ms": 2.295, "fwd_bwd_ms": 9.282, "out_err": 0.00048828125, "grad_err": [ 0.0009765625, 0.0009765625, 0.0009765625 ], "status": "ok" }, { "label": "h100_torch_A", "host": "uc3n106.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "mla_1k756", "B": 4, "Lq": 262144, "Lkv": 1024, "H": 1, "D": 128, "masked": true, "impl": "flex:mosaic", "fwd_ms": 1.293, "fwd_bwd_ms": 10.127, "out_err": 0.00390625, "grad_err": [ 0.0078125, 0.03125, 0.0625 ], "status": "ok" }, { "label": "h100_torch_B", "host": "uc3n106.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "selfm_d128_t16384", "B": 4, "Lq": 16384, "Lkv": 16384, "H": 8, "D": 128, "masked": true, "impl": "flex:mosaic", "fwd_ms": 7.393, "fwd_bwd_ms": 34.3, "out_err": 0.00048828125, "grad_err": [ 0.0009765625, 0.0009765625, 0.0009765625 ], "status": "ok" }, { "label": "h100_torch_B", "host": "uc3n106.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "cross_d128_256x16384", "B": 4, "Lq": 256, "Lkv": 16384, "H": 8, "D": 128, "masked": true, "impl": "flex:mosaic", "fwd_ms": 0.277, "fwd_bwd_ms": 0.914, "out_err": 0.00048828125, "grad_err": [ 0.00048828125, 0.00048828125, 0.000244140625 ], "status": "ok" }, { "label": "h100_torch_A", "host": "uc3n106.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "isab_ind_rows", "B": 56, "Lq": 256, "Lkv": 32768, "H": 8, "D": 32, "masked": true, "impl": "flex:mosaic", "fwd_ms": 5.448, "fwd_bwd_ms": 16.557, "out_err": 0.00048828125, "grad_err": [ 0.0009765625, 0.0009765625, 0.00048828125 ], "status": "ok" }, { "label": "h100_torch_B", "host": "uc3n106.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "cross_d64_16384x256", "B": 4, "Lq": 16384, "Lkv": 256, "H": 8, "D": 64, "masked": true, "impl": "flex:mosaic", "fwd_ms": 0.304, "fwd_bwd_ms": 0.943, "out_err": 0.001953125, "grad_err": [ 0.00390625, 0.015625, 0.03125 ], "status": "ok" }, { "label": "h100_torch_A", "host": "uc3n106.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "isab_1k_ind", "B": 1008, "Lq": 1024, "Lkv": 256, "H": 8, "D": 32, "masked": false, "impl": "flex:mosaic", "fwd_ms": 5.002, "fwd_bwd_ms": 12.958, "out_err": 0.0078125, "grad_err": [ 0.015625, 0.015625, 0.009765625 ], "status": "ok" }, { "label": "h100_torch_A", "host": "uc3n106.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "sq_isab_256x131072", "B": 28, "Lq": 256, "Lkv": 131072, "H": 8, "D": 32, "masked": true, "impl": "flex:mosaic", "fwd_ms": 10.797, "fwd_bwd_ms": 33.948, "out_err": 0.000244140625, "grad_err": [ 0.000244140625, 0.000244140625, 0.0001220703125 ], "status": "ok" }, { "label": "h100_torch_A", "host": "uc3n106.localdomain", "gpu": "NVIDIA H100", "cc": "9.0", "jax": "0.11.2", "shape": "sq_128x262144_d64", "B": 4, "Lq": 128, "Lkv": 262144, "H": 8, "D": 64, "masked": true, "impl": "flex:mosaic", "fwd_ms": 2.135, "fwd_bwd_ms": 9.151, "out_err": 0.0001220703125, "grad_err": [ 6.103515625e-05, 1.52587890625e-05, 3.814697265625e-06 ], "status": "ok" } ] }