116 lines
5.2 KiB
Bash
Executable file
116 lines
5.2 KiB
Bash
Executable file
#!/bin/sh
|
|
set -eu
|
|
|
|
[ "$#" -eq 1 ] || [ "$#" -eq 2 ] || {
|
|
echo "usage: $0 PROFILE_JSON [OUTPUT_JSON]" >&2
|
|
exit 2
|
|
}
|
|
profile=$1
|
|
output=${2:-}
|
|
|
|
analyze() {
|
|
jq '
|
|
def sum_indices($xs; $indices): [$indices[] as $i | $xs[$i]] | add;
|
|
def mean($xs): $xs | add / length;
|
|
def sample_stddev($xs):
|
|
mean($xs) as $m |
|
|
(($xs | map((. - $m) * (. - $m)) | add) / (($xs | length) - 1)) | sqrt;
|
|
def cv_pct($xs): 100 * sample_stddev($xs) / mean($xs);
|
|
|
|
. as $p |
|
|
["attention_norm", "qkv_projections", "qkv_transform", "kv_attention",
|
|
"attention_output", "ffn_norm", "ffn_gate_up", "ffn_activation",
|
|
"ffn_down_output"] as $expected_names |
|
|
($p.stages.all.mean_seconds | add) as $layer_stage_total |
|
|
sum_indices($p.stages.all.mean_seconds; [1, 4, 6, 8]) as $projection_total |
|
|
$p.stages.all.mean_seconds[3] as $attention_total |
|
|
sum_indices($p.stages.all.mean_seconds; [0, 2, 5, 7]) as $elementwise_total |
|
|
$p.stages.sliding.mean_seconds[3] as $sliding_attention |
|
|
$p.stages.full.mean_seconds[3] as $full_attention |
|
|
if ($p.schema_version != 1 or $p.engine != "gem-native-c-metal" or
|
|
$p.workload != "long-prefill-stage-profile" or
|
|
$p.stage_names != $expected_names or
|
|
$p.layer_count != 60 or $p.sliding_layer_count != 50 or
|
|
$p.full_layer_count != 10 or $p.tile_tokens != 32 or
|
|
$p.repetitions < 3 or
|
|
($p.wall_seconds_samples | length) != $p.repetitions or
|
|
($p.gpu_seconds_samples | length) != $p.repetitions or
|
|
($p.sampled_gpu_seconds_samples | length) != $p.repetitions or
|
|
($p.stages.all.mean_seconds | length) != 9 or
|
|
($p.stages.sliding.mean_seconds | length) != 9 or
|
|
($p.stages.full.mean_seconds | length) != 9 or
|
|
($p.stages.all.samples_seconds | length) != $p.repetitions or
|
|
($p.stages.sliding.samples_seconds | length) != $p.repetitions or
|
|
($p.stages.full.samples_seconds | length) != $p.repetitions or
|
|
($p.stages.all.samples_seconds | any(length != 9)) or
|
|
($p.stages.sliding.samples_seconds | any(length != 9)) or
|
|
($p.stages.full.samples_seconds | any(length != 9)))
|
|
then error("invalid production prefill-stage profile")
|
|
else
|
|
{
|
|
schema_version: 1,
|
|
scope: "production-graph prefill stage attribution",
|
|
source: {
|
|
model: $p.model,
|
|
attention_kernel: $p.attention_kernel,
|
|
kv_type: $p.kv_type,
|
|
prefix_mode: ($p.prefix_mode // "evaluated-prompt"),
|
|
prefix_tokens: $p.prefix_tokens,
|
|
tile_tokens: $p.tile_tokens,
|
|
repetitions: $p.repetitions,
|
|
layer_count: $p.layer_count,
|
|
sliding_layer_count: $p.sliding_layer_count,
|
|
full_layer_count: $p.full_layer_count
|
|
},
|
|
timing: {
|
|
mean_wall_seconds: mean($p.wall_seconds_samples),
|
|
mean_gpu_seconds: mean($p.gpu_seconds_samples),
|
|
mean_sampled_gpu_seconds: mean($p.sampled_gpu_seconds_samples),
|
|
sampled_gpu_fraction: (mean($p.sampled_gpu_seconds_samples) /
|
|
mean($p.gpu_seconds_samples)),
|
|
wall_cv_pct: cv_pct($p.wall_seconds_samples),
|
|
gpu_cv_pct: cv_pct($p.gpu_seconds_samples)
|
|
},
|
|
ranked_layer_stages: ([range(0; 9) as $i | {
|
|
name: $p.stage_names[$i],
|
|
mean_seconds: $p.stages.all.mean_seconds[$i],
|
|
layer_stage_fraction: ($p.stages.all.mean_seconds[$i] / $layer_stage_total),
|
|
sliding_seconds: $p.stages.sliding.mean_seconds[$i],
|
|
full_seconds: $p.stages.full.mean_seconds[$i]
|
|
}] | sort_by(-.mean_seconds)),
|
|
families: {
|
|
projections: {
|
|
mean_seconds: $projection_total,
|
|
sampled_gpu_fraction: ($projection_total / $p.mean_sampled_gpu_seconds)
|
|
},
|
|
kv_attention: {
|
|
mean_seconds: $attention_total,
|
|
sampled_gpu_fraction: ($attention_total / $p.mean_sampled_gpu_seconds),
|
|
sliding_seconds: $sliding_attention,
|
|
full_seconds: $full_attention
|
|
},
|
|
elementwise: {
|
|
mean_seconds: $elementwise_total,
|
|
sampled_gpu_fraction: ($elementwise_total / $p.mean_sampled_gpu_seconds)
|
|
},
|
|
embedding_and_logits: {
|
|
mean_seconds: ($p.mean_embedding_seconds + $p.mean_logits_seconds),
|
|
sampled_gpu_fraction: (($p.mean_embedding_seconds + $p.mean_logits_seconds) /
|
|
$p.mean_sampled_gpu_seconds)
|
|
}
|
|
},
|
|
selected_next_axis: (
|
|
if $projection_total > $attention_total then "q4-projection-kernels"
|
|
elif $sliding_attention > $full_attention then
|
|
"sliding-attention-preserving-rowwise-arithmetic"
|
|
else "full-attention-head-dim512" end),
|
|
screening: {
|
|
promotion_eligible: false,
|
|
reason: "stage attribution only; implement and benchmark one selected axis in the full graph"
|
|
}
|
|
}
|
|
end
|
|
' "$profile"
|
|
}
|
|
|
|
if [ -n "$output" ]; then analyze > "$output"; else analyze; fi
|