diff --git a/.github/actions/run-inference/action.yml b/.github/actions/run-inference/action.yml index cbb90495..7d659c04 100644 --- a/.github/actions/run-inference/action.yml +++ b/.github/actions/run-inference/action.yml @@ -3,7 +3,7 @@ description: Run one llama-tornado inference pass and write the metrics + sideca inputs: backend: - description: 'GPU backend (opencl, ptx, or cuda)' + description: 'GPU backend (opencl, cuda, or metal)' required: true model_file: description: 'Model filename inside $MODELS_DIR (e.g. Llama-3.2-1B-Instruct-F16.gguf)' @@ -43,7 +43,7 @@ runs: METRICS_FILE: ${{ inputs.metrics_file }} run: | # Run inference and emit raw metrics JSON via JAVA_TOOL_OPTIONS - # Backend (opencl/ptx/cuda) is auto-detected from $TORNADOVM_HOME/etc/tornado.backend + # Backend (opencl/cuda/metal) is auto-detected from $TORNADOVM_HOME/etc/tornado.backend ./llama-tornado --gpu \ --model $MODELS_DIR/${{ inputs.model_file }} \ --prompt "${{ inputs.prompt }}" \ diff --git a/.github/actions/setup-java/action.yml b/.github/actions/setup-java/action.yml index c461c1ff..2291624b 100644 --- a/.github/actions/setup-java/action.yml +++ b/.github/actions/setup-java/action.yml @@ -10,7 +10,10 @@ runs: using: composite steps: - name: Set up Java with SDKMAN - shell: bash + # macOS ships bash 3.2 (Apple keeps it frozen for GPLv3 reasons); SDKMAN's own + # scripts use bash 4+ syntax (${var^^}) and crash under it, so use Homebrew's + # bash on macOS runners instead. Requires `brew install bash` on the runner. + shell: ${{ runner.os == 'macOS' && '/opt/homebrew/bin/bash --noprofile --norc -eo pipefail {0}' || 'bash' }} run: | source "$HOME/.sdkman/bin/sdkman-init.sh" if ! sdk list java | grep -q "${{ inputs.java_version }}"; then diff --git a/.github/actions/setup-tornadovm/action.yml b/.github/actions/setup-tornadovm/action.yml index 62b99e53..218bfefc 100644 --- a/.github/actions/setup-tornadovm/action.yml +++ b/.github/actions/setup-tornadovm/action.yml @@ -3,22 +3,27 @@ description: Build TornadoVM develop once per backend and reuse across runs via inputs: backend: - description: 'TornadoVM backend to build (opencl, ptx, or cuda)' + description: 'TornadoVM backend to build (opencl, cuda, or metal)' required: true +# macOS ships bash 3.2 (Apple keeps it frozen for GPLv3 reasons); SDKMAN's own scripts use +# bash 4+ syntax (${var^^}) and crash under it, so every step below uses Homebrew's bash on +# macOS runners instead (requires `brew install bash` on the runner) and plain "bash" on Linux. +# GitHub's action-manifest parser doesn't support YAML anchors, so the expression is repeated +# per step rather than shared via &shell/*shell. runs: using: composite steps: - name: Get TornadoVM develop SHA id: tornado_sha - shell: bash + shell: ${{ runner.os == 'macOS' && '/opt/homebrew/bin/bash --noprofile --norc -eo pipefail {0}' || 'bash' }} run: | SHA=$(git ls-remote https://github.com/beehive-lab/TornadoVM refs/heads/develop | cut -f1) echo "sha=$SHA" >> $GITHUB_OUTPUT - name: Check local build sentinel id: sentinel - shell: bash + shell: ${{ runner.os == 'macOS' && '/opt/homebrew/bin/bash --noprofile --norc -eo pipefail {0}' || 'bash' }} run: | SENTINEL="$TORNADO_ROOT/.built-sha" if [ -f "$SENTINEL" ] && [ "$(cat $SENTINEL)" = "${{ steps.tornado_sha.outputs.sha }}" ]; then @@ -29,7 +34,7 @@ runs: - name: Clone TornadoVM develop if: steps.sentinel.outputs.up-to-date != 'true' - shell: bash + shell: ${{ runner.os == 'macOS' && '/opt/homebrew/bin/bash --noprofile --norc -eo pipefail {0}' || 'bash' }} run: | rm -rf $TORNADO_ROOT git clone --depth 1 --branch develop \ @@ -38,7 +43,7 @@ runs: - name: Set up Python venv for TornadoVM if: steps.sentinel.outputs.up-to-date != 'true' - shell: bash + shell: ${{ runner.os == 'macOS' && '/opt/homebrew/bin/bash --noprofile --norc -eo pipefail {0}' || 'bash' }} run: | python3 -m venv $TORNADO_ROOT/venv source $TORNADO_ROOT/venv/bin/activate @@ -46,7 +51,7 @@ runs: - name: Build TornadoVM if: steps.sentinel.outputs.up-to-date != 'true' - shell: bash + shell: ${{ runner.os == 'macOS' && '/opt/homebrew/bin/bash --noprofile --norc -eo pipefail {0}' || 'bash' }} run: | cd $TORNADO_ROOT mkdir -p graalJars && cp $GRAAL_JARS/* graalJars/ @@ -57,7 +62,7 @@ runs: - name: Find TornadoVM SDK directory id: find_sdk - shell: bash + shell: ${{ runner.os == 'macOS' && '/opt/homebrew/bin/bash --noprofile --norc -eo pipefail {0}' || 'bash' }} run: | SDK_DIR=$(find $TORNADO_ROOT/dist -type d -maxdepth 3 -path "*/tornadovm-*-${{ inputs.backend }}" | head -n 1) if [ -z "$SDK_DIR" ]; then @@ -70,7 +75,7 @@ runs: # Runs on both fresh build and sentinel hit — sets TORNADOVM_HOME and PATH # for all subsequent steps in the calling workflow. - name: Configure TornadoVM environment - shell: bash + shell: ${{ runner.os == 'macOS' && '/opt/homebrew/bin/bash --noprofile --norc -eo pipefail {0}' || 'bash' }} run: | FULL_SDK="${{ steps.find_sdk.outputs.sdk_dir }}" echo "Detected TornadoVM SDK: $FULL_SDK" diff --git a/.github/workflows/build-and-run.yml b/.github/workflows/build-and-run.yml index af61e397..c18cbdcb 100644 --- a/.github/workflows/build-and-run.yml +++ b/.github/workflows/build-and-run.yml @@ -18,7 +18,7 @@ env: jobs: code-quality: if: github.repository == 'beehive-lab/GPULlama3.java' - runs-on: [self-hosted] + runs-on: [self-hosted, Linux] timeout-minutes: 30 steps: @@ -31,21 +31,26 @@ jobs: # ./mvnw -T12C -Pspotless spotless:check # Build: TornadoVM → GPULlama3 → Quarkus LangChain4j - # max-parallel: 1 ensures the opencl, ptx and cuda variants run sequentially so - # there are no workspace conflicts between matrix jobs. + # No max-parallel cap: opencl/cuda still serialize against each other since they share + # the same self-hosted Linux runner, but that's the runner queue doing the serializing, + # not the matrix — the independent Mac runner must be free to start immediately rather + # than wait for a slot in a shared parallelism budget it doesn't actually contend for. + # fail-fast: false so a Linux failure doesn't cancel the independent Mac runner's job (and vice versa). build: if: github.repository == 'beehive-lab/GPULlama3.java' - runs-on: [self-hosted] + runs-on: ${{ matrix.backend.os == 'macos' && fromJSON('["self-hosted","macOS"]') || fromJSON('["self-hosted","Linux"]') }} needs: code-quality timeout-minutes: 30 strategy: - fail-fast: true - max-parallel: 1 + fail-fast: false matrix: backend: - name: opencl - - name: ptx + os: linux - name: cuda + os: linux + - name: metal + os: macos steps: - name: Checkout GPULlama3 @@ -75,7 +80,7 @@ jobs: fi echo "TornadoVM Maven version: $TORNADO_VERSION" # Strip any pre-existing -SNAPSHOT suffix before appending, making this step idempotent - # across sequential matrix variants (ptx runs after opencl on the same workspace). + # across sequential matrix variants (cuda runs after opencl on the same Linux workspace). BASE_VERSION=$(./mvnw help:evaluate -Dexpression=project.version -q -DforceStdout | sed 's/-SNAPSHOT$//') GPULLAMA3_VERSION="${BASE_VERSION}-SNAPSHOT" echo "GPULlama3 version: $GPULLAMA3_VERSION" @@ -92,22 +97,28 @@ jobs: - name: Build Quarkus LangChain4j run: | cd ${{ github.workspace }}/quarkus-langchain4j - sed -i 's/.*<\/gpu-llama3\.version>/'$GPULLAMA3_VERSION'<\/gpu-llama3.version>/' pom.xml + # -i.bak (extension glued to the flag) is the portable form: GNU sed (Linux) accepts it, + # and BSD sed (macOS) requires it — a bare `-i ''` is parsed as the script on BSD sed. + sed -i.bak 's/.*<\/gpu-llama3\.version>/'$GPULLAMA3_VERSION'<\/gpu-llama3.version>/' pom.xml + rm -f pom.xml.bak # -Dtornado activates the TornadoVM profile; -am builds only the gpu-llama3 module + deps mvn clean install -pl integration-tests/gpu-llama3 -am -DskipTests -Dtornado standalone-inference: if: github.repository == 'beehive-lab/GPULlama3.java' - runs-on: [self-hosted] + runs-on: ${{ matrix.backend.os == 'macos' && fromJSON('["self-hosted","macOS"]') || fromJSON('["self-hosted","Linux"]') }} needs: build timeout-minutes: 30 strategy: - fail-fast: true + fail-fast: false matrix: backend: - name: opencl - - name: ptx + os: linux - name: cuda + os: linux + - name: metal + os: macos steps: - name: Checkout GPULlama3 @@ -161,31 +172,6 @@ jobs: flags: --with-prefill-decode --batch-prefill-size 32 metrics_file: ${{ runner.temp }}/metrics-${{ matrix.backend.name }}-llama-1b-f16-batch-prefill-decode.json - # PTX-only: CUDA-graph variants - - name: PTX - FP16 - Run Llama-3.2-1B-Instruct-F16.gguf - Prefill-Decode-CUDA-Graphs - if: matrix.backend.name == 'ptx' - uses: ./.github/actions/run-inference - with: - backend: ${{ matrix.backend.name }} - model_file: Llama-3.2-1B-Instruct-F16.gguf - model: Llama-3.2-1B-Instruct - quantization: F16 - configuration: prefill-decode-cuda-graphs - flags: --with-prefill-decode --cuda-graphs - metrics_file: ${{ runner.temp }}/metrics-ptx-llama-1b-f16-prefill-decode-cuda-graphs.json - - - name: PTX - FP16 - Run Llama-3.2-1B-Instruct-F16.gguf - Batch-Prefill-Decode-CUDA-Graphs - if: matrix.backend.name == 'ptx' - uses: ./.github/actions/run-inference - with: - backend: ${{ matrix.backend.name }} - model_file: Llama-3.2-1B-Instruct-F16.gguf - model: Llama-3.2-1B-Instruct - quantization: F16 - configuration: batch-prefill-decode-cuda-graphs - flags: --with-prefill-decode --batch-prefill-size 32 --cuda-graphs - metrics_file: ${{ runner.temp }}/metrics-ptx-llama-1b-f16-batch-prefill-decode-cuda-graphs.json - - name: FP16 - Run Qwen3-0.6B-f16.gguf uses: ./.github/actions/run-inference with: @@ -218,32 +204,10 @@ jobs: flags: --with-prefill-decode --batch-prefill-size 32 metrics_file: ${{ runner.temp }}/metrics-${{ matrix.backend.name }}-qwen3-0.6b-f16-batch-prefill-decode.json - # PTX-only: CUDA-graph variants - - name: PTX - FP16 - Run Qwen3-0.6B-f16.gguf - Prefill-Decode-CUDA-Graphs - if: matrix.backend.name == 'ptx' - uses: ./.github/actions/run-inference - with: - backend: ${{ matrix.backend.name }} - model_file: Qwen3-0.6B-f16.gguf - model: Qwen3-0.6B - quantization: F16 - configuration: prefill-decode-cuda-graphs - flags: --with-prefill-decode --cuda-graphs - metrics_file: ${{ runner.temp }}/metrics-ptx-qwen3-0.6b-f16-prefill-decode-cuda-graphs.json - - - name: PTX - FP16 - Run Qwen3-0.6B-f16.gguf - Batch-Prefill-Decode-CUDA-Graphs - if: matrix.backend.name == 'ptx' - uses: ./.github/actions/run-inference - with: - backend: ${{ matrix.backend.name }} - model_file: Qwen3-0.6B-f16.gguf - model: Qwen3-0.6B - quantization: F16 - configuration: batch-prefill-decode-cuda-graphs - flags: --with-prefill-decode --batch-prefill-size 32 --cuda-graphs - metrics_file: ${{ runner.temp }}/metrics-ptx-qwen3-0.6b-f16-batch-prefill-decode-cuda-graphs.json - - name: FP16 - Run Mistral-7B-Instruct-v0.3.fp16.gguf + # TODO: re-enable once the Mistral-7B-on-Metal timeout is root-caused; see the + # Metal TornadoOutOfMemoryException also hit by the tiny Llama-3.2-1B quarkus job. + if: matrix.backend.name != 'metal' uses: ./.github/actions/run-inference with: backend: ${{ matrix.backend.name }} @@ -274,6 +238,8 @@ jobs: metrics_file: ${{ runner.temp }}/metrics-${{ matrix.backend.name }}-phi3-mini-fp16-standard.json - name: FP16 - Run Granite-3.2-2b-instruct-f16.gguf + # TODO: not yet provisioned under /opt/models on the Mac runner; re-enable once it is. + if: matrix.backend.name != 'metal' uses: ./.github/actions/run-inference with: backend: ${{ matrix.backend.name }} @@ -351,59 +317,6 @@ jobs: "flags=--with-prefill-decode --batch-prefill-size 32" \ prompt="Say hello" - # ── PTX-only: CUDA-graph variants ──────────────────────────────────────── - - name: PTX - Q8 - Run Llama-3.2-1B-Instruct-Q8_0.gguf - Prefill-Decode-CUDA-Graphs - if: matrix.backend.name == 'ptx' - env: - JAVA_TOOL_OPTIONS: >- - -Dllama.metrics.format=json - -Dllama.metrics.output=file - -Dllama.metrics.file=${{ runner.temp }}/metrics-ptx-llama-1b-q8-prefill-decode-cuda-graphs.json - run: | - cd ${{ github.workspace }} - export PATH="$TORNADOVM_HOME/bin:$JAVA_HOME/bin:$PATH" - ./llama-tornado --gpu \ - --model $MODELS_DIR/Llama-3.2-1B-Instruct-Q8_0.gguf \ - --prompt "Say hello" \ - --with-prefill-decode \ - --cuda-graphs - python3 scripts/write_metrics_sidecar.py \ - --out "${{ runner.temp }}/metrics-ptx-llama-1b-q8-prefill-decode-cuda-graphs.meta.json" \ - backend=ptx \ - task=llama-inference \ - model_file=Llama-3.2-1B-Instruct-Q8_0.gguf \ - model=Llama-3.2-1B-Instruct \ - quantization=Q8_0 \ - configuration=prefill-decode-cuda-graphs \ - "flags=--with-prefill-decode --cuda-graphs" \ - prompt="Say hello" - - - name: PTX - Q8 - Run Llama-3.2-1B-Instruct-Q8_0.gguf - Batch-Prefill-Decode-CUDA-Graphs - if: matrix.backend.name == 'ptx' - env: - JAVA_TOOL_OPTIONS: >- - -Dllama.metrics.format=json - -Dllama.metrics.output=file - -Dllama.metrics.file=${{ runner.temp }}/metrics-ptx-llama-1b-q8-batch-prefill-decode-cuda-graphs.json - run: | - cd ${{ github.workspace }} - export PATH="$TORNADOVM_HOME/bin:$JAVA_HOME/bin:$PATH" - ./llama-tornado --gpu \ - --model $MODELS_DIR/Llama-3.2-1B-Instruct-Q8_0.gguf \ - --prompt "Say hello" \ - --with-prefill-decode --batch-prefill-size 32 \ - --cuda-graphs - python3 scripts/write_metrics_sidecar.py \ - --out "${{ runner.temp }}/metrics-ptx-llama-1b-q8-batch-prefill-decode-cuda-graphs.meta.json" \ - backend=ptx \ - task=llama-inference \ - model_file=Llama-3.2-1B-Instruct-Q8_0.gguf \ - model=Llama-3.2-1B-Instruct \ - quantization=Q8_0 \ - configuration=batch-prefill-decode-cuda-graphs \ - "flags=--with-prefill-decode --batch-prefill-size 32 --cuda-graphs" \ - prompt="Say hello" - - name: Q8 - Run Qwen3-0.6B-Q8_0.gguf uses: ./.github/actions/run-inference with: @@ -436,31 +349,6 @@ jobs: flags: --with-prefill-decode --batch-prefill-size 32 metrics_file: ${{ runner.temp }}/metrics-${{ matrix.backend.name }}-qwen3-0-6b-q8-batch-prefill-decode.json - # PTX-only: CUDA-graph variants - - name: PTX - Q8 - Run Qwen3-0.6B-Q8_0.gguf - Prefill-Decode-CUDA-Graphs - if: matrix.backend.name == 'ptx' - uses: ./.github/actions/run-inference - with: - backend: ${{ matrix.backend.name }} - model_file: Qwen3-0.6B-Q8_0.gguf - model: Qwen3-0.6B - quantization: Q8_0 - configuration: prefill-decode-cuda-graphs - flags: --with-prefill-decode --cuda-graphs - metrics_file: ${{ runner.temp }}/metrics-ptx-qwen3-0-6b-q8-prefill-decode-cuda-graphs.json - - - name: PTX - Q8 - Run Qwen3-0.6B-Q8_0.gguf - Batch-Prefill-Decode-CUDA-Graphs - if: matrix.backend.name == 'ptx' - uses: ./.github/actions/run-inference - with: - backend: ${{ matrix.backend.name }} - model_file: Qwen3-0.6B-Q8_0.gguf - model: Qwen3-0.6B - quantization: Q8_0 - configuration: batch-prefill-decode-cuda-graphs - flags: --with-prefill-decode --batch-prefill-size 32 --cuda-graphs - metrics_file: ${{ runner.temp }}/metrics-ptx-qwen3-0-6b-q8-batch-prefill-decode-cuda-graphs.json - - name: Q8 - Run Phi-3-mini-4k-instruct-Q8_0.gguf uses: ./.github/actions/run-inference with: @@ -482,6 +370,8 @@ jobs: metrics_file: ${{ runner.temp }}/metrics-${{ matrix.backend.name }}-qwen2-5-1-5b-q8-standard.json - name: Q8 - Run Mistral-7B-Instruct-v0.3.Q8_0.gguf + # TODO: re-enable once Mistral-7B-on-Metal is root-caused; see the FP16 step above. + if: matrix.backend.name != 'metal' uses: ./.github/actions/run-inference with: backend: ${{ matrix.backend.name }} @@ -523,16 +413,19 @@ jobs: # Test integration with Quarkus-langchain4j quarkus-langchain4j-integration: if: github.repository == 'beehive-lab/GPULlama3.java' - runs-on: [self-hosted] + runs-on: ${{ matrix.backend.os == 'macos' && fromJSON('["self-hosted","macOS"]') || fromJSON('["self-hosted","Linux"]') }} needs: build timeout-minutes: 10 strategy: - fail-fast: true + fail-fast: false matrix: backend: - name: opencl - - name: ptx + os: linux - name: cuda + os: linux + - name: metal + os: macos steps: - name: Checkout GPULlama3 @@ -610,10 +503,20 @@ jobs: - name: Start Quarkus Application run: | cd ${{ github.workspace }}/quarkus-langchain4j/integration-tests/gpu-llama3 + APP_LOG="${{ runner.temp }}/quarkus-app.log" + echo "APP_LOG=$APP_LOG" >> $GITHUB_ENV + # Background process output stops being captured by the Actions log once this + # step ends, so redirect it to a file we can dump later if a request fails. + # + # -Dtornado.device.memory here is dead: the gpu-llama3 extension's ChatModelConfig + # defaults device-memory to "4GB" and calls System.setProperty("tornado.device.memory", ...) + # itself during model init, overwriting whatever we pass as a bare JVM flag. On Metal + # that 4GB budget isn't enough even for Llama-3.2-1B in batch-prefill-decode mode + # (TornadoOutOfMemoryException), so override the actual Quarkus config property instead. java @"$TORNADOVM_HOME/tornado-argfile" \ - -Dtornado.device.memory=8GB \ + -Dquarkus.langchain4j.gpu-llama3.chat-model.device-memory=8GB \ -Dquarkus.http.port=$QUARKUS_PORT \ - -jar target/quarkus-app/quarkus-run.jar & + -jar target/quarkus-app/quarkus-run.jar > "$APP_LOG" 2>&1 & APP_PID=$! echo "APP_PID=$APP_PID" >> $GITHUB_ENV for i in {1..30}; do @@ -650,7 +553,9 @@ jobs: run: | for attempt in 1 2 3; do echo "Attempt $attempt of 3 for streaming endpoint..." - HTTP_CODE=$(timeout 10s curl -s -o /dev/null -w "%{http_code}" http://localhost:$QUARKUS_PORT/chat/streaming) + # curl's own --max-time (not the external `timeout` command, which stock + # macOS doesn't ship) so this works the same on both runner OSes. + HTTP_CODE=$(curl -s --max-time 10 -o /dev/null -w "%{http_code}" http://localhost:$QUARKUS_PORT/chat/streaming) if [ "$HTTP_CODE" = "200" ]; then echo "SUCCESS: HTTP $HTTP_CODE" break @@ -660,6 +565,10 @@ jobs: [ $attempt -eq 3 ] && { echo "::error::Streaming endpoint failed after 3 attempts"; exit 1; } done + - name: Dump Quarkus Application Log + if: failure() + run: cat "$APP_LOG" || true + - name: Cleanup & Shutdown if: always() run: | @@ -675,7 +584,7 @@ jobs: github.event_name == 'push' && github.ref == 'refs/heads/main' - runs-on: [self-hosted] + runs-on: [self-hosted, Linux] needs: standalone-inference timeout-minutes: 15 diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/Qwen3FP16FFNLayers.java b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/Qwen3FP16FFNLayers.java index f0da8b54..e6a205f3 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/Qwen3FP16FFNLayers.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/Qwen3FP16FFNLayers.java @@ -34,9 +34,11 @@ public class Qwen3FP16FFNLayers extends AbstractTransformerLayerTaskGraphs wrapAttSplit. - unifiedLayer.task("attention", - TransformerComputeKernelsLayered::processHeadsFlashAttentionSplitKV, - context, - qwen3State.wrapQ, // query vectors - qwen3State.wrapKeyCache, // key cache - qwen3State.wrapValueCache, // value cache - qwen3State.wrapAttSplit, // scratch: per-head split partials (compact layout) - config.numberOfHeads(), // nHeads - nEmbdHead, // headSize - nEmbdGqa, // kvDim - gqa, // kvMul (nHeads / nHeadKv) - qwen3State.positionHolder, // position - layerIndex, // layer index - config.contextLength(), // context length - attentionSplits); // number of KV splits per head - // Phase 2: combine the per-head split partials into the final attention output -> wrapXb. - unifiedLayer.task("attention_combine", - TransformerComputeKernelsLayered::combineSplitKVAttention, - context, - qwen3State.wrapAttSplit, // scratch: per-head split partials (compact layout) - qwen3State.wrapXb, // output: attention result - config.numberOfHeads(), // nHeads - nEmbdHead, // headSize - attentionSplits); // number of KV splits per head + if (isMetalBackend) { + // Metal: single-workgroup-per-head online-softmax attention, writing directly to wrapXb. + // No combine phase needed (TornadoVM fails to JIT the multi-workgroup split-KV kernel here). + unifiedLayer.task("attention", + TransformerComputeKernelsLayered::processHeadsFlashAttention, + context, + qwen3State.wrapQ, // query vectors + qwen3State.wrapKeyCache, // key cache + qwen3State.wrapValueCache, // value cache + qwen3State.wrapXb, // output: attention result + config.numberOfHeads(), // nHeads + nEmbdHead, // headSize + nEmbdGqa, // kvDim + gqa, // kvMul (nHeads / nHeadKv) + qwen3State.positionHolder, // position + layerIndex, // layer index + config.contextLength()); // context length + } else { + // Split-KV (flash-decoding) attention. + // Phase 1: split each head's KV range across attentionSplits workgroups; partials -> wrapAttSplit. + unifiedLayer.task("attention", + TransformerComputeKernelsLayered::processHeadsFlashAttentionSplitKV, + context, + qwen3State.wrapQ, // query vectors + qwen3State.wrapKeyCache, // key cache + qwen3State.wrapValueCache, // value cache + qwen3State.wrapAttSplit, // scratch: per-head split partials (compact layout) + config.numberOfHeads(), // nHeads + nEmbdHead, // headSize + nEmbdGqa, // kvDim + gqa, // kvMul (nHeads / nHeadKv) + qwen3State.positionHolder, // position + layerIndex, // layer index + config.contextLength(), // context length + attentionSplits); // number of KV splits per head + // Phase 2: combine the per-head split partials into the final attention output -> wrapXb. + unifiedLayer.task("attention_combine", + TransformerComputeKernelsLayered::combineSplitKVAttention, + context, + qwen3State.wrapAttSplit, // scratch: per-head split partials (compact layout) + qwen3State.wrapXb, // output: attention result + config.numberOfHeads(), // nHeads + nEmbdHead, // headSize + attentionSplits); // number of KV splits per head + } // Output Projection with Residual if (useWarpMatmul) { diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/Qwen3Q8_0FFNLayers.java b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/Qwen3Q8_0FFNLayers.java index 5b184894..cd5bf4b8 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/Qwen3Q8_0FFNLayers.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/Qwen3Q8_0FFNLayers.java @@ -39,9 +39,11 @@ public class Qwen3Q8_0FFNLayers extends AbstractTransformerLayerTaskGraphs wrapAttSplit. - unifiedLayer.task("attention", - TransformerComputeKernelsLayered::processHeadsFlashAttentionSplitKV, - context, - qwen3State.wrapQ, // query vectors - qwen3State.wrapKeyCache, // key cache - qwen3State.wrapValueCache, // value cache - qwen3State.wrapAttSplit, // scratch: per-head split partials (compact layout) - config.numberOfHeads(), // nHeads - nEmbdHead, // headSize - nEmbdGqa, // kvDim - gqa, // kvMul (nHeads / nHeadKv) - qwen3State.positionHolder, // position - layerIndex, // layer index - config.contextLength(), // context length - attentionSplits); // number of KV splits per head - // Phase 2: combine the per-head split partials into the final attention output -> wrapXb. - unifiedLayer.task("attention_combine", - TransformerComputeKernelsLayered::combineSplitKVAttention, - context, - qwen3State.wrapAttSplit, // scratch: per-head split partials (compact layout) - qwen3State.wrapXb, // output: attention result - config.numberOfHeads(), // nHeads - nEmbdHead, // headSize - attentionSplits); // number of KV splits per head + if (isMetalBackend) { + // Metal: single-workgroup-per-head online-softmax attention, writing directly to wrapXb. + // No combine phase needed (TornadoVM fails to JIT the multi-workgroup split-KV kernel here). + unifiedLayer.task("attention", + TransformerComputeKernelsLayered::processHeadsFlashAttention, + context, + qwen3State.wrapQ, // query vectors + qwen3State.wrapKeyCache, // key cache + qwen3State.wrapValueCache, // value cache + qwen3State.wrapXb, // output: attention result + config.numberOfHeads(), // nHeads + nEmbdHead, // headSize + nEmbdGqa, // kvDim + gqa, // kvMul (nHeads / nHeadKv) + qwen3State.positionHolder, // position + layerIndex, // layer index + config.contextLength()); // context length + } else { + // Split-KV (flash-decoding) attention. + // Phase 1: split each head's KV range across attentionSplits workgroups; partials -> wrapAttSplit. + unifiedLayer.task("attention", + TransformerComputeKernelsLayered::processHeadsFlashAttentionSplitKV, + context, + qwen3State.wrapQ, // query vectors + qwen3State.wrapKeyCache, // key cache + qwen3State.wrapValueCache, // value cache + qwen3State.wrapAttSplit, // scratch: per-head split partials (compact layout) + config.numberOfHeads(), // nHeads + nEmbdHead, // headSize + nEmbdGqa, // kvDim + gqa, // kvMul (nHeads / nHeadKv) + qwen3State.positionHolder, // position + layerIndex, // layer index + config.contextLength(), // context length + attentionSplits); // number of KV splits per head + // Phase 2: combine the per-head split partials into the final attention output -> wrapXb. + unifiedLayer.task("attention_combine", + TransformerComputeKernelsLayered::combineSplitKVAttention, + context, + qwen3State.wrapAttSplit, // scratch: per-head split partials (compact layout) + qwen3State.wrapXb, // output: attention result + config.numberOfHeads(), // nHeads + nEmbdHead, // headSize + attentionSplits); // number of KV splits per head + } // Output Projection with Residual if (useWarpMatmul) { diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/scheduling/SchedulerDetectionService.java b/src/main/java/org/beehive/gpullama3/tornadovm/scheduling/SchedulerDetectionService.java index d777cb78..ce4021a0 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/scheduling/SchedulerDetectionService.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/scheduling/SchedulerDetectionService.java @@ -26,6 +26,18 @@ public static boolean isWarpShuffleSupported() { // return backendType == TornadoVMBackendType.PTX || backendType == TornadoVMBackendType.CUDA; } + /** + * Whether the active TornadoVM backend is Metal. The Metal backend fails to JIT the + * multi-workgroup split-KV flash-decoding attention kernel ({@code processHeadsFlashAttentionSplitKV}), + * so Qwen3 layers must fall back to the single-workgroup-per-head online-softmax kernel there. + */ + public static boolean isMetalBackend() { + TornadoVMBackendType backendType = TornadoRuntimeProvider.getTornadoRuntime() + .getBackend(0) + .getBackendType(); + return backendType == TornadoVMBackendType.METAL; + } + // @formatter:off public static SchedulerType determineSchedulerType(Model model) { TornadoRuntime tornadoRuntime = TornadoRuntimeProvider.getTornadoRuntime();