freesolo-flash 0.2.26__tar.gz → 0.2.30__tar.gz

freesolo_flash-0.2.30/.github/workflows/auto-rebake.yml

@@ -0,0 +1,245 @@
+name: auto-rebake per-arch kernel cache
+
+# Keeps the per-arch ghcr.io/freesolo-co/flash-worker:cu128-sm* images fresh whenever the :cu128 base
+# is rebuilt. Each cu128-smXX is physically `FROM :cu128` + a baked mega-cache, so a new base leaves
+# the per-arch tags pointing at the OLD base until rebaked. "Out of date" has two flavors, keyed off
+# the fp_cache / fp_base labels worker-image.yml stamps on :cu128 (docker/kernel_fingerprint.py):
+#
+#   * fp_cache changed -> the baked Triton/Inductor mega-cache is INVALID -> a real GPU re-warm is
+#     required (bake-kernel-cache.yml). PAID, and AUTO-fired by this detection (no manual approval) --
+#     fp_cache only moves on a deliberate toolchain bump, so spend stays tied to real need.
+#   * fp_cache same but fp_base changed -> the cache is still valid, the base just moved (FA wheel /
+#     causal-conv1d / rp_handler / a non-kernel dep) -> a FREE re-layer grafts the existing cache onto
+#     the fresh base (docker/Dockerfile.kernelcache.relayer), no GPU.
+#   * both same -> skip (this is the pure-flash-code case the old reminder over-fired on).
+#
+# The gate READS the published labels (authoritative for what was actually built, incl. a manual FA3
+# override) rather than recomputing from the checkout, and aborts unless the base's revision label ==
+# the build's head_sha (so a superseded base build or a `tag=...` dispatch that never touched :cu128
+# can't trigger a stale rebake).
+
+on:
+  workflow_run:
+    workflows: ["flash worker image"]
+    types: [completed]
+
+permissions:
+  contents: read
+  packages: write
+  actions: read
+  issues: write
+
+# Scope the group to the triggering base build's commit so a no-op completion (e.g. a tag=cu128-test
+# worker-image run, whose head_sha differs) can't cancel an in-flight real bake. The revision guard
+# already neutralizes superseded runs, so a global cancel-in-progress isn't needed.
+concurrency:
+  group: auto-rebake-${{ github.event.workflow_run.head_sha }}
+  cancel-in-progress: true
+
+env:
+  BASE: ghcr.io/freesolo-co/flash-worker:cu128
+
+jobs:
+  gate:
+    if: github.event.workflow_run.conclusion == 'success'
+    runs-on: ubuntu-latest
+    outputs:
+      proceed: ${{ steps.base.outputs.proceed }}
+      relayer_sms: ${{ steps.classify.outputs.relayer_sms }}
+      gpu_sms: ${{ steps.classify.outputs.gpu_sms }}
+      fp_cache: ${{ steps.base.outputs.fp_cache }}
+      fp_base: ${{ steps.base.outputs.fp_base }}
+      base_ref: ${{ steps.base.outputs.base_ref }}
+      base_revision: ${{ steps.base.outputs.revision }}
+    steps:
+      - uses: actions/checkout@v6
+        with:
+          ref: ${{ github.event.workflow_run.head_sha }}
+
+      - name: Install crane
+        run: |
+          curl -sSL https://github.com/google/go-containerregistry/releases/download/v0.20.2/go-containerregistry_Linux_x86_64.tar.gz \
+            | sudo tar -xz -C /usr/local/bin crane
+          crane version
+
+      - name: Log in to GHCR
+        uses: docker/login-action@v3
+        with:
+          registry: ghcr.io
+          username: ${{ github.actor }}
+          password: ${{ secrets.GITHUB_TOKEN }}
+
+      - name: Read base labels
+        id: base
+        env:
+          HEAD_SHA: ${{ github.event.workflow_run.head_sha }}
+        run: |
+          # canonical digest ref (repo@sha256, no tag) so every tool (crane/docker/runpod) accepts it.
+          digest=$(crane digest "$BASE")
+          repo=${BASE%:*}
+          cfg=$(crane config "$repo@$digest")
+          fpc=$(echo "$cfg" | jq -r '.config.Labels["co.freesolo.flash.kernelcache.fp_cache"] // ""')
+          fpb=$(echo "$cfg" | jq -r '.config.Labels["co.freesolo.flash.kernelcache.fp_base"] // ""')
+          rev=$(echo "$cfg" | jq -r '.config.Labels["org.opencontainers.image.revision"] // ""')
+          echo "base_ref=$repo@$digest" >> "$GITHUB_OUTPUT"
+          echo "fp_cache=$fpc" >> "$GITHUB_OUTPUT"
+          echo "fp_base=$fpb" >> "$GITHUB_OUTPUT"
+          echo "revision=$rev" >> "$GITHUB_OUTPUT"
+          # Only proceed if this base IS the build that triggered us: its revision label must match the
+          # build's head_sha. Guards (a) a newer base that already superseded this run, and (b) a
+          # `tag=cu128-test` dispatch that never updated :cu128 (the workflow_run payload can't reveal
+          # the dispatch tag). An unlabeled base (pre-rollout) also no-ops here.
+          if [ -n "$fpc" ] && [ "$rev" = "$HEAD_SHA" ]; then
+            echo "proceed=true" >> "$GITHUB_OUTPUT"
+            echo "base $BASE@$digest is rev $rev (== head_sha) fp_cache=$fpc fp_base=$fpb"
+          else
+            echo "proceed=false" >> "$GITHUB_OUTPUT"
+            echo "::notice::auto-rebake no-op: base rev '$rev' != head_sha '$HEAD_SHA' (or base unlabeled)"
+          fi
+
+      - name: Classify arches
+        id: classify
+        if: steps.base.outputs.proceed == 'true'
+        env:
+          FP_CACHE: ${{ steps.base.outputs.fp_cache }}
+          FP_BASE: ${{ steps.base.outputs.fp_base }}
+          REV: ${{ steps.base.outputs.revision }}
+        run: |
+          # arch list is the single source of truth in deps.py (mirrored by the bake matrix; a test
+          # pins that mirror). textual parse -> no flash import / uv sync needed.
+          arches=$(python3 -c "import re,pathlib; t=pathlib.Path('flash/providers/runpod/train/deps.py').read_text(); m=re.search(r'BAKED_PER_SM_ARCHES\s*=\s*frozenset\(\{([^}]+)\}\)', t); print(' '.join(re.findall(r'sm\d+', m.group(1))))")
+          relayer=(); gpu=()
+          summary="| arch | published fp_cache | published fp_base | action |\n|---|---|---|---|\n"
+          for sm in $arches; do
+            img="ghcr.io/freesolo-co/flash-worker:cu128-$sm"
+            if cfg=$(crane config "$img" 2>/dev/null); then
+              pc=$(echo "$cfg" | jq -r '.config.Labels["co.freesolo.flash.kernelcache.fp_cache"] // ""')
+              pb=$(echo "$cfg" | jq -r '.config.Labels["co.freesolo.flash.kernelcache.fp_base"] // ""')
+            else
+              pc="(missing)"; pb="(missing)"
+            fi
+            if [ "$pc" != "$FP_CACHE" ]; then
+              gpu+=("$sm"); action="GPU re-warm (cache invalid)"
+            elif [ "$pb" != "$FP_BASE" ]; then
+              relayer+=("$sm"); action="re-layer (free)"
+            else
+              action="up to date"
+            fi
+            echo "$sm: pc=$pc pb=$pb -> $action"
+            summary="$summary| $sm | \`$pc\` | \`$pb\` | $action |\n"
+          done
+          if [ ${#relayer[@]} -gt 0 ]; then
+            relayer_json=$(printf '%s\n' "${relayer[@]}" | jq -R . | jq -cs .)
+          else
+            relayer_json='[]'
+          fi
+          gpu_csv=$(IFS=,; echo "${gpu[*]}")
+          echo "relayer_sms=$relayer_json" >> "$GITHUB_OUTPUT"
+          echo "gpu_sms=$gpu_csv" >> "$GITHUB_OUTPUT"
+          {
+            echo "### auto-rebake plan (base rev \`$REV\`, fp_cache \`$FP_CACHE\` fp_base \`$FP_BASE\`)"
+            echo ""
+            echo -e "$summary"
+            echo ""
+            echo "re-layer (free): \`$relayer_json\` | GPU re-warm (paid, auto-started): \`${gpu_csv:-none}\`"
+          } >> "$GITHUB_STEP_SUMMARY"
+
+      # FYI notification (not an action item -- the bake auto-runs). Keeps a heads-up + cost trail of
+      # which arches a paid GPU re-warm fired for.
+      - name: Notify GPU re-bake auto-started
+        if: steps.base.outputs.proceed == 'true' && steps.classify.outputs.gpu_sms != ''
+        env:
+          GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+          GPU_SMS: ${{ steps.classify.outputs.gpu_sms }}
+          REV: ${{ steps.base.outputs.revision }}
+          RUN_URL: ${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}
+        run: |
+          title="kernel-cache: per-arch GPU re-bake auto-started"
+          body=$(printf 'The kernel toolchain changed (fp_cache), so a PAID GPU re-warm auto-started for these per-arch images (no approval needed):\n\n- arches: `%s`\n- base rev: `%s`\n\nTrack it in the run: %s' "$GPU_SMS" "$REV" "$RUN_URL")
+          # best-effort: a labeling/permission hiccup must never fail the gate. self-heal the label so
+          # no manual setup is needed, then reuse one open issue (comment) instead of spamming new ones.
+          gh label create kernel-rebake-pending --color B60205 --description "per-arch kernel cache GPU re-bake auto-started" 2>/dev/null || true
+          existing=$(gh issue list --label kernel-rebake-pending --state open --json number -q '.[0].number' 2>/dev/null || true)
+          if [ -n "${existing:-}" ]; then
+            gh issue comment "$existing" --body "$body" || true
+          else
+            gh issue create --title "$title" --body "$body" --label kernel-rebake-pending \
+              || gh issue create --title "$title" --body "$body" || true
+          fi
+
+  # FREE path: graft the existing valid cache onto the fresh base, no GPU. Runs automatically.
+  relayer:
+    needs: gate
+    if: needs.gate.outputs.proceed == 'true' && needs.gate.outputs.relayer_sms != '[]' && needs.gate.outputs.relayer_sms != ''
+    runs-on: ubuntu-24.04-8core
+    strategy:
+      fail-fast: false
+      matrix:
+        sm: ${{ fromJson(needs.gate.outputs.relayer_sms || '[]') }}
+    env:
+      BASE_REF: ${{ needs.gate.outputs.base_ref }}
+      FP_CACHE: ${{ needs.gate.outputs.fp_cache }}
+      FP_BASE: ${{ needs.gate.outputs.fp_base }}
+      BASE_REVISION: ${{ needs.gate.outputs.base_revision }}
+    steps:
+      - uses: actions/checkout@v6
+        with:
+          ref: ${{ github.event.workflow_run.head_sha }}
+
+      # The graft pulls two ~20GB images (old per-arch tag + fresh base); reclaim runner disk first.
+      - name: Free disk space
+        run: |
+          sudo rm -rf /usr/share/dotnet /opt/ghc /usr/local/.ghcup \
+            /usr/local/lib/android /opt/hostedtoolcache /usr/share/swift \
+            /usr/local/share/boost /usr/local/lib/node_modules \
+            /usr/local/share/powershell /usr/share/miniconda || true
+          sudo apt-get clean || true
+          sudo docker image prune -af || true
+          df -h /
+
+      - uses: docker/setup-buildx-action@v3
+
+      - name: Log in to GHCR
+        uses: docker/login-action@v3
+        with:
+          registry: ghcr.io
+          username: ${{ github.actor }}
+          password: ${{ secrets.GITHUB_TOKEN }}
+
+      # Pin OLDCACHE by digest so a concurrent overwrite of the source tag can't graft a wrong cache.
+      - name: Resolve old per-arch cache digest
+        id: old
+        run: |
+          img="ghcr.io/freesolo-co/flash-worker:cu128-${{ matrix.sm }}"
+          digest=$(docker buildx imagetools inspect "$img" --format '{{.Manifest.Digest}}')
+          # canonical digest ref (repo@sha256, no tag) for the FROM in the relayer build.
+          echo "ref=${img%:*}@$digest" >> "$GITHUB_OUTPUT"
+
+      - name: Re-layer cache onto the fresh base (no GPU)
+        uses: docker/build-push-action@v6
+        with:
+          context: .
+          file: docker/Dockerfile.kernelcache.relayer
+          push: true
+          tags: ghcr.io/freesolo-co/flash-worker:cu128-${{ matrix.sm }}
+          provenance: false
+          build-args: |
+            BASE=${{ env.BASE_REF }}
+            OLDCACHE=${{ steps.old.outputs.ref }}
+            FP_CACHE=${{ env.FP_CACHE }}
+            FP_BASE=${{ env.FP_BASE }}
+            BASE_REVISION=${{ env.BASE_REVISION }}
+
+  # PAID path: re-warm the cache on a real GPU per arch. Runs automatically -- the fingerprint gate
+  # above is the only thing that decides it fires (only when fp_cache changed), so no manual approval.
+  gpu-warm:
+    needs: gate
+    if: needs.gate.outputs.proceed == 'true' && needs.gate.outputs.gpu_sms != ''
+    uses: ./.github/workflows/bake-kernel-cache.yml
+    with:
+      sms: ${{ needs.gate.outputs.gpu_sms }}
+      fp_cache: ${{ needs.gate.outputs.fp_cache }}
+      fp_base: ${{ needs.gate.outputs.fp_base }}
+      base_ref: ${{ needs.gate.outputs.base_ref }}
+      base_revision: ${{ needs.gate.outputs.base_revision }}
+    secrets: inherit

{freesolo_flash-0.2.26 → freesolo_flash-0.2.30}/.github/workflows/bake-kernel-cache.yml

@@ -10,8 +10,12 @@ name: flash worker kernel-cache (per-arch)
 # Depends on the base ghcr.io/freesolo-co/flash-worker:cu128 image existing (worker-image.yml) — the
 # warmup runs INSIDE it so the cache matches the image's pinned torch/triton/fla/liger toolchain.
 #
-# Requires the RUNPOD_API_KEY repo secret (the warmup GPU pod) + the default GITHUB_TOKEN (GHCR push)
-# + HF_TOKEN (the temp dataset that ferries code + the cache artifact). Manual trigger; heavy + paid.
+# Secrets: RUNPOD_API_KEY (the warmup GPU pod) + HF_TOKEN (the temp dataset that ferries code + the
+# cache artifact) are native repo Actions secrets, synced from Infisical (prod:/flash-ci) by the
+# github-ci-flash Secret Sync; the GHCR push uses the default GITHUB_TOKEN. Heavy + paid,
+# but UNGATED: the fingerprint detection in auto-rebake.yml decides WHEN this runs (only on a real
+# kernel-toolchain change), so there is no manual approval step. Invoked by hand (workflow_dispatch)
+# OR automatically by auto-rebake.yml (workflow_call) when the kernel toolchain changed.
 
 on:
   workflow_dispatch:
@@ -20,12 +24,40 @@ on:
         description: "comma-separated sm list to bake; default is all validated arches"
         default: "sm80,sm86,sm89,sm90,sm120"
         required: false
+  # auto-rebake.yml calls this for the EXPENSIVE path (the cache is invalid and must be re-warmed on a
+  # real GPU). It passes the arch subset, the fingerprints to stamp, and a digest-pinned base so the
+  # warm runs inside the exact :cu128 the per-arch tag will sit on.
+  workflow_call:
+    inputs:
+      sms:
+        type: string
+        required: true
+      fp_cache:
+        type: string
+        required: false
+        default: ""
+      fp_base:
+        type: string
+        required: false
+        default: ""
+      base_ref:
+        type: string
+        required: false
+        default: "ghcr.io/freesolo-co/flash-worker:cu128"
+      base_revision:
+        type: string
+        required: false
+        default: ""
 
 permissions:
   contents: read
   packages: write
 
 jobs:
+  # Ungated: the PAID RunPod bake runs as soon as it's called/dispatched. The decision of WHEN to run
+  # lives in auto-rebake.yml's fingerprint gate (it only invokes this when fp_cache changed), so a
+  # second manual approval here would be redundant. Re-add an `approve` job with `environment: gpu-bake`
+  # (and `needs: approve`) if you ever want a human checkpoint back.
   bake:
     runs-on: ubuntu-24.04-8core
     strategy:
@@ -77,19 +109,63 @@ jobs:
               allowed_cuda: "13.0",
             }
     steps:
-      - uses: actions/checkout@v6
-
-      # skip arches not requested in the dispatch input
+      # Gate FIRST (no checkout needed) so skipped arches do no crane / checkout / build work.
       - name: Should bake ${{ matrix.sm }}?
         id: gate
         run: |
-          req="${{ github.event.inputs.sms || 'sm80,sm86,sm89,sm90,sm120' }}"
+          req="${{ inputs.sms || 'sm80,sm86,sm89,sm90,sm120' }}"
           req="${req// /}"   # tolerate spaces, e.g. "sm80, sm86"
           case ",$req," in
             *,${{ matrix.sm }},*) echo "run=true" >> "$GITHUB_OUTPUT" ;;
             *) echo "run=false" >> "$GITHUB_OUTPUT"; echo "skipping ${{ matrix.sm }}" ;;
           esac
 
+      - name: Install crane
+        if: steps.gate.outputs.run == 'true'
+        run: |
+          curl -sSL https://github.com/google/go-containerregistry/releases/download/v0.20.2/go-containerregistry_Linux_x86_64.tar.gz \
+            | sudo tar -xz -C /usr/local/bin crane
+          crane version
+
+      - name: Log in to GHCR
+        if: steps.gate.outputs.run == 'true'
+        uses: docker/login-action@v3
+        with:
+          registry: ghcr.io
+          username: ${{ github.actor }}
+          password: ${{ secrets.GITHUB_TOKEN }}
+
+      # Resolve the base + fingerprints BEFORE checkout, so the checkout (hence the warmed flash code)
+      # is pinned to the EXACT commit the base was built from. auto-rebake passes everything (fast path);
+      # a manual dispatch READS the labels off the published base, so the per-arch image always matches
+      # the base it extends (same source of truth as the gate) instead of recomputing from the checkout.
+      - name: Resolve base + fingerprints
+        id: base
+        if: steps.gate.outputs.run == 'true'
+        run: |
+          if [ -n "${{ inputs.fp_cache }}" ]; then
+            echo "ref=${{ inputs.base_ref }}" >> "$GITHUB_OUTPUT"
+            echo "fp_cache=${{ inputs.fp_cache }}" >> "$GITHUB_OUTPUT"
+            echo "fp_base=${{ inputs.fp_base }}" >> "$GITHUB_OUTPUT"
+            echo "revision=${{ inputs.base_revision }}" >> "$GITHUB_OUTPUT"
+          else
+            base="${{ inputs.base_ref || 'ghcr.io/freesolo-co/flash-worker:cu128' }}"
+            digest=$(crane digest "$base"); repo=${base%:*}
+            cfg=$(crane config "$repo@$digest")
+            echo "ref=$repo@$digest" >> "$GITHUB_OUTPUT"
+            echo "fp_cache=$(echo "$cfg" | jq -r '.config.Labels["co.freesolo.flash.kernelcache.fp_cache"] // ""')" >> "$GITHUB_OUTPUT"
+            echo "fp_base=$(echo "$cfg" | jq -r '.config.Labels["co.freesolo.flash.kernelcache.fp_base"] // ""')" >> "$GITHUB_OUTPUT"
+            echo "revision=$(echo "$cfg" | jq -r '.config.Labels["org.opencontainers.image.revision"] // ""')" >> "$GITHUB_OUTPUT"
+          fi
+
+      # Pin the checkout to the base's source commit so the warmed flash code matches the base exactly
+      # (the cache is content-keyed by the code that compiled it). Falls back to the event sha if the
+      # base predates the revision label.
+      - uses: actions/checkout@v6
+        if: steps.gate.outputs.run == 'true'
+        with:
+          ref: ${{ steps.base.outputs.revision || github.sha }}
+
       - name: Install uv
         if: steps.gate.outputs.run == 'true'
         uses: astral-sh/setup-uv@v5
@@ -103,17 +179,24 @@ jobs:
         if: steps.gate.outputs.run == 'true'
         run: uv sync --extra server
 
+      # RUNPOD_API_KEY (the warmup GPU pod) + HF_TOKEN (the temp dataset that ferries code/cache) are
+      # native repo Actions secrets, synced from Infisical (prod:/flash-ci) by the github-ci-flash
+      # Secret Sync. No Infisical CLI/login at job time -- Infisical stays the source of truth and
+      # rotations propagate via the sync. FLASH_CHALK_SPEC isn't synced, so bake_kernel_cache.py uses
+      # the repo-default DEFAULT_CHALK_SPEC (the value the fingerprint hashed) automatically. BASE_IMAGE
+      # is digest-pinned so the warm runs inside the EXACT base the per-arch tag will sit on.
       - name: Warm kernels on a RunPod ${{ matrix.sm }} GPU -> build/kernel_cache
         if: steps.gate.outputs.run == 'true'
         env:
           RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
           HF_TOKEN: ${{ secrets.HF_TOKEN }}
+          BASE_IMAGE: ${{ steps.base.outputs.ref }}
         run: |
           uv run python docker/bake_kernel_cache.py \
             --arch "${{ matrix.arch }}" --sm "${{ matrix.sm }}" \
             --gpu-type-id "${{ matrix.gpu_type_id }}" \
             --allowed-cuda "${{ matrix.allowed_cuda }}" \
-            --image ghcr.io/freesolo-co/flash-worker:cu128 \
+            --image "$BASE_IMAGE" \
             --out build/kernel_cache
 
       - name: Free disk space
@@ -130,17 +213,9 @@ jobs:
       - uses: docker/setup-buildx-action@v3
         if: steps.gate.outputs.run == 'true'
 
-      - name: Log in to GHCR
-        if: steps.gate.outputs.run == 'true'
-        uses: docker/login-action@v3
-        with:
-          registry: ghcr.io
-          username: ${{ github.actor }}
-          password: ${{ secrets.GITHUB_TOKEN }}
-
       # thin layer FROM the SAME base the warmup ran inside -> the cache toolchain always matches the
       # shipped image (rebuilding the full Dockerfile.worker from the checkout could drift from the
-      # published base and bake a cache that silently won't load).
+      # published base and bake a cache that silently won't load). GHCR auth from the login step above.
       - name: Build + push baked per-sm image (thin layer on the warmed base)
         if: steps.gate.outputs.run == 'true'
         uses: docker/build-push-action@v6
@@ -151,4 +226,7 @@ jobs:
           tags: ghcr.io/freesolo-co/flash-worker:cu128-${{ matrix.sm }}
           provenance: false
           build-args: |
-            BASE=ghcr.io/freesolo-co/flash-worker:cu128
+            BASE=${{ steps.base.outputs.ref }}
+            FP_CACHE=${{ steps.base.outputs.fp_cache }}
+            FP_BASE=${{ steps.base.outputs.fp_base }}
+            BASE_REVISION=${{ steps.base.outputs.revision }}

{freesolo_flash-0.2.26 → freesolo_flash-0.2.30}/.github/workflows/worker-image.yml

@@ -36,17 +36,40 @@ on:
       # so handler changes (e.g. the weight-cache preload branch) MUST rebuild the image. flash/**
       # covers the handler source (endpoints.py); the generator lives outside flash/, so watch it too.
       - docker/make_rp_handler.py
+      # The fp_cache/fp_base labels are computed by this script at build time, so a change to the
+      # fingerprint logic must rebuild :cu128 to re-stamp them (else auto-rebake never reclassifies).
+      - docker/kernel_fingerprint.py
 
 permissions:
   contents: read
   packages: write
 
+# A newer base build supersedes an older one: cancel the in-flight run so auto-rebake (which triggers
+# on this workflow's completion) only ever reasons about the latest :cu128. The tag is in the group so
+# a non-prod `tag=cu128-test` dispatch can't cancel an in-flight prod `cu128` push on the same ref
+# (which would leave :cu128 unpublished for that commit and auto-rebake blind to it).
+concurrency:
+  group: worker-image-${{ github.ref }}-${{ github.event.inputs.tag || 'cu128' }}
+  cancel-in-progress: true
+
 jobs:
   build:
     runs-on: ubuntu-24.04-8core
+    # Single source for the FlashAttention wheels: used BOTH as build-args and as fingerprint inputs
+    # (so a wheel bump moves fp_base). Keep FA3_SPEC in sync with Dockerfile.worker's ARG default.
+    env:
+      FA2_SPEC: https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.9.0/flash_attn-2.8.3%2Bcu128torch2.10-cp312-cp312-linux_x86_64.whl
+      FA3_SPEC: ${{ github.event.inputs.flash_attn_3_spec || 'https://github.com/windreamer/flash-attention3-wheels/releases/download/2026.03.19-850211f/flash_attn_3-3.0.0%2B20260318.cu128torch2100cxx11abitrue.8afc61-cp39-abi3-linux_x86_64.whl' }}
     steps:
       - uses: actions/checkout@v6
 
+      # Compute the kernel-cache fingerprints (stdlib-only, no uv sync) and stamp them as labels on
+      # :cu128 below, so auto-rebake.yml can read what this base was built from. Fails the build if a
+      # pin can't be parsed -- a silent miss would let a stale per-arch cache ship as a cold-JIT.
+      - name: Compute kernel fingerprints
+        id: fp
+        run: python3 docker/kernel_fingerprint.py --format github --fa2-spec "$FA2_SPEC" --fa3-spec "$FA3_SPEC" >> "$GITHUB_OUTPUT"
+
       - name: Free disk space
         run: |
           sudo rm -rf /usr/share/dotnet /opt/ghc /usr/local/.ghcup \
@@ -77,21 +100,20 @@ jobs:
           # so a fork run must not push to its own owner namespace and leave the real image stale.
           tags: ghcr.io/freesolo-co/flash-worker:${{ github.event.inputs.tag || 'cu128' }}
           provenance: false
+          # Stamp the kernel-cache fingerprints + the source commit so auto-rebake.yml can read what
+          # this base was built from (docker/kernel_fingerprint.py is the source of truth).
+          labels: |
+            co.freesolo.flash.kernelcache.fp_cache=${{ steps.fp.outputs.fp_cache }}
+            co.freesolo.flash.kernelcache.fp_base=${{ steps.fp.outputs.fp_base }}
+            org.opencontainers.image.revision=${{ github.sha }}
           # Install a PREBUILT flash-attn wheel instead of compiling from source (a multi-arch
           # source compile is huge and OOM-kills the runner). The wheel matches the base exactly
-          # (cu128 / torch2.10 / cp312).
+          # (cu128 / torch2.10 / cp312). FA specs come from the job-level env (single source).
           build-args: |
-            FLASH_ATTN_SPEC=https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.9.0/flash_attn-2.8.3%2Bcu128torch2.10-cp312-cp312-linux_x86_64.whl
-            FLASH_ATTN_3_SPEC=${{ github.event.inputs.flash_attn_3_spec || 'https://github.com/windreamer/flash-attention3-wheels/releases/download/2026.03.19-850211f/flash_attn_3-3.0.0%2B20260318.cu128torch2100cxx11abitrue.8afc61-cp39-abi3-linux_x86_64.whl' }}
+            FLASH_ATTN_SPEC=${{ env.FA2_SPEC }}
+            FLASH_ATTN_3_SPEC=${{ env.FA3_SPEC }}
 
-      # The per-arch baked images (cu128-<sm>) are rebaked MANUALLY (bake-kernel-cache.yml). When this
-      # rebuilds :cu128, those tags go stale -> prompt a rebake. Only when the cu128 tag was (re)built.
-      - name: Remind to rebake per-SM kernel-cache images
-        if: ${{ success() && (github.event.inputs.tag || 'cu128') == 'cu128' }}
-        run: |
-          echo "::warning title=Rebake per-SM images::A new :cu128 base was published. If FLASH_WORKER_IMAGE_PER_SM=1 is enabled, run bake-kernel-cache.yml to refresh cu128-sm{80,86,89,90,120}, else workers run stale baked deps."
-          {
-            echo "### ⚠️ Per-SM kernel-cache images may now be stale"
-            echo ""
-            echo "A fresh \`:cu128\` base was just published. If \`FLASH_WORKER_IMAGE_PER_SM=1\` is enabled, re-run **bake-kernel-cache.yml** to refresh the per-arch \`cu128-sm*\` tags, or workers will run the previous baked deps."
-          } >> "$GITHUB_STEP_SUMMARY"
+      # When :cu128 is republished the per-arch cu128-sm* tags go stale. auto-rebake.yml now handles
+      # that automatically on this workflow's completion: a free re-layer when only the base moved, or
+      # a one-click GPU re-warm when the kernel toolchain changed (it reads the fp_cache/fp_base labels
+      # stamped above vs each published per-arch tag). The old manual reminder is gone.

{freesolo_flash-0.2.26 → freesolo_flash-0.2.30}/Dockerfile.worker

@@ -25,9 +25,22 @@ ENV DEBIAN_FRONTEND=noninteractive \
 # JIT on a cold worker writes here too. These ENVs are always set (cheap); whether the dir is
 # pre-populated is what the opt-in bake controls. engine.worker._load_kernel_cache_if_present reads
 # /opt/flash/kernelcache at boot.
+# FLASHINFER_CUBIN_DIR / FLASHINFER_CACHE_DIR / FLASHINFER_WORKSPACE_BASE: FlashInfer's attention-backend
+# cubins (the Blackwell sm120 / B200 rollout attention backend) are NOT bundled in the wheel — flashinfer
+# downloads them from NVIDIA Artifactory on first use. Bake the cubin store + JIT cache into the same
+# persistent kernelcache tree so
+# a baked/volume cache serves them and a cold worker never re-fetches from Artifactory (kernel_warmup
+# also sets these). FLASHINFER_WORKSPACE_BASE is the var flashinfer 0.6.x actually reads to relocate the
+# JIT host-wrapper cache (FLASHINFER_CACHE_DIR is a DERIVED constant there, not an env knob), so set it
+# too — otherwise a B200 path that touches flashinfer before kernel_warmup runs lands the wrapper under
+# the default home workspace and re-JITs instead of using the baked cache.
 ENV TRITON_CACHE_DIR=/opt/flash/kernelcache/triton \
-    TORCHINDUCTOR_CACHE_DIR=/opt/flash/kernelcache/inductor
-RUN mkdir -p /opt/flash/kernelcache/triton /opt/flash/kernelcache/inductor
+    TORCHINDUCTOR_CACHE_DIR=/opt/flash/kernelcache/inductor \
+    FLASHINFER_CUBIN_DIR=/opt/flash/kernelcache/flashinfer_cubin \
+    FLASHINFER_CACHE_DIR=/opt/flash/kernelcache/flashinfer \
+    FLASHINFER_WORKSPACE_BASE=/opt/flash/kernelcache/flashinfer
+RUN mkdir -p /opt/flash/kernelcache/triton /opt/flash/kernelcache/inductor \
+    /opt/flash/kernelcache/flashinfer_cubin /opt/flash/kernelcache/flashinfer
 
 # build-essential for Triton/Inductor (Liger) + any source builds; git for pip VCS installs.
 RUN apt-get update && apt-get install -y --no-install-recommends build-essential git curl \
@@ -39,6 +52,7 @@ RUN pip install --no-cache-dir \
     "trl>=1.6,<1.7" \
     "peft>=0.19" \
     "vllm==0.19.1" \
+    "flashinfer-python==0.6.6" \
     "bitsandbytes>=0.49" \
     "datasets>=4.7,<6" \
     "huggingface_hub>=0.25" \
@@ -187,10 +201,12 @@ fa_ok = False; \
 exec('try:\n import flash_attn\n fa_ok = True\nexcept Exception:\n fa_ok = False'); \
 fa3_ok = False; \
 exec('try:\n import flash_attn_interface\n fa3_ok = True\nexcept Exception:\n fa3_ok = False'); \
+fi_ok = False; \
+exec('try:\n import flashinfer\n fi_ok = True\nexcept Exception:\n fi_ok = False'); \
 cudnn = torch.backends.cudnn.version(); assert cudnn, 'cuDNN not compiled into torch'; \
 ver = lambda n: M.version(n); \
-print('worker stack OK | torch', torch.__version__, '| cudnn', cudnn, '| triton', triton.__version__, '| flash_attn', 'present' if fa_ok else 'ABSENT (SDPA fallback)', '| flash_attn_3', 'present (Hopper FA3)' if fa3_ok else 'ABSENT (SDPA on Hopper)'); \
-[print('  ', d, ver(d)) for d in ['transformers','trl','peft','vllm','datasets','huggingface-hub', \
+print('worker stack OK | torch', torch.__version__, '| cudnn', cudnn, '| triton', triton.__version__, '| flash_attn', 'present' if fa_ok else 'ABSENT (SDPA fallback)', '| flash_attn_3', 'present (Hopper FA3)' if fa3_ok else 'ABSENT (SDPA on Hopper)', '| flashinfer', 'present (B200 MoE/attn)' if fi_ok else 'ABSENT (Triton fused-MoE fallback)'); \
+[print('  ', d, ver(d)) for d in ['transformers','trl','peft','vllm','flashinfer-python','datasets','huggingface-hub', \
 'accelerate','wandb','liger-kernel','flash-linear-attention','bitsandbytes', \
 'hf-transfer','freesolo','runpod']]"
 

{freesolo_flash-0.2.26 → freesolo_flash-0.2.30}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: freesolo-flash
-Version: 0.2.26
+Version: 0.2.30
 Summary: Flash — managed LoRA post-training (SFT/GRPO) for Freesolo environments, driven by the `flash` CLI
 Project-URL: Homepage, https://github.com/freesolo-co/flash
 Project-URL: Repository, https://github.com/freesolo-co/flash

{freesolo_flash-0.2.26 → freesolo_flash-0.2.30}/docker/Dockerfile.kernelcache

@@ -16,3 +16,13 @@ ENV TRITON_CACHE_DIR=/opt/flash/kernelcache/triton \
     TORCHINDUCTOR_CACHE_DIR=/opt/flash/kernelcache/inductor
 
 COPY build/kernel_cache/ /opt/flash/kernelcache/
+
+# Fingerprint + base-revision stamps so auto-rebake.yml can read what this image was baked from and
+# decide skip / re-layer / re-warm (docker/kernel_fingerprint.py is the source of truth for the
+# values). Default to "unknown" so a hand build still works; CI always passes the real ones.
+ARG FP_CACHE=unknown
+ARG FP_BASE=unknown
+ARG BASE_REVISION=unknown
+LABEL co.freesolo.flash.kernelcache.fp_cache=$FP_CACHE \
+      co.freesolo.flash.kernelcache.fp_base=$FP_BASE \
+      org.opencontainers.image.revision=$BASE_REVISION

freesolo_flash-0.2.30/docker/Dockerfile.kernelcache.relayer

@@ -0,0 +1,27 @@
+# Cheap per-arch RE-LAYER: graft the EXISTING baked mega-cache onto a freshly rebuilt :cu128 base,
+# with NO GPU. Used by auto-rebake.yml when fp_cache is unchanged (the cache is still valid) but the
+# base moved (fp_base changed) -- e.g. an FA2/FA3/causal-conv1d/rp_handler bump. The valid cache
+# already lives inside the currently-published per-arch tag, so OLDCACHE IS the cache store; we just
+# re-stack it on the new base. This is what makes a base-only change free (vs a paid GPU re-warm).
+#
+# BASE and OLDCACHE are passed DIGEST-PINNED by the gate so a concurrent base/tag overwrite can't
+# graft the cache onto the wrong base.
+ARG BASE=ghcr.io/freesolo-co/flash-worker:cu128
+ARG OLDCACHE=ghcr.io/freesolo-co/flash-worker:cu128-sm89
+FROM ${OLDCACHE} AS oldcache
+FROM ${BASE}
+
+# Match Dockerfile.worker / Dockerfile.kernelcache so the layer is self-contained even on an older base.
+ENV TRITON_CACHE_DIR=/opt/flash/kernelcache/triton \
+    TORCHINDUCTOR_CACHE_DIR=/opt/flash/kernelcache/inductor
+
+# Pull the warmed cache straight out of the old per-arch image -- no GPU, no external cache store.
+COPY --from=oldcache /opt/flash/kernelcache /opt/flash/kernelcache
+
+# Re-stamp with the CURRENT fingerprints (fp_cache unchanged, fp_base + revision are the new base's).
+ARG FP_CACHE=unknown
+ARG FP_BASE=unknown
+ARG BASE_REVISION=unknown
+LABEL co.freesolo.flash.kernelcache.fp_cache=$FP_CACHE \
+      co.freesolo.flash.kernelcache.fp_base=$FP_BASE \
+      org.opencontainers.image.revision=$BASE_REVISION