#!/usr/bin/env bash
# Creates a patched copy of vllm/utils/flashinfer.py that forces the FP8 GEMM
# backend to cuBLAS instead of "auto" (which selects cuDNN on GB10 / sm_121a,
# see https://github.com/flashinfer-ai/flashinfer/issues/3566).
#
# Valid ONLY for the image below. Regenerate after every image change and check
# first whether the FlashInfer issue is fixed in the new image.
#
# Usage:  ./make-flashinfer-cublas-patch.sh [target-dir]
#         default target-dir: /srv/vllm-patches  (needs write access)
set -euo pipefail

IMG="ghcr.io/timothystewart6/vllm-gb10:v0.28.0-gb10.2"
SRC="/usr/local/lib/python3.12/dist-packages/vllm/utils/flashinfer.py"
TARGET_DIR="${1:-/srv/vllm-patches}"
OUT="$TARGET_DIR/flashinfer_fp8_cublas.py"

WORK="$(mktemp -d)"
CID=""
cleanup() {
  [ -n "$CID" ] && docker rm -f "$CID" >/dev/null 2>&1 || true
  rm -rf "$WORK"
}
trap cleanup EXIT

CID="$(docker create "$IMG")"
docker cp "$CID:$SRC" "$WORK/flashinfer.py"
cp "$WORK/flashinfer.py" "$WORK/flashinfer_fp8_cublas.py"

# The file must contain exactly two lines of this form: the backend argument
# of the two FP8 GEMM calls. Anything else means a different vLLM version.
count="$(grep -c '^        "auto",$' "$WORK/flashinfer_fp8_cublas.py" || true)"
if [ "$count" != "2" ]; then
  echo "Expected 2 matching lines, found $count. Not patching." >&2
  exit 1
fi
sed -i 's/^        "auto",$/        "cublas",/' "$WORK/flashinfer_fp8_cublas.py"

echo "Changes against the original:"
diff "$WORK/flashinfer.py" "$WORK/flashinfer_fp8_cublas.py" || true

mkdir -p "$TARGET_DIR"
install -m 0644 "$WORK/flashinfer_fp8_cublas.py" "$OUT"
echo "Written: $OUT"
echo "Mount it read-only into the container:"
echo "  -v $OUT:$SRC:ro"