diff --git a/.github/actions/run-inference/action.yml b/.github/actions/run-inference/action.yml index 314a8be5..fe3d574b 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 or ptx)' + description: 'GPU backend (opencl, ptx, or cuda)' required: true model_file: description: 'Model filename inside $MODELS_DIR (e.g. Llama-3.2-1B-Instruct-F16.gguf)' diff --git a/.github/actions/setup-tornadovm/action.yml b/.github/actions/setup-tornadovm/action.yml index f280959a..62b99e53 100644 --- a/.github/actions/setup-tornadovm/action.yml +++ b/.github/actions/setup-tornadovm/action.yml @@ -3,7 +3,7 @@ description: Build TornadoVM develop once per backend and reuse across runs via inputs: backend: - description: 'TornadoVM backend to build (opencl or ptx)' + description: 'TornadoVM backend to build (opencl, ptx, or cuda)' required: true runs: diff --git a/.github/workflows/build-and-run.yml b/.github/workflows/build-and-run.yml index aa246c19..65515711 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] timeout-minutes: 30 steps: @@ -31,7 +31,7 @@ jobs: # ./mvnw -T12C -Pspotless spotless:check # Build: TornadoVM → GPULlama3 → Quarkus LangChain4j - # max-parallel: 1 ensures the opencl and ptx variants run sequentially so + # max-parallel: 1 ensures the opencl, ptx and cuda variants run sequentially so # there are no workspace conflicts between matrix jobs. build: if: github.repository == 'beehive-lab/GPULlama3.java' @@ -45,6 +45,7 @@ jobs: backend: - name: opencl - name: ptx + - name: cuda steps: - name: Checkout GPULlama3 @@ -106,6 +107,7 @@ jobs: backend: - name: opencl - name: ptx + - name: cuda steps: - name: Checkout GPULlama3 @@ -530,6 +532,7 @@ jobs: backend: - name: opencl - name: ptx + - name: cuda steps: - name: Checkout GPULlama3 diff --git a/README.md b/README.md index 97d05ce6..e7a99102 100644 --- a/README.md +++ b/README.md @@ -66,7 +66,8 @@ GPULlama3ChatModel model = GPULlama3ChatModel.builder() Ensure you have the following installed and configured: - **Java 21**: Required for Vector API support & TornadoVM. -- [TornadoVM](https://github.com/beehive-lab/TornadoVM) with OpenCL or PTX backends. +- [TornadoVM](https://github.com/beehive-lab/TornadoVM) with OpenCL, PTX, or CUDA backends. + - The `--cuda` backend requires a TornadoVM build that includes the CUDA backend from [TornadoVM PR #861](https://github.com/beehive-lab/TornadoVM/pull/861). This project currently builds against TornadoVM `5.0.0-jdk21-dev`. - GCC/G++ 13 or newer: Required to build and run TornadoVM native components. ### Install, Build, and Run @@ -305,6 +306,12 @@ Run a model with a text prompt: ./llama-tornado --gpu --verbose-init --opencl --model beehive-llama-3.2-1b-instruct-fp16.gguf --prompt "Explain the benefits of GPU acceleration." ``` +Select a backend explicitly with `--opencl`, `--ptx`, or `--cuda` (NVIDIA), or `--metal` (Apple Silicon). For example, to run on the CUDA backend: + +```bash +./llama-tornado --gpu --cuda --model beehive-llama-3.2-1b-instruct-fp16.gguf --prompt "Explain the benefits of GPU acceleration." +``` + #### GPU Execution (FP16 Model) Enable GPU acceleration with Q8_0 quantization: ```bash @@ -393,7 +400,7 @@ Supported command-line options include: ```bash cmd ➜ llama-tornado --help usage: llama-tornado [-h] --model MODEL_PATH [--prompt PROMPT] [-sp SYSTEM_PROMPT] [--temperature TEMPERATURE] [--top-p TOP_P] [--seed SEED] [-n MAX_TOKENS] - [--stream STREAM] [--echo ECHO] [-i] [--instruct] [--gpu] [--opencl] [--ptx] [--gpu-memory GPU_MEMORY] [--heap-min HEAP_MIN] [--heap-max HEAP_MAX] + [--stream STREAM] [--echo ECHO] [-i] [--instruct] [--gpu] [--opencl] [--ptx] [--cuda] [--metal] [--gpu-memory GPU_MEMORY] [--heap-min HEAP_MIN] [--heap-max HEAP_MAX] [--debug] [--profiler] [--profiler-dump-dir PROFILER_DUMP_DIR] [--print-bytecodes] [--print-threads] [--print-kernel] [--full-dump] [--show-command] [--execute-after-show] [--opencl-flags OPENCL_FLAGS] [--max-wait-events MAX_WAIT_EVENTS] [--verbose] @@ -424,7 +431,9 @@ Mode Selection: Hardware Configuration: --gpu Enable GPU acceleration (default: False) --opencl Use OpenCL backend (default) (default: None) - --ptx Use PTX/CUDA backend (default: None) + --ptx Use PTX backend (default: None) + --cuda Use CUDA backend (requires TornadoVM built with the CUDA backend) (default: None) + --metal Use Apple Metal backend (macOS only) (default: None) --gpu-memory GPU_MEMORY GPU memory allocation (default: 7GB) --heap-min HEAP_MIN Minimum JVM heap size (default: 20g) @@ -480,9 +489,9 @@ View TornadoVM's internal behavior: - **Support for GGUF format models** with full FP16 and partial support for Q8_0 and Q4_0 quantization. - **Instruction-following and chat modes** for various use cases. - **Interactive CLI** with `--interactive` and `--instruct` modes. - - **Flexible backend switching** - choose OpenCL or PTX at runtime (need to build TornadoVM with both enabled). + - **Flexible backend switching** - choose OpenCL, PTX, or CUDA at runtime (need to build TornadoVM with the chosen backends enabled). - **Cross-platform compatibility**: - - ✅ NVIDIA GPUs (OpenCL & PTX ) + - ✅ NVIDIA GPUs (OpenCL, PTX & CUDA) - ✅ Intel GPUs (OpenCL) - ✅ Apple GPUs (OpenCL) diff --git a/llama-tornado b/llama-tornado index 1d6c3d23..78388295 100755 --- a/llama-tornado +++ b/llama-tornado @@ -1,7 +1,7 @@ #!/usr/bin/env python3 """ llama-tornado: GPU-accelerated Java LLM runner with TornadoVM -Run LLM models using either OpenCL or PTX backends. +Run LLM models using OpenCL, PTX, CUDA, or Metal backends. """ import argparse @@ -19,6 +19,7 @@ from enum import Enum class Backend(Enum): OPENCL = "opencl" PTX = "ptx" + CUDA = "cuda" METAL = "metal" @@ -178,6 +179,14 @@ class LlamaRunner: "ALL-SYSTEM,jdk.incubator.vector,tornado.runtime,tornado.annotation,tornado.drivers.common,tornado.drivers.ptx", ] ) + elif args.backend == Backend.CUDA: + module_config.extend( + [ + f"@{self.tornado_sdk}/etc/exportLists/cuda-exports", + "--add-modules", + "ALL-SYSTEM,jdk.incubator.vector,tornado.runtime,tornado.annotation,tornado.drivers.common,tornado.drivers.cuda", + ] + ) elif args.backend == Backend.METAL: module_config.extend( [ @@ -426,7 +435,14 @@ def create_parser() -> argparse.ArgumentParser: dest="backend", action="store_const", const=Backend.PTX, - help="Use PTX/CUDA backend", + help="Use PTX backend", + ) + hw_group.add_argument( + "--cuda", + dest="backend", + action="store_const", + const=Backend.CUDA, + help="Use CUDA backend (requires TornadoVM built with the CUDA backend)", ) hw_group.add_argument( "--metal", diff --git a/pom.xml b/pom.xml index 1bddf721..f5dd084f 100644 --- a/pom.xml +++ b/pom.xml @@ -39,8 +39,9 @@ 0.5.0 - 4.0.1 + 5.0.0 -jdk21 + ${tornadovm.base.version}${jdk.version.suffix} 25 diff --git a/src/main/java/org/beehive/gpullama3/inference/state/State.java b/src/main/java/org/beehive/gpullama3/inference/state/State.java index 0807b756..1f884b98 100644 --- a/src/main/java/org/beehive/gpullama3/inference/state/State.java +++ b/src/main/java/org/beehive/gpullama3/inference/state/State.java @@ -85,6 +85,16 @@ public abstract class State { public final FloatArray attnScaleBatch; // B (per-token RMS scale, attn) public final FloatArray ffnScaleBatch; // B (per-token RMS scale, FFN) public final IntArray batchStartPosHolder; // 1 (start position of chunk) + public final HalfFloatArray normedXFFNFP16; + public final FloatArray ffnGateResult; + public final FloatArray ffnUpResult; + public final HalfFloatArray xbFP16Batch; + public final HalfFloatArray attnOutFP16; + public final FloatArray woOut; + public final HalfFloatArray wrapHbFP16Batch; + public final FloatArray w2Out; + public final FloatArray qkvResultBatch; // B × (dim + 2*kvDim), packed [q|k|v] rows + public final FloatArray gateUpResultBatch; // B × 2*hiddenDim, packed [gate|up] rows protected State(Configuration config, int batchsize) { this.batchsize = batchsize; @@ -135,11 +145,16 @@ protected State(Configuration config, int batchsize) { int gpuBatchSize = Integer.getInteger("llama.prefillBatchSize", 1); if (gpuBatchSize > 1) { + // The tensor-core GEMM kernels operate on full 128-row M tiles + // (BM = 128). Pad the GEMM-adjacent activation buffers so any + // batch size launches whole tiles; rows >= gpuBatchSize hold + // garbage and are never read by the non-GEMM kernels. + int paddedGpuBatch = (gpuBatchSize + 127) & ~127; int qDim = batchQDim(config); int kvDim = batchKvDim(config); this.embeddingXBatch = new HalfFloatArray(gpuBatchSize * config.dim()); this.wrapXBatch = new FloatArray(gpuBatchSize * config.dim()); - this.wrapXbFP16Batch = new HalfFloatArray(gpuBatchSize * config.dim()); + this.wrapXbFP16Batch = new HalfFloatArray(paddedGpuBatch * config.dim()); this.wrapQBatch = new FloatArray(gpuBatchSize * qDim); this.wrapKBatch = new FloatArray(gpuBatchSize * kvDim); this.wrapVBatch = new FloatArray(gpuBatchSize * kvDim); @@ -148,6 +163,17 @@ protected State(Configuration config, int batchsize) { this.attnScaleBatch = new FloatArray(gpuBatchSize); this.ffnScaleBatch = new FloatArray(gpuBatchSize); this.batchStartPosHolder = new IntArray(1); + this.normedXFFNFP16 = new HalfFloatArray(paddedGpuBatch * config.dim()); + this.ffnGateResult = new FloatArray(gpuBatchSize * config.hiddenDim()); + this.ffnUpResult = new FloatArray(gpuBatchSize * config.hiddenDim()); + + this.xbFP16Batch = new HalfFloatArray(gpuBatchSize * config.dim()); + this.attnOutFP16 = new HalfFloatArray(paddedGpuBatch * qDim); // qDim == dim for Llama; qDim = nHeads*headDim for Qwen3 + this.woOut = new FloatArray(paddedGpuBatch * config.dim()); + this.wrapHbFP16Batch = new HalfFloatArray(paddedGpuBatch * config.hiddenDim()); + this.w2Out = new FloatArray(paddedGpuBatch * config.dim()); + this.qkvResultBatch = new FloatArray(paddedGpuBatch * (qDim + 2 * kvDim)); + this.gateUpResultBatch = new FloatArray(paddedGpuBatch * 2 * config.hiddenDim()); } else { this.embeddingXBatch = null; this.wrapXBatch = null; @@ -160,6 +186,16 @@ protected State(Configuration config, int batchsize) { this.attnScaleBatch = null; this.ffnScaleBatch = null; this.batchStartPosHolder = null; + this.normedXFFNFP16 = null; + this.ffnGateResult = null; + this.ffnUpResult = null; + this.xbFP16Batch = null; + this.attnOutFP16 = null; + this.woOut = null; + this.wrapHbFP16Batch = null; + this.w2Out = null; + this.qkvResultBatch = null; + this.gateUpResultBatch = null; } } diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/TensorCoreSupport.java b/src/main/java/org/beehive/gpullama3/tornadovm/TensorCoreSupport.java new file mode 100644 index 00000000..0383e24b --- /dev/null +++ b/src/main/java/org/beehive/gpullama3/tornadovm/TensorCoreSupport.java @@ -0,0 +1,20 @@ +package org.beehive.gpullama3.tornadovm; + +import uk.ac.manchester.tornado.api.enums.TornadoVMBackendType; +import uk.ac.manchester.tornado.api.runtime.TornadoRuntimeProvider; + +/** + * Detects whether the active TornadoVM backend can execute the tensor-core + * (MMA) batch-prefill kernels. TornadoVM lowers the MMA intrinsics + * ({@code mmaLoadA/B}, {@code mma}, {@code mmaStore}) only on the NVIDIA + * CUDA backend. + */ +public final class TensorCoreSupport { + + public static boolean isTensorCoreCapableBackend() { + TornadoVMBackendType backendType = TornadoRuntimeProvider.getTornadoRuntime() + .getBackend(0) + .getBackendType(); + return backendType == TornadoVMBackendType.CUDA; + } +} diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/kernels/Qwen3Kernels.java b/src/main/java/org/beehive/gpullama3/tornadovm/kernels/Qwen3Kernels.java index 512c5774..a5546f1a 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/kernels/Qwen3Kernels.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/kernels/Qwen3Kernels.java @@ -1381,5 +1381,137 @@ public static void batchedRopeWithKVCacheQwen3( wrapValueCache.set(cacheOff + ic + halfEmbdHead, wrapVBatch.get(kHeadBase + ic + halfEmbdHead)); } } + + // ── Packed-QKV variants for the tensor-core batch prefill path ─────────── + // These mirror batchedFusedQKRmsNorm / batchedRopeWithKVCacheQwen3 but read + // and write the packed [q | k | v] buffer produced by gemmMMAQKV + // (row stride qDim + 2*kvDim), so no separate Q/K/V buffers are needed. + + public static void batchedFusedQKRmsNormPacked( + KernelContext context, + FloatArray qkvBatch, + FloatArray qWeights, + FloatArray kWeights, + int nHeads, + int nHeadKv, + int nEmbdHead, + int qDim, + int kvDim, + float rmsNormEps) { + + int groupId = context.globalIdx / nEmbdHead; + int localId = context.localIdx; + int localSize = context.localGroupSizeX; + int totalHeadsPerBatch = nHeads + nHeadKv; + int qkvStride = qDim + 2 * kvDim; + + int batchIdx = groupId / totalHeadsPerBatch; + int headSlot = groupId % totalHeadsPerBatch; + + float[] localSum = context.allocateFloatLocalArray(nEmbdHead); + + if (headSlot < nHeads) { + // Q head (packed offset 0) + int headOffset = batchIdx * qkvStride + headSlot * nEmbdHead; + float partialSum = 0.0f; + for (int i = localId; i < nEmbdHead; i += localSize) { + float val = qkvBatch.get(headOffset + i); + partialSum += val * val; + } + localSum[localId] = partialSum; + context.localBarrier(); + for (int stride = localSize / 2; stride > 0; stride >>= 1) { + if (localId < stride) { + localSum[localId] += localSum[localId + stride]; + } + context.localBarrier(); + } + float ss = localSum[0] / nEmbdHead + rmsNormEps; + ss = 1.0f / TornadoMath.sqrt(ss); + context.localBarrier(); + for (int i = localId; i < nEmbdHead; i += localSize) { + qkvBatch.set(headOffset + i, qWeights.get(i) * ss * qkvBatch.get(headOffset + i)); + } + } else { + // K head (packed offset qDim) + int kHeadIdx = headSlot - nHeads; + int headOffset = batchIdx * qkvStride + qDim + kHeadIdx * nEmbdHead; + float partialSum = 0.0f; + for (int i = localId; i < nEmbdHead; i += localSize) { + float val = qkvBatch.get(headOffset + i); + partialSum += val * val; + } + localSum[localId] = partialSum; + context.localBarrier(); + for (int stride = localSize / 2; stride > 0; stride >>= 1) { + if (localId < stride) { + localSum[localId] += localSum[localId + stride]; + } + context.localBarrier(); + } + float ss = localSum[0] / nEmbdHead + rmsNormEps; + ss = 1.0f / TornadoMath.sqrt(ss); + context.localBarrier(); + for (int i = localId; i < nEmbdHead; i += localSize) { + qkvBatch.set(headOffset + i, kWeights.get(i) * ss * qkvBatch.get(headOffset + i)); + } + } + } + + public static void batchedRopeWithKVCacheQwen3Packed( + KernelContext context, + IntArray batchStartPosHolder, + FloatArray qkvBatch, + FloatArray wrapKeyCache, + FloatArray wrapValueCache, + int kvDim, + int nEmbdHead, + int layerIndex, + int contextLength, + int qDim) { + + int globalIdx = context.globalIdx; + int halfQDim = qDim / 2; + int batchIdx = globalIdx / halfQDim; + int pairIdx = globalIdx % halfQDim; + int qkvStride = qDim + 2 * kvDim; + + int pos = batchStartPosHolder.get(0) + batchIdx; + + // Qwen3 uses split-half RoPE: pair element ic with ic + nEmbdHead/2 within each head. + int halfEmbdHead = nEmbdHead / 2; + int ic = pairIdx % halfEmbdHead; + int headIdx = pairIdx / halfEmbdHead; + + float freq = 1.0f / TornadoMath.pow(1000000.0f, 2.0f * ic / (float) nEmbdHead); + float val = pos * freq; + float fcr = TornadoMath.cos(val); + float fci = TornadoMath.sin(val); + + // Rotate Q in place (packed offset 0) + int qHeadBase = batchIdx * qkvStride + headIdx * nEmbdHead; + float v0q = qkvBatch.get(qHeadBase + ic); + float v1q = qkvBatch.get(qHeadBase + ic + halfEmbdHead); + qkvBatch.set(qHeadBase + ic, v0q * fcr - v1q * fci); + qkvBatch.set(qHeadBase + ic + halfEmbdHead, v0q * fci + v1q * fcr); + + // Rotate K (packed offset qDim) and write K,V to cache + if (pairIdx < kvDim / 2) { + int kHeadIdx = pairIdx / halfEmbdHead; + int kHeadBase = batchIdx * qkvStride + qDim + kHeadIdx * nEmbdHead; + int vHeadBase = batchIdx * qkvStride + qDim + kvDim + kHeadIdx * nEmbdHead; + float v0k = qkvBatch.get(kHeadBase + ic); + float v1k = qkvBatch.get(kHeadBase + ic + halfEmbdHead); + float rotK0 = v0k * fcr - v1k * fci; + float rotK1 = v0k * fci + v1k * fcr; + + int cacheOff = layerIndex * contextLength * kvDim + pos * kvDim + kHeadIdx * nEmbdHead; + wrapKeyCache.set(cacheOff + ic, rotK0); + wrapKeyCache.set(cacheOff + ic + halfEmbdHead, rotK1); + wrapValueCache.set(cacheOff + ic, qkvBatch.get(vHeadBase + ic)); + wrapValueCache.set(cacheOff + ic + halfEmbdHead, qkvBatch.get(vHeadBase + ic + halfEmbdHead)); + } + } + } // @formatter:on diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/kernels/TransformerBatchPrefillKernels.java b/src/main/java/org/beehive/gpullama3/tornadovm/kernels/TransformerBatchPrefillKernels.java index 057e9f87..9d85c936 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/kernels/TransformerBatchPrefillKernels.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/kernels/TransformerBatchPrefillKernels.java @@ -1,6 +1,7 @@ package org.beehive.gpullama3.tornadovm.kernels; import uk.ac.manchester.tornado.api.KernelContext; +import uk.ac.manchester.tornado.api.enums.MMAShape; import uk.ac.manchester.tornado.api.math.TornadoMath; import uk.ac.manchester.tornado.api.types.HalfFloat; import uk.ac.manchester.tornado.api.types.arrays.ByteArray; @@ -718,5 +719,1277 @@ public static void batchedFusedRmsNormFFNGateUpQ8(KernelContext context, } } + /** + * RMS-apply for FFN, writing FP16. Mirrors batchedRmsApplyFP16 but pulls + * scale from ffnScaleBatch. Output is the A operand for the W1/W3 MMA tasks. + * + * Worker: B*dim global threads, localSize=256. + */ + public static void batchedFFNRmsApplyFP16(KernelContext context, + HalfFloatArray normedXFFNFP16, + FloatArray wrapXBatch, + FloatArray rmsFFNWeights, + FloatArray ffnScaleBatch, + int dim) { + int gid = context.globalIdx; + int b = gid / dim; + int i = gid % dim; + float scale = ffnScaleBatch.get(b); + float result = rmsFFNWeights.get(i) * scale * wrapXBatch.get(gid); + normedXFFNFP16.set(gid, new HalfFloat(result)); + } + + /** + * Fused SiLU(gate) * up after the two FFN matmuls. + * Operates on FP32 inputs (MMA writes FP32). + * + * Worker: B*hiddenDim global threads, localSize=256. + */ + public static void batchedFFNSwiGLU(KernelContext context, + FloatArray wrapHbBatch, + FloatArray ffnGateResult, + FloatArray ffnUpResult, + int hiddenDim) { + int gid = context.globalIdx; + float g = ffnGateResult.get(gid); + float u = ffnUpResult.get(gid); + float silu = g / (1.0f + TornadoMath.exp(-g)); + wrapHbBatch.set(gid, silu * u); + } + + private static final int WARP_SIZE = 32; + private static final int BM = 128, BN = 128, BK = 16; + private static final int WARPS_M = 4, WARPS_N = 2; + private static final int WARPS_PER_BLOCK = WARPS_M * WARPS_N; + private static final int THREADS_PER_BLOCK = WARPS_PER_BLOCK * WARP_SIZE; + private static final int WM = BM / WARPS_M; + private static final int WN = BN / WARPS_N; + private static final int B_SUBTILE_BYTES = 256; + + /** + * Packs two consecutive FP16 values into one int (lo | hi<<16) for the + * shared-memory ldmatrix tiles. Leaf helper; inlined by the Tornado JIT. + */ + private static int packHalves(HalfFloatArray src, int idxLo, int idxHi) { + int lo = src.get(idxLo).getHalfFloatValue() & 0xFFFF; + int hi = src.get(idxHi).getHalfFloatValue() & 0xFFFF; + return lo | (hi << 16); + } + + /** + * Tensor-core GEMM: C[M,N] (FP32) = A[M,K] (FP16, row-major) × B[N,K] (FP16, row-major). + * + *

Software-pipelined: each thread stages the NEXT K-step's A/B elements in + * registers while the CURRENT step's ldmatrix+MMA execute, so global-memory + * latency is hidden behind tensor-core compute. Shared memory stays + * single-buffered; the two block barriers per step preserve correctness + * (read-complete before overwrite, write-complete before next ldmatrix).

+ * + *

Requires M % 128 == 0, N % 128 == 0, K % 16 == 0, SM 8.0+.

+ * + * Worker: WorkerGrid2D((M/128)*256, N/128), local (256,1,1). + */ + public static void gemmMMA(KernelContext ctx, + HalfFloatArray A, HalfFloatArray B, FloatArray C, + int M, int N, int K) { + int tid = ctx.localIdx; + int warpId = tid / WARP_SIZE; + int warpM = warpId / WARPS_N; + int warpN = warpId % WARPS_N; + int blockRow = BM * ctx.groupIdx; + int blockCol = BN * ctx.groupIdy; + + int[] aTile = ctx.allocateIntLocalArray(BM * BK / 2); + int[] bTile = ctx.allocateIntLocalArray(BK * BN / 2); + + float[] c00 = ctx.mmaFragment(0.0f); float[] c01 = ctx.mmaFragment(0.0f); + float[] c02 = ctx.mmaFragment(0.0f); float[] c03 = ctx.mmaFragment(0.0f); + float[] c04 = ctx.mmaFragment(0.0f); float[] c05 = ctx.mmaFragment(0.0f); + float[] c06 = ctx.mmaFragment(0.0f); float[] c07 = ctx.mmaFragment(0.0f); + float[] c10 = ctx.mmaFragment(0.0f); float[] c11 = ctx.mmaFragment(0.0f); + float[] c12 = ctx.mmaFragment(0.0f); float[] c13 = ctx.mmaFragment(0.0f); + float[] c14 = ctx.mmaFragment(0.0f); float[] c15 = ctx.mmaFragment(0.0f); + float[] c16 = ctx.mmaFragment(0.0f); float[] c17 = ctx.mmaFragment(0.0f); + + // ── Per-thread staging index math (constant across K-steps) ────────── + // A tile: BM*BK/2 = 1024 ints; layout idx = m_row*(BK/2) + k_pair. + // m_row = idx >>> 3, k = (idx & 7)*2; global A element (blockRow+m_row, kBase+k). + int aIdx0 = tid; int gA0 = (blockRow + (aIdx0 >>> 3)) * K + ((aIdx0 & 7) << 1); + int aIdx1 = tid + 256; int gA1 = (blockRow + (aIdx1 >>> 3)) * K + ((aIdx1 & 7) << 1); + int aIdx2 = tid + 512; int gA2 = (blockRow + (aIdx2 >>> 3)) * K + ((aIdx2 & 7) << 1); + int aIdx3 = tid + 768; int gA3 = (blockRow + (aIdx3 >>> 3)) * K + ((aIdx3 & 7) << 1); + // B tile: BK*BN/2 = 1024 ints; subTileId = idx >>> 6, k_row = (idx & 63) >>> 2, + // col = subTileId*8 + (idx & 3)*2; B[col, k] at col*K + k (pair at +K). + int bIdx0 = tid; int gB0 = (blockCol + ((bIdx0 >>> 6) << 3) + ((bIdx0 & 3) << 1)) * K + ((bIdx0 & 63) >>> 2); + int bIdx1 = tid + 256; int gB1 = (blockCol + ((bIdx1 >>> 6) << 3) + ((bIdx1 & 3) << 1)) * K + ((bIdx1 & 63) >>> 2); + int bIdx2 = tid + 512; int gB2 = (blockCol + ((bIdx2 >>> 6) << 3) + ((bIdx2 & 3) << 1)) * K + ((bIdx2 & 63) >>> 2); + int bIdx3 = tid + 768; int gB3 = (blockCol + ((bIdx3 >>> 6) << 3) + ((bIdx3 & 3) << 1)) * K + ((bIdx3 & 63) >>> 2); + + // ── Prologue: stage K-step 0 ───────────────────────────────────────── + int aReg0 = packHalves(A, gA0, gA0 + 1); + int aReg1 = packHalves(A, gA1, gA1 + 1); + int aReg2 = packHalves(A, gA2, gA2 + 1); + int aReg3 = packHalves(A, gA3, gA3 + 1); + int bReg0 = packHalves(B, gB0, gB0 + K); + int bReg1 = packHalves(B, gB1, gB1 + K); + int bReg2 = packHalves(B, gB2, gB2 + K); + int bReg3 = packHalves(B, gB3, gB3 + K); + aTile[aIdx0] = aReg0; aTile[aIdx1] = aReg1; aTile[aIdx2] = aReg2; aTile[aIdx3] = aReg3; + bTile[bIdx0] = bReg0; bTile[bIdx1] = bReg1; bTile[bIdx2] = bReg2; bTile[bIdx3] = bReg3; + ctx.localBarrier(); + + int numKSteps = K / BK; + for (int kStep = 0; kStep < numKSteps; kStep++) { + // Issue next step's global loads FIRST: independent of the MMAs below, + // so their latency overlaps ldmatrix + tensor-core compute. + if (kStep + 1 < numKSteps) { + int kOff = (kStep + 1) * BK; + aReg0 = packHalves(A, gA0 + kOff, gA0 + kOff + 1); + aReg1 = packHalves(A, gA1 + kOff, gA1 + kOff + 1); + aReg2 = packHalves(A, gA2 + kOff, gA2 + kOff + 1); + aReg3 = packHalves(A, gA3 + kOff, gA3 + kOff + 1); + bReg0 = packHalves(B, gB0 + kOff, gB0 + kOff + K); + bReg1 = packHalves(B, gB1 + kOff, gB1 + kOff + K); + bReg2 = packHalves(B, gB2 + kOff, gB2 + kOff + K); + bReg3 = packHalves(B, gB3 + kOff, gB3 + kOff + K); + } + + int aOff0 = warpM * 1024; + int aOff1 = warpM * 1024 + 512; + HalfFloat[] a0 = ctx.mmaLoadA(aTile, BK, aOff0); + HalfFloat[] a1 = ctx.mmaLoadA(aTile, BK, aOff1); + int bBase = warpN * 8; + HalfFloat[] b0 = ctx.mmaLoadB(bTile, BK, (bBase + 0) * B_SUBTILE_BYTES); + HalfFloat[] b1 = ctx.mmaLoadB(bTile, BK, (bBase + 1) * B_SUBTILE_BYTES); + HalfFloat[] b2 = ctx.mmaLoadB(bTile, BK, (bBase + 2) * B_SUBTILE_BYTES); + HalfFloat[] b3 = ctx.mmaLoadB(bTile, BK, (bBase + 3) * B_SUBTILE_BYTES); + HalfFloat[] b4 = ctx.mmaLoadB(bTile, BK, (bBase + 4) * B_SUBTILE_BYTES); + HalfFloat[] b5 = ctx.mmaLoadB(bTile, BK, (bBase + 5) * B_SUBTILE_BYTES); + HalfFloat[] b6 = ctx.mmaLoadB(bTile, BK, (bBase + 6) * B_SUBTILE_BYTES); + HalfFloat[] b7 = ctx.mmaLoadB(bTile, BK, (bBase + 7) * B_SUBTILE_BYTES); + ctx.localBarrier(); // all shared reads for this K-step complete + + // Overwrite shared tiles for the next step; overlaps the MMAs below. + if (kStep + 1 < numKSteps) { + aTile[aIdx0] = aReg0; aTile[aIdx1] = aReg1; aTile[aIdx2] = aReg2; aTile[aIdx3] = aReg3; + bTile[bIdx0] = bReg0; bTile[bIdx1] = bReg1; bTile[bIdx2] = bReg2; bTile[bIdx3] = bReg3; + } + + c00 = ctx.mma(a0, b0, c00, MMAShape.M16N8K16); + c01 = ctx.mma(a0, b1, c01, MMAShape.M16N8K16); + c02 = ctx.mma(a0, b2, c02, MMAShape.M16N8K16); + c03 = ctx.mma(a0, b3, c03, MMAShape.M16N8K16); + c04 = ctx.mma(a0, b4, c04, MMAShape.M16N8K16); + c05 = ctx.mma(a0, b5, c05, MMAShape.M16N8K16); + c06 = ctx.mma(a0, b6, c06, MMAShape.M16N8K16); + c07 = ctx.mma(a0, b7, c07, MMAShape.M16N8K16); + c10 = ctx.mma(a1, b0, c10, MMAShape.M16N8K16); + c11 = ctx.mma(a1, b1, c11, MMAShape.M16N8K16); + c12 = ctx.mma(a1, b2, c12, MMAShape.M16N8K16); + c13 = ctx.mma(a1, b3, c13, MMAShape.M16N8K16); + c14 = ctx.mma(a1, b4, c14, MMAShape.M16N8K16); + c15 = ctx.mma(a1, b5, c15, MMAShape.M16N8K16); + c16 = ctx.mma(a1, b6, c16, MMAShape.M16N8K16); + c17 = ctx.mma(a1, b7, c17, MMAShape.M16N8K16); + ctx.localBarrier(); // shared writes visible before next step's ldmatrix + } + + int rBase = blockRow + warpM * WM; + int cBase = blockCol + warpN * WN; + ctx.mmaStore(c00, C, rBase + 0, cBase + 0, N); + ctx.mmaStore(c01, C, rBase + 0, cBase + 8, N); + ctx.mmaStore(c02, C, rBase + 0, cBase + 16, N); + ctx.mmaStore(c03, C, rBase + 0, cBase + 24, N); + ctx.mmaStore(c04, C, rBase + 0, cBase + 32, N); + ctx.mmaStore(c05, C, rBase + 0, cBase + 40, N); + ctx.mmaStore(c06, C, rBase + 0, cBase + 48, N); + ctx.mmaStore(c07, C, rBase + 0, cBase + 56, N); + ctx.mmaStore(c10, C, rBase + 16, cBase + 0, N); + ctx.mmaStore(c11, C, rBase + 16, cBase + 8, N); + ctx.mmaStore(c12, C, rBase + 16, cBase + 16, N); + ctx.mmaStore(c13, C, rBase + 16, cBase + 24, N); + ctx.mmaStore(c14, C, rBase + 16, cBase + 32, N); + ctx.mmaStore(c15, C, rBase + 16, cBase + 40, N); + ctx.mmaStore(c16, C, rBase + 16, cBase + 48, N); + ctx.mmaStore(c17, C, rBase + 16, cBase + 56, N); + } + + // ── Residual add (FP32) ─────────────────────────────────────────────── + // gemmMMA overwrites C, but Wo and W2 both need x = x + W·a. + // Worker: B*dim global threads (valid rows only), localSize=256. + public static void batchedResidualAddFP32(KernelContext context, + FloatArray residual, // x (in/out) + FloatArray delta) { // GEMM output + int gid = context.globalIdx; + residual.set(gid, residual.get(gid) + delta.get(gid)); + } + + // ── SwiGLU emitting FP16 ────────────────────────────────────────────── + // Replaces batchedFFNSwiGLU. Output is the A operand for the W2 GEMM. + // Worker: B*hiddenDim global threads, localSize=256. + public static void batchedFFNSwiGLUFP16(KernelContext context, + HalfFloatArray wrapHbFP16Batch, + FloatArray ffnGateResult, + FloatArray ffnUpResult, + int hiddenDim) { + int gid = context.globalIdx; + float g = ffnGateResult.get(gid); + float u = ffnUpResult.get(gid); + float silu = g / (1.0f + TornadoMath.exp(-g)); + wrapHbFP16Batch.set(gid, new HalfFloat(silu * u)); + } + + // ── FP32 → FP16 cast (Option B only, see Wo below) ──────────────────── + // Worker: B*dim global threads, localSize=256. + public static void batchedConvertFP32toFP16(KernelContext context, + FloatArray in, + HalfFloatArray out) { + int gid = context.globalIdx; + out.set(gid, new HalfFloat(in.get(gid))); + } + + + // ── Fused MMA projections ───────────────────────────────────────────────── + // + // Q, K, V (and gate/up) share the same A operand and the same K dimension, + // so they are fused into ONE kernel launch each. The N grid spans the packed + // output [dim | kvDim | kvDim] (resp. [hidDim | hidDim]); each thread block + // selects its weight matrix from groupIdy — a block-uniform branch, so there + // is zero divergence and no weight duplication in memory. This restores the + // A-reuse of the old fused matvec kernels AND fixes grid starvation for the + // skinny GQA projections (kvDim/128 blocks alone cannot fill the SMs). + + /** + * Fused QKV tensor-core GEMM into a PACKED output: + * qkvOut[M, dim+2*kvDim] = A[M,K] × [Wq | Wk | Wv] (each [N_i, K] row-major). + * + *

Layout of a row of qkvOut: [ q(0..dim) | k(0..kvDim) | v(0..kvDim) ]. + * Requires dim % 128 == 0 and kvDim % 128 == 0.

+ * + * Worker: WorkerGrid2D((M/128)*256, (dim+2*kvDim)/128), local (256,1,1). + */ + public static void gemmMMAQKV(KernelContext ctx, + HalfFloatArray A, + HalfFloatArray wq, HalfFloatArray wk, HalfFloatArray wv, + FloatArray qkvOut, + int M, int dim, int kvDim, int K) { + int tid = ctx.localIdx; + int warpId = tid / WARP_SIZE; + int warpM = warpId / WARPS_N; + int warpN = warpId % WARPS_N; + int blockRow = BM * ctx.groupIdx; + int blockCol = BN * ctx.groupIdy; // column in the packed output + int qkvStride = dim + 2 * kvDim; + + // Column base inside the segment's own weight matrix (block-uniform). + int wColBase = blockCol; + if (blockCol >= dim) wColBase -= dim; + if (blockCol >= dim + kvDim) wColBase -= kvDim; + + int[] aTile = ctx.allocateIntLocalArray(BM * BK / 2); + int[] bTile = ctx.allocateIntLocalArray(BK * BN / 2); + + float[] c00 = ctx.mmaFragment(0.0f); float[] c01 = ctx.mmaFragment(0.0f); + float[] c02 = ctx.mmaFragment(0.0f); float[] c03 = ctx.mmaFragment(0.0f); + float[] c04 = ctx.mmaFragment(0.0f); float[] c05 = ctx.mmaFragment(0.0f); + float[] c06 = ctx.mmaFragment(0.0f); float[] c07 = ctx.mmaFragment(0.0f); + float[] c10 = ctx.mmaFragment(0.0f); float[] c11 = ctx.mmaFragment(0.0f); + float[] c12 = ctx.mmaFragment(0.0f); float[] c13 = ctx.mmaFragment(0.0f); + float[] c14 = ctx.mmaFragment(0.0f); float[] c15 = ctx.mmaFragment(0.0f); + float[] c16 = ctx.mmaFragment(0.0f); float[] c17 = ctx.mmaFragment(0.0f); + + int aIdx0 = tid; int gA0 = (blockRow + (aIdx0 >>> 3)) * K + ((aIdx0 & 7) << 1); + int aIdx1 = tid + 256; int gA1 = (blockRow + (aIdx1 >>> 3)) * K + ((aIdx1 & 7) << 1); + int aIdx2 = tid + 512; int gA2 = (blockRow + (aIdx2 >>> 3)) * K + ((aIdx2 & 7) << 1); + int aIdx3 = tid + 768; int gA3 = (blockRow + (aIdx3 >>> 3)) * K + ((aIdx3 & 7) << 1); + int bIdx0 = tid; int gB0 = (wColBase + ((bIdx0 >>> 6) << 3) + ((bIdx0 & 3) << 1)) * K + ((bIdx0 & 63) >>> 2); + int bIdx1 = tid + 256; int gB1 = (wColBase + ((bIdx1 >>> 6) << 3) + ((bIdx1 & 3) << 1)) * K + ((bIdx1 & 63) >>> 2); + int bIdx2 = tid + 512; int gB2 = (wColBase + ((bIdx2 >>> 6) << 3) + ((bIdx2 & 3) << 1)) * K + ((bIdx2 & 63) >>> 2); + int bIdx3 = tid + 768; int gB3 = (wColBase + ((bIdx3 >>> 6) << 3) + ((bIdx3 & 3) << 1)) * K + ((bIdx3 & 63) >>> 2); + + // ── Prologue: stage K-step 0 ── + int aReg0 = packHalves(A, gA0, gA0 + 1); + int aReg1 = packHalves(A, gA1, gA1 + 1); + int aReg2 = packHalves(A, gA2, gA2 + 1); + int aReg3 = packHalves(A, gA3, gA3 + 1); + int bReg0; int bReg1; int bReg2; int bReg3; + if (blockCol < dim) { + bReg0 = packHalves(wq, gB0, gB0 + K); + bReg1 = packHalves(wq, gB1, gB1 + K); + bReg2 = packHalves(wq, gB2, gB2 + K); + bReg3 = packHalves(wq, gB3, gB3 + K); + } else if (blockCol < dim + kvDim) { + bReg0 = packHalves(wk, gB0, gB0 + K); + bReg1 = packHalves(wk, gB1, gB1 + K); + bReg2 = packHalves(wk, gB2, gB2 + K); + bReg3 = packHalves(wk, gB3, gB3 + K); + } else { + bReg0 = packHalves(wv, gB0, gB0 + K); + bReg1 = packHalves(wv, gB1, gB1 + K); + bReg2 = packHalves(wv, gB2, gB2 + K); + bReg3 = packHalves(wv, gB3, gB3 + K); + } + aTile[aIdx0] = aReg0; aTile[aIdx1] = aReg1; aTile[aIdx2] = aReg2; aTile[aIdx3] = aReg3; + bTile[bIdx0] = bReg0; bTile[bIdx1] = bReg1; bTile[bIdx2] = bReg2; bTile[bIdx3] = bReg3; + ctx.localBarrier(); + + int numKSteps = K / BK; + for (int kStep = 0; kStep < numKSteps; kStep++) { + if (kStep + 1 < numKSteps) { + int kOff = (kStep + 1) * BK; + aReg0 = packHalves(A, gA0 + kOff, gA0 + kOff + 1); + aReg1 = packHalves(A, gA1 + kOff, gA1 + kOff + 1); + aReg2 = packHalves(A, gA2 + kOff, gA2 + kOff + 1); + aReg3 = packHalves(A, gA3 + kOff, gA3 + kOff + 1); + if (blockCol < dim) { + bReg0 = packHalves(wq, gB0 + kOff, gB0 + kOff + K); + bReg1 = packHalves(wq, gB1 + kOff, gB1 + kOff + K); + bReg2 = packHalves(wq, gB2 + kOff, gB2 + kOff + K); + bReg3 = packHalves(wq, gB3 + kOff, gB3 + kOff + K); + } else if (blockCol < dim + kvDim) { + bReg0 = packHalves(wk, gB0 + kOff, gB0 + kOff + K); + bReg1 = packHalves(wk, gB1 + kOff, gB1 + kOff + K); + bReg2 = packHalves(wk, gB2 + kOff, gB2 + kOff + K); + bReg3 = packHalves(wk, gB3 + kOff, gB3 + kOff + K); + } else { + bReg0 = packHalves(wv, gB0 + kOff, gB0 + kOff + K); + bReg1 = packHalves(wv, gB1 + kOff, gB1 + kOff + K); + bReg2 = packHalves(wv, gB2 + kOff, gB2 + kOff + K); + bReg3 = packHalves(wv, gB3 + kOff, gB3 + kOff + K); + } + } + + int aOff0 = warpM * 1024; + int aOff1 = warpM * 1024 + 512; + HalfFloat[] a0 = ctx.mmaLoadA(aTile, BK, aOff0); + HalfFloat[] a1 = ctx.mmaLoadA(aTile, BK, aOff1); + int bBase = warpN * 8; + HalfFloat[] b0 = ctx.mmaLoadB(bTile, BK, (bBase + 0) * B_SUBTILE_BYTES); + HalfFloat[] b1 = ctx.mmaLoadB(bTile, BK, (bBase + 1) * B_SUBTILE_BYTES); + HalfFloat[] b2 = ctx.mmaLoadB(bTile, BK, (bBase + 2) * B_SUBTILE_BYTES); + HalfFloat[] b3 = ctx.mmaLoadB(bTile, BK, (bBase + 3) * B_SUBTILE_BYTES); + HalfFloat[] b4 = ctx.mmaLoadB(bTile, BK, (bBase + 4) * B_SUBTILE_BYTES); + HalfFloat[] b5 = ctx.mmaLoadB(bTile, BK, (bBase + 5) * B_SUBTILE_BYTES); + HalfFloat[] b6 = ctx.mmaLoadB(bTile, BK, (bBase + 6) * B_SUBTILE_BYTES); + HalfFloat[] b7 = ctx.mmaLoadB(bTile, BK, (bBase + 7) * B_SUBTILE_BYTES); + ctx.localBarrier(); + + if (kStep + 1 < numKSteps) { + aTile[aIdx0] = aReg0; aTile[aIdx1] = aReg1; aTile[aIdx2] = aReg2; aTile[aIdx3] = aReg3; + bTile[bIdx0] = bReg0; bTile[bIdx1] = bReg1; bTile[bIdx2] = bReg2; bTile[bIdx3] = bReg3; + } + + c00 = ctx.mma(a0, b0, c00, MMAShape.M16N8K16); + c01 = ctx.mma(a0, b1, c01, MMAShape.M16N8K16); + c02 = ctx.mma(a0, b2, c02, MMAShape.M16N8K16); + c03 = ctx.mma(a0, b3, c03, MMAShape.M16N8K16); + c04 = ctx.mma(a0, b4, c04, MMAShape.M16N8K16); + c05 = ctx.mma(a0, b5, c05, MMAShape.M16N8K16); + c06 = ctx.mma(a0, b6, c06, MMAShape.M16N8K16); + c07 = ctx.mma(a0, b7, c07, MMAShape.M16N8K16); + c10 = ctx.mma(a1, b0, c10, MMAShape.M16N8K16); + c11 = ctx.mma(a1, b1, c11, MMAShape.M16N8K16); + c12 = ctx.mma(a1, b2, c12, MMAShape.M16N8K16); + c13 = ctx.mma(a1, b3, c13, MMAShape.M16N8K16); + c14 = ctx.mma(a1, b4, c14, MMAShape.M16N8K16); + c15 = ctx.mma(a1, b5, c15, MMAShape.M16N8K16); + c16 = ctx.mma(a1, b6, c16, MMAShape.M16N8K16); + c17 = ctx.mma(a1, b7, c17, MMAShape.M16N8K16); + ctx.localBarrier(); + } + + // Stores are uniform: packed column == blockCol-relative column, + // stride is the packed row width. + int rBase = blockRow + warpM * WM; + int cBase = blockCol + warpN * WN; + ctx.mmaStore(c00, qkvOut, rBase + 0, cBase + 0, qkvStride); + ctx.mmaStore(c01, qkvOut, rBase + 0, cBase + 8, qkvStride); + ctx.mmaStore(c02, qkvOut, rBase + 0, cBase + 16, qkvStride); + ctx.mmaStore(c03, qkvOut, rBase + 0, cBase + 24, qkvStride); + ctx.mmaStore(c04, qkvOut, rBase + 0, cBase + 32, qkvStride); + ctx.mmaStore(c05, qkvOut, rBase + 0, cBase + 40, qkvStride); + ctx.mmaStore(c06, qkvOut, rBase + 0, cBase + 48, qkvStride); + ctx.mmaStore(c07, qkvOut, rBase + 0, cBase + 56, qkvStride); + ctx.mmaStore(c10, qkvOut, rBase + 16, cBase + 0, qkvStride); + ctx.mmaStore(c11, qkvOut, rBase + 16, cBase + 8, qkvStride); + ctx.mmaStore(c12, qkvOut, rBase + 16, cBase + 16, qkvStride); + ctx.mmaStore(c13, qkvOut, rBase + 16, cBase + 24, qkvStride); + ctx.mmaStore(c14, qkvOut, rBase + 16, cBase + 32, qkvStride); + ctx.mmaStore(c15, qkvOut, rBase + 16, cBase + 40, qkvStride); + ctx.mmaStore(c16, qkvOut, rBase + 16, cBase + 48, qkvStride); + ctx.mmaStore(c17, qkvOut, rBase + 16, cBase + 56, qkvStride); + } + + /** + * Fused W1/W3 (gate/up) tensor-core GEMM into a PACKED output: + * gateUpOut[M, 2*hidDim] = A[M,K] × [W1 | W3] (each [hidDim, K] row-major). + * + *

Layout of a row: [ gate(0..hidDim) | up(0..hidDim) ]. + * Requires hidDim % 128 == 0.

+ * + * Worker: WorkerGrid2D((M/128)*256, (2*hidDim)/128), local (256,1,1). + */ + public static void gemmMMAGateUp(KernelContext ctx, + HalfFloatArray A, + HalfFloatArray w1, HalfFloatArray w3, + FloatArray gateUpOut, + int M, int hidDim, int K) { + int tid = ctx.localIdx; + int warpId = tid / WARP_SIZE; + int warpM = warpId / WARPS_N; + int warpN = warpId % WARPS_N; + int blockRow = BM * ctx.groupIdx; + int blockCol = BN * ctx.groupIdy; // column in the packed output + int outStride = 2 * hidDim; + + int wColBase = (blockCol < hidDim) ? blockCol : (blockCol - hidDim); + + int[] aTile = ctx.allocateIntLocalArray(BM * BK / 2); + int[] bTile = ctx.allocateIntLocalArray(BK * BN / 2); + + float[] c00 = ctx.mmaFragment(0.0f); float[] c01 = ctx.mmaFragment(0.0f); + float[] c02 = ctx.mmaFragment(0.0f); float[] c03 = ctx.mmaFragment(0.0f); + float[] c04 = ctx.mmaFragment(0.0f); float[] c05 = ctx.mmaFragment(0.0f); + float[] c06 = ctx.mmaFragment(0.0f); float[] c07 = ctx.mmaFragment(0.0f); + float[] c10 = ctx.mmaFragment(0.0f); float[] c11 = ctx.mmaFragment(0.0f); + float[] c12 = ctx.mmaFragment(0.0f); float[] c13 = ctx.mmaFragment(0.0f); + float[] c14 = ctx.mmaFragment(0.0f); float[] c15 = ctx.mmaFragment(0.0f); + float[] c16 = ctx.mmaFragment(0.0f); float[] c17 = ctx.mmaFragment(0.0f); + + int aIdx0 = tid; int gA0 = (blockRow + (aIdx0 >>> 3)) * K + ((aIdx0 & 7) << 1); + int aIdx1 = tid + 256; int gA1 = (blockRow + (aIdx1 >>> 3)) * K + ((aIdx1 & 7) << 1); + int aIdx2 = tid + 512; int gA2 = (blockRow + (aIdx2 >>> 3)) * K + ((aIdx2 & 7) << 1); + int aIdx3 = tid + 768; int gA3 = (blockRow + (aIdx3 >>> 3)) * K + ((aIdx3 & 7) << 1); + int bIdx0 = tid; int gB0 = (wColBase + ((bIdx0 >>> 6) << 3) + ((bIdx0 & 3) << 1)) * K + ((bIdx0 & 63) >>> 2); + int bIdx1 = tid + 256; int gB1 = (wColBase + ((bIdx1 >>> 6) << 3) + ((bIdx1 & 3) << 1)) * K + ((bIdx1 & 63) >>> 2); + int bIdx2 = tid + 512; int gB2 = (wColBase + ((bIdx2 >>> 6) << 3) + ((bIdx2 & 3) << 1)) * K + ((bIdx2 & 63) >>> 2); + int bIdx3 = tid + 768; int gB3 = (wColBase + ((bIdx3 >>> 6) << 3) + ((bIdx3 & 3) << 1)) * K + ((bIdx3 & 63) >>> 2); + + int aReg0 = packHalves(A, gA0, gA0 + 1); + int aReg1 = packHalves(A, gA1, gA1 + 1); + int aReg2 = packHalves(A, gA2, gA2 + 1); + int aReg3 = packHalves(A, gA3, gA3 + 1); + int bReg0; int bReg1; int bReg2; int bReg3; + if (blockCol < hidDim) { + bReg0 = packHalves(w1, gB0, gB0 + K); + bReg1 = packHalves(w1, gB1, gB1 + K); + bReg2 = packHalves(w1, gB2, gB2 + K); + bReg3 = packHalves(w1, gB3, gB3 + K); + } else { + bReg0 = packHalves(w3, gB0, gB0 + K); + bReg1 = packHalves(w3, gB1, gB1 + K); + bReg2 = packHalves(w3, gB2, gB2 + K); + bReg3 = packHalves(w3, gB3, gB3 + K); + } + aTile[aIdx0] = aReg0; aTile[aIdx1] = aReg1; aTile[aIdx2] = aReg2; aTile[aIdx3] = aReg3; + bTile[bIdx0] = bReg0; bTile[bIdx1] = bReg1; bTile[bIdx2] = bReg2; bTile[bIdx3] = bReg3; + ctx.localBarrier(); + + int numKSteps = K / BK; + for (int kStep = 0; kStep < numKSteps; kStep++) { + if (kStep + 1 < numKSteps) { + int kOff = (kStep + 1) * BK; + aReg0 = packHalves(A, gA0 + kOff, gA0 + kOff + 1); + aReg1 = packHalves(A, gA1 + kOff, gA1 + kOff + 1); + aReg2 = packHalves(A, gA2 + kOff, gA2 + kOff + 1); + aReg3 = packHalves(A, gA3 + kOff, gA3 + kOff + 1); + if (blockCol < hidDim) { + bReg0 = packHalves(w1, gB0 + kOff, gB0 + kOff + K); + bReg1 = packHalves(w1, gB1 + kOff, gB1 + kOff + K); + bReg2 = packHalves(w1, gB2 + kOff, gB2 + kOff + K); + bReg3 = packHalves(w1, gB3 + kOff, gB3 + kOff + K); + } else { + bReg0 = packHalves(w3, gB0 + kOff, gB0 + kOff + K); + bReg1 = packHalves(w3, gB1 + kOff, gB1 + kOff + K); + bReg2 = packHalves(w3, gB2 + kOff, gB2 + kOff + K); + bReg3 = packHalves(w3, gB3 + kOff, gB3 + kOff + K); + } + } + + int aOff0 = warpM * 1024; + int aOff1 = warpM * 1024 + 512; + HalfFloat[] a0 = ctx.mmaLoadA(aTile, BK, aOff0); + HalfFloat[] a1 = ctx.mmaLoadA(aTile, BK, aOff1); + int bBase = warpN * 8; + HalfFloat[] b0 = ctx.mmaLoadB(bTile, BK, (bBase + 0) * B_SUBTILE_BYTES); + HalfFloat[] b1 = ctx.mmaLoadB(bTile, BK, (bBase + 1) * B_SUBTILE_BYTES); + HalfFloat[] b2 = ctx.mmaLoadB(bTile, BK, (bBase + 2) * B_SUBTILE_BYTES); + HalfFloat[] b3 = ctx.mmaLoadB(bTile, BK, (bBase + 3) * B_SUBTILE_BYTES); + HalfFloat[] b4 = ctx.mmaLoadB(bTile, BK, (bBase + 4) * B_SUBTILE_BYTES); + HalfFloat[] b5 = ctx.mmaLoadB(bTile, BK, (bBase + 5) * B_SUBTILE_BYTES); + HalfFloat[] b6 = ctx.mmaLoadB(bTile, BK, (bBase + 6) * B_SUBTILE_BYTES); + HalfFloat[] b7 = ctx.mmaLoadB(bTile, BK, (bBase + 7) * B_SUBTILE_BYTES); + ctx.localBarrier(); + + if (kStep + 1 < numKSteps) { + aTile[aIdx0] = aReg0; aTile[aIdx1] = aReg1; aTile[aIdx2] = aReg2; aTile[aIdx3] = aReg3; + bTile[bIdx0] = bReg0; bTile[bIdx1] = bReg1; bTile[bIdx2] = bReg2; bTile[bIdx3] = bReg3; + } + + c00 = ctx.mma(a0, b0, c00, MMAShape.M16N8K16); + c01 = ctx.mma(a0, b1, c01, MMAShape.M16N8K16); + c02 = ctx.mma(a0, b2, c02, MMAShape.M16N8K16); + c03 = ctx.mma(a0, b3, c03, MMAShape.M16N8K16); + c04 = ctx.mma(a0, b4, c04, MMAShape.M16N8K16); + c05 = ctx.mma(a0, b5, c05, MMAShape.M16N8K16); + c06 = ctx.mma(a0, b6, c06, MMAShape.M16N8K16); + c07 = ctx.mma(a0, b7, c07, MMAShape.M16N8K16); + c10 = ctx.mma(a1, b0, c10, MMAShape.M16N8K16); + c11 = ctx.mma(a1, b1, c11, MMAShape.M16N8K16); + c12 = ctx.mma(a1, b2, c12, MMAShape.M16N8K16); + c13 = ctx.mma(a1, b3, c13, MMAShape.M16N8K16); + c14 = ctx.mma(a1, b4, c14, MMAShape.M16N8K16); + c15 = ctx.mma(a1, b5, c15, MMAShape.M16N8K16); + c16 = ctx.mma(a1, b6, c16, MMAShape.M16N8K16); + c17 = ctx.mma(a1, b7, c17, MMAShape.M16N8K16); + ctx.localBarrier(); + } + + int rBase = blockRow + warpM * WM; + int cBase = blockCol + warpN * WN; + ctx.mmaStore(c00, gateUpOut, rBase + 0, cBase + 0, outStride); + ctx.mmaStore(c01, gateUpOut, rBase + 0, cBase + 8, outStride); + ctx.mmaStore(c02, gateUpOut, rBase + 0, cBase + 16, outStride); + ctx.mmaStore(c03, gateUpOut, rBase + 0, cBase + 24, outStride); + ctx.mmaStore(c04, gateUpOut, rBase + 0, cBase + 32, outStride); + ctx.mmaStore(c05, gateUpOut, rBase + 0, cBase + 40, outStride); + ctx.mmaStore(c06, gateUpOut, rBase + 0, cBase + 48, outStride); + ctx.mmaStore(c07, gateUpOut, rBase + 0, cBase + 56, outStride); + ctx.mmaStore(c10, gateUpOut, rBase + 16, cBase + 0, outStride); + ctx.mmaStore(c11, gateUpOut, rBase + 16, cBase + 8, outStride); + ctx.mmaStore(c12, gateUpOut, rBase + 16, cBase + 16, outStride); + ctx.mmaStore(c13, gateUpOut, rBase + 16, cBase + 24, outStride); + ctx.mmaStore(c14, gateUpOut, rBase + 16, cBase + 32, outStride); + ctx.mmaStore(c15, gateUpOut, rBase + 16, cBase + 40, outStride); + ctx.mmaStore(c16, gateUpOut, rBase + 16, cBase + 48, outStride); + ctx.mmaStore(c17, gateUpOut, rBase + 16, cBase + 56, outStride); + } + + // ── Parallel RMS reductions ─────────────────────────────────────────────── + // Replace the localSize=1 sequential reductions (one thread walking `dim` + // elements alone) with one 256-thread workgroup per token and a shared-memory + // tree reduction. + + /** + * Parallel RMS square-sum reduction. One workgroup per batch token. + * + * Worker: B workgroups × localSize threads (localSize=256). + */ + public static void batchedRmsReduceParallel(KernelContext context, + FloatArray wrapXBatch, + FloatArray scaleBatch, + int dim, float eps, int localSize) { + int tid = context.localIdx; + int b = context.groupIdx; + int localSz = context.localGroupSizeX; + float[] partial = context.allocateFloatLocalArray(localSize); + + int base = b * dim; + float ss = 0.0f; + for (int i = tid; i < dim; i += localSz) { + float v = wrapXBatch.get(base + i); + ss += v * v; + } + partial[tid] = ss; + context.localBarrier(); + for (int s = localSz / 2; s > 0; s >>= 1) { + if (tid < s) { + partial[tid] += partial[tid + s]; + } + context.localBarrier(); + } + if (tid == 0) { + float m = partial[0] / dim + eps; + scaleBatch.set(b, 1.0f / TornadoMath.sqrt(m)); + } + } + + /** + * Parallel RMS reduction FUSED with the pending residual add: + * x[b,i] += delta[b,i] first, then square-sum over the updated row. + * Replaces the separate woResid task + FFN RMS reduce (each element is + * visited exactly once, so the in-place update is race-free). + * + * Worker: B workgroups × localSize threads (localSize=256). + */ + public static void batchedRmsReduceFusedResidual(KernelContext context, + FloatArray wrapXBatch, + FloatArray delta, + FloatArray scaleBatch, + int dim, float eps, int localSize) { + int tid = context.localIdx; + int b = context.groupIdx; + int localSz = context.localGroupSizeX; + float[] partial = context.allocateFloatLocalArray(localSize); + + int base = b * dim; + float ss = 0.0f; + for (int i = tid; i < dim; i += localSz) { + float v = wrapXBatch.get(base + i) + delta.get(base + i); + wrapXBatch.set(base + i, v); + ss += v * v; + } + partial[tid] = ss; + context.localBarrier(); + for (int s = localSz / 2; s > 0; s >>= 1) { + if (tid < s) { + partial[tid] += partial[tid + s]; + } + context.localBarrier(); + } + if (tid == 0) { + float m = partial[0] / dim + eps; + scaleBatch.set(b, 1.0f / TornadoMath.sqrt(m)); + } + } + + // ── RoPE + KV cache over the packed QKV buffer ──────────────────────────── + + /** + * Fused batched RoPE rotation + KV cache write, reading/writing the PACKED + * QKV buffer produced by {@link #gemmMMAQKV}. Row layout: [ q | k | v ]. + * Q is rotated in place (consumed from the packed buffer by attention). + * + * Worker: B*(dim/2) global threads, localSize=512 (or less). + */ + public static void batchedRopeWithKVCachePacked(KernelContext context, + IntArray batchStartPosHolder, + FloatArray qkvBatch, + FloatArray wrapKeyCache, + FloatArray wrapValueCache, + int kvDim, int headSize, + int layerIndex, int contextLength, int dim) { + int globalIdx = context.globalIdx; + int halfDim = dim / 2; + int batchIdx = globalIdx / halfDim; + int pairIdx = globalIdx % halfDim; + int i = pairIdx * 2; + int qkvStride = dim + 2 * kvDim; + + int pos = batchStartPosHolder.get(0) + batchIdx; + int qOffset = batchIdx * qkvStride; + int kOffset = batchIdx * qkvStride + dim; + int vOffset = batchIdx * qkvStride + dim + kvDim; + + if (i + 1 < dim) { + int head_dim = i % headSize; + float freq = 1.0f / TornadoMath.pow(50000.0f, head_dim / (float) headSize); + float val = pos * freq; + float fcr = TornadoMath.cos(val); + float fci = TornadoMath.sin(val); + + // Rotate Q in place + float v0q = qkvBatch.get(qOffset + i); + float v1q = qkvBatch.get(qOffset + i + 1); + qkvBatch.set(qOffset + i, v0q * fcr - v1q * fci); + qkvBatch.set(qOffset + i + 1, v0q * fci + v1q * fcr); + + // Rotate K and write K,V to cache + if (i + 1 < kvDim) { + float v0k = qkvBatch.get(kOffset + i); + float v1k = qkvBatch.get(kOffset + i + 1); + float rotK0 = v0k * fcr - v1k * fci; + float rotK1 = v0k * fci + v1k * fcr; + + int cacheOff = layerIndex * contextLength * kvDim + pos * kvDim; + wrapKeyCache.set(cacheOff + i, rotK0); + wrapKeyCache.set(cacheOff + i + 1, rotK1); + wrapValueCache.set(cacheOff + i, qkvBatch.get(vOffset + i)); + wrapValueCache.set(cacheOff + i + 1, qkvBatch.get(vOffset + i + 1)); + } + } + } + + // ── Flash attention (fixed accumulation, FP16 output) ──────────────────── + + /** + * Batched causal flash attention over the packed QKV buffer, writing FP16 + * directly (the A operand of the Wo GEMM — eliminates the attnCast pass). + * + *

Fixes the redundant accumulation of the previous version: each thread + * now OWNS output dims {tid, tid+localSz} and accumulates them in registers, + * instead of every thread redundantly computing the full headSize output + * vector into a (spilled) private array. K/V tile loads are flattened over + * (t, d) so consecutive threads issue coalesced reads.

+ * + *

Requires headSize <= 2*localSz (localSz = min(headSize, 128)).

+ * + * Worker: B*nHeads workgroups × min(headSize,128) threads. + */ + public static void batchedFlashAttentionFP16Out(KernelContext context, + IntArray batchStartPosHolder, + FloatArray qkvBatch, + FloatArray wrapKeyCache, + FloatArray wrapValueCache, + HalfFloatArray attnOutFP16, + int nHeads, int headSize, + int kvDim, int kvMul, + int layerIndex, int contextLength, int dim) { + int tid = context.localIdx; + int groupId = context.groupIdx; + int localSz = context.localGroupSizeX; + + int batchIdx = groupId / nHeads; + int h = groupId % nHeads; + int pos = batchStartPosHolder.get(0) + batchIdx; + int loff = layerIndex * contextLength * kvDim; + int kvHeadIdx = h / kvMul; + int BLOCK_C = 16; + int qkvStride = dim + 2 * kvDim; + + float[] qShared = context.allocateFloatLocalArray(headSize); + float[] kTile = context.allocateFloatLocalArray(BLOCK_C * headSize); + float[] vTile = context.allocateFloatLocalArray(BLOCK_C * headSize); + float[] sTile = context.allocateFloatLocalArray(BLOCK_C); + + // Load Q (rotated, from the packed QKV buffer) into shared memory + int qOffset = batchIdx * qkvStride + h * headSize; + for (int i = tid; i < headSize; i += localSz) { + qShared[i] = qkvBatch.get(qOffset + i); + } + context.localBarrier(); + + float maxScore = Float.NEGATIVE_INFINITY; + float sumExp = 0.0f; + // Each thread owns output dims d0 = tid and (if headSize > localSz) d1. + float acc0 = 0.0f; + float acc1 = 0.0f; + int d1 = tid + localSz; + + for (int tileC = 0; tileC <= pos; tileC += BLOCK_C) { + int tileEnd = Math.min(tileC + BLOCK_C - 1, pos); + int tileLen = tileEnd - tileC + 1; + + // Load K/V tile — flattened over (t, d) for coalescing + for (int idx = tid; idx < tileLen * headSize; idx += localSz) { + int tInTile = idx / headSize; + int d = idx % headSize; + int kvOff = loff + (tileC + tInTile) * kvDim + kvHeadIdx * headSize + d; + kTile[tInTile * headSize + d] = wrapKeyCache.get(kvOff); + vTile[tInTile * headSize + d] = wrapValueCache.get(kvOff); + } + context.localBarrier(); + + // Scores: one thread per key position in the tile + for (int t = tileC + tid; t <= tileEnd; t += localSz) { + int tInTile = t - tileC; + float score = 0.0f; + for (int d = 0; d < headSize; d++) { + score += qShared[d] * kTile[tInTile * headSize + d]; + } + sTile[tInTile] = score / TornadoMath.sqrt(headSize); + } + context.localBarrier(); + + // Tile max: redundant per-thread scan over <= 16 shared values — + // deterministic across the workgroup, no broadcast needed. + float tileMax = Float.NEGATIVE_INFINITY; + for (int t = 0; t < tileLen; t++) { + if (sTile[t] > tileMax) { + tileMax = sTile[t]; + } + } + + float newMax = Math.max(maxScore, tileMax); + if (maxScore != Float.NEGATIVE_INFINITY && newMax != maxScore) { + float corr = TornadoMath.exp(maxScore - newMax); + sumExp *= corr; + acc0 *= corr; + acc1 *= corr; + } + maxScore = newMax; + + for (int t = 0; t < tileLen; t++) { + float p = TornadoMath.exp(sTile[t] - maxScore); + sumExp += p; + acc0 += p * vTile[t * headSize + tid]; + if (d1 < headSize) { + acc1 += p * vTile[t * headSize + d1]; + } + } + context.localBarrier(); + } + + float norm = (sumExp > 0.0f) ? (1.0f / sumExp) : 0.0f; + int outOffset = batchIdx * dim + h * headSize; + attnOutFP16.set(outOffset + tid, new HalfFloat(acc0 * norm)); + if (d1 < headSize) { + attnOutFP16.set(outOffset + d1, new HalfFloat(acc1 * norm)); + } + } + + // ── SwiGLU over the packed gate/up buffer, emitting FP16 ───────────────── + + /** + * Fused SiLU(gate) * up over the PACKED [gate | up] GEMM output, + * emitting FP16 (the A operand of the W2 GEMM). + * + * Worker: B*hiddenDim global threads, localSize=256. + */ + public static void batchedFFNSwiGLUFP16Packed(KernelContext context, + HalfFloatArray wrapHbFP16Batch, + FloatArray gateUpResult, + int hiddenDim) { + int gid = context.globalIdx; + int b = gid / hiddenDim; + int i = gid % hiddenDim; + int rowBase = b * 2 * hiddenDim; + float g = gateUpResult.get(rowBase + i); + float u = gateUpResult.get(rowBase + hiddenDim + i); + float silu = g / (1.0f + TornadoMath.exp(-g)); + wrapHbFP16Batch.set(gid, new HalfFloat(silu * u)); + } + + + // ── Q8_0 tensor-core GEMMs (W8A16) ─────────────────────────────────────── + // + // Q8_0 weights stay quantized in global memory (34-byte GGUF blocks: + // FP16 scale + 32 int8 quants) and are dequantized to FP16 *in the + // register-staging step* of the software pipeline, then flow through the + // identical ldmatrix + m16n8k16 path as the FP16 GEMMs. This halves the + // weight-side memory traffic relative to FP16 while reusing the proven + // FP16 tensor-core pipeline. A true INT8 (m16n8k32) path would need + // per-k-block accumulator rescaling — blocked on fragment-level access + // in the TornadoVM intrinsics; see paper future work. + // + // Note: BK = 16, so a K-step never straddles a Q8_0 block boundary + // (32 | K), and each staged pair reads exactly one scale per column. + + /** + * Dequantizes and packs two vertically-adjacent (col, col+1) Q8_0 weight + * elements at depth k into one int of two FP16 values, matching the + * bTile layout expected by mmaLoadB. Leaf helper; inlined by the JIT. + */ + private static int packQ8Halves(ByteArray w, int col, int k, int blocksPerRow) { + int kBlock = k >>> 5; // k / 32 + int kIn = k & 31; // k % 32 + int off0 = (col * blocksPerRow + kBlock) * 34; + int off1 = off0 + blocksPerRow * 34; // column col+1, same k-block + float v0 = w.getHalfFloat(off0).getFloat32() * w.get(off0 + 2 + kIn); + float v1 = w.getHalfFloat(off1).getFloat32() * w.get(off1 + 2 + kIn); + int lo = new HalfFloat(v0).getHalfFloatValue() & 0xFFFF; + int hi = new HalfFloat(v1).getHalfFloatValue() & 0xFFFF; + return lo | (hi << 16); + } + + /** + * Tensor-core GEMM with Q8_0 weights: + * C[M,N] (FP32) = A[M,K] (FP16) × B[N,K] (Q8_0 blocks, row-major). + * Same tiling, pipeline, and constraints as {@link #gemmMMA}. + */ + public static void gemmMMAQ8(KernelContext ctx, + HalfFloatArray A, ByteArray B, FloatArray C, + int M, int N, int K) { + int tid = ctx.localIdx; + int warpId = tid / WARP_SIZE; + int warpM = warpId / WARPS_N; + int warpN = warpId % WARPS_N; + int blockRow = BM * ctx.groupIdx; + int blockCol = BN * ctx.groupIdy; + int blocksPerRow = K / 32; + + int[] aTile = ctx.allocateIntLocalArray(BM * BK / 2); + int[] bTile = ctx.allocateIntLocalArray(BK * BN / 2); + + float[] c00 = ctx.mmaFragment(0.0f); float[] c01 = ctx.mmaFragment(0.0f); + float[] c02 = ctx.mmaFragment(0.0f); float[] c03 = ctx.mmaFragment(0.0f); + float[] c04 = ctx.mmaFragment(0.0f); float[] c05 = ctx.mmaFragment(0.0f); + float[] c06 = ctx.mmaFragment(0.0f); float[] c07 = ctx.mmaFragment(0.0f); + float[] c10 = ctx.mmaFragment(0.0f); float[] c11 = ctx.mmaFragment(0.0f); + float[] c12 = ctx.mmaFragment(0.0f); float[] c13 = ctx.mmaFragment(0.0f); + float[] c14 = ctx.mmaFragment(0.0f); float[] c15 = ctx.mmaFragment(0.0f); + float[] c16 = ctx.mmaFragment(0.0f); float[] c17 = ctx.mmaFragment(0.0f); + + int aIdx0 = tid; int gA0 = (blockRow + (aIdx0 >>> 3)) * K + ((aIdx0 & 7) << 1); + int aIdx1 = tid + 256; int gA1 = (blockRow + (aIdx1 >>> 3)) * K + ((aIdx1 & 7) << 1); + int aIdx2 = tid + 512; int gA2 = (blockRow + (aIdx2 >>> 3)) * K + ((aIdx2 & 7) << 1); + int aIdx3 = tid + 768; int gA3 = (blockRow + (aIdx3 >>> 3)) * K + ((aIdx3 & 7) << 1); + // B staging keeps (col, k) coordinates explicit for block-offset math. + int bIdx0 = tid; int bCol0 = blockCol + ((bIdx0 >>> 6) << 3) + ((bIdx0 & 3) << 1); int bK0 = (bIdx0 & 63) >>> 2; + int bIdx1 = tid + 256; int bCol1 = blockCol + ((bIdx1 >>> 6) << 3) + ((bIdx1 & 3) << 1); int bK1 = (bIdx1 & 63) >>> 2; + int bIdx2 = tid + 512; int bCol2 = blockCol + ((bIdx2 >>> 6) << 3) + ((bIdx2 & 3) << 1); int bK2 = (bIdx2 & 63) >>> 2; + int bIdx3 = tid + 768; int bCol3 = blockCol + ((bIdx3 >>> 6) << 3) + ((bIdx3 & 3) << 1); int bK3 = (bIdx3 & 63) >>> 2; + + int aReg0 = packHalves(A, gA0, gA0 + 1); + int aReg1 = packHalves(A, gA1, gA1 + 1); + int aReg2 = packHalves(A, gA2, gA2 + 1); + int aReg3 = packHalves(A, gA3, gA3 + 1); + int bReg0 = packQ8Halves(B, bCol0, bK0, blocksPerRow); + int bReg1 = packQ8Halves(B, bCol1, bK1, blocksPerRow); + int bReg2 = packQ8Halves(B, bCol2, bK2, blocksPerRow); + int bReg3 = packQ8Halves(B, bCol3, bK3, blocksPerRow); + aTile[aIdx0] = aReg0; aTile[aIdx1] = aReg1; aTile[aIdx2] = aReg2; aTile[aIdx3] = aReg3; + bTile[bIdx0] = bReg0; bTile[bIdx1] = bReg1; bTile[bIdx2] = bReg2; bTile[bIdx3] = bReg3; + ctx.localBarrier(); + + int numKSteps = K / BK; + for (int kStep = 0; kStep < numKSteps; kStep++) { + if (kStep + 1 < numKSteps) { + int kOff = (kStep + 1) * BK; + aReg0 = packHalves(A, gA0 + kOff, gA0 + kOff + 1); + aReg1 = packHalves(A, gA1 + kOff, gA1 + kOff + 1); + aReg2 = packHalves(A, gA2 + kOff, gA2 + kOff + 1); + aReg3 = packHalves(A, gA3 + kOff, gA3 + kOff + 1); + bReg0 = packQ8Halves(B, bCol0, kOff + bK0, blocksPerRow); + bReg1 = packQ8Halves(B, bCol1, kOff + bK1, blocksPerRow); + bReg2 = packQ8Halves(B, bCol2, kOff + bK2, blocksPerRow); + bReg3 = packQ8Halves(B, bCol3, kOff + bK3, blocksPerRow); + } + + int aOff0 = warpM * 1024; + int aOff1 = warpM * 1024 + 512; + HalfFloat[] a0 = ctx.mmaLoadA(aTile, BK, aOff0); + HalfFloat[] a1 = ctx.mmaLoadA(aTile, BK, aOff1); + int bBase = warpN * 8; + HalfFloat[] b0 = ctx.mmaLoadB(bTile, BK, (bBase + 0) * B_SUBTILE_BYTES); + HalfFloat[] b1 = ctx.mmaLoadB(bTile, BK, (bBase + 1) * B_SUBTILE_BYTES); + HalfFloat[] b2 = ctx.mmaLoadB(bTile, BK, (bBase + 2) * B_SUBTILE_BYTES); + HalfFloat[] b3 = ctx.mmaLoadB(bTile, BK, (bBase + 3) * B_SUBTILE_BYTES); + HalfFloat[] b4 = ctx.mmaLoadB(bTile, BK, (bBase + 4) * B_SUBTILE_BYTES); + HalfFloat[] b5 = ctx.mmaLoadB(bTile, BK, (bBase + 5) * B_SUBTILE_BYTES); + HalfFloat[] b6 = ctx.mmaLoadB(bTile, BK, (bBase + 6) * B_SUBTILE_BYTES); + HalfFloat[] b7 = ctx.mmaLoadB(bTile, BK, (bBase + 7) * B_SUBTILE_BYTES); + ctx.localBarrier(); + + if (kStep + 1 < numKSteps) { + aTile[aIdx0] = aReg0; aTile[aIdx1] = aReg1; aTile[aIdx2] = aReg2; aTile[aIdx3] = aReg3; + bTile[bIdx0] = bReg0; bTile[bIdx1] = bReg1; bTile[bIdx2] = bReg2; bTile[bIdx3] = bReg3; + } + + c00 = ctx.mma(a0, b0, c00, MMAShape.M16N8K16); + c01 = ctx.mma(a0, b1, c01, MMAShape.M16N8K16); + c02 = ctx.mma(a0, b2, c02, MMAShape.M16N8K16); + c03 = ctx.mma(a0, b3, c03, MMAShape.M16N8K16); + c04 = ctx.mma(a0, b4, c04, MMAShape.M16N8K16); + c05 = ctx.mma(a0, b5, c05, MMAShape.M16N8K16); + c06 = ctx.mma(a0, b6, c06, MMAShape.M16N8K16); + c07 = ctx.mma(a0, b7, c07, MMAShape.M16N8K16); + c10 = ctx.mma(a1, b0, c10, MMAShape.M16N8K16); + c11 = ctx.mma(a1, b1, c11, MMAShape.M16N8K16); + c12 = ctx.mma(a1, b2, c12, MMAShape.M16N8K16); + c13 = ctx.mma(a1, b3, c13, MMAShape.M16N8K16); + c14 = ctx.mma(a1, b4, c14, MMAShape.M16N8K16); + c15 = ctx.mma(a1, b5, c15, MMAShape.M16N8K16); + c16 = ctx.mma(a1, b6, c16, MMAShape.M16N8K16); + c17 = ctx.mma(a1, b7, c17, MMAShape.M16N8K16); + ctx.localBarrier(); + } + + int rBase = blockRow + warpM * WM; + int cBase = blockCol + warpN * WN; + ctx.mmaStore(c00, C, rBase + 0, cBase + 0, N); + ctx.mmaStore(c01, C, rBase + 0, cBase + 8, N); + ctx.mmaStore(c02, C, rBase + 0, cBase + 16, N); + ctx.mmaStore(c03, C, rBase + 0, cBase + 24, N); + ctx.mmaStore(c04, C, rBase + 0, cBase + 32, N); + ctx.mmaStore(c05, C, rBase + 0, cBase + 40, N); + ctx.mmaStore(c06, C, rBase + 0, cBase + 48, N); + ctx.mmaStore(c07, C, rBase + 0, cBase + 56, N); + ctx.mmaStore(c10, C, rBase + 16, cBase + 0, N); + ctx.mmaStore(c11, C, rBase + 16, cBase + 8, N); + ctx.mmaStore(c12, C, rBase + 16, cBase + 16, N); + ctx.mmaStore(c13, C, rBase + 16, cBase + 24, N); + ctx.mmaStore(c14, C, rBase + 16, cBase + 32, N); + ctx.mmaStore(c15, C, rBase + 16, cBase + 40, N); + ctx.mmaStore(c16, C, rBase + 16, cBase + 48, N); + ctx.mmaStore(c17, C, rBase + 16, cBase + 56, N); + } + + /** + * Fused QKV tensor-core GEMM with Q8_0 weights into the PACKED output: + * qkvOut[M, dim+2*kvDim] = A[M,K] (FP16) × [Wq | Wk | Wv] (Q8_0, [N_i,K] row-major). + * Same layout, fusion, and constraints as {@link #gemmMMAQKV}. + */ + public static void gemmMMAQKVQ8(KernelContext ctx, + HalfFloatArray A, + ByteArray wq, ByteArray wk, ByteArray wv, + FloatArray qkvOut, + int M, int dim, int kvDim, int K) { + int tid = ctx.localIdx; + int warpId = tid / WARP_SIZE; + int warpM = warpId / WARPS_N; + int warpN = warpId % WARPS_N; + int blockRow = BM * ctx.groupIdx; + int blockCol = BN * ctx.groupIdy; + int qkvStride = dim + 2 * kvDim; + int blocksPerRow = K / 32; + + int wColBase = blockCol; + if (blockCol >= dim) wColBase -= dim; + if (blockCol >= dim + kvDim) wColBase -= kvDim; + + int[] aTile = ctx.allocateIntLocalArray(BM * BK / 2); + int[] bTile = ctx.allocateIntLocalArray(BK * BN / 2); + + float[] c00 = ctx.mmaFragment(0.0f); float[] c01 = ctx.mmaFragment(0.0f); + float[] c02 = ctx.mmaFragment(0.0f); float[] c03 = ctx.mmaFragment(0.0f); + float[] c04 = ctx.mmaFragment(0.0f); float[] c05 = ctx.mmaFragment(0.0f); + float[] c06 = ctx.mmaFragment(0.0f); float[] c07 = ctx.mmaFragment(0.0f); + float[] c10 = ctx.mmaFragment(0.0f); float[] c11 = ctx.mmaFragment(0.0f); + float[] c12 = ctx.mmaFragment(0.0f); float[] c13 = ctx.mmaFragment(0.0f); + float[] c14 = ctx.mmaFragment(0.0f); float[] c15 = ctx.mmaFragment(0.0f); + float[] c16 = ctx.mmaFragment(0.0f); float[] c17 = ctx.mmaFragment(0.0f); + + int aIdx0 = tid; int gA0 = (blockRow + (aIdx0 >>> 3)) * K + ((aIdx0 & 7) << 1); + int aIdx1 = tid + 256; int gA1 = (blockRow + (aIdx1 >>> 3)) * K + ((aIdx1 & 7) << 1); + int aIdx2 = tid + 512; int gA2 = (blockRow + (aIdx2 >>> 3)) * K + ((aIdx2 & 7) << 1); + int aIdx3 = tid + 768; int gA3 = (blockRow + (aIdx3 >>> 3)) * K + ((aIdx3 & 7) << 1); + int bIdx0 = tid; int bCol0 = wColBase + ((bIdx0 >>> 6) << 3) + ((bIdx0 & 3) << 1); int bK0 = (bIdx0 & 63) >>> 2; + int bIdx1 = tid + 256; int bCol1 = wColBase + ((bIdx1 >>> 6) << 3) + ((bIdx1 & 3) << 1); int bK1 = (bIdx1 & 63) >>> 2; + int bIdx2 = tid + 512; int bCol2 = wColBase + ((bIdx2 >>> 6) << 3) + ((bIdx2 & 3) << 1); int bK2 = (bIdx2 & 63) >>> 2; + int bIdx3 = tid + 768; int bCol3 = wColBase + ((bIdx3 >>> 6) << 3) + ((bIdx3 & 3) << 1); int bK3 = (bIdx3 & 63) >>> 2; + + int aReg0 = packHalves(A, gA0, gA0 + 1); + int aReg1 = packHalves(A, gA1, gA1 + 1); + int aReg2 = packHalves(A, gA2, gA2 + 1); + int aReg3 = packHalves(A, gA3, gA3 + 1); + int bReg0; int bReg1; int bReg2; int bReg3; + if (blockCol < dim) { + bReg0 = packQ8Halves(wq, bCol0, bK0, blocksPerRow); + bReg1 = packQ8Halves(wq, bCol1, bK1, blocksPerRow); + bReg2 = packQ8Halves(wq, bCol2, bK2, blocksPerRow); + bReg3 = packQ8Halves(wq, bCol3, bK3, blocksPerRow); + } else if (blockCol < dim + kvDim) { + bReg0 = packQ8Halves(wk, bCol0, bK0, blocksPerRow); + bReg1 = packQ8Halves(wk, bCol1, bK1, blocksPerRow); + bReg2 = packQ8Halves(wk, bCol2, bK2, blocksPerRow); + bReg3 = packQ8Halves(wk, bCol3, bK3, blocksPerRow); + } else { + bReg0 = packQ8Halves(wv, bCol0, bK0, blocksPerRow); + bReg1 = packQ8Halves(wv, bCol1, bK1, blocksPerRow); + bReg2 = packQ8Halves(wv, bCol2, bK2, blocksPerRow); + bReg3 = packQ8Halves(wv, bCol3, bK3, blocksPerRow); + } + aTile[aIdx0] = aReg0; aTile[aIdx1] = aReg1; aTile[aIdx2] = aReg2; aTile[aIdx3] = aReg3; + bTile[bIdx0] = bReg0; bTile[bIdx1] = bReg1; bTile[bIdx2] = bReg2; bTile[bIdx3] = bReg3; + ctx.localBarrier(); + + int numKSteps = K / BK; + for (int kStep = 0; kStep < numKSteps; kStep++) { + if (kStep + 1 < numKSteps) { + int kOff = (kStep + 1) * BK; + aReg0 = packHalves(A, gA0 + kOff, gA0 + kOff + 1); + aReg1 = packHalves(A, gA1 + kOff, gA1 + kOff + 1); + aReg2 = packHalves(A, gA2 + kOff, gA2 + kOff + 1); + aReg3 = packHalves(A, gA3 + kOff, gA3 + kOff + 1); + if (blockCol < dim) { + bReg0 = packQ8Halves(wq, bCol0, kOff + bK0, blocksPerRow); + bReg1 = packQ8Halves(wq, bCol1, kOff + bK1, blocksPerRow); + bReg2 = packQ8Halves(wq, bCol2, kOff + bK2, blocksPerRow); + bReg3 = packQ8Halves(wq, bCol3, kOff + bK3, blocksPerRow); + } else if (blockCol < dim + kvDim) { + bReg0 = packQ8Halves(wk, bCol0, kOff + bK0, blocksPerRow); + bReg1 = packQ8Halves(wk, bCol1, kOff + bK1, blocksPerRow); + bReg2 = packQ8Halves(wk, bCol2, kOff + bK2, blocksPerRow); + bReg3 = packQ8Halves(wk, bCol3, kOff + bK3, blocksPerRow); + } else { + bReg0 = packQ8Halves(wv, bCol0, kOff + bK0, blocksPerRow); + bReg1 = packQ8Halves(wv, bCol1, kOff + bK1, blocksPerRow); + bReg2 = packQ8Halves(wv, bCol2, kOff + bK2, blocksPerRow); + bReg3 = packQ8Halves(wv, bCol3, kOff + bK3, blocksPerRow); + } + } + + int aOff0 = warpM * 1024; + int aOff1 = warpM * 1024 + 512; + HalfFloat[] a0 = ctx.mmaLoadA(aTile, BK, aOff0); + HalfFloat[] a1 = ctx.mmaLoadA(aTile, BK, aOff1); + int bBase = warpN * 8; + HalfFloat[] b0 = ctx.mmaLoadB(bTile, BK, (bBase + 0) * B_SUBTILE_BYTES); + HalfFloat[] b1 = ctx.mmaLoadB(bTile, BK, (bBase + 1) * B_SUBTILE_BYTES); + HalfFloat[] b2 = ctx.mmaLoadB(bTile, BK, (bBase + 2) * B_SUBTILE_BYTES); + HalfFloat[] b3 = ctx.mmaLoadB(bTile, BK, (bBase + 3) * B_SUBTILE_BYTES); + HalfFloat[] b4 = ctx.mmaLoadB(bTile, BK, (bBase + 4) * B_SUBTILE_BYTES); + HalfFloat[] b5 = ctx.mmaLoadB(bTile, BK, (bBase + 5) * B_SUBTILE_BYTES); + HalfFloat[] b6 = ctx.mmaLoadB(bTile, BK, (bBase + 6) * B_SUBTILE_BYTES); + HalfFloat[] b7 = ctx.mmaLoadB(bTile, BK, (bBase + 7) * B_SUBTILE_BYTES); + ctx.localBarrier(); + + if (kStep + 1 < numKSteps) { + aTile[aIdx0] = aReg0; aTile[aIdx1] = aReg1; aTile[aIdx2] = aReg2; aTile[aIdx3] = aReg3; + bTile[bIdx0] = bReg0; bTile[bIdx1] = bReg1; bTile[bIdx2] = bReg2; bTile[bIdx3] = bReg3; + } + + c00 = ctx.mma(a0, b0, c00, MMAShape.M16N8K16); + c01 = ctx.mma(a0, b1, c01, MMAShape.M16N8K16); + c02 = ctx.mma(a0, b2, c02, MMAShape.M16N8K16); + c03 = ctx.mma(a0, b3, c03, MMAShape.M16N8K16); + c04 = ctx.mma(a0, b4, c04, MMAShape.M16N8K16); + c05 = ctx.mma(a0, b5, c05, MMAShape.M16N8K16); + c06 = ctx.mma(a0, b6, c06, MMAShape.M16N8K16); + c07 = ctx.mma(a0, b7, c07, MMAShape.M16N8K16); + c10 = ctx.mma(a1, b0, c10, MMAShape.M16N8K16); + c11 = ctx.mma(a1, b1, c11, MMAShape.M16N8K16); + c12 = ctx.mma(a1, b2, c12, MMAShape.M16N8K16); + c13 = ctx.mma(a1, b3, c13, MMAShape.M16N8K16); + c14 = ctx.mma(a1, b4, c14, MMAShape.M16N8K16); + c15 = ctx.mma(a1, b5, c15, MMAShape.M16N8K16); + c16 = ctx.mma(a1, b6, c16, MMAShape.M16N8K16); + c17 = ctx.mma(a1, b7, c17, MMAShape.M16N8K16); + ctx.localBarrier(); + } + + int rBase = blockRow + warpM * WM; + int cBase = blockCol + warpN * WN; + ctx.mmaStore(c00, qkvOut, rBase + 0, cBase + 0, qkvStride); + ctx.mmaStore(c01, qkvOut, rBase + 0, cBase + 8, qkvStride); + ctx.mmaStore(c02, qkvOut, rBase + 0, cBase + 16, qkvStride); + ctx.mmaStore(c03, qkvOut, rBase + 0, cBase + 24, qkvStride); + ctx.mmaStore(c04, qkvOut, rBase + 0, cBase + 32, qkvStride); + ctx.mmaStore(c05, qkvOut, rBase + 0, cBase + 40, qkvStride); + ctx.mmaStore(c06, qkvOut, rBase + 0, cBase + 48, qkvStride); + ctx.mmaStore(c07, qkvOut, rBase + 0, cBase + 56, qkvStride); + ctx.mmaStore(c10, qkvOut, rBase + 16, cBase + 0, qkvStride); + ctx.mmaStore(c11, qkvOut, rBase + 16, cBase + 8, qkvStride); + ctx.mmaStore(c12, qkvOut, rBase + 16, cBase + 16, qkvStride); + ctx.mmaStore(c13, qkvOut, rBase + 16, cBase + 24, qkvStride); + ctx.mmaStore(c14, qkvOut, rBase + 16, cBase + 32, qkvStride); + ctx.mmaStore(c15, qkvOut, rBase + 16, cBase + 40, qkvStride); + ctx.mmaStore(c16, qkvOut, rBase + 16, cBase + 48, qkvStride); + ctx.mmaStore(c17, qkvOut, rBase + 16, cBase + 56, qkvStride); + } + + /** + * Fused W1/W3 (gate/up) tensor-core GEMM with Q8_0 weights into the PACKED + * output gateUpOut[M, 2*hidDim]. Same layout and constraints as + * {@link #gemmMMAGateUp}. + */ + public static void gemmMMAGateUpQ8(KernelContext ctx, + HalfFloatArray A, + ByteArray w1, ByteArray w3, + FloatArray gateUpOut, + int M, int hidDim, int K) { + int tid = ctx.localIdx; + int warpId = tid / WARP_SIZE; + int warpM = warpId / WARPS_N; + int warpN = warpId % WARPS_N; + int blockRow = BM * ctx.groupIdx; + int blockCol = BN * ctx.groupIdy; + int outStride = 2 * hidDim; + int blocksPerRow = K / 32; + + int wColBase = (blockCol < hidDim) ? blockCol : (blockCol - hidDim); + + int[] aTile = ctx.allocateIntLocalArray(BM * BK / 2); + int[] bTile = ctx.allocateIntLocalArray(BK * BN / 2); + + float[] c00 = ctx.mmaFragment(0.0f); float[] c01 = ctx.mmaFragment(0.0f); + float[] c02 = ctx.mmaFragment(0.0f); float[] c03 = ctx.mmaFragment(0.0f); + float[] c04 = ctx.mmaFragment(0.0f); float[] c05 = ctx.mmaFragment(0.0f); + float[] c06 = ctx.mmaFragment(0.0f); float[] c07 = ctx.mmaFragment(0.0f); + float[] c10 = ctx.mmaFragment(0.0f); float[] c11 = ctx.mmaFragment(0.0f); + float[] c12 = ctx.mmaFragment(0.0f); float[] c13 = ctx.mmaFragment(0.0f); + float[] c14 = ctx.mmaFragment(0.0f); float[] c15 = ctx.mmaFragment(0.0f); + float[] c16 = ctx.mmaFragment(0.0f); float[] c17 = ctx.mmaFragment(0.0f); + + int aIdx0 = tid; int gA0 = (blockRow + (aIdx0 >>> 3)) * K + ((aIdx0 & 7) << 1); + int aIdx1 = tid + 256; int gA1 = (blockRow + (aIdx1 >>> 3)) * K + ((aIdx1 & 7) << 1); + int aIdx2 = tid + 512; int gA2 = (blockRow + (aIdx2 >>> 3)) * K + ((aIdx2 & 7) << 1); + int aIdx3 = tid + 768; int gA3 = (blockRow + (aIdx3 >>> 3)) * K + ((aIdx3 & 7) << 1); + int bIdx0 = tid; int bCol0 = wColBase + ((bIdx0 >>> 6) << 3) + ((bIdx0 & 3) << 1); int bK0 = (bIdx0 & 63) >>> 2; + int bIdx1 = tid + 256; int bCol1 = wColBase + ((bIdx1 >>> 6) << 3) + ((bIdx1 & 3) << 1); int bK1 = (bIdx1 & 63) >>> 2; + int bIdx2 = tid + 512; int bCol2 = wColBase + ((bIdx2 >>> 6) << 3) + ((bIdx2 & 3) << 1); int bK2 = (bIdx2 & 63) >>> 2; + int bIdx3 = tid + 768; int bCol3 = wColBase + ((bIdx3 >>> 6) << 3) + ((bIdx3 & 3) << 1); int bK3 = (bIdx3 & 63) >>> 2; + + int aReg0 = packHalves(A, gA0, gA0 + 1); + int aReg1 = packHalves(A, gA1, gA1 + 1); + int aReg2 = packHalves(A, gA2, gA2 + 1); + int aReg3 = packHalves(A, gA3, gA3 + 1); + int bReg0; int bReg1; int bReg2; int bReg3; + if (blockCol < hidDim) { + bReg0 = packQ8Halves(w1, bCol0, bK0, blocksPerRow); + bReg1 = packQ8Halves(w1, bCol1, bK1, blocksPerRow); + bReg2 = packQ8Halves(w1, bCol2, bK2, blocksPerRow); + bReg3 = packQ8Halves(w1, bCol3, bK3, blocksPerRow); + } else { + bReg0 = packQ8Halves(w3, bCol0, bK0, blocksPerRow); + bReg1 = packQ8Halves(w3, bCol1, bK1, blocksPerRow); + bReg2 = packQ8Halves(w3, bCol2, bK2, blocksPerRow); + bReg3 = packQ8Halves(w3, bCol3, bK3, blocksPerRow); + } + aTile[aIdx0] = aReg0; aTile[aIdx1] = aReg1; aTile[aIdx2] = aReg2; aTile[aIdx3] = aReg3; + bTile[bIdx0] = bReg0; bTile[bIdx1] = bReg1; bTile[bIdx2] = bReg2; bTile[bIdx3] = bReg3; + ctx.localBarrier(); + + int numKSteps = K / BK; + for (int kStep = 0; kStep < numKSteps; kStep++) { + if (kStep + 1 < numKSteps) { + int kOff = (kStep + 1) * BK; + aReg0 = packHalves(A, gA0 + kOff, gA0 + kOff + 1); + aReg1 = packHalves(A, gA1 + kOff, gA1 + kOff + 1); + aReg2 = packHalves(A, gA2 + kOff, gA2 + kOff + 1); + aReg3 = packHalves(A, gA3 + kOff, gA3 + kOff + 1); + if (blockCol < hidDim) { + bReg0 = packQ8Halves(w1, bCol0, kOff + bK0, blocksPerRow); + bReg1 = packQ8Halves(w1, bCol1, kOff + bK1, blocksPerRow); + bReg2 = packQ8Halves(w1, bCol2, kOff + bK2, blocksPerRow); + bReg3 = packQ8Halves(w1, bCol3, kOff + bK3, blocksPerRow); + } else { + bReg0 = packQ8Halves(w3, bCol0, kOff + bK0, blocksPerRow); + bReg1 = packQ8Halves(w3, bCol1, kOff + bK1, blocksPerRow); + bReg2 = packQ8Halves(w3, bCol2, kOff + bK2, blocksPerRow); + bReg3 = packQ8Halves(w3, bCol3, kOff + bK3, blocksPerRow); + } + } + + int aOff0 = warpM * 1024; + int aOff1 = warpM * 1024 + 512; + HalfFloat[] a0 = ctx.mmaLoadA(aTile, BK, aOff0); + HalfFloat[] a1 = ctx.mmaLoadA(aTile, BK, aOff1); + int bBase = warpN * 8; + HalfFloat[] b0 = ctx.mmaLoadB(bTile, BK, (bBase + 0) * B_SUBTILE_BYTES); + HalfFloat[] b1 = ctx.mmaLoadB(bTile, BK, (bBase + 1) * B_SUBTILE_BYTES); + HalfFloat[] b2 = ctx.mmaLoadB(bTile, BK, (bBase + 2) * B_SUBTILE_BYTES); + HalfFloat[] b3 = ctx.mmaLoadB(bTile, BK, (bBase + 3) * B_SUBTILE_BYTES); + HalfFloat[] b4 = ctx.mmaLoadB(bTile, BK, (bBase + 4) * B_SUBTILE_BYTES); + HalfFloat[] b5 = ctx.mmaLoadB(bTile, BK, (bBase + 5) * B_SUBTILE_BYTES); + HalfFloat[] b6 = ctx.mmaLoadB(bTile, BK, (bBase + 6) * B_SUBTILE_BYTES); + HalfFloat[] b7 = ctx.mmaLoadB(bTile, BK, (bBase + 7) * B_SUBTILE_BYTES); + ctx.localBarrier(); + + if (kStep + 1 < numKSteps) { + aTile[aIdx0] = aReg0; aTile[aIdx1] = aReg1; aTile[aIdx2] = aReg2; aTile[aIdx3] = aReg3; + bTile[bIdx0] = bReg0; bTile[bIdx1] = bReg1; bTile[bIdx2] = bReg2; bTile[bIdx3] = bReg3; + } + + c00 = ctx.mma(a0, b0, c00, MMAShape.M16N8K16); + c01 = ctx.mma(a0, b1, c01, MMAShape.M16N8K16); + c02 = ctx.mma(a0, b2, c02, MMAShape.M16N8K16); + c03 = ctx.mma(a0, b3, c03, MMAShape.M16N8K16); + c04 = ctx.mma(a0, b4, c04, MMAShape.M16N8K16); + c05 = ctx.mma(a0, b5, c05, MMAShape.M16N8K16); + c06 = ctx.mma(a0, b6, c06, MMAShape.M16N8K16); + c07 = ctx.mma(a0, b7, c07, MMAShape.M16N8K16); + c10 = ctx.mma(a1, b0, c10, MMAShape.M16N8K16); + c11 = ctx.mma(a1, b1, c11, MMAShape.M16N8K16); + c12 = ctx.mma(a1, b2, c12, MMAShape.M16N8K16); + c13 = ctx.mma(a1, b3, c13, MMAShape.M16N8K16); + c14 = ctx.mma(a1, b4, c14, MMAShape.M16N8K16); + c15 = ctx.mma(a1, b5, c15, MMAShape.M16N8K16); + c16 = ctx.mma(a1, b6, c16, MMAShape.M16N8K16); + c17 = ctx.mma(a1, b7, c17, MMAShape.M16N8K16); + ctx.localBarrier(); + } + + int rBase = blockRow + warpM * WM; + int cBase = blockCol + warpN * WN; + ctx.mmaStore(c00, gateUpOut, rBase + 0, cBase + 0, outStride); + ctx.mmaStore(c01, gateUpOut, rBase + 0, cBase + 8, outStride); + ctx.mmaStore(c02, gateUpOut, rBase + 0, cBase + 16, outStride); + ctx.mmaStore(c03, gateUpOut, rBase + 0, cBase + 24, outStride); + ctx.mmaStore(c04, gateUpOut, rBase + 0, cBase + 32, outStride); + ctx.mmaStore(c05, gateUpOut, rBase + 0, cBase + 40, outStride); + ctx.mmaStore(c06, gateUpOut, rBase + 0, cBase + 48, outStride); + ctx.mmaStore(c07, gateUpOut, rBase + 0, cBase + 56, outStride); + ctx.mmaStore(c10, gateUpOut, rBase + 16, cBase + 0, outStride); + ctx.mmaStore(c11, gateUpOut, rBase + 16, cBase + 8, outStride); + ctx.mmaStore(c12, gateUpOut, rBase + 16, cBase + 16, outStride); + ctx.mmaStore(c13, gateUpOut, rBase + 16, cBase + 24, outStride); + ctx.mmaStore(c14, gateUpOut, rBase + 16, cBase + 32, outStride); + ctx.mmaStore(c15, gateUpOut, rBase + 16, cBase + 40, outStride); + ctx.mmaStore(c16, gateUpOut, rBase + 16, cBase + 48, outStride); + ctx.mmaStore(c17, gateUpOut, rBase + 16, cBase + 56, outStride); + } + // @formatter:on } diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/layers/BatchPrefillTransformerLayerTaskGraphs.java b/src/main/java/org/beehive/gpullama3/tornadovm/layers/BatchPrefillTransformerLayerTaskGraphs.java index cf406091..9818eaa5 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/layers/BatchPrefillTransformerLayerTaskGraphs.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/layers/BatchPrefillTransformerLayerTaskGraphs.java @@ -8,7 +8,7 @@ /** * Interface for a group of N batched-prefill transformer-layer TornadoVM TaskGraphs. * - *

Implemented by {@code LlamaFP16LayersBatchPrefill} and {@code LlamaQ8_0LayersBatchPrefill}.

+ *

Implemented by {@code LlamaFP16LayersBatchPrefillMMA}, {@code LlamaFP16LayersBatchPrefill}, {@code LlamaQ8_0LayersBatchPrefillMMA} and {@code LlamaQ8_0LayersBatchPrefill}.

*/ public interface BatchPrefillTransformerLayerTaskGraphs { List getLayerImmutableTaskGraphs(); diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/prefill/LlamaFP16LayersBatchPrefillMMA.java b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/prefill/LlamaFP16LayersBatchPrefillMMA.java new file mode 100644 index 00000000..7d04117c --- /dev/null +++ b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/prefill/LlamaFP16LayersBatchPrefillMMA.java @@ -0,0 +1,293 @@ +package org.beehive.gpullama3.tornadovm.layers.type.fp16.prefill; + +import org.beehive.gpullama3.inference.state.LlamaState; +import org.beehive.gpullama3.inference.weights.tornado.LlamaTornadoWeights; +import org.beehive.gpullama3.model.llama.LlamaConfiguration; +import org.beehive.gpullama3.tornadovm.kernels.TransformerBatchPrefillKernels; +import org.beehive.gpullama3.tornadovm.scheduling.WorkerGridFactory; +import org.beehive.gpullama3.tornadovm.layers.BatchPrefillTransformerLayerTaskGraphs; +import uk.ac.manchester.tornado.api.GridScheduler; +import uk.ac.manchester.tornado.api.ImmutableTaskGraph; +import uk.ac.manchester.tornado.api.KernelContext; +import uk.ac.manchester.tornado.api.TaskGraph; +import uk.ac.manchester.tornado.api.WorkerGrid; +import uk.ac.manchester.tornado.api.WorkerGrid1D; +import uk.ac.manchester.tornado.api.WorkerGrid2D; +import uk.ac.manchester.tornado.api.enums.DataTransferMode; + +import java.util.List; +import java.util.stream.IntStream; + +/** + * Batched-prefill transformer-layer TaskGraphs for the unified batched prefill-decode plan + * ({@link org.beehive.gpullama3.tornadovm.TornadoVMMasterPlanBatchPrefillDecode}). + * + *

One {@link ImmutableTaskGraph} per transformer layer, each processing + * {@code batchSize} tokens simultaneously via {@link TransformerBatchPrefillKernels}.

+ * + *

Tensor-core layer pipeline (12 tasks, all GEMMs on MMA):

+ * 
+ *   batch_attn_rms        parallel RMS square-sum reduction (256 thr/token)
+ *   batch_attn_rms_apply  RMS apply + FP16 quantize → wrapXbFP16Batch
+ *   qkvProj               ONE fused MMA GEMM → packed qkvResultBatch [q|k|v]
+ *   batch_rope_kv         RoPE over the packed buffer + KV cache write
+ *   batch_attention       flash attention (register-partitioned P·V) → attnOutFP16
+ *   woProj                MMA GEMM → woOut
+ *   batch_ffn_rms         parallel RMS reduce FUSED with x += woOut
+ *   batch_ffn_rms_apply   RMS apply + FP16 quantize → normedXFFNFP16
+ *   gateUpProj            ONE fused MMA GEMM → packed gateUpResultBatch [gate|up]
+ *   swiglu                SiLU(gate)*up over packed buffer → wrapHbFP16Batch
+ *   w2Proj                MMA GEMM → w2Out
+ *   w2Resid               x += w2Out
+ *
+ * + *

KV cache ({@code wrapKeyCache}, {@code wrapValueCache}) is persisted on device + * after every layer so the subsequent single-token decode layers can consume it.

+ */ +public class LlamaFP16LayersBatchPrefillMMA implements BatchPrefillTransformerLayerTaskGraphs { + + // Local size for the parallel RMS reductions (one workgroup per token). + static final int RMS_LOCAL_SIZE = 256; + + private final LlamaState state; + private final LlamaTornadoWeights weights; + private final LlamaConfiguration config; + private final KernelContext context = new KernelContext(); + private final int batchSize; + // GEMM M dimension rounded up to whole 128-row tiles (BM). The GEMM-adjacent + // buffers in State are allocated at this padded size; rows >= batchSize are + // computed but never consumed. All non-GEMM kernels use the true batchSize. + private final int paddedBatch; + private final List layerITGs; + private String lastLayerTaskGraphID; + + public LlamaFP16LayersBatchPrefillMMA(LlamaState state, LlamaTornadoWeights weights, + LlamaConfiguration config, int batchSize) { + this.state = state; + this.weights = weights; + this.config = config; + this.batchSize = batchSize; + this.paddedBatch = (batchSize + 127) & ~127; + if (batchSize % 128 != 0) { + System.out.printf("[GPULlama3] prefill batch %d padded to %d for tensor-core tiles; " + + "GEMM efficiency is %d/%d — use a multiple of 128 for best throughput.%n", + batchSize, paddedBatch, batchSize, paddedBatch); + } + this.layerITGs = IntStream.range(0, config.numberOfLayers()) + .mapToObj(this::createBatchPrefillLayerTaskGraph) + .map(TaskGraph::snapshot) + .toList(); + } + + // @formatter:off + private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) { + String graphName = "batchPrefillLayer_" + layerIndex; + if (layerIndex == config.numberOfLayers() - 1) lastLayerTaskGraphID = graphName; + + TaskGraph batchPrefillLayer = new TaskGraph(graphName); + + // ── Data Transfers ───────────────────────────────────────────────────── + if (layerIndex == 0) { + // batchStartPosHolder is set by host before each chunk → EVERY_EXECUTION + batchPrefillLayer.transferToDevice(DataTransferMode.EVERY_EXECUTION, state.batchStartPosHolder); + // Allocate persistent GPU-side intermediates once + batchPrefillLayer.transferToDevice(DataTransferMode.FIRST_EXECUTION, + context, + state.attnScaleBatch, state.ffnScaleBatch, + state.wrapXbFP16Batch, + state.qkvResultBatch, + state.wrapKeyCache, state.wrapValueCache, + state.normedXFFNFP16, state.gateUpResultBatch, + state.attnOutFP16, state.woOut, state.wrapHbFP16Batch, state.w2Out); + // wrapXBatch produced by the prefillActivation graph and persists in device memory + // to consume it from there we should use the explicit uniqueTaskGraph name + // the no-arg form would use current graph name, which causes NPE without CUDA Graphs + batchPrefillLayer.consumeFromDevice("prefillActivation", state.wrapXBatch); + } else { + // for the same reasons as above, we should use the explicit uniqueTaskGraph name to consume + String pred = "batchPrefillLayer_" + (layerIndex - 1); + batchPrefillLayer.consumeFromDevice(pred, + context, + state.wrapXBatch, + state.batchStartPosHolder, + state.attnScaleBatch, state.ffnScaleBatch, + state.wrapXbFP16Batch, + state.qkvResultBatch, + state.wrapKeyCache, state.wrapValueCache, + state.normedXFFNFP16, state.gateUpResultBatch, + state.attnOutFP16, state.woOut, state.wrapHbFP16Batch, state.w2Out); + } + + // Per-layer weights: upload once + batchPrefillLayer.transferToDevice(DataTransferMode.FIRST_EXECUTION, + weights.rms_att_weightLayered[layerIndex].asFloatArray(), + weights.wqLayered[layerIndex].asHalfFloatArray(), + weights.wkLayered[layerIndex].asHalfFloatArray(), + weights.wvLayered[layerIndex].asHalfFloatArray(), + weights.woLayered[layerIndex].asHalfFloatArray(), + weights.rms_ffn_weightLayered[layerIndex].asFloatArray(), + weights.w1Layered[layerIndex].asHalfFloatArray(), + weights.w2Layered[layerIndex].asHalfFloatArray(), + weights.w3Layered[layerIndex].asHalfFloatArray()); + + int dim = config.dim(); + int kvDim = config.kvDim(); + int hidDim = config.hiddenDim(); + + // ── Attention Block ──────────────────────────────────────────────────── + batchPrefillLayer.task("batch_attn_rms", + TransformerBatchPrefillKernels::batchedRmsReduceParallel, + context, state.wrapXBatch, state.attnScaleBatch, + dim, config.rmsNormEps(), RMS_LOCAL_SIZE); + + batchPrefillLayer.task("batch_attn_rms_apply", + TransformerBatchPrefillKernels::batchedRmsApplyFP16, + context, state.wrapXbFP16Batch, state.wrapXBatch, + weights.rms_att_weightLayered[layerIndex].asFloatArray(), + state.attnScaleBatch, dim); + + // Q, K, V in ONE tensor-core launch → packed [q|k|v] rows. + // Grid spans (dim + 2*kvDim)/128 N-blocks: no grid starvation on the + // skinny GQA projections, and the A operand is read once, not thrice. + batchPrefillLayer.task("qkvProj", + TransformerBatchPrefillKernels::gemmMMAQKV, + context, state.wrapXbFP16Batch, + weights.wqLayered[layerIndex].asHalfFloatArray(), + weights.wkLayered[layerIndex].asHalfFloatArray(), + weights.wvLayered[layerIndex].asHalfFloatArray(), + state.qkvResultBatch, paddedBatch, dim, kvDim, dim); + + batchPrefillLayer.task("batch_rope_kv", + TransformerBatchPrefillKernels::batchedRopeWithKVCachePacked, + context, state.batchStartPosHolder, + state.qkvResultBatch, + state.wrapKeyCache, state.wrapValueCache, + kvDim, config.headSize(), layerIndex, config.contextLength(), dim); + + // Register-partitioned P·V accumulation + direct FP16 emission + // (replaces batchedFlashAttention + attnCast). + batchPrefillLayer.task("batch_attention", + TransformerBatchPrefillKernels::batchedFlashAttentionFP16Out, + context, state.batchStartPosHolder, + state.qkvResultBatch, state.wrapKeyCache, state.wrapValueCache, + state.attnOutFP16, + config.numberOfHeads(), config.headSize(), + kvDim, config.kvMul(), layerIndex, config.contextLength(), dim); + + batchPrefillLayer.task("woProj", TransformerBatchPrefillKernels::gemmMMA, + context, state.attnOutFP16, + weights.woLayered[layerIndex].asHalfFloatArray(), + state.woOut, paddedBatch, dim, dim); + + // ── FFN Block ────────────────────────────────────────────────────────── + // x += woOut is fused into the FFN RMS reduction (drops the woResid task). + batchPrefillLayer.task("batch_ffn_rms", + TransformerBatchPrefillKernels::batchedRmsReduceFusedResidual, + context, state.wrapXBatch, state.woOut, state.ffnScaleBatch, + dim, config.rmsNormEps(), RMS_LOCAL_SIZE); + + batchPrefillLayer.task("batch_ffn_rms_apply", + TransformerBatchPrefillKernels::batchedFFNRmsApplyFP16, + context, state.normedXFFNFP16, state.wrapXBatch, + weights.rms_ffn_weightLayered[layerIndex].asFloatArray(), + state.ffnScaleBatch, dim); + + // W1 and W3 in ONE tensor-core launch → packed [gate|up] rows. + batchPrefillLayer.task("gateUpProj", + TransformerBatchPrefillKernels::gemmMMAGateUp, + context, state.normedXFFNFP16, + weights.w1Layered[layerIndex].asHalfFloatArray(), + weights.w3Layered[layerIndex].asHalfFloatArray(), + state.gateUpResultBatch, paddedBatch, hidDim, dim); + + batchPrefillLayer.task("swiglu", + TransformerBatchPrefillKernels::batchedFFNSwiGLUFP16Packed, + context, state.wrapHbFP16Batch, state.gateUpResultBatch, hidDim) + .task("w2Proj", TransformerBatchPrefillKernels::gemmMMA, + context, state.wrapHbFP16Batch, + weights.w2Layered[layerIndex].asHalfFloatArray(), + state.w2Out, batchSize, dim, hidDim) + .task("w2Resid", TransformerBatchPrefillKernels::batchedResidualAddFP32, + context, state.wrapXBatch, state.w2Out); + + // Persist wrapXBatch for the next layer, and KV cache so the decode + // layers can consume it via the activation graph pass-through. + batchPrefillLayer.persistOnDevice(state.wrapXBatch, state.wrapKeyCache, state.wrapValueCache); + + return batchPrefillLayer; + } + // @formatter:on + + // gemmMMA family: 256 threads/block (1D within block), grid over M- and N-blocks. + static WorkerGrid mmaGrid(int paddedM, int N) { + int mBlocks = paddedM / 128; // BM + int nBlocks = N / 128; // BN + WorkerGrid2D g = new WorkerGrid2D(mBlocks * 256, nBlocks); + g.setLocalWork(256, 1, 1); // groupIdx∈[0,mBlocks), groupIdy∈[0,nBlocks) + return g; + } + + static WorkerGrid elementwiseGrid(int n) { // n must be a multiple of 256 + WorkerGrid1D g = new WorkerGrid1D(n); + g.setLocalWork(256, 1, 1); + return g; + } + + /** Registers all batch layer workers in the shared {@link GridScheduler}. */ + public void updateGridScheduler(GridScheduler scheduler) { + int dim = config.dim(); + int kvDim = config.kvDim(); + int hidDim = config.hiddenDim(); + int nHeads = config.numberOfHeads(); + int headSz = config.headSize(); + + // Parallel RMS reductions: one 256-thread workgroup per batch token + WorkerGrid rmsWorker = WorkerGridFactory.genericWorker( + batchSize * RMS_LOCAL_SIZE, RMS_LOCAL_SIZE); + + // RMS apply: B*dim threads, local=256 (dim is always a multiple of 256 for LLaMA) + WorkerGrid rmsApplyWorker = WorkerGridFactory.genericWorker(batchSize * dim, 256); + WorkerGrid ffnRmsApplyWorker = WorkerGridFactory.genericWorker(batchSize * dim, 256); + + // RoPE+KV cache: B*(dim/2) threads, local=512 + int ropeGlobal = batchSize * (dim / 2); + int ropeLocal = Math.min(512, ropeGlobal); + while (ropeLocal > 1 && ropeGlobal % ropeLocal != 0) ropeLocal--; + WorkerGrid ropeWorker = WorkerGridFactory.genericWorker(ropeGlobal, ropeLocal); + + // Attention (flash): B*nHeads workgroups × min(headSize,128) threads. + // The kernel requires headSize <= 2*localSize. + int attnLocal = Math.min(headSz, 128); + WorkerGrid attnWorker = WorkerGridFactory.genericWorker( + batchSize * nHeads * attnLocal, attnLocal); + + // MMA grids + WorkerGrid mmaQkvWorker = mmaGrid(paddedBatch, dim + 2 * kvDim); // fused QKV + WorkerGrid mmaDimWorker = mmaGrid(paddedBatch, dim); // woProj, w2Proj + WorkerGrid mmaGateUpWorker = mmaGrid(paddedBatch, 2 * hidDim); // fused W1/W3 + + // Elementwise grids (one thread per valid element; dim & hidDim are mult. of 256) + WorkerGrid ewDimWorker = elementwiseGrid(batchSize * dim); // w2Resid + WorkerGrid ewHidWorker = elementwiseGrid(batchSize * hidDim); // swiglu + + for (int i = 0; i < config.numberOfLayers(); i++) { + String p = "batchPrefillLayer_" + i + "."; + scheduler.addWorkerGrid(p + "batch_attn_rms", rmsWorker); + scheduler.addWorkerGrid(p + "batch_attn_rms_apply", rmsApplyWorker); + scheduler.addWorkerGrid(p + "qkvProj", mmaQkvWorker); + scheduler.addWorkerGrid(p + "batch_rope_kv", ropeWorker); + scheduler.addWorkerGrid(p + "batch_attention", attnWorker); + scheduler.addWorkerGrid(p + "woProj", mmaDimWorker); + scheduler.addWorkerGrid(p + "batch_ffn_rms", rmsWorker); + scheduler.addWorkerGrid(p + "batch_ffn_rms_apply", ffnRmsApplyWorker); + scheduler.addWorkerGrid(p + "gateUpProj", mmaGateUpWorker); + scheduler.addWorkerGrid(p + "swiglu", ewHidWorker); + scheduler.addWorkerGrid(p + "w2Proj", mmaDimWorker); + scheduler.addWorkerGrid(p + "w2Resid", ewDimWorker); + } + } + + public List getLayerImmutableTaskGraphs() { return layerITGs; } + public String getLastLayerTaskGraphID() { return lastLayerTaskGraphID; } + public KernelContext getContext() { return context; } +} diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/prefill/Qwen3FP16LayersBatchPrefillMMA.java b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/prefill/Qwen3FP16LayersBatchPrefillMMA.java new file mode 100644 index 00000000..f12b842d --- /dev/null +++ b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/prefill/Qwen3FP16LayersBatchPrefillMMA.java @@ -0,0 +1,302 @@ +package org.beehive.gpullama3.tornadovm.layers.type.fp16.prefill; + +import org.beehive.gpullama3.inference.state.Qwen3State; +import org.beehive.gpullama3.inference.weights.tornado.Qwen3TornadoWeights; +import org.beehive.gpullama3.model.qwen3.Qwen3Configuration; +import org.beehive.gpullama3.tornadovm.kernels.Qwen3Kernels; +import org.beehive.gpullama3.tornadovm.kernels.TransformerBatchPrefillKernels; +import org.beehive.gpullama3.tornadovm.scheduling.WorkerGridFactory; +import org.beehive.gpullama3.tornadovm.layers.BatchPrefillTransformerLayerTaskGraphs; +import uk.ac.manchester.tornado.api.GridScheduler; +import uk.ac.manchester.tornado.api.ImmutableTaskGraph; +import uk.ac.manchester.tornado.api.KernelContext; +import uk.ac.manchester.tornado.api.TaskGraph; +import uk.ac.manchester.tornado.api.WorkerGrid; +import uk.ac.manchester.tornado.api.WorkerGrid1D; +import uk.ac.manchester.tornado.api.WorkerGrid2D; +import uk.ac.manchester.tornado.api.enums.DataTransferMode; + +import java.util.List; +import java.util.stream.IntStream; + +/** + * Qwen3 batched-prefill transformer-layer TaskGraphs on the tensor-core (MMA) + * pipeline. Mirrors {@link LlamaFP16LayersBatchPrefillMMA} with the Qwen3 + * architectural additions: per-head Q/K RMS normalization between the QKV + * projection and RoPE, split-half RoPE pairing, and qDim-shaped attention + * (qDim = nHeads * headDim, which may differ from the model dim). + * + *

Tensor-core layer pipeline (13 tasks, all GEMMs on MMA):

+ * 
+ *   batch_attn_rms        parallel RMS square-sum reduction (256 thr/token)
+ *   batch_attn_rms_apply  RMS apply + FP16 quantize → wrapXbFP16Batch
+ *   qkvProj               ONE fused MMA GEMM → packed qkvResultBatch [q|k|v]
+ *   batch_qk_rmsnorm      per-head Q/K RMS norm over the packed buffer
+ *   batch_rope_kv         split-half RoPE over the packed buffer + KV cache write
+ *   batch_attention       flash attention (register-partitioned P·V) → attnOutFP16
+ *   woProj                MMA GEMM [dim × qDim] → woOut
+ *   batch_ffn_rms         parallel RMS reduce FUSED with x += woOut
+ *   batch_ffn_rms_apply   RMS apply + FP16 quantize → normedXFFNFP16
+ *   gateUpProj            ONE fused MMA GEMM → packed gateUpResultBatch [gate|up]
+ *   swiglu                SiLU(gate)*up over packed buffer → wrapHbFP16Batch
+ *   w2Proj                MMA GEMM → w2Out
+ *   w2Resid               x += w2Out
+ *
+ * + *

Requires dim, qDim, kvDim, and hidDim to be multiples of 128 + * (holds for all standard Qwen3 checkpoints).

+ */ +public class Qwen3FP16LayersBatchPrefillMMA implements BatchPrefillTransformerLayerTaskGraphs { + + // Local size for the parallel RMS reductions (one workgroup per token). + static final int RMS_LOCAL_SIZE = 256; + + private final Qwen3State state; + private final Qwen3TornadoWeights weights; + private final Qwen3Configuration config; + private final KernelContext context = new KernelContext(); + private final int batchSize; + // GEMM M dimension rounded up to whole 128-row tiles (BM); see the Llama + // planner for the padding rationale. Non-GEMM kernels use the true batchSize. + private final int paddedBatch; + private final int nHeadKv; + private final int nEmbdHead; + private final int qDim; + private final int kvDim; + private final int gqa; + private final List layerITGs; + private String lastLayerTaskGraphID; + + public Qwen3FP16LayersBatchPrefillMMA(Qwen3State state, Qwen3TornadoWeights weights, + Qwen3Configuration config, int batchSize) { + this.state = state; + this.weights = weights; + this.config = config; + this.batchSize = batchSize; + this.paddedBatch = (batchSize + 127) & ~127; + if (batchSize % 128 != 0) { + System.out.printf("[GPULlama3] prefill batch %d padded to %d for tensor-core tiles; " + + "GEMM efficiency is %d/%d — use a multiple of 128 for best throughput.%n", + batchSize, paddedBatch, batchSize, paddedBatch); + } + this.nHeadKv = config.numberOfKeyValueHeads(); + this.nEmbdHead = config.numberOfHeadsValue(); + this.qDim = config.numberOfHeadsKey() * config.numberOfHeads(); + this.kvDim = config.numberOfHeadsValue() * nHeadKv; + this.gqa = config.numberOfHeads() / nHeadKv; + this.layerITGs = IntStream.range(0, config.numberOfLayers()) + .mapToObj(this::createBatchPrefillLayerTaskGraph) + .map(TaskGraph::snapshot) + .toList(); + } + + // @formatter:off + private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) { + String graphName = "batchPrefillLayer_" + layerIndex; + if (layerIndex == config.numberOfLayers() - 1) lastLayerTaskGraphID = graphName; + + TaskGraph batchPrefillLayer = new TaskGraph(graphName); + int dim = config.dim(); + int hidDim = config.hiddenDim(); + + // ── Data Transfers ───────────────────────────────────────────────────── + if (layerIndex == 0) { + batchPrefillLayer.transferToDevice(DataTransferMode.EVERY_EXECUTION, state.batchStartPosHolder); + batchPrefillLayer.transferToDevice(DataTransferMode.FIRST_EXECUTION, + context, + state.attnScaleBatch, state.ffnScaleBatch, + state.wrapXbFP16Batch, + state.qkvResultBatch, + state.wrapKeyCache, state.wrapValueCache, + state.normedXFFNFP16, state.gateUpResultBatch, + state.attnOutFP16, state.woOut, state.wrapHbFP16Batch, state.w2Out); + batchPrefillLayer.consumeFromDevice("prefillActivation", state.wrapXBatch); + } else { + String pred = "batchPrefillLayer_" + (layerIndex - 1); + batchPrefillLayer.consumeFromDevice(pred, + context, + state.wrapXBatch, + state.batchStartPosHolder, + state.attnScaleBatch, state.ffnScaleBatch, + state.wrapXbFP16Batch, + state.qkvResultBatch, + state.wrapKeyCache, state.wrapValueCache, + state.normedXFFNFP16, state.gateUpResultBatch, + state.attnOutFP16, state.woOut, state.wrapHbFP16Batch, state.w2Out); + } + + // Per-layer weights: upload once + batchPrefillLayer.transferToDevice(DataTransferMode.FIRST_EXECUTION, + weights.rms_att_weightLayered[layerIndex].asFloatArray(), + weights.wqLayered[layerIndex].asHalfFloatArray(), + weights.wkLayered[layerIndex].asHalfFloatArray(), + weights.wvLayered[layerIndex].asHalfFloatArray(), + weights.woLayered[layerIndex].asHalfFloatArray(), + weights.rms_att_QNormLayered[layerIndex].asFloatArray(), + weights.rms_att_KNormLayered[layerIndex].asFloatArray(), + weights.rms_ffn_weightLayered[layerIndex].asFloatArray(), + weights.w1Layered[layerIndex].asHalfFloatArray(), + weights.w2Layered[layerIndex].asHalfFloatArray(), + weights.w3Layered[layerIndex].asHalfFloatArray()); + + // ── Attention Block ──────────────────────────────────────────────────── + batchPrefillLayer.task("batch_attn_rms", + TransformerBatchPrefillKernels::batchedRmsReduceParallel, + context, state.wrapXBatch, state.attnScaleBatch, + dim, config.rmsNormEps(), RMS_LOCAL_SIZE); + + batchPrefillLayer.task("batch_attn_rms_apply", + TransformerBatchPrefillKernels::batchedRmsApplyFP16, + context, state.wrapXbFP16Batch, state.wrapXBatch, + weights.rms_att_weightLayered[layerIndex].asFloatArray(), + state.attnScaleBatch, dim); + + // Q, K, V in ONE tensor-core launch → packed [q|k|v] rows (stride qDim+2*kvDim) + batchPrefillLayer.task("qkvProj", + TransformerBatchPrefillKernels::gemmMMAQKV, + context, state.wrapXbFP16Batch, + weights.wqLayered[layerIndex].asHalfFloatArray(), + weights.wkLayered[layerIndex].asHalfFloatArray(), + weights.wvLayered[layerIndex].asHalfFloatArray(), + state.qkvResultBatch, paddedBatch, qDim, kvDim, dim); + + // Qwen3: per-head RMS norm on Q and K before RoPE + batchPrefillLayer.task("batch_qk_rmsnorm", + Qwen3Kernels::batchedFusedQKRmsNormPacked, + context, state.qkvResultBatch, + weights.rms_att_QNormLayered[layerIndex].asFloatArray(), + weights.rms_att_KNormLayered[layerIndex].asFloatArray(), + config.numberOfHeads(), nHeadKv, nEmbdHead, + qDim, kvDim, config.rmsNormEps()); + + batchPrefillLayer.task("batch_rope_kv", + Qwen3Kernels::batchedRopeWithKVCacheQwen3Packed, + context, state.batchStartPosHolder, + state.qkvResultBatch, + state.wrapKeyCache, state.wrapValueCache, + kvDim, nEmbdHead, layerIndex, config.contextLength(), qDim); + + // Register-partitioned flash attention over the packed buffer. + // The 'dim' parameter doubles as the packed-Q stride base and the + // attnOutFP16 row width — both are qDim for Qwen3. + batchPrefillLayer.task("batch_attention", + TransformerBatchPrefillKernels::batchedFlashAttentionFP16Out, + context, state.batchStartPosHolder, + state.qkvResultBatch, state.wrapKeyCache, state.wrapValueCache, + state.attnOutFP16, + config.numberOfHeads(), nEmbdHead, + kvDim, gqa, layerIndex, config.contextLength(), qDim); + + // Output projection: [M=batch, N=dim, K=qDim] + batchPrefillLayer.task("woProj", TransformerBatchPrefillKernels::gemmMMA, + context, state.attnOutFP16, + weights.woLayered[layerIndex].asHalfFloatArray(), + state.woOut, paddedBatch, dim, qDim); + + // ── FFN Block ────────────────────────────────────────────────────────── + batchPrefillLayer.task("batch_ffn_rms", + TransformerBatchPrefillKernels::batchedRmsReduceFusedResidual, + context, state.wrapXBatch, state.woOut, state.ffnScaleBatch, + dim, config.rmsNormEps(), RMS_LOCAL_SIZE); + + batchPrefillLayer.task("batch_ffn_rms_apply", + TransformerBatchPrefillKernels::batchedFFNRmsApplyFP16, + context, state.normedXFFNFP16, state.wrapXBatch, + weights.rms_ffn_weightLayered[layerIndex].asFloatArray(), + state.ffnScaleBatch, dim); + + batchPrefillLayer.task("gateUpProj", + TransformerBatchPrefillKernels::gemmMMAGateUp, + context, state.normedXFFNFP16, + weights.w1Layered[layerIndex].asHalfFloatArray(), + weights.w3Layered[layerIndex].asHalfFloatArray(), + state.gateUpResultBatch, paddedBatch, hidDim, dim); + + batchPrefillLayer.task("swiglu", + TransformerBatchPrefillKernels::batchedFFNSwiGLUFP16Packed, + context, state.wrapHbFP16Batch, state.gateUpResultBatch, hidDim) + .task("w2Proj", TransformerBatchPrefillKernels::gemmMMA, + context, state.wrapHbFP16Batch, + weights.w2Layered[layerIndex].asHalfFloatArray(), + state.w2Out, paddedBatch, dim, hidDim) + .task("w2Resid", TransformerBatchPrefillKernels::batchedResidualAddFP32, + context, state.wrapXBatch, state.w2Out); + + batchPrefillLayer.persistOnDevice(state.wrapXBatch, state.wrapKeyCache, state.wrapValueCache); + + return batchPrefillLayer; + } + // @formatter:on + + // gemmMMA family: 256 threads/block (1D within block), grid over M- and N-blocks. + static WorkerGrid mmaGrid(int paddedM, int N) { + int mBlocks = paddedM / 128; // BM + int nBlocks = N / 128; // BN + WorkerGrid2D g = new WorkerGrid2D(mBlocks * 256, nBlocks); + g.setLocalWork(256, 1, 1); + return g; + } + + static WorkerGrid elementwiseGrid(int n) { // n must be a multiple of 256 + WorkerGrid1D g = new WorkerGrid1D(n); + g.setLocalWork(256, 1, 1); + return g; + } + + /** Registers all batch layer workers in the shared {@link GridScheduler}. */ + public void updateGridScheduler(GridScheduler scheduler) { + int dim = config.dim(); + int hidDim = config.hiddenDim(); + int nHeads = config.numberOfHeads(); + + WorkerGrid rmsWorker = WorkerGridFactory.genericWorker( + batchSize * RMS_LOCAL_SIZE, RMS_LOCAL_SIZE); + + WorkerGrid rmsApplyWorker = WorkerGridFactory.genericWorker(batchSize * dim, 256); + WorkerGrid ffnRmsApplyWorker = WorkerGridFactory.genericWorker(batchSize * dim, 256); + + // Q/K per-head RMS norm: one nEmbdHead-thread workgroup per (token, head) + WorkerGrid qkRmsNormWorker = WorkerGridFactory.genericWorker( + batchSize * (nHeads + nHeadKv) * nEmbdHead, nEmbdHead); + + // Split-half RoPE: B*(qDim/2) threads + int ropeGlobal = batchSize * (qDim / 2); + int ropeLocal = Math.min(512, ropeGlobal); + while (ropeLocal > 1 && ropeGlobal % ropeLocal != 0) ropeLocal--; + WorkerGrid ropeWorker = WorkerGridFactory.genericWorker(ropeGlobal, ropeLocal); + + // Attention: B*nHeads workgroups × min(nEmbdHead,128) threads + int attnLocal = Math.min(nEmbdHead, 128); + WorkerGrid attnWorker = WorkerGridFactory.genericWorker( + batchSize * nHeads * attnLocal, attnLocal); + + // MMA grids + WorkerGrid mmaQkvWorker = mmaGrid(paddedBatch, qDim + 2 * kvDim); // fused QKV + WorkerGrid mmaDimWorker = mmaGrid(paddedBatch, dim); // woProj, w2Proj + WorkerGrid mmaGateUpWorker = mmaGrid(paddedBatch, 2 * hidDim); // fused W1/W3 + + WorkerGrid ewDimWorker = elementwiseGrid(batchSize * dim); // w2Resid + WorkerGrid ewHidWorker = elementwiseGrid(batchSize * hidDim); // swiglu + + for (int i = 0; i < config.numberOfLayers(); i++) { + String p = "batchPrefillLayer_" + i + "."; + scheduler.addWorkerGrid(p + "batch_attn_rms", rmsWorker); + scheduler.addWorkerGrid(p + "batch_attn_rms_apply", rmsApplyWorker); + scheduler.addWorkerGrid(p + "qkvProj", mmaQkvWorker); + scheduler.addWorkerGrid(p + "batch_qk_rmsnorm", qkRmsNormWorker); + scheduler.addWorkerGrid(p + "batch_rope_kv", ropeWorker); + scheduler.addWorkerGrid(p + "batch_attention", attnWorker); + scheduler.addWorkerGrid(p + "woProj", mmaDimWorker); + scheduler.addWorkerGrid(p + "batch_ffn_rms", rmsWorker); + scheduler.addWorkerGrid(p + "batch_ffn_rms_apply", ffnRmsApplyWorker); + scheduler.addWorkerGrid(p + "gateUpProj", mmaGateUpWorker); + scheduler.addWorkerGrid(p + "swiglu", ewHidWorker); + scheduler.addWorkerGrid(p + "w2Proj", mmaDimWorker); + scheduler.addWorkerGrid(p + "w2Resid", ewDimWorker); + } + } + + public List getLayerImmutableTaskGraphs() { return layerITGs; } + public String getLastLayerTaskGraphID() { return lastLayerTaskGraphID; } + public KernelContext getContext() { return context; } +} diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/prefill/LlamaQ8_0LayersBatchPrefillMMA.java b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/prefill/LlamaQ8_0LayersBatchPrefillMMA.java new file mode 100644 index 00000000..24264400 --- /dev/null +++ b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/prefill/LlamaQ8_0LayersBatchPrefillMMA.java @@ -0,0 +1,294 @@ +package org.beehive.gpullama3.tornadovm.layers.type.q8_0.prefill; + +import org.beehive.gpullama3.inference.state.LlamaState; +import org.beehive.gpullama3.inference.weights.tornado.LlamaTornadoWeights; +import org.beehive.gpullama3.model.llama.LlamaConfiguration; +import org.beehive.gpullama3.tornadovm.kernels.TransformerBatchPrefillKernels; +import org.beehive.gpullama3.tornadovm.scheduling.WorkerGridFactory; +import org.beehive.gpullama3.tornadovm.layers.BatchPrefillTransformerLayerTaskGraphs; +import uk.ac.manchester.tornado.api.GridScheduler; +import uk.ac.manchester.tornado.api.ImmutableTaskGraph; +import uk.ac.manchester.tornado.api.KernelContext; +import uk.ac.manchester.tornado.api.TaskGraph; +import uk.ac.manchester.tornado.api.WorkerGrid; +import uk.ac.manchester.tornado.api.WorkerGrid1D; +import uk.ac.manchester.tornado.api.WorkerGrid2D; +import uk.ac.manchester.tornado.api.enums.DataTransferMode; + +import java.util.List; +import java.util.stream.IntStream; + +/** + * Batched-prefill transformer-layer TaskGraphs for the unified batched prefill-decode plan + * ({@link org.beehive.gpullama3.tornadovm.TornadoVMMasterPlanBatchPrefillDecode}). + * + *

One {@link ImmutableTaskGraph} per transformer layer, each processing + * {@code batchSize} tokens simultaneously via {@link TransformerBatchPrefillKernels}.

+ * + *

Tensor-core layer pipeline (12 tasks, all GEMMs on MMA, Q8_0 weights + * dequantized to FP16 in the GEMM staging registers — W8A16):

+ * 
+ *   batch_attn_rms        parallel RMS square-sum reduction (256 thr/token)
+ *   batch_attn_rms_apply  RMS apply + FP16 quantize → wrapXbFP16Batch
+ *   qkvProj               ONE fused MMA GEMM → packed qkvResultBatch [q|k|v]
+ *   batch_rope_kv         RoPE over the packed buffer + KV cache write
+ *   batch_attention       flash attention (register-partitioned P·V) → attnOutFP16
+ *   woProj                MMA GEMM → woOut
+ *   batch_ffn_rms         parallel RMS reduce FUSED with x += woOut
+ *   batch_ffn_rms_apply   RMS apply + FP16 quantize → normedXFFNFP16
+ *   gateUpProj            ONE fused MMA GEMM → packed gateUpResultBatch [gate|up]
+ *   swiglu                SiLU(gate)*up over packed buffer → wrapHbFP16Batch
+ *   w2Proj                MMA GEMM → w2Out
+ *   w2Resid               x += w2Out
+ *
+ * + *

KV cache ({@code wrapKeyCache}, {@code wrapValueCache}) is persisted on device + * after every layer so the subsequent single-token decode layers can consume it.

+ */ +public class LlamaQ8_0LayersBatchPrefillMMA implements BatchPrefillTransformerLayerTaskGraphs { + + // Local size for the parallel RMS reductions (one workgroup per token). + static final int RMS_LOCAL_SIZE = 256; + + private final LlamaState state; + private final LlamaTornadoWeights weights; + private final LlamaConfiguration config; + private final KernelContext context = new KernelContext(); + private final int batchSize; + // GEMM M dimension rounded up to whole 128-row tiles (BM). The GEMM-adjacent + // buffers in State are allocated at this padded size; rows >= batchSize are + // computed but never consumed. All non-GEMM kernels use the true batchSize. + private final int paddedBatch; + private final List layerITGs; + private String lastLayerTaskGraphID; + + public LlamaQ8_0LayersBatchPrefillMMA(LlamaState state, LlamaTornadoWeights weights, + LlamaConfiguration config, int batchSize) { + this.state = state; + this.weights = weights; + this.config = config; + this.batchSize = batchSize; + this.paddedBatch = (batchSize + 127) & ~127; + if (batchSize % 128 != 0) { + System.out.printf("[GPULlama3] prefill batch %d padded to %d for tensor-core tiles; " + + "GEMM efficiency is %d/%d — use a multiple of 128 for best throughput.%n", + batchSize, paddedBatch, batchSize, paddedBatch); + } + this.layerITGs = IntStream.range(0, config.numberOfLayers()) + .mapToObj(this::createBatchPrefillLayerTaskGraph) + .map(TaskGraph::snapshot) + .toList(); + } + + // @formatter:off + private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) { + String graphName = "batchPrefillLayer_" + layerIndex; + if (layerIndex == config.numberOfLayers() - 1) lastLayerTaskGraphID = graphName; + + TaskGraph batchPrefillLayer = new TaskGraph(graphName); + + // ── Data Transfers ───────────────────────────────────────────────────── + if (layerIndex == 0) { + // batchStartPosHolder is set by host before each chunk → EVERY_EXECUTION + batchPrefillLayer.transferToDevice(DataTransferMode.EVERY_EXECUTION, state.batchStartPosHolder); + // Allocate persistent GPU-side intermediates once + batchPrefillLayer.transferToDevice(DataTransferMode.FIRST_EXECUTION, + context, + state.attnScaleBatch, state.ffnScaleBatch, + state.wrapXbFP16Batch, + state.qkvResultBatch, + state.wrapKeyCache, state.wrapValueCache, + state.normedXFFNFP16, state.gateUpResultBatch, + state.attnOutFP16, state.woOut, state.wrapHbFP16Batch, state.w2Out); + // wrapXBatch produced by the prefillActivation graph and persists in device memory + // to consume it from there we should use the explicit uniqueTaskGraph name + // the no-arg form would use current graph name, which causes NPE without CUDA Graphs + batchPrefillLayer.consumeFromDevice("prefillActivation", state.wrapXBatch); + } else { + // for the same reasons as above, we should use the explicit uniqueTaskGraph name to consume + String pred = "batchPrefillLayer_" + (layerIndex - 1); + batchPrefillLayer.consumeFromDevice(pred, + context, + state.wrapXBatch, + state.batchStartPosHolder, + state.attnScaleBatch, state.ffnScaleBatch, + state.wrapXbFP16Batch, + state.qkvResultBatch, + state.wrapKeyCache, state.wrapValueCache, + state.normedXFFNFP16, state.gateUpResultBatch, + state.attnOutFP16, state.woOut, state.wrapHbFP16Batch, state.w2Out); + } + + // Per-layer weights: upload once + batchPrefillLayer.transferToDevice(DataTransferMode.FIRST_EXECUTION, + weights.rms_att_weightLayered[layerIndex].asFloatArray(), + weights.wqLayered[layerIndex].asByteArray(), + weights.wkLayered[layerIndex].asByteArray(), + weights.wvLayered[layerIndex].asByteArray(), + weights.woLayered[layerIndex].asByteArray(), + weights.rms_ffn_weightLayered[layerIndex].asFloatArray(), + weights.w1Layered[layerIndex].asByteArray(), + weights.w2Layered[layerIndex].asByteArray(), + weights.w3Layered[layerIndex].asByteArray()); + + int dim = config.dim(); + int kvDim = config.kvDim(); + int hidDim = config.hiddenDim(); + + // ── Attention Block ──────────────────────────────────────────────────── + batchPrefillLayer.task("batch_attn_rms", + TransformerBatchPrefillKernels::batchedRmsReduceParallel, + context, state.wrapXBatch, state.attnScaleBatch, + dim, config.rmsNormEps(), RMS_LOCAL_SIZE); + + batchPrefillLayer.task("batch_attn_rms_apply", + TransformerBatchPrefillKernels::batchedRmsApplyFP16, + context, state.wrapXbFP16Batch, state.wrapXBatch, + weights.rms_att_weightLayered[layerIndex].asFloatArray(), + state.attnScaleBatch, dim); + + // Q, K, V in ONE tensor-core launch → packed [q|k|v] rows. + // Grid spans (dim + 2*kvDim)/128 N-blocks: no grid starvation on the + // skinny GQA projections, and the A operand is read once, not thrice. + batchPrefillLayer.task("qkvProj", + TransformerBatchPrefillKernels::gemmMMAQKVQ8, + context, state.wrapXbFP16Batch, + weights.wqLayered[layerIndex].asByteArray(), + weights.wkLayered[layerIndex].asByteArray(), + weights.wvLayered[layerIndex].asByteArray(), + state.qkvResultBatch, paddedBatch, dim, kvDim, dim); + + batchPrefillLayer.task("batch_rope_kv", + TransformerBatchPrefillKernels::batchedRopeWithKVCachePacked, + context, state.batchStartPosHolder, + state.qkvResultBatch, + state.wrapKeyCache, state.wrapValueCache, + kvDim, config.headSize(), layerIndex, config.contextLength(), dim); + + // Register-partitioned P·V accumulation + direct FP16 emission + // (replaces batchedFlashAttention + attnCast). + batchPrefillLayer.task("batch_attention", + TransformerBatchPrefillKernels::batchedFlashAttentionFP16Out, + context, state.batchStartPosHolder, + state.qkvResultBatch, state.wrapKeyCache, state.wrapValueCache, + state.attnOutFP16, + config.numberOfHeads(), config.headSize(), + kvDim, config.kvMul(), layerIndex, config.contextLength(), dim); + + batchPrefillLayer.task("woProj", TransformerBatchPrefillKernels::gemmMMAQ8, + context, state.attnOutFP16, + weights.woLayered[layerIndex].asByteArray(), + state.woOut, paddedBatch, dim, dim); + + // ── FFN Block ────────────────────────────────────────────────────────── + // x += woOut is fused into the FFN RMS reduction (drops the woResid task). + batchPrefillLayer.task("batch_ffn_rms", + TransformerBatchPrefillKernels::batchedRmsReduceFusedResidual, + context, state.wrapXBatch, state.woOut, state.ffnScaleBatch, + dim, config.rmsNormEps(), RMS_LOCAL_SIZE); + + batchPrefillLayer.task("batch_ffn_rms_apply", + TransformerBatchPrefillKernels::batchedFFNRmsApplyFP16, + context, state.normedXFFNFP16, state.wrapXBatch, + weights.rms_ffn_weightLayered[layerIndex].asFloatArray(), + state.ffnScaleBatch, dim); + + // W1 and W3 in ONE tensor-core launch → packed [gate|up] rows. + batchPrefillLayer.task("gateUpProj", + TransformerBatchPrefillKernels::gemmMMAGateUpQ8, + context, state.normedXFFNFP16, + weights.w1Layered[layerIndex].asByteArray(), + weights.w3Layered[layerIndex].asByteArray(), + state.gateUpResultBatch, paddedBatch, hidDim, dim); + + batchPrefillLayer.task("swiglu", + TransformerBatchPrefillKernels::batchedFFNSwiGLUFP16Packed, + context, state.wrapHbFP16Batch, state.gateUpResultBatch, hidDim) + .task("w2Proj", TransformerBatchPrefillKernels::gemmMMAQ8, + context, state.wrapHbFP16Batch, + weights.w2Layered[layerIndex].asByteArray(), + state.w2Out, batchSize, dim, hidDim) + .task("w2Resid", TransformerBatchPrefillKernels::batchedResidualAddFP32, + context, state.wrapXBatch, state.w2Out); + + // Persist wrapXBatch for the next layer, and KV cache so the decode + // layers can consume it via the activation graph pass-through. + batchPrefillLayer.persistOnDevice(state.wrapXBatch, state.wrapKeyCache, state.wrapValueCache); + + return batchPrefillLayer; + } + // @formatter:on + + // gemmMMA family: 256 threads/block (1D within block), grid over M- and N-blocks. + static WorkerGrid mmaGrid(int paddedM, int N) { + int mBlocks = paddedM / 128; // BM + int nBlocks = N / 128; // BN + WorkerGrid2D g = new WorkerGrid2D(mBlocks * 256, nBlocks); + g.setLocalWork(256, 1, 1); // groupIdx∈[0,mBlocks), groupIdy∈[0,nBlocks) + return g; + } + + static WorkerGrid elementwiseGrid(int n) { // n must be a multiple of 256 + WorkerGrid1D g = new WorkerGrid1D(n); + g.setLocalWork(256, 1, 1); + return g; + } + + /** Registers all batch layer workers in the shared {@link GridScheduler}. */ + public void updateGridScheduler(GridScheduler scheduler) { + int dim = config.dim(); + int kvDim = config.kvDim(); + int hidDim = config.hiddenDim(); + int nHeads = config.numberOfHeads(); + int headSz = config.headSize(); + + // Parallel RMS reductions: one 256-thread workgroup per batch token + WorkerGrid rmsWorker = WorkerGridFactory.genericWorker( + batchSize * RMS_LOCAL_SIZE, RMS_LOCAL_SIZE); + + // RMS apply: B*dim threads, local=256 (dim is always a multiple of 256 for LLaMA) + WorkerGrid rmsApplyWorker = WorkerGridFactory.genericWorker(batchSize * dim, 256); + WorkerGrid ffnRmsApplyWorker = WorkerGridFactory.genericWorker(batchSize * dim, 256); + + // RoPE+KV cache: B*(dim/2) threads, local=512 + int ropeGlobal = batchSize * (dim / 2); + int ropeLocal = Math.min(512, ropeGlobal); + while (ropeLocal > 1 && ropeGlobal % ropeLocal != 0) ropeLocal--; + WorkerGrid ropeWorker = WorkerGridFactory.genericWorker(ropeGlobal, ropeLocal); + + // Attention (flash): B*nHeads workgroups × min(headSize,128) threads. + // The kernel requires headSize <= 2*localSize. + int attnLocal = Math.min(headSz, 128); + WorkerGrid attnWorker = WorkerGridFactory.genericWorker( + batchSize * nHeads * attnLocal, attnLocal); + + // MMA grids + WorkerGrid mmaQkvWorker = mmaGrid(paddedBatch, dim + 2 * kvDim); // fused QKV + WorkerGrid mmaDimWorker = mmaGrid(paddedBatch, dim); // woProj, w2Proj + WorkerGrid mmaGateUpWorker = mmaGrid(paddedBatch, 2 * hidDim); // fused W1/W3 + + // Elementwise grids (one thread per valid element; dim & hidDim are mult. of 256) + WorkerGrid ewDimWorker = elementwiseGrid(batchSize * dim); // w2Resid + WorkerGrid ewHidWorker = elementwiseGrid(batchSize * hidDim); // swiglu + + for (int i = 0; i < config.numberOfLayers(); i++) { + String p = "batchPrefillLayer_" + i + "."; + scheduler.addWorkerGrid(p + "batch_attn_rms", rmsWorker); + scheduler.addWorkerGrid(p + "batch_attn_rms_apply", rmsApplyWorker); + scheduler.addWorkerGrid(p + "qkvProj", mmaQkvWorker); + scheduler.addWorkerGrid(p + "batch_rope_kv", ropeWorker); + scheduler.addWorkerGrid(p + "batch_attention", attnWorker); + scheduler.addWorkerGrid(p + "woProj", mmaDimWorker); + scheduler.addWorkerGrid(p + "batch_ffn_rms", rmsWorker); + scheduler.addWorkerGrid(p + "batch_ffn_rms_apply", ffnRmsApplyWorker); + scheduler.addWorkerGrid(p + "gateUpProj", mmaGateUpWorker); + scheduler.addWorkerGrid(p + "swiglu", ewHidWorker); + scheduler.addWorkerGrid(p + "w2Proj", mmaDimWorker); + scheduler.addWorkerGrid(p + "w2Resid", ewDimWorker); + } + } + + public List getLayerImmutableTaskGraphs() { return layerITGs; } + public String getLastLayerTaskGraphID() { return lastLayerTaskGraphID; } + public KernelContext getContext() { return context; } +} diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/prefill/Qwen3Q8_0LayersBatchPrefillMMA.java b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/prefill/Qwen3Q8_0LayersBatchPrefillMMA.java new file mode 100644 index 00000000..1cd02ee4 --- /dev/null +++ b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/q8_0/prefill/Qwen3Q8_0LayersBatchPrefillMMA.java @@ -0,0 +1,304 @@ +package org.beehive.gpullama3.tornadovm.layers.type.q8_0.prefill; + +import org.beehive.gpullama3.inference.state.Qwen3State; +import org.beehive.gpullama3.inference.weights.tornado.Qwen3TornadoWeights; +import org.beehive.gpullama3.model.qwen3.Qwen3Configuration; +import org.beehive.gpullama3.tornadovm.kernels.Qwen3Kernels; +import org.beehive.gpullama3.tornadovm.kernels.TransformerBatchPrefillKernels; +import org.beehive.gpullama3.tornadovm.layers.type.fp16.prefill.LlamaFP16LayersBatchPrefillMMA; +import org.beehive.gpullama3.tornadovm.scheduling.WorkerGridFactory; +import org.beehive.gpullama3.tornadovm.layers.BatchPrefillTransformerLayerTaskGraphs; +import uk.ac.manchester.tornado.api.GridScheduler; +import uk.ac.manchester.tornado.api.ImmutableTaskGraph; +import uk.ac.manchester.tornado.api.KernelContext; +import uk.ac.manchester.tornado.api.TaskGraph; +import uk.ac.manchester.tornado.api.WorkerGrid; +import uk.ac.manchester.tornado.api.WorkerGrid1D; +import uk.ac.manchester.tornado.api.WorkerGrid2D; +import uk.ac.manchester.tornado.api.enums.DataTransferMode; + +import java.util.List; +import java.util.stream.IntStream; + +/** + * Qwen3 batched-prefill transformer-layer TaskGraphs on the tensor-core (MMA) + * pipeline. Mirrors {@link LlamaFP16LayersBatchPrefillMMA} with the Qwen3 + * architectural additions: per-head Q/K RMS normalization between the QKV + * projection and RoPE, split-half RoPE pairing, and qDim-shaped attention + * (qDim = nHeads * headDim, which may differ from the model dim). + * + *

Tensor-core layer pipeline (13 tasks, all GEMMs on MMA, Q8_0 weights dequantized + * to FP16 in the GEMM staging registers — W8A16):

+ * 
+ *   batch_attn_rms        parallel RMS square-sum reduction (256 thr/token)
+ *   batch_attn_rms_apply  RMS apply + FP16 quantize → wrapXbFP16Batch
+ *   qkvProj               ONE fused MMA GEMM → packed qkvResultBatch [q|k|v]
+ *   batch_qk_rmsnorm      per-head Q/K RMS norm over the packed buffer
+ *   batch_rope_kv         split-half RoPE over the packed buffer + KV cache write
+ *   batch_attention       flash attention (register-partitioned P·V) → attnOutFP16
+ *   woProj                MMA GEMM [dim × qDim] → woOut
+ *   batch_ffn_rms         parallel RMS reduce FUSED with x += woOut
+ *   batch_ffn_rms_apply   RMS apply + FP16 quantize → normedXFFNFP16
+ *   gateUpProj            ONE fused MMA GEMM → packed gateUpResultBatch [gate|up]
+ *   swiglu                SiLU(gate)*up over packed buffer → wrapHbFP16Batch
+ *   w2Proj                MMA GEMM → w2Out
+ *   w2Resid               x += w2Out
+ *
+ * + *

Requires dim, qDim, kvDim, and hidDim to be multiples of 128 + * (holds for all standard Qwen3 checkpoints).

+ */ +public class Qwen3Q8_0LayersBatchPrefillMMA implements BatchPrefillTransformerLayerTaskGraphs { + + // Local size for the parallel RMS reductions (one workgroup per token). + static final int RMS_LOCAL_SIZE = 256; + + private final Qwen3State state; + private final Qwen3TornadoWeights weights; + private final Qwen3Configuration config; + private final KernelContext context = new KernelContext(); + private final int batchSize; + // GEMM M dimension rounded up to whole 128-row tiles (BM); see the Llama + // planner for the padding rationale. Non-GEMM kernels use the true batchSize. + private final int paddedBatch; + private final int nHeadKv; + private final int nEmbdHead; + private final int qDim; + private final int kvDim; + private final int gqa; + private final List layerITGs; + private String lastLayerTaskGraphID; + + public Qwen3Q8_0LayersBatchPrefillMMA(Qwen3State state, Qwen3TornadoWeights weights, + Qwen3Configuration config, int batchSize) { + this.state = state; + this.weights = weights; + this.config = config; + this.batchSize = batchSize; + this.paddedBatch = (batchSize + 127) & ~127; + if (batchSize % 128 != 0) { + System.out.printf("[GPULlama3] prefill batch %d padded to %d for tensor-core tiles; " + + "GEMM efficiency is %d/%d — use a multiple of 128 for best throughput.%n", + batchSize, paddedBatch, batchSize, paddedBatch); + } + this.nHeadKv = config.numberOfKeyValueHeads(); + this.nEmbdHead = config.numberOfHeadsValue(); + this.qDim = config.numberOfHeadsKey() * config.numberOfHeads(); + this.kvDim = config.numberOfHeadsValue() * nHeadKv; + this.gqa = config.numberOfHeads() / nHeadKv; + this.layerITGs = IntStream.range(0, config.numberOfLayers()) + .mapToObj(this::createBatchPrefillLayerTaskGraph) + .map(TaskGraph::snapshot) + .toList(); + } + + // @formatter:off + private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) { + String graphName = "batchPrefillLayer_" + layerIndex; + if (layerIndex == config.numberOfLayers() - 1) lastLayerTaskGraphID = graphName; + + TaskGraph batchPrefillLayer = new TaskGraph(graphName); + int dim = config.dim(); + int hidDim = config.hiddenDim(); + + // ── Data Transfers ───────────────────────────────────────────────────── + if (layerIndex == 0) { + batchPrefillLayer.transferToDevice(DataTransferMode.EVERY_EXECUTION, state.batchStartPosHolder); + batchPrefillLayer.transferToDevice(DataTransferMode.FIRST_EXECUTION, + context, + state.attnScaleBatch, state.ffnScaleBatch, + state.wrapXbFP16Batch, + state.qkvResultBatch, + state.wrapKeyCache, state.wrapValueCache, + state.normedXFFNFP16, state.gateUpResultBatch, + state.attnOutFP16, state.woOut, state.wrapHbFP16Batch, state.w2Out); + batchPrefillLayer.consumeFromDevice("prefillActivation", state.wrapXBatch); + } else { + String pred = "batchPrefillLayer_" + (layerIndex - 1); + batchPrefillLayer.consumeFromDevice(pred, + context, + state.wrapXBatch, + state.batchStartPosHolder, + state.attnScaleBatch, state.ffnScaleBatch, + state.wrapXbFP16Batch, + state.qkvResultBatch, + state.wrapKeyCache, state.wrapValueCache, + state.normedXFFNFP16, state.gateUpResultBatch, + state.attnOutFP16, state.woOut, state.wrapHbFP16Batch, state.w2Out); + } + + // Per-layer weights: upload once + batchPrefillLayer.transferToDevice(DataTransferMode.FIRST_EXECUTION, + weights.rms_att_weightLayered[layerIndex].asFloatArray(), + weights.wqLayered[layerIndex].asByteArray(), + weights.wkLayered[layerIndex].asByteArray(), + weights.wvLayered[layerIndex].asByteArray(), + weights.woLayered[layerIndex].asByteArray(), + weights.rms_att_QNormLayered[layerIndex].asFloatArray(), + weights.rms_att_KNormLayered[layerIndex].asFloatArray(), + weights.rms_ffn_weightLayered[layerIndex].asFloatArray(), + weights.w1Layered[layerIndex].asByteArray(), + weights.w2Layered[layerIndex].asByteArray(), + weights.w3Layered[layerIndex].asByteArray()); + + // ── Attention Block ──────────────────────────────────────────────────── + batchPrefillLayer.task("batch_attn_rms", + TransformerBatchPrefillKernels::batchedRmsReduceParallel, + context, state.wrapXBatch, state.attnScaleBatch, + dim, config.rmsNormEps(), RMS_LOCAL_SIZE); + + batchPrefillLayer.task("batch_attn_rms_apply", + TransformerBatchPrefillKernels::batchedRmsApplyFP16, + context, state.wrapXbFP16Batch, state.wrapXBatch, + weights.rms_att_weightLayered[layerIndex].asFloatArray(), + state.attnScaleBatch, dim); + + // Q, K, V in ONE tensor-core launch → packed [q|k|v] rows (stride qDim+2*kvDim) + batchPrefillLayer.task("qkvProj", + TransformerBatchPrefillKernels::gemmMMAQKVQ8, + context, state.wrapXbFP16Batch, + weights.wqLayered[layerIndex].asByteArray(), + weights.wkLayered[layerIndex].asByteArray(), + weights.wvLayered[layerIndex].asByteArray(), + state.qkvResultBatch, paddedBatch, qDim, kvDim, dim); + + // Qwen3: per-head RMS norm on Q and K before RoPE + batchPrefillLayer.task("batch_qk_rmsnorm", + Qwen3Kernels::batchedFusedQKRmsNormPacked, + context, state.qkvResultBatch, + weights.rms_att_QNormLayered[layerIndex].asFloatArray(), + weights.rms_att_KNormLayered[layerIndex].asFloatArray(), + config.numberOfHeads(), nHeadKv, nEmbdHead, + qDim, kvDim, config.rmsNormEps()); + + batchPrefillLayer.task("batch_rope_kv", + Qwen3Kernels::batchedRopeWithKVCacheQwen3Packed, + context, state.batchStartPosHolder, + state.qkvResultBatch, + state.wrapKeyCache, state.wrapValueCache, + kvDim, nEmbdHead, layerIndex, config.contextLength(), qDim); + + // Register-partitioned flash attention over the packed buffer. + // The 'dim' parameter doubles as the packed-Q stride base and the + // attnOutFP16 row width — both are qDim for Qwen3. + batchPrefillLayer.task("batch_attention", + TransformerBatchPrefillKernels::batchedFlashAttentionFP16Out, + context, state.batchStartPosHolder, + state.qkvResultBatch, state.wrapKeyCache, state.wrapValueCache, + state.attnOutFP16, + config.numberOfHeads(), nEmbdHead, + kvDim, gqa, layerIndex, config.contextLength(), qDim); + + // Output projection: [M=batch, N=dim, K=qDim] + batchPrefillLayer.task("woProj", TransformerBatchPrefillKernels::gemmMMAQ8, + context, state.attnOutFP16, + weights.woLayered[layerIndex].asByteArray(), + state.woOut, paddedBatch, dim, qDim); + + // ── FFN Block ────────────────────────────────────────────────────────── + batchPrefillLayer.task("batch_ffn_rms", + TransformerBatchPrefillKernels::batchedRmsReduceFusedResidual, + context, state.wrapXBatch, state.woOut, state.ffnScaleBatch, + dim, config.rmsNormEps(), RMS_LOCAL_SIZE); + + batchPrefillLayer.task("batch_ffn_rms_apply", + TransformerBatchPrefillKernels::batchedFFNRmsApplyFP16, + context, state.normedXFFNFP16, state.wrapXBatch, + weights.rms_ffn_weightLayered[layerIndex].asFloatArray(), + state.ffnScaleBatch, dim); + + batchPrefillLayer.task("gateUpProj", + TransformerBatchPrefillKernels::gemmMMAGateUpQ8, + context, state.normedXFFNFP16, + weights.w1Layered[layerIndex].asByteArray(), + weights.w3Layered[layerIndex].asByteArray(), + state.gateUpResultBatch, paddedBatch, hidDim, dim); + + batchPrefillLayer.task("swiglu", + TransformerBatchPrefillKernels::batchedFFNSwiGLUFP16Packed, + context, state.wrapHbFP16Batch, state.gateUpResultBatch, hidDim) + .task("w2Proj", TransformerBatchPrefillKernels::gemmMMAQ8, + context, state.wrapHbFP16Batch, + weights.w2Layered[layerIndex].asByteArray(), + state.w2Out, paddedBatch, dim, hidDim) + .task("w2Resid", TransformerBatchPrefillKernels::batchedResidualAddFP32, + context, state.wrapXBatch, state.w2Out); + + batchPrefillLayer.persistOnDevice(state.wrapXBatch, state.wrapKeyCache, state.wrapValueCache); + + return batchPrefillLayer; + } + // @formatter:on + + // gemmMMA family: 256 threads/block (1D within block), grid over M- and N-blocks. + static WorkerGrid mmaGrid(int paddedM, int N) { + int mBlocks = paddedM / 128; // BM + int nBlocks = N / 128; // BN + WorkerGrid2D g = new WorkerGrid2D(mBlocks * 256, nBlocks); + g.setLocalWork(256, 1, 1); + return g; + } + + static WorkerGrid elementwiseGrid(int n) { // n must be a multiple of 256 + WorkerGrid1D g = new WorkerGrid1D(n); + g.setLocalWork(256, 1, 1); + return g; + } + + /** Registers all batch layer workers in the shared {@link GridScheduler}. */ + public void updateGridScheduler(GridScheduler scheduler) { + int dim = config.dim(); + int hidDim = config.hiddenDim(); + int nHeads = config.numberOfHeads(); + + WorkerGrid rmsWorker = WorkerGridFactory.genericWorker( + batchSize * RMS_LOCAL_SIZE, RMS_LOCAL_SIZE); + + WorkerGrid rmsApplyWorker = WorkerGridFactory.genericWorker(batchSize * dim, 256); + WorkerGrid ffnRmsApplyWorker = WorkerGridFactory.genericWorker(batchSize * dim, 256); + + // Q/K per-head RMS norm: one nEmbdHead-thread workgroup per (token, head) + WorkerGrid qkRmsNormWorker = WorkerGridFactory.genericWorker( + batchSize * (nHeads + nHeadKv) * nEmbdHead, nEmbdHead); + + // Split-half RoPE: B*(qDim/2) threads + int ropeGlobal = batchSize * (qDim / 2); + int ropeLocal = Math.min(512, ropeGlobal); + while (ropeLocal > 1 && ropeGlobal % ropeLocal != 0) ropeLocal--; + WorkerGrid ropeWorker = WorkerGridFactory.genericWorker(ropeGlobal, ropeLocal); + + // Attention: B*nHeads workgroups × min(nEmbdHead,128) threads + int attnLocal = Math.min(nEmbdHead, 128); + WorkerGrid attnWorker = WorkerGridFactory.genericWorker( + batchSize * nHeads * attnLocal, attnLocal); + + // MMA grids + WorkerGrid mmaQkvWorker = mmaGrid(paddedBatch, qDim + 2 * kvDim); // fused QKV + WorkerGrid mmaDimWorker = mmaGrid(paddedBatch, dim); // woProj, w2Proj + WorkerGrid mmaGateUpWorker = mmaGrid(paddedBatch, 2 * hidDim); // fused W1/W3 + + WorkerGrid ewDimWorker = elementwiseGrid(batchSize * dim); // w2Resid + WorkerGrid ewHidWorker = elementwiseGrid(batchSize * hidDim); // swiglu + + for (int i = 0; i < config.numberOfLayers(); i++) { + String p = "batchPrefillLayer_" + i + "."; + scheduler.addWorkerGrid(p + "batch_attn_rms", rmsWorker); + scheduler.addWorkerGrid(p + "batch_attn_rms_apply", rmsApplyWorker); + scheduler.addWorkerGrid(p + "qkvProj", mmaQkvWorker); + scheduler.addWorkerGrid(p + "batch_qk_rmsnorm", qkRmsNormWorker); + scheduler.addWorkerGrid(p + "batch_rope_kv", ropeWorker); + scheduler.addWorkerGrid(p + "batch_attention", attnWorker); + scheduler.addWorkerGrid(p + "woProj", mmaDimWorker); + scheduler.addWorkerGrid(p + "batch_ffn_rms", rmsWorker); + scheduler.addWorkerGrid(p + "batch_ffn_rms_apply", ffnRmsApplyWorker); + scheduler.addWorkerGrid(p + "gateUpProj", mmaGateUpWorker); + scheduler.addWorkerGrid(p + "swiglu", ewHidWorker); + scheduler.addWorkerGrid(p + "w2Proj", mmaDimWorker); + scheduler.addWorkerGrid(p + "w2Resid", ewDimWorker); + } + } + + public List getLayerImmutableTaskGraphs() { return layerITGs; } + public String getLastLayerTaskGraphID() { return lastLayerTaskGraphID; } + public KernelContext getContext() { return context; } +} diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/plan/components/fp16/LlamaFP16PlanComponents.java b/src/main/java/org/beehive/gpullama3/tornadovm/plan/components/fp16/LlamaFP16PlanComponents.java index e5166844..a4490d88 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/plan/components/fp16/LlamaFP16PlanComponents.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/plan/components/fp16/LlamaFP16PlanComponents.java @@ -14,7 +14,9 @@ import org.beehive.gpullama3.tornadovm.layers.type.fp16.decode.LlamaFP16FFNLayersDecode; import org.beehive.gpullama3.tornadovm.layers.type.fp16.decode.LlamaFP16FFNLayersPrefillDecode; import org.beehive.gpullama3.tornadovm.layers.type.fp16.decode.LogitsFP16LayerDecode; +import org.beehive.gpullama3.tornadovm.layers.type.fp16.prefill.LlamaFP16LayersBatchPrefillMMA; import org.beehive.gpullama3.tornadovm.layers.type.fp16.prefill.LlamaFP16LayersBatchPrefill; +import org.beehive.gpullama3.tornadovm.TensorCoreSupport; import org.beehive.gpullama3.tornadovm.plan.components.BatchPrefillDecodeForwardPlanComponents; import org.beehive.gpullama3.tornadovm.plan.components.activation.BatchDecodeActivation; import org.beehive.gpullama3.tornadovm.plan.components.activation.BatchPrefillActivation; @@ -81,6 +83,9 @@ public TransformerLayerTaskGraphs batchDecodeTransformerLayers() { @Override public BatchPrefillTransformerLayerTaskGraphs batchPrefillTransformerLayers(int batchSize) { + if (TensorCoreSupport.isTensorCoreCapableBackend()) { + return new LlamaFP16LayersBatchPrefillMMA(state, weights, config, batchSize); + } return new LlamaFP16LayersBatchPrefill(state, weights, config, batchSize); } diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/plan/components/fp16/Qwen3FP16PlanComponents.java b/src/main/java/org/beehive/gpullama3/tornadovm/plan/components/fp16/Qwen3FP16PlanComponents.java index 4cb7a0aa..7a8008f8 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/plan/components/fp16/Qwen3FP16PlanComponents.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/plan/components/fp16/Qwen3FP16PlanComponents.java @@ -14,7 +14,9 @@ import org.beehive.gpullama3.tornadovm.layers.type.fp16.decode.LogitsFP16LayerDecode; import org.beehive.gpullama3.tornadovm.layers.type.fp16.decode.Qwen3FP16FFNLayersDecode; import org.beehive.gpullama3.tornadovm.layers.type.fp16.decode.Qwen3FP16FFNLayersPrefillDecode; +import org.beehive.gpullama3.tornadovm.layers.type.fp16.prefill.Qwen3FP16LayersBatchPrefillMMA; import org.beehive.gpullama3.tornadovm.layers.type.fp16.prefill.Qwen3FP16LayersBatchPrefill; +import org.beehive.gpullama3.tornadovm.TensorCoreSupport; import org.beehive.gpullama3.tornadovm.plan.components.BatchPrefillDecodeForwardPlanComponents; import org.beehive.gpullama3.tornadovm.plan.components.activation.BatchDecodeActivation; import org.beehive.gpullama3.tornadovm.plan.components.activation.BatchPrefillActivation; @@ -76,6 +78,9 @@ public TransformerLayerTaskGraphs batchDecodeTransformerLayers() { @Override public BatchPrefillTransformerLayerTaskGraphs batchPrefillTransformerLayers(int batchSize) { + if (TensorCoreSupport.isTensorCoreCapableBackend()) { + return new Qwen3FP16LayersBatchPrefillMMA(state, weights, config, batchSize); + } return new Qwen3FP16LayersBatchPrefill(state, weights, config, batchSize); } diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/plan/components/q8_0/LlamaQ8_0PlanComponents.java b/src/main/java/org/beehive/gpullama3/tornadovm/plan/components/q8_0/LlamaQ8_0PlanComponents.java index 536a8789..bd7402f2 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/plan/components/q8_0/LlamaQ8_0PlanComponents.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/plan/components/q8_0/LlamaQ8_0PlanComponents.java @@ -14,7 +14,9 @@ import org.beehive.gpullama3.tornadovm.layers.type.q8_0.decode.LlamaQ8_0FFNLayersDecode; import org.beehive.gpullama3.tornadovm.layers.type.q8_0.decode.LlamaQ8_0FFNLayersPrefillDecode; import org.beehive.gpullama3.tornadovm.layers.type.q8_0.decode.LogitsQ8_0LayerDecode; +import org.beehive.gpullama3.tornadovm.layers.type.q8_0.prefill.LlamaQ8_0LayersBatchPrefillMMA; import org.beehive.gpullama3.tornadovm.layers.type.q8_0.prefill.LlamaQ8_0LayersBatchPrefill; +import org.beehive.gpullama3.tornadovm.TensorCoreSupport; import org.beehive.gpullama3.tornadovm.plan.components.BatchPrefillDecodeForwardPlanComponents; import org.beehive.gpullama3.tornadovm.plan.components.activation.BatchDecodeActivation; import org.beehive.gpullama3.tornadovm.plan.components.activation.BatchPrefillActivation; @@ -85,6 +87,9 @@ public TransformerLayerTaskGraphs batchDecodeTransformerLayers() { @Override public BatchPrefillTransformerLayerTaskGraphs batchPrefillTransformerLayers(int batchSize) { + if (TensorCoreSupport.isTensorCoreCapableBackend()) { + return new LlamaQ8_0LayersBatchPrefillMMA(state, weights, config, batchSize); + } return new LlamaQ8_0LayersBatchPrefill(state, weights, config, batchSize); } diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/plan/components/q8_0/Qwen3Q8_0PlanComponents.java b/src/main/java/org/beehive/gpullama3/tornadovm/plan/components/q8_0/Qwen3Q8_0PlanComponents.java index 6f067d1f..e4d3b6e2 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/plan/components/q8_0/Qwen3Q8_0PlanComponents.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/plan/components/q8_0/Qwen3Q8_0PlanComponents.java @@ -14,7 +14,9 @@ import org.beehive.gpullama3.tornadovm.layers.type.q8_0.decode.LogitsQ8_0LayerDecode; import org.beehive.gpullama3.tornadovm.layers.type.q8_0.decode.Qwen3Q8_0FFNLayersDecode; import org.beehive.gpullama3.tornadovm.layers.type.q8_0.decode.Qwen3Q8_0FFNLayersPrefillDecode; +import org.beehive.gpullama3.tornadovm.layers.type.q8_0.prefill.Qwen3Q8_0LayersBatchPrefillMMA; import org.beehive.gpullama3.tornadovm.layers.type.q8_0.prefill.Qwen3Q8_0LayersBatchPrefill; +import org.beehive.gpullama3.tornadovm.TensorCoreSupport; import org.beehive.gpullama3.tornadovm.plan.components.BatchPrefillDecodeForwardPlanComponents; import org.beehive.gpullama3.tornadovm.plan.components.activation.BatchDecodeActivation; import org.beehive.gpullama3.tornadovm.plan.components.activation.BatchPrefillActivation; @@ -76,6 +78,9 @@ public TransformerLayerTaskGraphs batchDecodeTransformerLayers() { @Override public BatchPrefillTransformerLayerTaskGraphs batchPrefillTransformerLayers(int batchSize) { + if (TensorCoreSupport.isTensorCoreCapableBackend()) { + return new Qwen3Q8_0LayersBatchPrefillMMA(state, weights, config, batchSize); + } return new Qwen3Q8_0LayersBatchPrefill(state, weights, config, batchSize); }