82 lines
3.5 KiB
Bash
Executable file
82 lines
3.5 KiB
Bash
Executable file
#!/bin/sh
|
|
set -eu
|
|
|
|
[ "$#" -eq 4 ] || [ "$#" -eq 5 ] || {
|
|
echo "usage: $0 BASELINE_A1 CANDIDATE_B1 CANDIDATE_B2 BASELINE_A2 [OUTPUT_JSON]" >&2
|
|
exit 2
|
|
}
|
|
a1=$1
|
|
b1=$2
|
|
b2=$3
|
|
a2=$4
|
|
output=${5:-}
|
|
script_dir=$(CDPATH= cd -- "$(dirname -- "$0")" && pwd)
|
|
temporary=$(mktemp -d /private/tmp/gem-decode-attention-abba.XXXXXX)
|
|
trap 'rm -rf "$temporary"' EXIT INT TERM HUP
|
|
|
|
merge_pair() {
|
|
first=$1
|
|
second=$2
|
|
destination=$3
|
|
jq -s '
|
|
def controls: {
|
|
schema_version, engine, workload, shape, score_token_tile, value_kernel,
|
|
value_dimension_tile, value_threadgroups, query_heads, kv_heads, head_dimension,
|
|
tokens, live_tokens, cache_capacity, window, kv_cache_bytes, score_scratch_bytes,
|
|
value_partition_scratch_bytes, value_partitions, iterations, warmup, dispatches
|
|
};
|
|
.[0] as $x | .[1] as $y |
|
|
if ($x | controls) != ($y | controls) or
|
|
$x.repetitions < 3 or $y.repetitions < 3 or
|
|
($x.gpu_seconds_samples | length) != $x.repetitions or
|
|
($y.gpu_seconds_samples | length) != $y.repetitions or
|
|
($x.stages.score.samples_seconds | length) != $x.repetitions or
|
|
($y.stages.score.samples_seconds | length) != $y.repetitions
|
|
then error("ABBA arms do not preserve their tuning controls") else
|
|
$x |
|
|
.repetitions = ($x.repetitions + $y.repetitions) |
|
|
.gpu_seconds_samples = ($x.gpu_seconds_samples + $y.gpu_seconds_samples) |
|
|
.stages.score.samples_seconds =
|
|
($x.stages.score.samples_seconds + $y.stages.score.samples_seconds)
|
|
end
|
|
' "$first" "$second" > "$destination"
|
|
}
|
|
|
|
merge_pair "$a1" "$a2" "$temporary/baseline.json"
|
|
merge_pair "$b1" "$b2" "$temporary/candidate.json"
|
|
"$script_dir/analyze-decode-attention.sh" "$temporary/baseline.json" \
|
|
"$temporary/candidate.json" "$temporary/pooled.json"
|
|
|
|
analyze() {
|
|
jq --slurpfile a1 "$a1" --slurpfile b1 "$b1" --slurpfile b2 "$b2" \
|
|
--slurpfile a2 "$a2" '
|
|
def mean($xs): $xs | add / length;
|
|
def gain($before; $after): 100 * (1 - mean($after) / mean($before));
|
|
def drift($first; $second): 100 * (mean($second) / mean($first) - 1);
|
|
. + {bracket: {
|
|
order: ["baseline-a1", "candidate-b1", "candidate-b2", "baseline-a2"],
|
|
arm_repetitions: $a1[0].repetitions,
|
|
total_gpu: {
|
|
baseline_a1_mean_seconds: mean($a1[0].gpu_seconds_samples),
|
|
candidate_b1_mean_seconds: mean($b1[0].gpu_seconds_samples),
|
|
candidate_b2_mean_seconds: mean($b2[0].gpu_seconds_samples),
|
|
baseline_a2_mean_seconds: mean($a2[0].gpu_seconds_samples),
|
|
opening_improvement_pct: gain($a1[0].gpu_seconds_samples; $b1[0].gpu_seconds_samples),
|
|
closing_improvement_pct: gain($a2[0].gpu_seconds_samples; $b2[0].gpu_seconds_samples),
|
|
baseline_drift_pct: drift($a1[0].gpu_seconds_samples; $a2[0].gpu_seconds_samples),
|
|
candidate_drift_pct: drift($b1[0].gpu_seconds_samples; $b2[0].gpu_seconds_samples),
|
|
every_candidate_sample_faster:
|
|
(([$b1[0].gpu_seconds_samples[], $b2[0].gpu_seconds_samples[]] | max) <
|
|
([$a1[0].gpu_seconds_samples[], $a2[0].gpu_seconds_samples[]] | min))
|
|
},
|
|
score_stage: {
|
|
opening_improvement_pct:
|
|
gain($a1[0].stages.score.samples_seconds; $b1[0].stages.score.samples_seconds),
|
|
closing_improvement_pct:
|
|
gain($a2[0].stages.score.samples_seconds; $b2[0].stages.score.samples_seconds)
|
|
}
|
|
}}
|
|
' "$temporary/pooled.json"
|
|
}
|
|
|
|
if [ -n "$output" ]; then analyze > "$output"; else analyze; fi
|