Skip to content
Open
Show file tree
Hide file tree
Changes from 4 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
83 changes: 53 additions & 30 deletions .buildkite/scripts/generate_support_matrices.sh
Original file line number Diff line number Diff line change
Expand Up @@ -29,22 +29,31 @@ FEATURE_STAGES=("CorrectnessTest" "PerformanceTest")
FEATURE_STAGES_QUANTIZATION=("QuantizationMethods" "RecommendedTPUGenerations" "CorrectnessTest" "PerformanceTest")
FEATURE_STAGES_MICROBENCHMARKS=("CorrectnessTest" "PerformanceTest")

declare -A TPU_GENERATIONS=(
["INT8 W8A8"]="\"v5, v6\""
["INT4 W4A16"]="\"v5, v6\""
["FP8 W8A8"]="v7"
["FP8 W8A16"]="v7"
["FP4 W4A16"]="v7"
["AWQ INT4"]="\"v5, v6\""
)
declare -A QUANTIZATION_METHODS=(
["INT8 W8A8"]="compressed-tensor"
["INT4 W4A16"]="awq"
["FP8 W8A8"]="compressed-tensor"
["FP8 W8A16"]="compressed-tensor"
["FP4 W4A16"]="mxfp4"
["AWQ INT4"]=""
)
get_tpu_generation() {
local key="$1"
case "$key" in
"INT8 W8A8") echo "\"v5, v6\"" ;;
"INT4 W4A16") echo "\"v5, v6\"" ;;
"FP8 W8A8") echo "v7" ;;
"FP8 W8A16") echo "v7" ;;
"FP4 W4A16") echo "v7" ;;
"AWQ INT4") echo "\"v5, v6\"" ;;
*) echo "N/A" ;;
esac
}

get_quantization_method() {
local key="$1"
case "$key" in
"INT8 W8A8") echo "compressed-tensor" ;;
"INT4 W4A16") echo "awq" ;;
"FP8 W8A8") echo "compressed-tensor" ;;
"FP8 W8A16") echo "compressed-tensor" ;;
"FP4 W4A16") echo "mxfp4" ;;
"AWQ INT4") echo "" ;;
*) echo "N/A" ;;
esac
}
declare -a model_csv_files=()
declare -a feature_csv_files=()
declare -a default_feature_names=()
Expand Down Expand Up @@ -111,10 +120,10 @@ process_models() {
result="Text"
fi
else
result=$(buildkite-agent meta-data get "${TPU_METADATA_PREFIX}${model}:${stage}" --default "N/A")
result=$(buildkite-agent meta-data get "${TPU_METADATA_PREFIX}${model}:${stage}" --default "N/A")
fi
row="$row,$result"
if [ "$stage" != "Type" ] && [ "${result}" != "✅" ] && [ "${result}" != "N/A" ] && [ "${result}" != "unverified" ]; then
if [ "$stage" != "Type" ] && [ "${result}" != "✅ Passing" ] && [ "${result}" != "N/A" ] && [ "${result}" != "❓ Untested" ]; then
ANY_FAILED=true
fi
done
Expand Down Expand Up @@ -164,20 +173,32 @@ process_features() {
local result
if [ "$is_quantization_matrix" = true ] && [ "$stage" == "RecommendedTPUGenerations" ]; then
# If it's the quantization matrix, hardcode the TPU generation
result="${TPU_GENERATIONS["$feature"]:-N/A}"
result="$(get_tpu_generation "$feature")"
elif [ "$is_quantization_matrix" = true ] && [ "$stage" == "QuantizationMethods" ]; then
# If it's the quantization matrix, hardcode the quantization methods
result="${QUANTIZATION_METHODS["$feature"]:-N/A}"
result="$(get_quantization_method "$feature")"

elif [[ "$mode" == "DEFAULT" ]]; then
result="✅"
result="✅ Passing"
else
result=$(buildkite-agent meta-data get "${TPU_METADATA_PREFIX}${feature}:${stage}" --default "N/A")
result=$(buildkite-agent meta-data get "${TPU_METADATA_PREFIX}${feature}:${stage}" --default "⚪ N/A")
# Format any remaining custom strings from upstream configs
local result_lower
result_lower="$(echo "$result" | tr '[:upper:]' '[:lower:]')"
if [[ "$result_lower" == "beta" ]]; then
result="⚠️ Beta"
elif [[ "$result_lower" == "experimental" ]]; then
result="🧪 Experimental"
elif [[ "$result_lower" == "planned" ]]; then
result="📝 Planned"
elif [[ "$result_lower" == "unplanned" ]]; then
result="⛔️ Unplanned"
fi
fi

row="$row,$result"

# Check for failure (exclude the hardcoded TPU generation column and Quantization Methods column)
if [ "$stage" != "QuantizationMethods" ] && [ "$stage" != "RecommendedTPUGenerations" ] && [ "${result}" != "✅" ] && [ "${result}" != "N/A" ] && [ "${result}" != "unverified" ]; then
if [ "$stage" != "QuantizationMethods" ] && [ "$stage" != "RecommendedTPUGenerations" ] && [[ "${result}" != "✅ Passing" && "${result}" != "N/A" && "${result}" != "❓ Untested" && "${result}" != "⚠️ Beta" && "${result}" != "🧪 Experimental" && "${result}" != "📝 Planned" && "${result}" != "⛔️ Unplanned" ]]; then
ANY_FAILED=true
fi

Expand All @@ -187,6 +208,7 @@ process_features() {
done
}

# Pivot Logic (Microbenchmarks)
process_kernel_matrix_to_pivot() {
local input_csv="${TPU_DIR}/kernel_support_matrix_microbenchmarks.csv"
local output_file="${TPU_DIR}/kernel_support_matrix-microbenchmarks.csv"
Expand All @@ -197,12 +219,13 @@ process_kernel_matrix_to_pivot() {
fi

# Define Headers for Display
local header="Kernel,W16 A16 (Correctness),W16 A16 (Performance),W8 A8 (Correctness),W8 A8 (Performance),W8 A16 (Correctness),W8 A16 (Performance),W4 A4 (Correctness),W4 A4 (Performance),W4 A8 (Correctness),W4 A8 (Performance),W4 A16 (Correctness),W4 A16 (Performance)"
local header="Kernel,W16 A16 (Corr),W16 A16 (Perf),W8 A8 (Corr),W8 A8 (Perf),W8 A16 (Corr),W8 A16 (Perf),W4 A4 (Corr),W4 A4 (Perf),W4 A8 (Corr),W4 A8 (Perf),W4 A16 (Corr),W4 A16 (Perf)"
echo "$header" > "$output_file"

# Define the quantization order to match the header
local quant_cols_list="w16a16 w8a8 w8a16 w4a4 w4a8 w4a16"


# Awk Script for Pivoting (Data Rows)
awk -v AWK_QUANT_COLS="$quant_cols_list" '
BEGIN {
Expand All @@ -221,7 +244,7 @@ process_kernel_matrix_to_pivot() {
}

# Store original Correctness ($2) and Performance ($3)
matrix[base_kernel_key][quant_type] = $2 OFS $3;
matrix[base_kernel_key, quant_type] = $2 OFS $3;

if (! (base_kernel_key in kernels)) {
kernels[base_kernel_key] = 1;
Expand All @@ -245,17 +268,17 @@ process_kernel_matrix_to_pivot() {
row = out_name;
for (j=1; j<=6; j++) {
q = q_order[j];
# Use original N/A,N/A if data is missing
data = (matrix[k][q] == "") ? "N/A,N/A" : matrix[k][q];
# Use formatted N/A if data is missing
data = (matrix[k, q] == "") ? "N/A,N/A" : matrix[k, q];
row = row OFS data;
}
print row >> "'"$output_file"'";
}
}
' "$input_csv"

# Display the pivoted result
echo "--- $output_file ---"
# Upload the newly created pivot table
echo "--- Uploading Pivoted Kernel Matrix: $output_file ---"
cat "$output_file"
buildkite-agent artifact upload "$output_file"
}
Expand Down
8 changes: 4 additions & 4 deletions .buildkite/scripts/record_step_result.sh
Original file line number Diff line number Diff line change
Expand Up @@ -44,16 +44,16 @@ message=""

case $OUTCOME in
"passed")
message="✅"
message="✅ Passing"
;;
"skipped")
message="N/A"
message="N/A"
;;
"unverified")
message="unverified"
message="❓ Untested"
;;
*)
message="❌"
message="❌ Failing"
;;
esac

Expand Down
132 changes: 132 additions & 0 deletions tests/test_generate_support_matrices.sh
Original file line number Diff line number Diff line change
@@ -0,0 +1,132 @@
#!/bin/bash
# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

set -euo pipefail

echo "Running Support Matrices Generator Tests..."

export TPU_VERSION="tpu6e"
export BUILDKITE_TAG="test_version"
TEST_DIR=$(mktemp -d)
export TEST_DIR
trap 'rm -rf "$TEST_DIR" v6' EXIT

cd "$(dirname "$0")/.."

# Mock buildkite-agent
buildkite-agent() {
local cmd="$1"
local subcmd="$2"
if [[ "$cmd" == "meta-data" && "$subcmd" == "get" ]]; then
local key="$3"
if [[ "$key" == "model-list" ]]; then
echo -e "test_model"
elif [[ "$key" == "feature-list" ]]; then
echo -e "kernel support matrix microbenchmarks"
elif [[ "$key" == "v6test_model_category" ]]; then
echo "text-only"
elif [[ "$key" == "v6test_model:UnitTest" ]]; then
echo "✅ Passing"
elif [[ "$key" == "v6test_model:Accuracy/Correctness" ]]; then
echo "❌ Failing"
elif [[ "$key" == "v6test_model:Benchmark" ]]; then
echo "⚪ N/A"
elif [[ "$key" == "v6kernel support matrix microbenchmarks_category" ]]; then
echo "kernel support matrix microbenchmarks"
else
echo "❓ Untested"
fi
elif [[ "$cmd" == "meta-data" && "$subcmd" == "set" ]]; then
echo "$3=$4" >> "$TEST_DIR/metadata.txt"
elif [[ "$cmd" == "artifact" && "$subcmd" == "upload" ]]; then
cp "$3" "$TEST_DIR/"
fi
}
export -f buildkite-agent

# Mock mapfile for MacOS compatibility (Bash 3.2 doesn't have it)
mapfile() {
local array_name=""
for arg in "$@"; do
if [[ "$arg" != "-t" ]]; then
array_name="$arg"
fi
done
local lines=()
while IFS= read -r line || [ -n "$line" ]; do
lines+=("$line")
done
if [[ -n "$array_name" ]]; then
eval "$array_name=(\"\${lines[@]}\")"
fi
}
export -f mapfile

# Set up mock microbenchmark input
mkdir -p v6
cat << 'EOF' > v6/kernel_support_matrix_microbenchmarks.csv
kernels,CorrectnessTest,PerformanceTest,TPU Versions
fused_moe-w8a8,❓ Untested,❓ Untested,v6
fused_moe-w8a16,✅ Passing,✅ Passing,v6
fused_moe-w16a16,⚪ N/A,⚪ N/A,v6
mla,✅ Passing,❓ Untested,v6
generic ragged paged attention v3,✅ Passing,✅ Passing,v6
EOF

# Execute the script
bash .buildkite/scripts/generate_support_matrices.sh > /dev/null

# Assertions
echo "Verifying Outputs..."
if ! grep -q "CI_TESTS_FAILED=true" "$TEST_DIR/metadata.txt"; then
echo "❌ ERROR: CI_TESTS_FAILED was not set to true despite test_model failing."
exit 1
fi

PIVOT_CSV="$TEST_DIR/kernel_support_matrix-microbenchmarks.csv"
if [[ ! -f "$PIVOT_CSV" ]]; then
echo "❌ ERROR: Pivoted Microbenchmark CSV not uploaded."
exit 1
fi

# Verify header
if ! grep -q "Kernel,W16 A16 (Corr),W16 A16 (Perf),W8 A8 (Corr),W8 A8 (Perf),W8 A16 (Corr),W8 A16 (Perf),W4 A4 (Corr),W4 A4 (Perf),W4 A8 (Corr),W4 A8 (Perf),W4 A16 (Corr),W4 A16 (Perf)" "$PIVOT_CSV"; then
echo "❌ ERROR: Incorrect Pivot Header."
cat "$PIVOT_CSV"
exit 1
fi

# Verify mla outputs correctly
if ! grep -q "\"mla\*\",✅ Passing,❓ Untested,⚪ N/A,⚪ N/A,⚪ N/A,⚪ N/A,⚪ N/A,⚪ N/A,⚪ N/A,⚪ N/A,⚪ N/A,⚪ N/A" "$PIVOT_CSV"; then
echo "❌ ERROR: mla parsing failed or default values are incorrect."
cat "$PIVOT_CSV"
exit 1
fi

# Verify fused_moe outputs correctly
if ! grep -q "\"fused_moe\",⚪ N/A,⚪ N/A,❓ Untested,❓ Untested,✅ Passing,✅ Passing,⚪ N/A,⚪ N/A,⚪ N/A,⚪ N/A,⚪ N/A,⚪ N/A" "$PIVOT_CSV"; then
echo "❌ ERROR: fused_moe parsing failed."
cat "$PIVOT_CSV"
exit 1
fi

# Verify string substitution correctness
if ! grep -q "\"generic ragged paged<br>attention v3\*\",✅ Passing,✅ Passing,⚪ N/A,⚪ N/A,⚪ N/A,⚪ N/A,⚪ N/A,⚪ N/A,⚪ N/A,⚪ N/A,⚪ N/A,⚪ N/A" "$PIVOT_CSV"; then
echo "❌ ERROR: generic ragged paged attention v3 substitution failed."
cat "$PIVOT_CSV"
exit 1
fi

echo "✅ All tests passed successfully!"
Loading