| { |
| "op": "ai.onnx.Split", |
| "cases": [ |
| { |
| "name": "qkv_like", |
| "attrs": { "axis": 2 }, |
| "inputs": { "input": { "dtype": "float32", "shape": [1, 128, 2304] } }, |
| "outputs": { |
| "y0": { "dtype": "float32", "shape": [1, 128, 768] }, |
| "y1": { "dtype": "float32", "shape": [1, 128, 1536] } |
| }, |
| "preset": "smoke", |
| "vars": { "count": 294912, "dtype": "float32" }, |
| "bench": { |
| "primary": true, |
| "metrics": [{ "type": "bandwidth", "value": "args.count * dtypeBytes(args.dtype) * 2" }] |
| } |
| }, |
| { |
| "name": "three-way-equal-axis1-m1500-n1152", |
| "preset": "model", |
| "attrs": { "axis": 1 }, |
| "vars": { "count": 1728000, "dtype": "float32" }, |
| "inputs": { "input": { "shape": [1500, 1152], "dtype": "float32", "dist": "normal", "seed": 976, "scale": 0.5 } }, |
| "outputs": { |
| "y0": { "shape": [1500, 384], "dtype": "float32" }, |
| "y1": { "shape": [1500, 384], "dtype": "float32" }, |
| "y2": { "shape": [1500, 384], "dtype": "float32" } |
| }, |
| "bench": { |
| "primary": true, |
| "metrics": [{ "type": "bandwidth", "value": "args.count * dtypeBytes(args.dtype) * 2" }] |
| } |
| }, |
| { |
| "name": "split-block-axis0-f32-8m", |
| "preset": "smoke", |
| "attrs": { "axis": 0 }, |
| "vars": { "count": 8388608, "dtype": "float32" }, |
| "inputs": { "input": { "shape": [8192, 1024], "dtype": "float32", "dist": "normal", "seed": 970, "scale": 1 } }, |
| "outputs": { |
| "y0": { "shape": [4096, 1024], "dtype": "float32" }, |
| "y1": { "shape": [4096, 1024], "dtype": "float32" } |
| }, |
| "bench": { "metrics": [{ "type": "bandwidth", "value": "args.count * dtypeBytes(args.dtype) * 2" }] } |
| }, |
| { |
| "name": "split-scalar-cliff-last-axis-unaligned", |
| "attrs": { "axis": 2 }, |
| "inputs": { "input": { "dtype": "float32", "shape": [1, 128, 2306] } }, |
| "outputs": { |
| "y0": { "dtype": "float32", "shape": [1, 128, 770] }, |
| "y1": { "dtype": "float32", "shape": [1, 128, 1536] } |
| }, |
| "preset": "smoke", |
| "vars": { "count": 295168, "dtype": "float32" }, |
| "bench": { "metrics": [{ "type": "bandwidth", "value": "args.count * dtypeBytes(args.dtype) * 2" }] } |
| }, |
| { |
| "name": "qkv_like_f16", |
| "attrs": { "axis": 2 }, |
| "inputs": { "input": { "dtype": "float16", "shape": [1, 128, 2304] } }, |
| "outputs": { |
| "y0": { "dtype": "float16", "shape": [1, 128, 768] }, |
| "y1": { "dtype": "float16", "shape": [1, 128, 1536] } |
| }, |
| "preset": "smoke", |
| "vars": { "count": 294912, "dtype": "float16" }, |
| "bench": { "metrics": [{ "type": "bandwidth", "value": "args.count * dtypeBytes(args.dtype) * 2" }] } |
| }, |
| { |
| "name": "split-five-outputs-scalar-8m", |
| "attrs": { "axis": 1 }, |
| "inputs": { "input": { "dtype": "float32", "shape": [8, 1310720], "dist": "normal", "seed": 971, "scale": 1 } }, |
| "outputs": { |
| "y0": { "dtype": "float32", "shape": [8, 262144] }, |
| "y1": { "dtype": "float32", "shape": [8, 262144] }, |
| "y2": { "dtype": "float32", "shape": [8, 262144] }, |
| "y3": { "dtype": "float32", "shape": [8, 262144] }, |
| "y4": { "dtype": "float32", "shape": [8, 262144] } |
| }, |
| "preset": "smoke", |
| "vars": { "count": 10485760, "dtype": "float32" }, |
| "bench": { "metrics": [{ "type": "bandwidth", "value": "args.count * dtypeBytes(args.dtype) * 2" }] } |
| }, |
| { |
| "name": "split-scalar-rank7-middle-axis-2m", |
| "preset": "stress", |
| "attrs": { "axis": 2 }, |
| "vars": { "count": 1990656, "dtype": "float32" }, |
| "inputs": { |
| "input": { "dtype": "float32", "shape": [64, 64, 6, 3, 3, 3, 3], "dist": "normal", "seed": 973, "scale": 1 } |
| }, |
| "outputs": { |
| "y0": { "dtype": "float32", "shape": [64, 64, 3, 3, 3, 3, 3] }, |
| "y1": { "dtype": "float32", "shape": [64, 64, 3, 3, 3, 3, 3] } |
| }, |
| "bench": { "metrics": [{ "type": "bandwidth", "value": "args.count * dtypeBytes(args.dtype) * 2" }] } |
| }, |
| { |
| "name": "split-scalar-four-outputs-unaligned-runs-2m", |
| "preset": "stress", |
| "attrs": { "axis": 1 }, |
| "vars": { "count": 2211840, "dtype": "float32" }, |
| "inputs": { |
| "input": { "dtype": "float32", "shape": [8192, 10, 3, 3, 3], "dist": "normal", "seed": 974, "scale": 1 } |
| }, |
| "outputs": { |
| "y0": { "dtype": "float32", "shape": [8192, 3, 3, 3, 3] }, |
| "y1": { "dtype": "float32", "shape": [8192, 3, 3, 3, 3] }, |
| "y2": { "dtype": "float32", "shape": [8192, 2, 3, 3, 3] }, |
| "y3": { "dtype": "float32", "shape": [8192, 2, 3, 3, 3] } |
| }, |
| "bench": { "metrics": [{ "type": "bandwidth", "value": "args.count * dtypeBytes(args.dtype) * 2" }] } |
| }, |
| { |
| "name": "split-three-output-q2050-kv512-odd-segment-pathology", |
| "preset": "stress", |
| "attrs": { "axis": 2 }, |
| "vars": { "count": 1573888, "dtype": "float32" }, |
| "inputs": { "input": { "dtype": "float32", "shape": [1, 512, 3074], "dist": "normal", "seed": 975, "scale": 1 } }, |
| "outputs": { |
| "y0": { "dtype": "float32", "shape": [1, 512, 2050] }, |
| "y1": { "dtype": "float32", "shape": [1, 512, 512] }, |
| "y2": { "dtype": "float32", "shape": [1, 512, 512] } |
| }, |
| "bench": { |
| "primary": true, |
| "metrics": [{ "type": "bandwidth", "value": "args.count * dtypeBytes(args.dtype) * 2" }] |
| } |
| } |
| ], |
| "tunableSpace": { "WORKGROUP_SIZE": [64, 128, 256] } |
| } |
|
|
