diff --git a/.github/workflows/build.yaml b/.github/workflows/build.yaml
index 1f5aa0b8..543d64fe 100644
--- a/.github/workflows/build.yaml
+++ b/.github/workflows/build.yaml
@@ -18,161 +18,160 @@ on:
         type: boolean
 
 jobs:
-  # build-and-push:
-  #   outputs:
-  #     docker_image: ${{ steps.final.outputs.docker_image }}
-  #     docker_devices: ${{ steps.final.outputs.docker_devices }}
-  #     runs_on: ${{ steps.final.outputs.runs_on }}
-  #     label: ${{ steps.final.outputs.label }}
-  #     extra_pytest: ${{ steps.final.outputs.extra_pytest }}
-  #   concurrency:
-  #     group: ${{ github.workflow }}-build-and-push-image-${{ inputs.hardware }}-${{ github.head_ref || github.run_id }}
-  #     cancel-in-progress: true
-  #   runs-on:
-  #     group: aws-highmemory-32-plus-priv
-  #   permissions:
-  #     contents: write
-  #     packages: write
-  #   steps:
-  #     - name: Checkout repository
-  #       uses: actions/checkout@v4
-  #     - name: Inject slug/short variables
-  #       uses: rlespinasse/github-slug-action@v4.4.1
-  #     - name: Construct harware variables
-  #       shell: bash
-  #       run: |
-  #         case ${{ inputs.hardware }} in
-  #           cuda)
-  #               export dockerfile="Dockerfile"
-  #               export label_extension=""
-  #               export docker_devices=""
-  #               export runs_on="aws-g6-12xl-plus-priv-cache"
-  #               export platform=""
-  #               export extra_pytest=""
-  #               ;;
-  #           rocm)
-  #               export dockerfile="Dockerfile_amd"
-  #               export label_extension="-rocm"
-  #               export docker_devices="/dev/kfd,/dev/dri"
-  #               export runs_on="amd-gpu-runners"
-  #               export platform=""
-  #               export extra_pytest="-k flash_llama_grammar_tools"
-  #               ;;
-  #           intel-xpu)
-  #               export dockerfile="Dockerfile_intel"
-  #               export label_extension="-intel-xpu"
-  #               export docker_devices=""
-  #               export runs_on="ubuntu-latest"
-  #               export platform="xpu"
-  #               export extra_pytest=""
-  #               ;;
-  #           intel-cpu)
-  #               export dockerfile="Dockerfile_intel"
-  #               export label_extension="-intel-cpu"
-  #               export docker_devices=""
-  #               export runs_on="ubuntu-latest"
-  #               export platform="cpu"
-  #               export extra_pytest=""
-  #               ;;
-  #         esac
-  #         echo $dockerfile
-  #         echo "Dockerfile=${dockerfile}"
-  #         echo $label_extension
-  #         echo $docker_devices
-  #         echo $runs_on
-  #         echo $platform
-  #         echo "DOCKERFILE=${dockerfile}" >> $GITHUB_ENV
-  #         echo "LABEL=${label_extension}" >> $GITHUB_ENV
-  #         echo "PLATFORM=${platform}" >> $GITHUB_ENV
-  #         echo "DOCKER_DEVICES=${docker_devices}" >> $GITHUB_ENV
-  #         echo "RUNS_ON=${runs_on}" >> $GITHUB_ENV
-  #         echo "EXTRA_PYTEST=${extra_pytest}" >> $GITHUB_ENV
-  #         echo REGISTRY_MIRROR=$REGISTRY_MIRROR >> $GITHUB_ENV
-  #     - name: Initialize Docker Buildx
-  #       uses: docker/setup-buildx-action@v3
-  #       with:
-  #         install: true
-  #         buildkitd-config: /tmp/buildkitd.toml
-  #     - name: Login to internal Container Registry
-  #       uses: docker/login-action@v3
-  #       with:
-  #         username: ${{ secrets.REGISTRY_USERNAME }}
-  #         password: ${{ secrets.REGISTRY_PASSWORD }}
-  #         registry: registry.internal.huggingface.tech
-  #     - name: Login to GitHub Container Registry
-  #       if: github.event_name != 'pull_request'
-  #       uses: docker/login-action@v3
-  #       with:
-  #         registry: ghcr.io
-  #         username: ${{ github.actor }}
-  #         password: ${{ secrets.GITHUB_TOKEN }}
-  #     - name: Login to Azure Container Registry
-  #       if: github.event_name != 'pull_request'
-  #       uses: docker/login-action@v3
-  #       with:
-  #         username: ${{ secrets.AZURE_DOCKER_USERNAME }}
-  #         password: ${{ secrets.AZURE_DOCKER_PASSWORD }}
-  #         registry: db4c2190dd824d1f950f5d1555fbadf0.azurecr.io
-  #     # If pull request
-  #     - name: Extract metadata (tags, labels) for Docker
-  #       if: ${{ github.event_name == 'pull_request' }}
-  #       id: meta-pr
-  #       uses: docker/metadata-action@v5
-  #       with:
-  #         images: |
-  #           registry.internal.huggingface.tech/api-inference/community/text-generation-inference
-  #         tags: |
-  #           type=raw,value=sha-${{ env.GITHUB_SHA_SHORT }}${{ env.LABEL }}
-  #     # If main, release or tag
-  #     - name: Extract metadata (tags, labels) for Docker
-  #       if: ${{ github.event_name != 'pull_request' }}
-  #       id: meta
-  #       uses: docker/metadata-action@v4.3.0
-  #       with:
-  #         flavor: |
-  #           latest=auto
-  #         images: |
-  #           registry.internal.huggingface.tech/api-inference/community/text-generation-inference
-  #           ghcr.io/huggingface/text-generation-inference
-  #           db4c2190dd824d1f950f5d1555fbadf0.azurecr.io/text-generation-inference
-  #         tags: |
-  #           type=semver,pattern={{version}}${{ env.LABEL }}
-  #           type=semver,pattern={{major}}.{{minor}}${{ env.LABEL }}
-  #           type=raw,value=latest${{ env.LABEL }},enable=${{ github.ref == format('refs/heads/{0}', github.event.repository.default_branch) }}
-  #           type=raw,value=sha-${{ env.GITHUB_SHA_SHORT }}${{ env.LABEL }}
-  #     - name: Build and push Docker image
-  #       id: build-and-push
-  #       uses: docker/build-push-action@v4
-  #       with:
-  #         context: .
-  #         file: ${{ env.DOCKERFILE }}
-  #         push: true
-  #         platforms: 'linux/amd64'
-  #         build-args: |
-  #           GIT_SHA=${{ env.GITHUB_SHA }}
-  #           DOCKER_LABEL=sha-${{ env.GITHUB_SHA_SHORT }}${{ env.LABEL }}
-  #           PLATFORM=${{ env.PLATFORM }}
-  #         tags: ${{ steps.meta.outputs.tags || steps.meta-pr.outputs.tags }}
-  #         labels: ${{ steps.meta.outputs.labels || steps.meta-pr.outputs.labels }}
-  #         cache-from: type=s3,region=us-east-1,bucket=ci-docker-buildx-cache,name=text-generation-inference-cache${{ env.LABEL }},mode=min,access_key_id=${{ secrets.S3_CI_DOCKER_BUILDX_CACHE_ACCESS_KEY_ID }},secret_access_key=${{ secrets.S3_CI_DOCKER_BUILDX_CACHE_SECRET_ACCESS_KEY }},mode=min
-  #         cache-to: type=s3,region=us-east-1,bucket=ci-docker-buildx-cache,name=text-generation-inference-cache${{ env.LABEL }},mode=min,access_key_id=${{ secrets.S3_CI_DOCKER_BUILDX_CACHE_ACCESS_KEY_ID }},secret_access_key=${{ secrets.S3_CI_DOCKER_BUILDX_CACHE_SECRET_ACCESS_KEY }},mode=min
-  #     - name: Final
-  #       id: final
-  #       run: |
-  #         echo "docker_image=registry.internal.huggingface.tech/api-inference/community/text-generation-inference:sha-${{ env.GITHUB_SHA_SHORT}}${{ env.LABEL }}" >> "$GITHUB_OUTPUT"
-  #         echo "docker_devices=${{ env.DOCKER_DEVICES }}" >> "$GITHUB_OUTPUT"
-  #         echo "runs_on=${{ env.RUNS_ON }}" >> "$GITHUB_OUTPUT"
-  #         echo "label=${{ env.LABEL }}" >> "$GITHUB_OUTPUT"
-  #         echo "extra_pytest=${{ env.EXTRA_PYTEST }}" >> "$GITHUB_OUTPUT"
+  build-and-push:
+    outputs:
+      docker_image: ${{ steps.final.outputs.docker_image }}
+      docker_devices: ${{ steps.final.outputs.docker_devices }}
+      runs_on: ${{ steps.final.outputs.runs_on }}
+      label: ${{ steps.final.outputs.label }}
+      extra_pytest: ${{ steps.final.outputs.extra_pytest }}
+    concurrency:
+      group: ${{ github.workflow }}-build-and-push-image-${{ inputs.hardware }}-${{ github.head_ref || github.run_id }}
+      cancel-in-progress: true
+    runs-on:
+      group: aws-highmemory-32-plus-priv
+    permissions:
+      contents: write
+      packages: write
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v4
+      - name: Inject slug/short variables
+        uses: rlespinasse/github-slug-action@v4.4.1
+      - name: Construct harware variables
+        shell: bash
+        run: |
+          case ${{ inputs.hardware }} in
+            cuda)
+                export dockerfile="Dockerfile"
+                export label_extension=""
+                export docker_devices=""
+                export runs_on="aws-g6-12xl-plus-priv-cache"
+                export platform=""
+                export extra_pytest=""
+                ;;
+            rocm)
+                export dockerfile="Dockerfile_amd"
+                export label_extension="-rocm"
+                export docker_devices="/dev/kfd,/dev/dri"
+                export runs_on="amd-gpu-runners"
+                export platform=""
+                export extra_pytest="-k test_flash_gemma_load"
+                ;;
+            intel-xpu)
+                export dockerfile="Dockerfile_intel"
+                export label_extension="-intel-xpu"
+                export docker_devices=""
+                export runs_on="ubuntu-latest"
+                export platform="xpu"
+                export extra_pytest=""
+                ;;
+            intel-cpu)
+                export dockerfile="Dockerfile_intel"
+                export label_extension="-intel-cpu"
+                export docker_devices=""
+                export runs_on="ubuntu-latest"
+                export platform="cpu"
+                export extra_pytest=""
+                ;;
+          esac
+          echo $dockerfile
+          echo "Dockerfile=${dockerfile}"
+          echo $label_extension
+          echo $docker_devices
+          echo $runs_on
+          echo $platform
+          echo "DOCKERFILE=${dockerfile}" >> $GITHUB_ENV
+          echo "LABEL=${label_extension}" >> $GITHUB_ENV
+          echo "PLATFORM=${platform}" >> $GITHUB_ENV
+          echo "DOCKER_DEVICES=${docker_devices}" >> $GITHUB_ENV
+          echo "RUNS_ON=${runs_on}" >> $GITHUB_ENV
+          echo "EXTRA_PYTEST=${extra_pytest}" >> $GITHUB_ENV
+          echo REGISTRY_MIRROR=$REGISTRY_MIRROR >> $GITHUB_ENV
+      - name: Initialize Docker Buildx
+        uses: docker/setup-buildx-action@v3
+        with:
+          install: true
+          buildkitd-config: /tmp/buildkitd.toml
+      - name: Login to internal Container Registry
+        uses: docker/login-action@v3
+        with:
+          username: ${{ secrets.REGISTRY_USERNAME }}
+          password: ${{ secrets.REGISTRY_PASSWORD }}
+          registry: registry.internal.huggingface.tech
+      - name: Login to GitHub Container Registry
+        if: github.event_name != 'pull_request'
+        uses: docker/login-action@v3
+        with:
+          registry: ghcr.io
+          username: ${{ github.actor }}
+          password: ${{ secrets.GITHUB_TOKEN }}
+      - name: Login to Azure Container Registry
+        if: github.event_name != 'pull_request'
+        uses: docker/login-action@v3
+        with:
+          username: ${{ secrets.AZURE_DOCKER_USERNAME }}
+          password: ${{ secrets.AZURE_DOCKER_PASSWORD }}
+          registry: db4c2190dd824d1f950f5d1555fbadf0.azurecr.io
+      # If pull request
+      - name: Extract metadata (tags, labels) for Docker
+        if: ${{ github.event_name == 'pull_request' }}
+        id: meta-pr
+        uses: docker/metadata-action@v5
+        with:
+          images: |
+            registry.internal.huggingface.tech/api-inference/community/text-generation-inference
+          tags: |
+            type=raw,value=sha-${{ env.GITHUB_SHA_SHORT }}${{ env.LABEL }}
+      # If main, release or tag
+      - name: Extract metadata (tags, labels) for Docker
+        if: ${{ github.event_name != 'pull_request' }}
+        id: meta
+        uses: docker/metadata-action@v4.3.0
+        with:
+          flavor: |
+            latest=auto
+          images: |
+            registry.internal.huggingface.tech/api-inference/community/text-generation-inference
+            ghcr.io/huggingface/text-generation-inference
+            db4c2190dd824d1f950f5d1555fbadf0.azurecr.io/text-generation-inference
+          tags: |
+            type=semver,pattern={{version}}${{ env.LABEL }}
+            type=semver,pattern={{major}}.{{minor}}${{ env.LABEL }}
+            type=raw,value=latest${{ env.LABEL }},enable=${{ github.ref == format('refs/heads/{0}', github.event.repository.default_branch) }}
+            type=raw,value=sha-${{ env.GITHUB_SHA_SHORT }}${{ env.LABEL }}
+      - name: Build and push Docker image
+        id: build-and-push
+        uses: docker/build-push-action@v4
+        with:
+          context: .
+          file: ${{ env.DOCKERFILE }}
+          push: true
+          platforms: 'linux/amd64'
+          build-args: |
+            GIT_SHA=${{ env.GITHUB_SHA }}
+            DOCKER_LABEL=sha-${{ env.GITHUB_SHA_SHORT }}${{ env.LABEL }}
+            PLATFORM=${{ env.PLATFORM }}
+          tags: ${{ steps.meta.outputs.tags || steps.meta-pr.outputs.tags }}
+          labels: ${{ steps.meta.outputs.labels || steps.meta-pr.outputs.labels }}
+          cache-from: type=s3,region=us-east-1,bucket=ci-docker-buildx-cache,name=text-generation-inference-cache${{ env.LABEL }},mode=min,access_key_id=${{ secrets.S3_CI_DOCKER_BUILDX_CACHE_ACCESS_KEY_ID }},secret_access_key=${{ secrets.S3_CI_DOCKER_BUILDX_CACHE_SECRET_ACCESS_KEY }},mode=min
+          cache-to: type=s3,region=us-east-1,bucket=ci-docker-buildx-cache,name=text-generation-inference-cache${{ env.LABEL }},mode=min,access_key_id=${{ secrets.S3_CI_DOCKER_BUILDX_CACHE_ACCESS_KEY_ID }},secret_access_key=${{ secrets.S3_CI_DOCKER_BUILDX_CACHE_SECRET_ACCESS_KEY }},mode=min
+      - name: Final
+        id: final
+        run: |
+          echo "docker_image=registry.internal.huggingface.tech/api-inference/community/text-generation-inference:sha-${{ env.GITHUB_SHA_SHORT}}${{ env.LABEL }}" >> "$GITHUB_OUTPUT"
+          echo "docker_devices=${{ env.DOCKER_DEVICES }}" >> "$GITHUB_OUTPUT"
+          echo "runs_on=${{ env.RUNS_ON }}" >> "$GITHUB_OUTPUT"
+          echo "label=${{ env.LABEL }}" >> "$GITHUB_OUTPUT"
+          echo "extra_pytest=${{ env.EXTRA_PYTEST }}" >> "$GITHUB_OUTPUT"
   integration_tests:
     concurrency:
       group: ${{ github.workflow }}-${{ github.job }}-${{ needs.build-and-push.outputs.label }}-${{ github.head_ref || github.run_id }}
       cancel-in-progress: true
-      # needs: build-and-push
+      needs: build-and-push
     runs-on:
-      # group: ${{ needs.build-and-push.outputs.runs_on }}
-      group: amd-gpu-runners
-      # if: needs.build-and-push.outputs.runs_on != 'ubuntu-latest'
+      group: ${{ needs.build-and-push.outputs.runs_on }}
+      if: needs.build-and-push.outputs.runs_on != 'ubuntu-latest'
     env:
       PYTEST_FLAGS: ${{ (startsWith(github.ref, 'refs/tags/') || github.ref == 'refs/heads/main' || inputs.release-tests == true) && '--release' || '--release' }}
     steps:
@@ -190,15 +189,9 @@ jobs:
       - name: Run tests
         run: |
           export DOCKER_VOLUME=/mnt2/
-          # export DOCKER_IMAGE=${{ needs.build-and-push.outputs.docker_image }}
-          export DOCKER_IMAGE=registry.internal.huggingface.tech/api-inference/community/text-generation-inference:sha-318bbc7-rocm
-          # export DOCKER_DEVICES=${{ needs.build-and-push.outputs.docker_devices }}
-          export DOCKER_DEVICES=/dev/kfd,/dev/dri
-          # export EXTRA_PYTEST="${{ needs.build-and-push.outputs.extra_pytest }}"
+          export DOCKER_IMAGE=${{ needs.build-and-push.outputs.docker_image }}
+          export DOCKER_DEVICES=${{ needs.build-and-push.outputs.docker_devices }}
+          export EXTRA_PYTEST="${{ needs.build-and-push.outputs.extra_pytest }}"
           export HF_TOKEN=${{ secrets.HF_TOKEN }}
           echo $DOCKER_IMAGE
-          python -c 'import docker; client = docker.from_env(); client.containers.run("registry.internal.huggingface.tech/api-inference/community/text-generation-inference:sha-318bbc7-rocm", devices=["/dev/kfd", "/dev/dri"], environment={"MODEL_ID": "gpt2", "NUM_SHARD": 2})'
-          # which pytest
-          pip freeze
-          # pytest -s -vv integration-tests -k flash_llama_grammar_tools_auto
-          # pytest -s -vv integration-tests ${PYTEST_FLAGS} ${EXTRA_PYTEST}
+          pytest -s -vv integration-tests ${PYTEST_FLAGS} ${EXTRA_PYTEST}