diff --git a/.github/workflows/docker.yaml b/.github/workflows/docker.yaml
index 838bd55df1..eb8c15c323 100644
--- a/.github/workflows/docker.yaml
+++ b/.github/workflows/docker.yaml
@@ -10,6 +10,9 @@ on:
       ref:
         description: 'The git ref (branch or tag) to build the Docker image from.'
         required: true
+      version:
+        description: 'Version tag for the Docker image (e.g., 0.3.3)'
+        required: true
 
 env:
   REGISTRY: ghcr.io
@@ -28,6 +31,9 @@ jobs:
     steps:
       - name: Checkout repository
         uses: actions/checkout@v5
+        with:
+          ref: ${{ github.event.inputs.ref || github.ref }}
+          path: trinity-${{ github.run_id }}
       # Uses the `docker/login-action` action to log in to the Container registry registry using the account and password that will publish the packages. Once published, the packages are scoped to the account defined here.
       - name: Log in to the Container registry
         uses: docker/login-action@v2
@@ -42,8 +48,10 @@ jobs:
         with:
           images: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}
           tags: |
-            type=match,pattern=\d.\d.\d
-            type=sha
+            type=match,pattern=\d.\d.\d,enable=${{ github.event_name == 'release' }}
+            type=sha,enable=${{ github.event_name == 'release' }}
+            type=raw,value=${{ github.event.inputs.version }},enable=${{ github.event_name == 'workflow_dispatch' }}
+            type=sha,enable=${{ github.event_name == 'workflow_dispatch' }}
       # This step uses the `docker/build-push-action` action to build the image, based on your repository's `Dockerfile`. If the build succeeds, it pushes the image to GitHub Packages.
       # It uses the `context` parameter to define the build's context as the set of files located in the specified path. For more information, see [Usage](https://github.com/docker/build-push-action#usage) in the README of the `docker/build-push-action` repository.
       # It uses the `tags` and `labels` parameters to tag and label the image with the output from the "meta" step.
@@ -51,9 +59,9 @@ jobs:
         id: push
         uses: docker/build-push-action@v4
         with:
-          context: .
+          context: trinity-${{ github.run_id }}
           push: true
-          file: scripts/docker/Dockerfile
+          file: trinity-${{ github.run_id }}/scripts/docker/Dockerfile
           shm-size: 64g
           tags: ${{ steps.meta.outputs.tags }}
           labels: ${{ steps.meta.outputs.labels }}
@@ -65,3 +73,8 @@ jobs:
           subject-name: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME}}
           subject-digest: ${{ steps.push.outputs.digest }}
           push-to-registry: true
+
+      - name: Cleanup workspace
+        if: always()
+        run: |
+          sudo rm -rf trinity-${{ github.run_id }} 2>/dev/null
diff --git a/scripts/docker/Dockerfile.uv b/scripts/docker/Dockerfile.uv
index ad653226bb..3d40a279c5 100644
--- a/scripts/docker/Dockerfile.uv
+++ b/scripts/docker/Dockerfile.uv
@@ -40,7 +40,7 @@ RUN . /opt/venv/bin/activate && \
 
 # Install flash_attn and Megatron
 RUN . /opt/venv/bin/activate && \
-    uv pip install -e.[flash_attn] && \
+    uv pip install flash_attn==2.8.1 --no-deps --no-cache-dir && \
     uv pip install -e .[megatron] && \
     NVCC_APPEND_FLAGS="--threads 4" APEX_PARALLEL_BUILD=8 \
     uv pip install -v --no-build-isolation \