From 52edca98f0d1b5121fb551ed8d5c5723eec1899f Mon Sep 17 00:00:00 2001 From: Orion Papadakis Date: Wed, 15 Jul 2026 20:47:02 +0300 Subject: [PATCH 1/5] Add flag-gated FP16 KV cache for Llama and Qwen3 single-token decode With -Dllama.kvcache.fp16=true (and the NVIDIA scheduler path), LlamaState and Qwen3State additionally allocate half-precision KV caches, and the standard single-token task graphs read/write those instead of the FP32 ones, halving decode KV-cache bandwidth. Requires TornadoVM packed half2 support (HalfFloatArray.getHalf2/setHalf2, Half2 conversion helpers). - Llama: ropeRotationWithCacheCopyFP16 packs the rotated K pair and the V pair with single 32-bit stores; processHeadsFlashAttentionFP16 reads the K/V tiles with packed half2 loads and expands to FP32 shared-memory tiles. All softmax/output arithmetic stays FP32. - Qwen3: ropeRotationWithCacheCopyFP16 (scalar half stores; Qwen3 rotation pairs sit a half-head apart, so pairs are not adjacent on the write side) and processHeadsFlashAttentionSplitKVFP16 (packed reads, FP32 online-softmax accumulation; combine phase unchanged). - processHeadsFlashAttentionFP16Scalar plus -Dllama.kvcache.fp16.scalar isolate the packed-load gain from the halved-footprint gain when benchmarking. - The prefill/decode graph variants share the KV cache with the batch prefill layers, so they override useFp16KVCache() to false; the flag currently applies to standard (single-token) mode only. FP32 buffers stay allocated, so all other paths are unaffected. --- .../gpullama3/inference/state/LlamaState.java | 7 + .../gpullama3/inference/state/Qwen3State.java | 7 + .../gpullama3/inference/state/State.java | 18 + .../tornadovm/kernels/Qwen3Kernels.java | 62 +++ .../TransformerComputeKernelsLayered.java | 383 ++++++++++++++++++ .../AbstractTransformerLayerTaskGraphs.java | 11 + .../layers/type/fp16/LlamaFP16FFNLayers.java | 92 ++++- .../layers/type/fp16/Qwen3FP16FFNLayers.java | 107 +++-- .../fp16/decode/LlamaFP16FFNLayersDecode.java | 9 + .../LlamaFP16FFNLayersPrefillDecode.java | 9 + .../fp16/decode/Qwen3FP16FFNLayersDecode.java | 9 + .../Qwen3FP16FFNLayersPrefillDecode.java | 9 + 12 files changed, 672 insertions(+), 51 deletions(-) diff --git a/src/main/java/org/beehive/gpullama3/inference/state/LlamaState.java b/src/main/java/org/beehive/gpullama3/inference/state/LlamaState.java index 6fbf5d97..baf8534d 100644 --- a/src/main/java/org/beehive/gpullama3/inference/state/LlamaState.java +++ b/src/main/java/org/beehive/gpullama3/inference/state/LlamaState.java @@ -3,6 +3,7 @@ import org.beehive.gpullama3.tensor.standard.ArrayFloatTensor; import org.beehive.gpullama3.tensor.standard.FloatTensor; import org.beehive.gpullama3.model.Configuration; +import uk.ac.manchester.tornado.api.types.HalfFloat; import uk.ac.manchester.tornado.api.types.arrays.FloatArray; import uk.ac.manchester.tornado.api.types.arrays.HalfFloatArray; import uk.ac.manchester.tornado.api.types.arrays.IntArray; @@ -70,6 +71,12 @@ protected StateFields createStateFields(Configuration config) { fields.wrapValueCache = new FloatArray(config.contextLength() * kvDim * config.numberOfLayers()); fields.wrapValueCache.init(0.f); fields.wrapKeyCache.init(0.f); + if (USE_FP16_KV) { + fields.wrapKeyCacheFP16 = new HalfFloatArray(config.contextLength() * kvDim * config.numberOfLayers()); + fields.wrapValueCacheFP16 = new HalfFloatArray(config.contextLength() * kvDim * config.numberOfLayers()); + fields.wrapKeyCacheFP16.init(new HalfFloat(0.f)); + fields.wrapValueCacheFP16.init(new HalfFloat(0.f)); + } fields.wrapAtt = new FloatArray(config.numberOfHeads() * config.contextLength()); fields.positionHolder = new IntArray(1); diff --git a/src/main/java/org/beehive/gpullama3/inference/state/Qwen3State.java b/src/main/java/org/beehive/gpullama3/inference/state/Qwen3State.java index 109310cf..b663e30b 100644 --- a/src/main/java/org/beehive/gpullama3/inference/state/Qwen3State.java +++ b/src/main/java/org/beehive/gpullama3/inference/state/Qwen3State.java @@ -4,6 +4,7 @@ import org.beehive.gpullama3.tensor.standard.FloatTensor; import org.beehive.gpullama3.model.Configuration; import org.beehive.gpullama3.model.qwen3.Qwen3Configuration; +import uk.ac.manchester.tornado.api.types.HalfFloat; import uk.ac.manchester.tornado.api.types.arrays.FloatArray; import uk.ac.manchester.tornado.api.types.arrays.HalfFloatArray; import uk.ac.manchester.tornado.api.types.arrays.IntArray; @@ -112,6 +113,12 @@ protected StateFields createStateFields(Configuration configuration) { fields.wrapValueCache = new FloatArray(config.contextLength() * nEmbdGqa * config.numberOfLayers()); fields.wrapValueCache.init(0.f); fields.wrapKeyCache.init(0.f); + if (USE_FP16_KV) { + fields.wrapKeyCacheFP16 = new HalfFloatArray(config.contextLength() * nEmbdGqa * config.numberOfLayers()); + fields.wrapValueCacheFP16 = new HalfFloatArray(config.contextLength() * nEmbdGqa * config.numberOfLayers()); + fields.wrapKeyCacheFP16.init(new HalfFloat(0.f)); + fields.wrapValueCacheFP16.init(new HalfFloat(0.f)); + } fields.wrapAtt = new FloatArray(config.numberOfHeads() * config.contextLength()); fields.positionHolder = new IntArray(1); 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 1f884b98..31e1f1f2 100644 --- a/src/main/java/org/beehive/gpullama3/inference/state/State.java +++ b/src/main/java/org/beehive/gpullama3/inference/state/State.java @@ -27,6 +27,19 @@ */ public abstract class State { + /** + * When set ({@code -Dllama.kvcache.fp16=true}), model states that support it additionally + * allocate half-precision KV caches, and the NVIDIA decode path reads/writes those instead of + * the FP32 ones (halving KV bandwidth; accumulation stays FP32). + */ + public static final boolean USE_FP16_KV = Boolean.getBoolean("llama.kvcache.fp16"); + + /** + * Evaluation aid: with the FP16 KV cache active, read it with scalar half loads instead of + * packed half2 loads ({@code -Dllama.kvcache.fp16.scalar=true}) to isolate the packed-load gain. + */ + public static final boolean FP16_KV_SCALAR = Boolean.getBoolean("llama.kvcache.fp16.scalar"); + // current wave of activations public final FloatTensor x; // activation at current time stamp (dim,) public final FloatTensor xb; // same, but inside a residual branch (dim,) @@ -58,6 +71,8 @@ public abstract class State { public final FloatArray wrapAtt; // FloatArray wrapper for the attention scores, optimized for TornadoVM. public final FloatArray wrapKeyCache; // FloatArray wrapper for the key cache, optimized for TornadoVM. public final FloatArray wrapValueCache; // FloatArray wrapper for the value cache, optimized for TornadoVM. + public final HalfFloatArray wrapKeyCacheFP16; // Optional half-precision key cache (see USE_FP16_KV); null unless enabled. + public final HalfFloatArray wrapValueCacheFP16; // Optional half-precision value cache (see USE_FP16_KV); null unless enabled. public final IntArray positionHolder; public TornadoNativeArray embeddingX; @@ -135,6 +150,8 @@ protected State(Configuration config, int batchsize) { // dim vs kvdim this.wrapKeyCache = fields.wrapKeyCache; this.wrapValueCache = fields.wrapValueCache; + this.wrapKeyCacheFP16 = fields.wrapKeyCacheFP16; + this.wrapValueCacheFP16 = fields.wrapValueCacheFP16; this.wrapAtt = fields.wrapAtt; this.positionHolder = fields.positionHolder; @@ -218,6 +235,7 @@ protected static class StateFields { public FloatTensor[] keyCache, valueCache; public FloatArray wrapX, wrapXb, wrapXb2, wrapHb, wrapHb2, wrapLogits; public FloatArray wrapQ, wrapK, wrapV, wrapAtt, wrapKeyCache, wrapValueCache; + public HalfFloatArray wrapKeyCacheFP16, wrapValueCacheFP16; public IntArray positionHolder; public FloatArray temp, tempFFN, tempLogits; public TornadoNativeArray embeddingX; 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 a5546f1a..ae768956 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/kernels/Qwen3Kernels.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/kernels/Qwen3Kernels.java @@ -360,6 +360,68 @@ public static void ropeRotationWithCacheCopy( } } + /** + * FP16 KV cache variant of {@link #ropeRotationWithCacheCopy}. Qwen3's rotation pairs sit a + * half-head apart (split-half layout), so the cache writes stay scalar half stores; the win is + * the halved cache footprint read back by the attention kernels. + */ + public static void ropeRotationWithCacheCopyFP16( + KernelContext context, + IntArray positionHolder, + FloatArray q, // Q vector (in/out) + FloatArray k, // K vector (in/out) + FloatArray v, // V vector (in only) + HalfFloatArray keyCache, // Key cache (out) + HalfFloatArray valueCache, // Value cache (out) + int numberOfKeyValueHeads, + int nEmbdHead, + int nEmbdGqa, + int layer, + int contextLength) { + + int h = context.globalIdx; + int ic = context.globalIdy; + + int pos = positionHolder.get(0); + int rotn = h < numberOfKeyValueHeads ? 2 : 1; + int poffset = h * nEmbdHead; + int nComplEmbdHead = nEmbdHead / 2; + + float theta = 1000000.0f; + int i = ic * 2; + float freq = 1.0f / TornadoMath.pow(theta, (float) i / (float) nEmbdHead); + + float val = pos * freq; + float fcr = TornadoMath.cos(val); + float fci = TornadoMath.sin(val); + + // Rotate Q (all heads) + float v0q = q.get(poffset + ic); + float v1q = q.get(poffset + ic + nComplEmbdHead); + q.set(poffset + ic, v0q * fcr - v1q * fci); + q.set(poffset + ic + nComplEmbdHead, v0q * fci + v1q * fcr); + + // Rotate K and copy K/V to cache (only for KV heads) + if (rotn > 1 && (poffset + ic + nComplEmbdHead) < k.getSize()) { + float v0k = k.get(poffset + ic); + float v1k = k.get(poffset + ic + nComplEmbdHead); + float rotatedK0 = v0k * fcr - v1k * fci; + float rotatedK1 = v0k * fci + v1k * fcr; + + k.set(poffset + ic, rotatedK0); + k.set(poffset + ic + nComplEmbdHead, rotatedK1); + + int cacheOffset = layer * contextLength * nEmbdGqa + pos * nEmbdGqa; + int kvIdx = h * nEmbdHead; + + keyCache.set(cacheOffset + kvIdx + ic, new HalfFloat(rotatedK0)); + keyCache.set(cacheOffset + kvIdx + ic + nComplEmbdHead, new HalfFloat(rotatedK1)); + + valueCache.set(cacheOffset + kvIdx + ic, new HalfFloat(v.get(poffset + ic))); + valueCache.set(cacheOffset + kvIdx + ic + nComplEmbdHead, new HalfFloat(v.get(poffset + ic + nComplEmbdHead))); + } + } + /** * Fused Q/K/V matrix-vector multiplication for Qwen3 GQA. * Q has full head dimension, K/V have reduced KV head dimension. diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/kernels/TransformerComputeKernelsLayered.java b/src/main/java/org/beehive/gpullama3/tornadovm/kernels/TransformerComputeKernelsLayered.java index 1b400e61..003b6f71 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/kernels/TransformerComputeKernelsLayered.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/kernels/TransformerComputeKernelsLayered.java @@ -8,6 +8,7 @@ import uk.ac.manchester.tornado.api.types.arrays.FloatArray; import uk.ac.manchester.tornado.api.types.arrays.HalfFloatArray; import uk.ac.manchester.tornado.api.types.arrays.Int8Array; +import uk.ac.manchester.tornado.api.types.vectors.Half2; import uk.ac.manchester.tornado.api.types.arrays.IntArray; public class TransformerComputeKernelsLayered { @@ -541,6 +542,47 @@ public static void ropeRotationWithCacheCopy(KernelContext context, IntArray pos } + /** + * FP16 KV cache variant of {@link #ropeRotationWithCacheCopy}: the rotated K pair and the V pair + * are packed to half precision and written to the caches with single 32-bit stores. + * Requires kvDim to be even (adjacent-pair indices stay even). + */ + public static void ropeRotationWithCacheCopyFP16(KernelContext context, IntArray positionHolder, FloatArray sq, FloatArray sk, FloatArray sv, HalfFloatArray keyCache, HalfFloatArray valueCache, + int kvDim, int headSize, int layer, int contextLength) { + + int i = context.globalIdx * 2; + int pos = positionHolder.get(0); + + if (i + 1 < sq.getSize()) { + 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 + float v0q = sq.get(i); + float v1q = sq.get(i + 1); + sq.set(i, v0q * fcr - v1q * fci); + sq.set(i + 1, v0q * fci + v1q * fcr); + + // Rotate K AND write to cache (only for kvDim elements) + if (i + 1 < kvDim) { + float v0k = sk.get(i); + float v1k = sk.get(i + 1); + float rotated0 = v0k * fcr - v1k * fci; + float rotated1 = v0k * fci + v1k * fcr; + + sk.set(i, rotated0); + sk.set(i + 1, rotated1); + + int cacheOffset = layer * contextLength * kvDim + pos * kvDim; + keyCache.setHalf2(cacheOffset + i, Half2.fromFloats(rotated0, rotated1)); + valueCache.setHalf2(cacheOffset + i, Half2.fromFloats(sv.get(i), sv.get(i + 1))); + } + } + } + /** * RoPE rotation using precomputed frequency tables (cos/sin) instead of on-the-fly computation. * Required for models with non-standard RoPE (e.g., YaRN scaling in Devstral 2). @@ -905,6 +947,227 @@ public static void processHeadsFlashAttention(KernelContext context, FloatArray } } + /** + * FP16 KV cache variant of {@link #processHeadsFlashAttention}. K/V tiles are read from the + * half-precision caches with packed 32-bit loads ({@code getHalf2}) and expanded to FP32 in + * shared memory; all score/softmax/output arithmetic stays FP32. Requires even headSize. + */ + public static void processHeadsFlashAttentionFP16(KernelContext context, FloatArray q, HalfFloatArray key_cache, HalfFloatArray value_cache, FloatArray xb, int nHeads, int headSize, int kvDim, + int kvMul, IntArray positionHolder, int layer, int contextLength) { + + int tid = context.localIdx; + int h = context.groupIdx; + int localSize = context.localGroupSizeX; + + if (h >= nHeads) { + return; + } + + int pos = positionHolder.get(0); + int loff = layer * contextLength * kvDim; + int kvHeadIdx = h / kvMul; + int BLOCK_SIZE_C = 16; + + float[] q_shared = context.allocateFloatLocalArray(headSize); + float[] k_tile = context.allocateFloatLocalArray(BLOCK_SIZE_C * headSize); + float[] v_tile = context.allocateFloatLocalArray(BLOCK_SIZE_C * headSize); + float[] s_tile = context.allocateFloatLocalArray(BLOCK_SIZE_C); + float[] shared_tile_max_holder = context.allocateFloatLocalArray(1); + + float maxScore = Float.NEGATIVE_INFINITY; + float sumExp = 0.0f; + + float[] output = new float[headSize]; + for (int i = 0; i < headSize; i++) { + output[i] = 0.0f; + } + + for (int i = tid; i < headSize; i += localSize) { + q_shared[i] = q.get(h * headSize + i); + } + + context.localBarrier(); + + for (int tileC = 0; tileC <= pos; tileC += BLOCK_SIZE_C) { + int tileEnd = Math.min(tileC + BLOCK_SIZE_C - 1, pos); + + // Packed 32-bit FP16 pair loads; expand to FP32 tiles in shared memory. + for (int tIdxInSeq = tileC + tid; tIdxInSeq <= tileEnd; tIdxInSeq += localSize) { + int tileMemOffset = (tIdxInSeq - tileC) * headSize; + int kvBase = loff + tIdxInSeq * kvDim + kvHeadIdx * headSize; + for (int d = 0; d < headSize; d += 2) { + Half2 kPair = key_cache.getHalf2(kvBase + d); + Half2 vPair = value_cache.getHalf2(kvBase + d); + k_tile[tileMemOffset + d] = Half2.lowFloat(kPair); + k_tile[tileMemOffset + d + 1] = Half2.highFloat(kPair); + v_tile[tileMemOffset + d] = Half2.lowFloat(vPair); + v_tile[tileMemOffset + d + 1] = Half2.highFloat(vPair); + } + } + + context.localBarrier(); + + for (int tIdxInSeq = tileC + tid; tIdxInSeq <= tileEnd; tIdxInSeq += localSize) { + int score_idx_in_tile = tIdxInSeq - tileC; + + float score = 0.0f; + for (int d = 0; d < headSize; d++) { + score += q_shared[d] * k_tile[score_idx_in_tile * headSize + d]; + } + score /= TornadoMath.sqrt(headSize); + s_tile[score_idx_in_tile] = score; + } + + context.localBarrier(); + + float tileLocalMax = Float.NEGATIVE_INFINITY; + for (int i = 0; i <= tileEnd - tileC; i++) { + if (s_tile[i] > tileLocalMax) { + tileLocalMax = s_tile[i]; + } + } + + if (tid == 0) { + shared_tile_max_holder[0] = tileLocalMax; + } + context.localBarrier(); + float currentTileMax = shared_tile_max_holder[0]; + + float newMax = Math.max(maxScore, currentTileMax); + if (newMax != maxScore && maxScore != Float.NEGATIVE_INFINITY) { + float scale = TornadoMath.exp(maxScore - newMax); + sumExp *= scale; + for (int d = 0; d < headSize; d++) { + output[d] *= scale; + } + } + maxScore = newMax; + + for (int t_idx_in_s_tile = 0; t_idx_in_s_tile <= tileEnd - tileC; t_idx_in_s_tile++) { + float expScore = TornadoMath.exp(s_tile[t_idx_in_s_tile] - maxScore); + sumExp += expScore; + + for (int d = 0; d < headSize; d++) { + output[d] += expScore * v_tile[t_idx_in_s_tile * headSize + d]; + } + } + context.localBarrier(); + } + + float normFactor = (sumExp > 0.0f) ? (1.0f / sumExp) : 0.0f; + for (int d = tid; d < headSize; d += localSize) { + xb.set(h * headSize + d, output[d] * normFactor); + } + } + + /** + * Scalar-read FP16 KV cache variant of {@link #processHeadsFlashAttentionFP16}: identical + * computation, but the K/V tiles are read one half element at a time. Used to isolate the + * benefit of the packed loads from the benefit of the halved cache footprint. + */ + public static void processHeadsFlashAttentionFP16Scalar(KernelContext context, FloatArray q, HalfFloatArray key_cache, HalfFloatArray value_cache, FloatArray xb, int nHeads, int headSize, + int kvDim, int kvMul, IntArray positionHolder, int layer, int contextLength) { + + int tid = context.localIdx; + int h = context.groupIdx; + int localSize = context.localGroupSizeX; + + if (h >= nHeads) { + return; + } + + int pos = positionHolder.get(0); + int loff = layer * contextLength * kvDim; + int kvHeadIdx = h / kvMul; + int BLOCK_SIZE_C = 16; + + float[] q_shared = context.allocateFloatLocalArray(headSize); + float[] k_tile = context.allocateFloatLocalArray(BLOCK_SIZE_C * headSize); + float[] v_tile = context.allocateFloatLocalArray(BLOCK_SIZE_C * headSize); + float[] s_tile = context.allocateFloatLocalArray(BLOCK_SIZE_C); + float[] shared_tile_max_holder = context.allocateFloatLocalArray(1); + + float maxScore = Float.NEGATIVE_INFINITY; + float sumExp = 0.0f; + + float[] output = new float[headSize]; + for (int i = 0; i < headSize; i++) { + output[i] = 0.0f; + } + + for (int i = tid; i < headSize; i += localSize) { + q_shared[i] = q.get(h * headSize + i); + } + + context.localBarrier(); + + for (int tileC = 0; tileC <= pos; tileC += BLOCK_SIZE_C) { + int tileEnd = Math.min(tileC + BLOCK_SIZE_C - 1, pos); + + for (int tIdxInSeq = tileC + tid; tIdxInSeq <= tileEnd; tIdxInSeq += localSize) { + int tileMemOffset = (tIdxInSeq - tileC) * headSize; + int kvBase = loff + tIdxInSeq * kvDim + kvHeadIdx * headSize; + for (int d = 0; d < headSize; d++) { + k_tile[tileMemOffset + d] = key_cache.get(kvBase + d).getFloat32(); + v_tile[tileMemOffset + d] = value_cache.get(kvBase + d).getFloat32(); + } + } + + context.localBarrier(); + + for (int tIdxInSeq = tileC + tid; tIdxInSeq <= tileEnd; tIdxInSeq += localSize) { + int score_idx_in_tile = tIdxInSeq - tileC; + + float score = 0.0f; + for (int d = 0; d < headSize; d++) { + score += q_shared[d] * k_tile[score_idx_in_tile * headSize + d]; + } + score /= TornadoMath.sqrt(headSize); + s_tile[score_idx_in_tile] = score; + } + + context.localBarrier(); + + float tileLocalMax = Float.NEGATIVE_INFINITY; + for (int i = 0; i <= tileEnd - tileC; i++) { + if (s_tile[i] > tileLocalMax) { + tileLocalMax = s_tile[i]; + } + } + + if (tid == 0) { + shared_tile_max_holder[0] = tileLocalMax; + } + context.localBarrier(); + float currentTileMax = shared_tile_max_holder[0]; + + float newMax = Math.max(maxScore, currentTileMax); + if (newMax != maxScore && maxScore != Float.NEGATIVE_INFINITY) { + float scale = TornadoMath.exp(maxScore - newMax); + sumExp *= scale; + for (int d = 0; d < headSize; d++) { + output[d] *= scale; + } + } + maxScore = newMax; + + for (int t_idx_in_s_tile = 0; t_idx_in_s_tile <= tileEnd - tileC; t_idx_in_s_tile++) { + float expScore = TornadoMath.exp(s_tile[t_idx_in_s_tile] - maxScore); + sumExp += expScore; + + for (int d = 0; d < headSize; d++) { + output[d] += expScore * v_tile[t_idx_in_s_tile * headSize + d]; + } + } + context.localBarrier(); + } + + float normFactor = (sumExp > 0.0f) ? (1.0f / sumExp) : 0.0f; + for (int d = tid; d < headSize; d += localSize) { + xb.set(h * headSize + d, output[d] * normFactor); + } + } + public static void processHeadsFlashAttentionOptV2(KernelContext context, FloatArray q, FloatArray key_cache, FloatArray value_cache, FloatArray xb, int nHeads, int headSize, // NOTE: Still used for logic, but not for allocation size int kvDim, int kvMul, IntArray positionHolder, int layer, int contextLength) { @@ -1356,6 +1619,126 @@ public static void processHeadsFlashAttentionSplitKV(KernelContext context, Floa } } + /** + * FP16 KV cache variant of {@link #processHeadsFlashAttentionSplitKV}. K/V rows are read from + * the half-precision caches with packed 32-bit loads and expanded to FP32; the online-softmax + * accumulation stays FP32. Pairs with {@link #combineSplitKVAttention} unchanged. + */ + public static void processHeadsFlashAttentionSplitKVFP16(KernelContext context, FloatArray q, HalfFloatArray key_cache, HalfFloatArray value_cache, FloatArray att, int nHeads, int headSize, + int kvDim, int kvMul, IntArray positionHolder, int layer, int contextLength, int nSplits) { + + final int MAX_HEAD_SIZE = 128; + final int MAX_LOCAL_SIZE = 64; + + int tid = context.localIdx; + int g = context.groupIdx; // 0 .. nHeads*nSplits - 1 + int localSize = context.localGroupSizeX; + int h = g / nSplits; + int s = g % nSplits; + + if (h >= nHeads) { + return; + } + + int pos = positionHolder.get(0); + int seqLen = pos + 1; + int chunk = (seqLen + nSplits - 1) / nSplits; + int startPos = s * chunk; + int endPos = Math.min(startPos + chunk, seqLen); // exclusive + + int loff = layer * contextLength * kvDim; + int kvHeadIdx = h / kvMul; + float invSqrt = 1.0f / TornadoMath.sqrt(headSize); + + float[] q_shared = context.allocateFloatLocalArray(MAX_HEAD_SIZE); + float[] accShared = context.allocateFloatLocalArray(MAX_LOCAL_SIZE * MAX_HEAD_SIZE); + float[] mShared = context.allocateFloatLocalArray(MAX_LOCAL_SIZE); + float[] lShared = context.allocateFloatLocalArray(MAX_LOCAL_SIZE); + float[] corrShared = context.allocateFloatLocalArray(MAX_LOCAL_SIZE); + float[] bcast = context.allocateFloatLocalArray(1); + + int headBase = h * nSplits * (headSize + 2); + int outBase = headBase + s * headSize; + int mBase = headBase + nSplits * headSize; + int lBase = mBase + nSplits; + + for (int i = tid; i < headSize; i += localSize) { + q_shared[i] = q.get(h * headSize + i); + } + int rowBase = tid * headSize; + for (int d = 0; d < headSize; d++) { + accShared[rowBase + d] = 0.0f; + } + context.localBarrier(); + + // Strided scan over this split's position chunk (no barriers). + float m = Float.NEGATIVE_INFINITY; + float l = 0.0f; + for (int p = startPos + tid; p < endPos; p += localSize) { + int base = loff + p * kvDim + kvHeadIdx * headSize; + float score = 0.0f; + for (int d = 0; d < headSize; d += 2) { + Half2 kPair = key_cache.getHalf2(base + d); + score += q_shared[d] * Half2.lowFloat(kPair); + score += q_shared[d + 1] * Half2.highFloat(kPair); + } + score *= invSqrt; + float newM = Math.max(m, score); + float corr = (m == Float.NEGATIVE_INFINITY) ? 0.0f : TornadoMath.exp(m - newM); + float e = TornadoMath.exp(score - newM); + for (int d = 0; d < headSize; d += 2) { + Half2 vPair = value_cache.getHalf2(base + d); + accShared[rowBase + d] = accShared[rowBase + d] * corr + e * Half2.lowFloat(vPair); + accShared[rowBase + d + 1] = accShared[rowBase + d + 1] * corr + e * Half2.highFloat(vPair); + } + l = l * corr + e; + m = newM; + } + mShared[tid] = m; + lShared[tid] = l; + context.localBarrier(); + + // Block max. + if (tid == 0) { + float blockMax = Float.NEGATIVE_INFINITY; + for (int t = 0; t < localSize; t++) { + if (mShared[t] > blockMax) { + blockMax = mShared[t]; + } + } + bcast[0] = blockMax; + } + context.localBarrier(); + float M = bcast[0]; + + corrShared[tid] = (mShared[tid] == Float.NEGATIVE_INFINITY) ? 0.0f : TornadoMath.exp(mShared[tid] - M); + context.localBarrier(); + + // Block sum L = Σ_t l_t · corr_t. + if (tid == 0) { + float blockSum = 0.0f; + for (int t = 0; t < localSize; t++) { + blockSum += lShared[t] * corrShared[t]; + } + bcast[0] = blockSum; + } + context.localBarrier(); + float L = bcast[0]; + + // Write UNNORMALIZED partial numerator (relative to block max M), plus M and L for the combine. + for (int d = tid; d < headSize; d += localSize) { + float acc = 0.0f; + for (int t = 0; t < localSize; t++) { + acc += corrShared[t] * accShared[t * headSize + d]; + } + att.set(outBase + d, acc); + } + if (tid == 0) { + att.set(mBase + s, M); + att.set(lBase + s, L); + } + } + /** * Qwen3-family decode attention, split-KV (flash-decoding) phase 2: combine. * diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/layers/AbstractTransformerLayerTaskGraphs.java b/src/main/java/org/beehive/gpullama3/tornadovm/layers/AbstractTransformerLayerTaskGraphs.java index 14a71e88..1b4ee532 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/layers/AbstractTransformerLayerTaskGraphs.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/layers/AbstractTransformerLayerTaskGraphs.java @@ -91,4 +91,15 @@ public String getLastFFNLayerTaskGraphID() { protected boolean shouldUseFinalNormalization() { return schedulerType == SchedulerType.NON_NVIDIA; } + + /** + * Whether this layer stack should use the half-precision KV cache: requested via + * {@code -Dllama.kvcache.fp16=true}, allocated by the model state, and running on the NVIDIA + * path (the FP16 kernels rely on packed half2 codegen in the CUDA backend). The packed + * accessors need even element indices, which holds for the (even) headSize/kvDim of the + * supported models. + */ + protected boolean useFp16KVCache() { + return State.USE_FP16_KV && state.wrapKeyCacheFP16 != null && schedulerType == SchedulerType.NVIDIA; + } } diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/LlamaFP16FFNLayers.java b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/LlamaFP16FFNLayers.java index 84379ee9..68e37cda 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/LlamaFP16FFNLayers.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/LlamaFP16FFNLayers.java @@ -203,19 +203,35 @@ protected TaskGraph createFFNLayerTaskGraph(int layerIndex) { LOCAL_WORK_GROUP_SIZE_ALLOC); // RoPE + KV Cache - unifiedLayer.task("rope_and_kv_cache", - TransformerComputeKernelsLayered::ropeRotationWithCacheCopy, - context, - state.positionHolder, - state.wrapQ, // Q (in/out) - state.wrapK, // K (in/out) - state.wrapV, // V (in only) - state.wrapKeyCache, // Key cache (out) - state.wrapValueCache, // Value cache (out) - config.kvDim(), - config.headSize(), - layerIndex, - config.contextLength()); + if (useFp16KVCache()) { + unifiedLayer.task("rope_and_kv_cache", + TransformerComputeKernelsLayered::ropeRotationWithCacheCopyFP16, + context, + state.positionHolder, + state.wrapQ, // Q (in/out) + state.wrapK, // K (in/out) + state.wrapV, // V (in only) + state.wrapKeyCacheFP16, // Key cache (out, FP16) + state.wrapValueCacheFP16, // Value cache (out, FP16) + config.kvDim(), + config.headSize(), + layerIndex, + config.contextLength()); + } else { + unifiedLayer.task("rope_and_kv_cache", + TransformerComputeKernelsLayered::ropeRotationWithCacheCopy, + context, + state.positionHolder, + state.wrapQ, // Q (in/out) + state.wrapK, // K (in/out) + state.wrapV, // V (in only) + state.wrapKeyCache, // Key cache (out) + state.wrapValueCache, // Value cache (out) + config.kvDim(), + config.headSize(), + layerIndex, + config.contextLength()); + } // Attention configureAttention(unifiedLayer, layerIndex); // Output Projection (Wo) with residual @@ -258,8 +274,13 @@ protected TaskGraph createFFNLayerTaskGraph(int layerIndex) { weights.w2Layered[layerIndex].asHalfFloatArray(), config.hiddenDim(), config.dim(), LOCAL_WORK_GROUP_SIZE_ALLOC); - unifiedLayer.persistOnDevice(state.wrapX, state.wrapKeyCache, - state.wrapValueCache); + if (useFp16KVCache()) { + unifiedLayer.persistOnDevice(state.wrapX, state.wrapKeyCacheFP16, + state.wrapValueCacheFP16); + } else { + unifiedLayer.persistOnDevice(state.wrapX, state.wrapKeyCache, + state.wrapValueCache); + } return unifiedLayer; } @@ -284,6 +305,8 @@ protected String predecessorGraphName(int layerIndex) { } protected TaskGraph configureLayerDataTransfers(TaskGraph unifiedLayer, int layerIndex) { + Object keyCache = useFp16KVCache() ? state.wrapKeyCacheFP16 : state.wrapKeyCache; + Object valueCache = useFp16KVCache() ? state.wrapValueCacheFP16 : state.wrapValueCache; if (layerIndex == 0) { // First layer: Transfer initial data to device (one-time transfer) unifiedLayer.transferToDevice(DataTransferMode.EVERY_EXECUTION, @@ -298,7 +321,7 @@ protected TaskGraph configureLayerDataTransfers(TaskGraph unifiedLayer, int laye // QKV vectors state.wrapQ, state.wrapK, state.wrapV, // KV cache - state.wrapKeyCache, state.wrapValueCache, + keyCache, valueCache, // Attention & FFN buffers state.wrapAtt, state.wrapHb, state.wrapXbFP16); } else { @@ -314,7 +337,7 @@ protected TaskGraph configureLayerDataTransfers(TaskGraph unifiedLayer, int laye // QKV vectors state.wrapQ, state.wrapK, state.wrapV, // KV cache - state.wrapKeyCache, state.wrapValueCache, + keyCache, valueCache, // Attention & FFN buffers state.wrapAtt, state.wrapHb, // Position & misc @@ -324,7 +347,40 @@ protected TaskGraph configureLayerDataTransfers(TaskGraph unifiedLayer, int laye } private TaskGraph configureAttention(TaskGraph unifiedLayer, int layerIndex) { - if (schedulerType == SchedulerType.NVIDIA) { + if (useFp16KVCache()) { + // Flash Attention over the half-precision KV cache (FP32 accumulation). + // The scalar-read variant is an evaluation aid; the packed variant is the default. + if (State.FP16_KV_SCALAR) { + return unifiedLayer.task("attention", + TransformerComputeKernelsLayered::processHeadsFlashAttentionFP16Scalar, + context, + state.wrapQ, + state.wrapKeyCacheFP16, + state.wrapValueCacheFP16, + state.wrapXb, + config.numberOfHeads(), + config.headSize(), + config.kvDim(), + config.kvMul(), + state.positionHolder, + layerIndex, + config.contextLength()); + } + return unifiedLayer.task("attention", + TransformerComputeKernelsLayered::processHeadsFlashAttentionFP16, + context, + state.wrapQ, + state.wrapKeyCacheFP16, + state.wrapValueCacheFP16, + state.wrapXb, + config.numberOfHeads(), + config.headSize(), + config.kvDim(), + config.kvMul(), + state.positionHolder, + layerIndex, + config.contextLength()); + } else if (schedulerType == SchedulerType.NVIDIA) { // Flash Attention (optimized for NVIDIA GPUs) return unifiedLayer.task("attention", TransformerComputeKernelsLayered::processHeadsFlashAttention, diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/Qwen3FP16FFNLayers.java b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/Qwen3FP16FFNLayers.java index f0da8b54..ddfadf66 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/Qwen3FP16FFNLayers.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/Qwen3FP16FFNLayers.java @@ -310,38 +310,73 @@ protected TaskGraph createFFNLayerTaskGraph(int layerIndex) { config.rmsNormEps()); // epsilon // Fused RoPE Rotation + KV Cache Write - unifiedLayer.task("rope_and_kv_cache", - Qwen3Kernels::ropeRotationWithCacheCopy, - context, - qwen3State.positionHolder, // current position - qwen3State.wrapQ, // Q vectors (in/out, rotated) - qwen3State.wrapK, // K vectors (in/out, rotated) - qwen3State.wrapV, // V vectors (in only) - qwen3State.wrapKeyCache, // key cache (out) - qwen3State.wrapValueCache, // value cache (out) - config.numberOfKeyValueHeads(), // nHeadKv - nEmbdHead, // head dimension - nEmbdGqa, // kvDim - layerIndex, // layer index for cache offset - config.contextLength()); // max sequence length + if (useFp16KVCache()) { + unifiedLayer.task("rope_and_kv_cache", + Qwen3Kernels::ropeRotationWithCacheCopyFP16, + context, + qwen3State.positionHolder, // current position + qwen3State.wrapQ, // Q vectors (in/out, rotated) + qwen3State.wrapK, // K vectors (in/out, rotated) + qwen3State.wrapV, // V vectors (in only) + qwen3State.wrapKeyCacheFP16, // key cache (out, FP16) + qwen3State.wrapValueCacheFP16, // value cache (out, FP16) + config.numberOfKeyValueHeads(), // nHeadKv + nEmbdHead, // head dimension + nEmbdGqa, // kvDim + layerIndex, // layer index for cache offset + config.contextLength()); // max sequence length + } else { + unifiedLayer.task("rope_and_kv_cache", + Qwen3Kernels::ropeRotationWithCacheCopy, + context, + qwen3State.positionHolder, // current position + qwen3State.wrapQ, // Q vectors (in/out, rotated) + qwen3State.wrapK, // K vectors (in/out, rotated) + qwen3State.wrapV, // V vectors (in only) + qwen3State.wrapKeyCache, // key cache (out) + qwen3State.wrapValueCache, // value cache (out) + config.numberOfKeyValueHeads(), // nHeadKv + nEmbdHead, // head dimension + nEmbdGqa, // kvDim + layerIndex, // layer index for cache offset + config.contextLength()); // max sequence length + } // Split-KV (flash-decoding) attention. // Phase 1: split each head's KV range across attentionSplits workgroups; partials -> wrapAttSplit. - unifiedLayer.task("attention", - TransformerComputeKernelsLayered::processHeadsFlashAttentionSplitKV, - context, - qwen3State.wrapQ, // query vectors - qwen3State.wrapKeyCache, // key cache - qwen3State.wrapValueCache, // value cache - qwen3State.wrapAttSplit, // scratch: per-head split partials (compact layout) - config.numberOfHeads(), // nHeads - nEmbdHead, // headSize - nEmbdGqa, // kvDim - gqa, // kvMul (nHeads / nHeadKv) - qwen3State.positionHolder, // position - layerIndex, // layer index - config.contextLength(), // context length - attentionSplits); // number of KV splits per head + if (useFp16KVCache()) { + unifiedLayer.task("attention", + TransformerComputeKernelsLayered::processHeadsFlashAttentionSplitKVFP16, + context, + qwen3State.wrapQ, // query vectors + qwen3State.wrapKeyCacheFP16, // key cache (FP16) + qwen3State.wrapValueCacheFP16, // value cache (FP16) + qwen3State.wrapAttSplit, // scratch: per-head split partials (compact layout) + config.numberOfHeads(), // nHeads + nEmbdHead, // headSize + nEmbdGqa, // kvDim + gqa, // kvMul (nHeads / nHeadKv) + qwen3State.positionHolder, // position + layerIndex, // layer index + config.contextLength(), // context length + attentionSplits); // number of KV splits per head + } else { + unifiedLayer.task("attention", + TransformerComputeKernelsLayered::processHeadsFlashAttentionSplitKV, + context, + qwen3State.wrapQ, // query vectors + qwen3State.wrapKeyCache, // key cache + qwen3State.wrapValueCache, // value cache + qwen3State.wrapAttSplit, // scratch: per-head split partials (compact layout) + config.numberOfHeads(), // nHeads + nEmbdHead, // headSize + nEmbdGqa, // kvDim + gqa, // kvMul (nHeads / nHeadKv) + qwen3State.positionHolder, // position + layerIndex, // layer index + config.contextLength(), // context length + attentionSplits); // number of KV splits per head + } // Phase 2: combine the per-head split partials into the final attention output -> wrapXb. unifiedLayer.task("attention_combine", TransformerComputeKernelsLayered::combineSplitKVAttention, @@ -452,7 +487,11 @@ protected TaskGraph createFFNLayerTaskGraph(int layerIndex) { config.dim(), // output dim LOCAL_WORK_GROUP_SIZE_ALLOC); } - unifiedLayer.persistOnDevice(qwen3State.wrapX, qwen3State.wrapKeyCache, qwen3State.wrapValueCache); + if (useFp16KVCache()) { + unifiedLayer.persistOnDevice(qwen3State.wrapX, qwen3State.wrapKeyCacheFP16, qwen3State.wrapValueCacheFP16); + } else { + unifiedLayer.persistOnDevice(qwen3State.wrapX, qwen3State.wrapKeyCache, qwen3State.wrapValueCache); + } return unifiedLayer; } @@ -477,6 +516,8 @@ protected String predecessorGraphName(int layerIndex) { * Configure data transfers for first and subsequent layers */ protected TaskGraph configureLayerDataTransfers(TaskGraph unifiedLayer, int layerIndex) { + Object keyCache = useFp16KVCache() ? qwen3State.wrapKeyCacheFP16 : qwen3State.wrapKeyCache; + Object valueCache = useFp16KVCache() ? qwen3State.wrapValueCacheFP16 : qwen3State.wrapValueCache; if (layerIndex == 0) { // First layer: Transfer temporary buffers and QKV state every execution unifiedLayer.transferToDevice(DataTransferMode.EVERY_EXECUTION, qwen3State.positionHolder); @@ -485,7 +526,7 @@ protected TaskGraph configureLayerDataTransfers(TaskGraph unifiedLayer, int laye unifiedLayer.transferToDevice(DataTransferMode.FIRST_EXECUTION, // context, qwen3State.wrapXb, qwen3State.wrapXb2, // qwen3State.wrapQ, qwen3State.wrapK, qwen3State.wrapV, // - qwen3State.wrapKeyCache, qwen3State.wrapValueCache, // + keyCache, valueCache, // qwen3State.wrapAtt, qwen3State.wrapHb ); unifiedLayer.transferToDevice(DataTransferMode.FIRST_EXECUTION, qwen3State.wrapAttSplit); } else { @@ -496,8 +537,8 @@ protected TaskGraph configureLayerDataTransfers(TaskGraph unifiedLayer, int laye String pred = "layer_" + (layerIndex - 1); unifiedLayer.consumeFromDevice(pred, context, qwen3State.wrapXb, qwen3State.wrapXb2, // qwen3State.wrapQ, qwen3State.wrapK, // - qwen3State.wrapV, qwen3State.wrapKeyCache, // - qwen3State.wrapValueCache, qwen3State.wrapAtt, // + qwen3State.wrapV, keyCache, // + valueCache, qwen3State.wrapAtt, // qwen3State.wrapHb, qwen3State.positionHolder); // unifiedLayer.consumeFromDevice(pred, qwen3State.wrapAttSplit); } diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/decode/LlamaFP16FFNLayersDecode.java b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/decode/LlamaFP16FFNLayersDecode.java index a8010433..1d09ef01 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/decode/LlamaFP16FFNLayersDecode.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/decode/LlamaFP16FFNLayersDecode.java @@ -37,6 +37,15 @@ public LlamaFP16FFNLayersDecode(String taskGraph, LlamaState state, * Must match the {@code TaskGraph} names used in * {@code buildDecodeActivationGraph()} and {@code createFFNLayerTaskGraph()}.

*/ + /** + * The prefill/decode graph variants share the FP32 KV cache with the batch-prefill layers, + * so the FP16 KV cache path (standard single-token mode only) is disabled here. + */ + @Override + protected boolean useFp16KVCache() { + return false; + } + @Override protected String predecessorGraphName(int layerIndex) { return (layerIndex == 0) ? "decodeActivation" : "layer_" + (layerIndex - 1); diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/decode/LlamaFP16FFNLayersPrefillDecode.java b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/decode/LlamaFP16FFNLayersPrefillDecode.java index 4a74e69b..2246ad0a 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/decode/LlamaFP16FFNLayersPrefillDecode.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/decode/LlamaFP16FFNLayersPrefillDecode.java @@ -43,6 +43,15 @@ public LlamaFP16FFNLayersPrefillDecode(String taskGraph, LlamaState state, * Layer 0 receives {@code wrapX} from the decode activation graph; * layers 1+ receive it from the previous decode layer. */ + /** + * The prefill/decode graph variants share the FP32 KV cache with the batch-prefill layers, + * so the FP16 KV cache path (standard single-token mode only) is disabled here. + */ + @Override + protected boolean useFp16KVCache() { + return false; + } + @Override protected String predecessorGraphName(int layerIndex) { return (layerIndex == 0) ? "decodeActivation" : "layer_" + (layerIndex - 1); diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/decode/Qwen3FP16FFNLayersDecode.java b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/decode/Qwen3FP16FFNLayersDecode.java index 1e32485e..e25ebd68 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/decode/Qwen3FP16FFNLayersDecode.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/decode/Qwen3FP16FFNLayersDecode.java @@ -25,6 +25,15 @@ public Qwen3FP16FFNLayersDecode(String taskGraph, Qwen3State state, super(taskGraph, state, weights, config, schedulerType); } + /** + * The prefill/decode graph variants share the FP32 KV cache with the batch-prefill layers, + * so the FP16 KV cache path (standard single-token mode only) is disabled here. + */ + @Override + protected boolean useFp16KVCache() { + return false; + } + @Override protected String predecessorGraphName(int layerIndex) { return (layerIndex == 0) ? "decodeActivation" : "layer_" + (layerIndex - 1); diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/decode/Qwen3FP16FFNLayersPrefillDecode.java b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/decode/Qwen3FP16FFNLayersPrefillDecode.java index f6cd60b0..f143b415 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/decode/Qwen3FP16FFNLayersPrefillDecode.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/decode/Qwen3FP16FFNLayersPrefillDecode.java @@ -25,6 +25,15 @@ public Qwen3FP16FFNLayersPrefillDecode(String taskGraph, Qwen3State state, super(taskGraph, state, weights, config, schedulerType); } + /** + * The prefill/decode graph variants share the FP32 KV cache with the batch-prefill layers, + * so the FP16 KV cache path (standard single-token mode only) is disabled here. + */ + @Override + protected boolean useFp16KVCache() { + return false; + } + @Override protected String predecessorGraphName(int layerIndex) { return (layerIndex == 0) ? "decodeActivation" : "layer_" + (layerIndex - 1); From b7ae254dddd40f0aa222aabd10a12a67644c955a Mon Sep 17 00:00:00 2001 From: Orion Papadakis Date: Wed, 22 Jul 2026 18:33:26 +0300 Subject: [PATCH 2/5] Consume FP16 KV cache in decode and batch-decode task graphs Route the persistOnDevice/consumeFromDevice calls in the decode and batch-decode graphs to the FP16 key/value buffers when the FP16 KV cache is active, so the flag-gated cache added for single-token decode is also honoured on the prefill-decode and batch paths. --- .../fp16/decode/LlamaFP16FFNLayersDecode.java | 22 +++++++++---------- .../fp16/decode/LogitsFP16LayerDecode.java | 17 ++++++++++---- .../fp16/decode/Qwen3FP16FFNLayersDecode.java | 16 ++++---------- .../activation/BatchDecodeActivation.java | 7 ++++-- 4 files changed, 33 insertions(+), 29 deletions(-) diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/decode/LlamaFP16FFNLayersDecode.java b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/decode/LlamaFP16FFNLayersDecode.java index 1d09ef01..7dfe49ca 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/decode/LlamaFP16FFNLayersDecode.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/decode/LlamaFP16FFNLayersDecode.java @@ -37,15 +37,6 @@ public LlamaFP16FFNLayersDecode(String taskGraph, LlamaState state, * Must match the {@code TaskGraph} names used in * {@code buildDecodeActivationGraph()} and {@code createFFNLayerTaskGraph()}.

*/ - /** - * The prefill/decode graph variants share the FP32 KV cache with the batch-prefill layers, - * so the FP16 KV cache path (standard single-token mode only) is disabled here. - */ - @Override - protected boolean useFp16KVCache() { - return false; - } - @Override protected String predecessorGraphName(int layerIndex) { return (layerIndex == 0) ? "decodeActivation" : "layer_" + (layerIndex - 1); @@ -53,6 +44,9 @@ protected String predecessorGraphName(int layerIndex) { @Override protected TaskGraph configureLayerDataTransfers(TaskGraph layer, int layerIndex) { + LlamaState llamaState = (LlamaState) state; + Object keyCache = useFp16KVCache() ? state.wrapKeyCacheFP16 : state.wrapKeyCache; + Object valueCache = useFp16KVCache() ? state.wrapValueCacheFP16 : state.wrapValueCache; if (layerIndex == 0) { // Same as parent layer 0, but wrapKeyCache/wrapValueCache come from device // (passed through by the decode activation graph, which relays them from @@ -64,8 +58,11 @@ protected TaskGraph configureLayerDataTransfers(TaskGraph layer, int layerIndex) state.wrapXb, state.wrapXb2, state.wrapQ, state.wrapK, state.wrapV, state.wrapAtt, state.wrapHb, state.wrapXbFP16); + if (splitKvAttentionEnabled()) { + layer.transferToDevice(DataTransferMode.FIRST_EXECUTION, llamaState.wrapAttSplit); + } // Explicit source — must match the TaskGraph name in buildDecodeActivationGraph(). - layer.consumeFromDevice("decodeActivation", state.wrapKeyCache, state.wrapValueCache); + layer.consumeFromDevice("decodeActivation", keyCache, valueCache); } else { // Layers 1+: use explicit predecessor name for ALL consumed objects. // Calling super here would use the no-arg form (source key = own graph name), @@ -75,9 +72,12 @@ protected TaskGraph configureLayerDataTransfers(TaskGraph layer, int layerIndex) context, state.wrapXb, state.wrapXb2, state.wrapQ, state.wrapK, state.wrapV, - state.wrapKeyCache, state.wrapValueCache, + keyCache, valueCache, state.wrapAtt, state.wrapHb, state.positionHolder, state.wrapXbFP16); + if (splitKvAttentionEnabled()) { + layer.consumeFromDevice(pred, llamaState.wrapAttSplit); + } } return layer; } diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/decode/LogitsFP16LayerDecode.java b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/decode/LogitsFP16LayerDecode.java index 7ee529d0..6780d355 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/decode/LogitsFP16LayerDecode.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/decode/LogitsFP16LayerDecode.java @@ -33,8 +33,17 @@ public LogitsFP16LayerDecode(String name, State state, Weights weights, Configur super(name, state, weights, config, lastTaskGraphID, schedulerType); } + /** The KV cache objects the decode layers actually persist (FP16 when that path is active). */ + private Object keyCache() { + return (State.USE_FP16_KV && state.wrapKeyCacheFP16 != null) ? state.wrapKeyCacheFP16 : state.wrapKeyCache; + } + + private Object valueCache() { + return (State.USE_FP16_KV && state.wrapValueCacheFP16 != null) ? state.wrapValueCacheFP16 : state.wrapValueCache; + } + /** - * Prepends {@code consumeFromDevice(lastTaskGraphID, wrapKeyCache, wrapValueCache)} before all tasks. + * Prepends {@code consumeFromDevice(lastTaskGraphID, keyCache, valueCache)} before all tasks. * *

Must use the named-source form so that {@code updatePersistedObjectState()} adds the KV cache * to the source-keyed map. Without the source name, the fallback in {@code updatePersistedObjectState} @@ -43,12 +52,12 @@ public LogitsFP16LayerDecode(String name, State state, Weights weights, Configur */ @Override protected void configureAdditionalConsumes(TaskGraph logits) { - logits.consumeFromDevice(lastTaskGraphID, state.wrapKeyCache, state.wrapValueCache); + logits.consumeFromDevice(lastTaskGraphID, keyCache(), valueCache()); } - /** Appends {@code persistOnDevice(wrapKeyCache, wrapValueCache)} after {@code transferToHost}. */ + /** Appends {@code persistOnDevice(keyCache, valueCache)} after {@code transferToHost}. */ @Override protected void configureAdditionalPersists(TaskGraph logits) { - logits.persistOnDevice(state.wrapKeyCache, state.wrapValueCache); + logits.persistOnDevice(keyCache(), valueCache()); } } diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/decode/Qwen3FP16FFNLayersDecode.java b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/decode/Qwen3FP16FFNLayersDecode.java index e25ebd68..9362c382 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/decode/Qwen3FP16FFNLayersDecode.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/decode/Qwen3FP16FFNLayersDecode.java @@ -25,15 +25,6 @@ public Qwen3FP16FFNLayersDecode(String taskGraph, Qwen3State state, super(taskGraph, state, weights, config, schedulerType); } - /** - * The prefill/decode graph variants share the FP32 KV cache with the batch-prefill layers, - * so the FP16 KV cache path (standard single-token mode only) is disabled here. - */ - @Override - protected boolean useFp16KVCache() { - return false; - } - @Override protected String predecessorGraphName(int layerIndex) { return (layerIndex == 0) ? "decodeActivation" : "layer_" + (layerIndex - 1); @@ -41,6 +32,8 @@ protected String predecessorGraphName(int layerIndex) { @Override protected TaskGraph configureLayerDataTransfers(TaskGraph layer, int layerIndex) { + Object keyCache = useFp16KVCache() ? qwen3State.wrapKeyCacheFP16 : qwen3State.wrapKeyCache; + Object valueCache = useFp16KVCache() ? qwen3State.wrapValueCacheFP16 : qwen3State.wrapValueCache; if (layerIndex == 0) { layer.transferToDevice(DataTransferMode.EVERY_EXECUTION, qwen3State.positionHolder, qwen3State.temp, qwen3State.tempFFN); @@ -51,15 +44,14 @@ protected TaskGraph configureLayerDataTransfers(TaskGraph layer, int layerIndex) qwen3State.wrapAtt, qwen3State.wrapHb); layer.transferToDevice(DataTransferMode.FIRST_EXECUTION, qwen3State.wrapAttSplit); // KV cache already allocated by batch prefill; relay from decode activation graph. - layer.consumeFromDevice("decodeActivation", - qwen3State.wrapKeyCache, qwen3State.wrapValueCache); + layer.consumeFromDevice("decodeActivation", keyCache, valueCache); } else { String pred = "layer_" + (layerIndex - 1); layer.consumeFromDevice(pred, context, qwen3State.wrapXb, qwen3State.wrapXb2, qwen3State.wrapQ, qwen3State.wrapK, qwen3State.wrapV, - qwen3State.wrapKeyCache, qwen3State.wrapValueCache, + keyCache, valueCache, qwen3State.wrapAtt, qwen3State.wrapHb, qwen3State.positionHolder, qwen3State.temp, qwen3State.tempFFN); diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/plan/components/activation/BatchDecodeActivation.java b/src/main/java/org/beehive/gpullama3/tornadovm/plan/components/activation/BatchDecodeActivation.java index ea9fb1fe..20021beb 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/plan/components/activation/BatchDecodeActivation.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/plan/components/activation/BatchDecodeActivation.java @@ -35,8 +35,11 @@ public BatchDecodeActivation(State state, Configuration config, String lastBatch // @formatter:off private TaskGraph buildGraph(KernelContext ctx, State state, String lastBatchLayerId, boolean isQ8) { + boolean fp16KV = State.USE_FP16_KV && state.wrapKeyCacheFP16 != null; + Object keyCache = fp16KV ? state.wrapKeyCacheFP16 : state.wrapKeyCache; + Object valueCache = fp16KV ? state.wrapValueCacheFP16 : state.wrapValueCache; TaskGraph tg = new TaskGraph("decodeActivation") - .consumeFromDevice(lastBatchLayerId, state.wrapKeyCache, state.wrapValueCache) + .consumeFromDevice(lastBatchLayerId, keyCache, valueCache) .transferToDevice(DataTransferMode.EVERY_EXECUTION, state.embeddingX); if (isQ8) { tg.task("updateX", TransformerComputeKernels::convertQ8_0toFP32, ctx, @@ -47,7 +50,7 @@ private TaskGraph buildGraph(KernelContext ctx, State state, (HalfFloatArray) state.embeddingX, state.wrapX); } - return tg.persistOnDevice(state.wrapX, state.wrapKeyCache, state.wrapValueCache); + return tg.persistOnDevice(state.wrapX, keyCache, valueCache); } // @formatter:on From dd0dae1b7f0797d935d75cc6c1db12631701784a Mon Sep 17 00:00:00 2001 From: Orion Papadakis Date: Wed, 22 Jul 2026 18:33:52 +0300 Subject: [PATCH 3/5] Add packed half2 split-KV attention (deep-half2) Introduce a -Dllama.attention.deepHalf2 path for the FP16 KV cache: stage Q once per workgroup as a __half2 local-memory tile and accumulate each K/V pair with a single __hfma2, converting to FP32 once per row (llama.cpp fattn-vec style) instead of per pair. Adds the packed single-token and batch-prefill attention kernels for Llama and Qwen3, and makes the split-KV count configurable via -Dllama.attention.splitKv.count. Requires the packed half2 local-array API (allocateHalf2LocalArray), so bump the TornadoVM dependency to 5.1.1-jdk21-dev. --- pom.xml | 6 +- .../gpullama3/inference/state/LlamaState.java | 7 + .../gpullama3/inference/state/Qwen3State.java | 2 +- .../gpullama3/inference/state/State.java | 8 + .../tornadovm/kernels/Qwen3Kernels.java | 112 ++++++- .../TransformerBatchPrefillKernels.java | 295 ++++++++++++++++++ .../TransformerComputeKernelsLayered.java | 197 ++++++++++-- .../layers/type/fp16/LlamaFP16FFNLayers.java | 66 +++- .../layers/type/fp16/Qwen3FP16FFNLayers.java | 5 +- .../LlamaFP16LayersBatchPrefillMMA.java | 69 ++-- .../Qwen3FP16LayersBatchPrefillMMA.java | 65 +++- 11 files changed, 755 insertions(+), 77 deletions(-) diff --git a/pom.xml b/pom.xml index f5dd084f..cfe25015 100644 --- a/pom.xml +++ b/pom.xml @@ -39,8 +39,8 @@ 0.5.0 - 5.0.0 - -jdk21 + 5.1.1 + -jdk21-dev ${tornadovm.base.version}${jdk.version.suffix} @@ -147,7 +147,7 @@ 21 21 - -jdk21 + -jdk21-dev ${tornadovm.base.version}${jdk.version.suffix} diff --git a/src/main/java/org/beehive/gpullama3/inference/state/LlamaState.java b/src/main/java/org/beehive/gpullama3/inference/state/LlamaState.java index baf8534d..98372c21 100644 --- a/src/main/java/org/beehive/gpullama3/inference/state/LlamaState.java +++ b/src/main/java/org/beehive/gpullama3/inference/state/LlamaState.java @@ -22,8 +22,15 @@ */ public final class LlamaState extends State { + /** Number of KV splits per head for opt-in split-KV decode attention. */ + public static final int SPLIT_KV = Integer.getInteger("llama.attention.splitKv.count", 8); + + // Split-KV attention scratch: per (head, split) partial numerator [headSize] plus block max/sum. + public final FloatArray wrapAttSplit; + public LlamaState(Configuration config, int batchsize) { super(config, batchsize); + this.wrapAttSplit = new FloatArray(config.numberOfHeads() * SPLIT_KV * (config.headSize() + 2)); } @Override diff --git a/src/main/java/org/beehive/gpullama3/inference/state/Qwen3State.java b/src/main/java/org/beehive/gpullama3/inference/state/Qwen3State.java index b663e30b..9e1be40f 100644 --- a/src/main/java/org/beehive/gpullama3/inference/state/Qwen3State.java +++ b/src/main/java/org/beehive/gpullama3/inference/state/Qwen3State.java @@ -34,7 +34,7 @@ public final class Qwen3State extends State { public FloatArray wrapAttSplit; /** Number of KV splits per head for split-KV (flash-decoding) decode attention. */ - public static final int SPLIT_KV = 8; + public static final int SPLIT_KV = Integer.getInteger("llama.attention.splitKv.count", 8); public Qwen3State(Configuration config, int batchsize) { super(config, batchsize); 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 31e1f1f2..a787d746 100644 --- a/src/main/java/org/beehive/gpullama3/inference/state/State.java +++ b/src/main/java/org/beehive/gpullama3/inference/state/State.java @@ -40,6 +40,14 @@ public abstract class State { */ public static final boolean FP16_KV_SCALAR = Boolean.getBoolean("llama.kvcache.fp16.scalar"); + /** + * With the FP16 KV cache and split-KV attention active, keep the K·Q score accumulation packed + * ({@code -Dllama.attention.deepHalf2=true}): Q is staged once per workgroup as a __half2 + * local-memory tile and each K pair is consumed with a single __hfma2, converting to FP32 only + * once per row (llama.cpp fattn-vec style) instead of per pair. + */ + public static final boolean ATTENTION_DEEP_HALF2 = Boolean.getBoolean("llama.attention.deepHalf2"); + // current wave of activations public final FloatTensor x; // activation at current time stamp (dim,) public final FloatTensor xb; // same, but inside a residual branch (dim,) 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 ae768956..c72bb35b 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/kernels/Qwen3Kernels.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/kernels/Qwen3Kernels.java @@ -8,6 +8,7 @@ import uk.ac.manchester.tornado.api.types.arrays.FloatArray; import uk.ac.manchester.tornado.api.types.arrays.HalfFloatArray; import uk.ac.manchester.tornado.api.types.arrays.IntArray; +import uk.ac.manchester.tornado.api.types.vectors.Half2; // @formatter:off public class Qwen3Kernels { @@ -456,8 +457,11 @@ public static void fusedQKVMatmul( int rowOffset = rowId * inputDim; float partialSum = 0.0f; - for (int j = localId; j < inputDim; j += localWorkGroupSize) { - partialSum += wq.get(rowOffset + j).getFloat32() * x.get(j); + // Packed FP16 pair loads (single 32-bit load per pair); inputDim is even. + for (int j = localId * 2; j < inputDim; j += localWorkGroupSize * 2) { + Half2 pair = wq.getHalf2(rowOffset + j); + partialSum += Half2.lowFloat(pair) * x.get(j); + partialSum += Half2.highFloat(pair) * x.get(j + 1); } localSum[localId] = partialSum; @@ -480,8 +484,11 @@ public static void fusedQKVMatmul( int rowOffset = kRow * inputDim; float partialSum = 0.0f; - for (int j = localId; j < inputDim; j += localWorkGroupSize) { - partialSum += wk.get(rowOffset + j).getFloat32() * x.get(j); + // Packed FP16 pair loads (single 32-bit load per pair); inputDim is even. + for (int j = localId * 2; j < inputDim; j += localWorkGroupSize * 2) { + Half2 pair = wk.getHalf2(rowOffset + j); + partialSum += Half2.lowFloat(pair) * x.get(j); + partialSum += Half2.highFloat(pair) * x.get(j + 1); } localSum[localId] = partialSum; @@ -504,8 +511,11 @@ public static void fusedQKVMatmul( int rowOffset = vRow * inputDim; float partialSum = 0.0f; - for (int j = localId; j < inputDim; j += localWorkGroupSize) { - partialSum += wv.get(rowOffset + j).getFloat32() * x.get(j); + // Packed FP16 pair loads (single 32-bit load per pair); inputDim is even. + for (int j = localId * 2; j < inputDim; j += localWorkGroupSize * 2) { + Half2 pair = wv.getHalf2(rowOffset + j); + partialSum += Half2.lowFloat(pair) * x.get(j); + partialSum += Half2.highFloat(pair) * x.get(j + 1); } localSum[localId] = partialSum; @@ -557,9 +567,13 @@ public static void fusedRmsNormQKVMatmul( int rowOffset = rowId * inputDim; float partialSum = 0.0f; - for (int j = localId; j < inputDim; j += localWorkGroupSize) { - float normalized = rmsWeights.get(j) * scale * x.get(j); - partialSum += wq.get(rowOffset + j).getFloat32() * normalized; + // Packed FP16 pair loads (single 32-bit load per pair); inputDim is even. + for (int j = localId * 2; j < inputDim; j += localWorkGroupSize * 2) { + float normalized0 = rmsWeights.get(j) * scale * x.get(j); + float normalized1 = rmsWeights.get(j + 1) * scale * x.get(j + 1); + Half2 pair = wq.getHalf2(rowOffset + j); + partialSum += Half2.lowFloat(pair) * normalized0; + partialSum += Half2.highFloat(pair) * normalized1; } localSum[localId] = partialSum; @@ -582,9 +596,13 @@ public static void fusedRmsNormQKVMatmul( int rowOffset = kRow * inputDim; float partialSum = 0.0f; - for (int j = localId; j < inputDim; j += localWorkGroupSize) { - float normalized = rmsWeights.get(j) * scale * x.get(j); - partialSum += wk.get(rowOffset + j).getFloat32() * normalized; + // Packed FP16 pair loads (single 32-bit load per pair); inputDim is even. + for (int j = localId * 2; j < inputDim; j += localWorkGroupSize * 2) { + float normalized0 = rmsWeights.get(j) * scale * x.get(j); + float normalized1 = rmsWeights.get(j + 1) * scale * x.get(j + 1); + Half2 pair = wk.getHalf2(rowOffset + j); + partialSum += Half2.lowFloat(pair) * normalized0; + partialSum += Half2.highFloat(pair) * normalized1; } localSum[localId] = partialSum; @@ -607,9 +625,13 @@ public static void fusedRmsNormQKVMatmul( int rowOffset = vRow * inputDim; float partialSum = 0.0f; - for (int j = localId; j < inputDim; j += localWorkGroupSize) { - float normalized = rmsWeights.get(j) * scale * x.get(j); - partialSum += wv.get(rowOffset + j).getFloat32() * normalized; + // Packed FP16 pair loads (single 32-bit load per pair); inputDim is even. + for (int j = localId * 2; j < inputDim; j += localWorkGroupSize * 2) { + float normalized0 = rmsWeights.get(j) * scale * x.get(j); + float normalized1 = rmsWeights.get(j + 1) * scale * x.get(j + 1); + Half2 pair = wv.getHalf2(rowOffset + j); + partialSum += Half2.lowFloat(pair) * normalized0; + partialSum += Half2.highFloat(pair) * normalized1; } localSum[localId] = partialSum; @@ -1575,5 +1597,65 @@ public static void batchedRopeWithKVCacheQwen3Packed( } } + /** + * {@link #batchedRopeWithKVCacheQwen3Packed} writing a half-precision KV cache. + * + *

Qwen3's split-half RoPE pairs element {@code ic} with {@code ic + nEmbdHead/2}, + * which are not adjacent, so the cache writes stay scalar.

+ */ + public static void batchedRopeWithKVCacheQwen3PackedFP16( + KernelContext context, + IntArray batchStartPosHolder, + FloatArray qkvBatch, + HalfFloatArray wrapKeyCache, + HalfFloatArray 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; + + 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 the half-precision 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, new HalfFloat(rotK0)); + wrapKeyCache.set(cacheOff + ic + halfEmbdHead, new HalfFloat(rotK1)); + wrapValueCache.set(cacheOff + ic, new HalfFloat(qkvBatch.get(vHeadBase + ic))); + wrapValueCache.set(cacheOff + ic + halfEmbdHead, new HalfFloat(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 9d85c936..3d4184e7 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/kernels/TransformerBatchPrefillKernels.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/kernels/TransformerBatchPrefillKernels.java @@ -8,6 +8,7 @@ import uk.ac.manchester.tornado.api.types.arrays.FloatArray; import uk.ac.manchester.tornado.api.types.arrays.HalfFloatArray; import uk.ac.manchester.tornado.api.types.arrays.IntArray; +import uk.ac.manchester.tornado.api.types.vectors.Half2; /** * GPU kernels for batched prefill. @@ -1518,6 +1519,300 @@ public static void batchedFlashAttentionFP16Out(KernelContext context, } } + // ── FP16 KV cache variants (batched prefill) ───────────────────────────── + + /** + * {@link #batchedRopeWithKVCachePacked} writing a half-precision KV cache. + * + *

K/V pairs (i, i+1) are adjacent and i is even, so each write is a single + * packed 32-bit store.

+ * + * Worker: B*(dim/2) global threads, localSize=512 (or less). + */ + public static void batchedRopeWithKVCachePackedFP16(KernelContext context, + IntArray batchStartPosHolder, + FloatArray qkvBatch, + HalfFloatArray wrapKeyCache, + HalfFloatArray 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 the half-precision 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.setHalf2(cacheOff + i, Half2.fromFloats(rotK0, rotK1)); + wrapValueCache.setHalf2(cacheOff + i, + Half2.fromFloats(qkvBatch.get(vOffset + i), qkvBatch.get(vOffset + i + 1))); + } + } + } + + /** + * {@link #batchedFlashAttentionFP16Out} reading a half-precision KV cache. + * + *

K/V tile loads are packed: each thread pulls two adjacent head dims per + * 32-bit load and unpacks them into the FP32 shared tiles, so accumulation and + * the online softmax are unchanged.

+ * + *

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

+ * + * Worker: B*nHeads workgroups × min(headSize,128) threads. + */ + public static void batchedFlashAttentionFP16OutKVFP16(KernelContext context, + IntArray batchStartPosHolder, + FloatArray qkvBatch, + HalfFloatArray wrapKeyCache, + HalfFloatArray 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; + int halfHead = headSize / 2; + + 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; + 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 — one packed 32-bit load per thread per pair of head dims + for (int idx = tid; idx < tileLen * halfHead; idx += localSz) { + int tInTile = idx / halfHead; + int dPair = (idx % halfHead) * 2; + int kvOff = loff + (tileC + tInTile) * kvDim + kvHeadIdx * headSize + dPair; + Half2 kPair = wrapKeyCache.getHalf2(kvOff); + Half2 vPair = wrapValueCache.getHalf2(kvOff); + int tileOff = tInTile * headSize + dPair; + kTile[tileOff] = Half2.lowFloat(kPair); + kTile[tileOff + 1] = Half2.highFloat(kPair); + vTile[tileOff] = Half2.lowFloat(vPair); + vTile[tileOff + 1] = Half2.highFloat(vPair); + } + 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(); + + 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)); + } + } + + /** + * {@link #batchedFlashAttentionFP16OutKVFP16} with packed __half2 shared-memory tiles. + * + *

K/V pairs stay packed from the 32-bit global load through the shared tile + * (halving the K/V tile footprint), Q is staged once as packed FP16 pairs, and the + * score loop consumes each pair with a single __hfma2 (llama.cpp fattn-tile style). + * The online softmax and the P·V accumulation stay FP32.

+ * + *

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

+ * + * Worker: B*nHeads workgroups × min(headSize,128) threads. + */ + public static void batchedFlashAttentionFP16OutKVFP16PackedTile(KernelContext context, + IntArray batchStartPosHolder, + FloatArray qkvBatch, + HalfFloatArray wrapKeyCache, + HalfFloatArray 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; + int halfHead = headSize / 2; + + Half2[] qPacked = context.allocateHalf2LocalArray(64); // headSize/2 <= 64 + Half2[] kTile = context.allocateHalf2LocalArray(16 * 64); // BLOCK_C * halfHead + Half2[] vTile = context.allocateHalf2LocalArray(16 * 64); + float[] sTile = context.allocateFloatLocalArray(BLOCK_C); + + // Load Q (rotated, from the packed QKV buffer) into shared memory as packed pairs + int qOffset = batchIdx * qkvStride + h * headSize; + for (int i = tid; i < halfHead; i += localSz) { + qPacked[i] = Half2.fromFloats(qkvBatch.get(qOffset + i * 2), qkvBatch.get(qOffset + i * 2 + 1)); + } + context.localBarrier(); + + float maxScore = Float.NEGATIVE_INFINITY; + float sumExp = 0.0f; + 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 — packed 32-bit load into a packed __half2 tile, no expansion + for (int idx = tid; idx < tileLen * halfHead; idx += localSz) { + int tInTile = idx / halfHead; + int dPair = idx % halfHead; + int kvOff = loff + (tileC + tInTile) * kvDim + kvHeadIdx * headSize + dPair * 2; + int tileOff = tInTile * halfHead + dPair; + kTile[tileOff] = wrapKeyCache.getHalf2(kvOff); + vTile[tileOff] = wrapValueCache.getHalf2(kvOff); + } + context.localBarrier(); + + // Scores: one thread per key position; one __hfma2 per pair, expanded once per row + for (int t = tileC + tid; t <= tileEnd; t += localSz) { + int tInTile = t - tileC; + Half2 scoreAcc = Half2.fromFloats(0.0f, 0.0f); + for (int dp = 0; dp < halfHead; dp++) { + scoreAcc = Half2.fma(kTile[tInTile * halfHead + dp], qPacked[dp], scoreAcc); + } + sTile[tInTile] = (Half2.lowFloat(scoreAcc) + Half2.highFloat(scoreAcc)) / TornadoMath.sqrt(headSize); + } + context.localBarrier(); + + 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; + + // P·V: each thread owns fixed head dims (tid, tid+localSz), so its lane within + // the packed pair is fixed; select low/high once per tile element. + int pair0 = tid >> 1; + boolean high0 = (tid & 1) == 1; + int pair1 = d1 >> 1; + boolean high1 = (d1 & 1) == 1; + for (int t = 0; t < tileLen; t++) { + float p = TornadoMath.exp(sTile[t] - maxScore); + sumExp += p; + Half2 v0 = vTile[t * halfHead + pair0]; + acc0 += p * (high0 ? Half2.highFloat(v0) : Half2.lowFloat(v0)); + if (d1 < headSize) { + Half2 v1 = vTile[t * halfHead + pair1]; + acc1 += p * (high1 ? Half2.highFloat(v1) : Half2.lowFloat(v1)); + } + } + 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 ───────────────── /** diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/kernels/TransformerComputeKernelsLayered.java b/src/main/java/org/beehive/gpullama3/tornadovm/kernels/TransformerComputeKernelsLayered.java index 003b6f71..2e027a69 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/kernels/TransformerComputeKernelsLayered.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/kernels/TransformerComputeKernelsLayered.java @@ -59,10 +59,15 @@ public static void fusedRmsNormFFNGateUp(KernelContext context, FloatArray x, int rowOffsetW3 = rowId * dim; // === W1 matmul with inline normalization === + // Weights are read as packed FP16 pairs (single 32-bit loads); dim is even for all + // supported models, so the pair indices stay even as getHalf2 requires. float sum1 = 0.0f; - for (int j = localId; j < dim; j += localWorkGroupSize) { - float normalized = rmsWeights.get(j) * scale * x.get(j); - sum1 += w1.get(rowOffsetW1 + j).getFloat32() * normalized; + for (int j = localId * 2; j < dim; j += localWorkGroupSize * 2) { + float normalized0 = rmsWeights.get(j) * scale * x.get(j); + float normalized1 = rmsWeights.get(j + 1) * scale * x.get(j + 1); + Half2 pairW1 = w1.getHalf2(rowOffsetW1 + j); + sum1 += Half2.lowFloat(pairW1) * normalized0; + sum1 += Half2.highFloat(pairW1) * normalized1; } localSum[localId] = sum1; @@ -78,9 +83,12 @@ public static void fusedRmsNormFFNGateUp(KernelContext context, FloatArray x, // === W3 matmul with inline normalization (same computation) === float sum3 = 0.0f; - for (int j = localId; j < dim; j += localWorkGroupSize) { - float normalized = rmsWeights.get(j) * scale * x.get(j); - sum3 += w3.get(rowOffsetW3 + j).getFloat32() * normalized; + for (int j = localId * 2; j < dim; j += localWorkGroupSize * 2) { + float normalized0 = rmsWeights.get(j) * scale * x.get(j); + float normalized1 = rmsWeights.get(j + 1) * scale * x.get(j + 1); + Half2 pairW3 = w3.getHalf2(rowOffsetW3 + j); + sum3 += Half2.lowFloat(pairW3) * normalized0; + sum3 += Half2.highFloat(pairW3) * normalized1; } localSum[localId] = sum3; @@ -1739,6 +1747,130 @@ public static void processHeadsFlashAttentionSplitKVFP16(KernelContext context, } } + /** + * Deep-half2 variant of {@link #processHeadsFlashAttentionSplitKVFP16}: the K·Q score loop + * stays packed end-to-end (llama.cpp fattn-vec style). Q is converted once per workgroup into + * a __half2 local-memory tile, each K pair is consumed with a single __hfma2 into a packed + * accumulator, and the pair sums are expanded to FP32 once per row instead of once per pair. + * The V/softmax state stays FP32 (it accumulates across the whole position chunk). + */ + public static void processHeadsFlashAttentionSplitKVFP16Packed(KernelContext context, FloatArray q, HalfFloatArray key_cache, HalfFloatArray value_cache, FloatArray att, int nHeads, + int headSize, int kvDim, int kvMul, IntArray positionHolder, int layer, int contextLength, int nSplits) { + + final int MAX_HEAD_SIZE = 128; + final int MAX_LOCAL_SIZE = 64; + + int tid = context.localIdx; + int g = context.groupIdx; // 0 .. nHeads*nSplits - 1 + int localSize = context.localGroupSizeX; + int h = g / nSplits; + int s = g % nSplits; + + if (h >= nHeads) { + return; + } + + int pos = positionHolder.get(0); + int seqLen = pos + 1; + int chunk = (seqLen + nSplits - 1) / nSplits; + int startPos = s * chunk; + int endPos = Math.min(startPos + chunk, seqLen); // exclusive + + int loff = layer * contextLength * kvDim; + int kvHeadIdx = h / kvMul; + float invSqrt = 1.0f / TornadoMath.sqrt(headSize); + + Half2[] qPacked = context.allocateHalf2LocalArray(MAX_HEAD_SIZE / 2); + float[] accShared = context.allocateFloatLocalArray(MAX_LOCAL_SIZE * MAX_HEAD_SIZE); + float[] mShared = context.allocateFloatLocalArray(MAX_LOCAL_SIZE); + float[] lShared = context.allocateFloatLocalArray(MAX_LOCAL_SIZE); + float[] corrShared = context.allocateFloatLocalArray(MAX_LOCAL_SIZE); + float[] bcast = context.allocateFloatLocalArray(1); + + int headBase = h * nSplits * (headSize + 2); + int outBase = headBase + s * headSize; + int mBase = headBase + nSplits * headSize; + int lBase = mBase + nSplits; + + // Stage Q once as packed FP16 pairs; the per-pair conversion cost is paid once per + // workgroup instead of the score loop paying two half->float expansions per K pair. + for (int i = tid; i < headSize / 2; i += localSize) { + qPacked[i] = Half2.fromFloats(q.get(h * headSize + i * 2), q.get(h * headSize + i * 2 + 1)); + } + int rowBase = tid * headSize; + for (int d = 0; d < headSize; d++) { + accShared[rowBase + d] = 0.0f; + } + context.localBarrier(); + + // Strided scan over this split's position chunk (no barriers). + float m = Float.NEGATIVE_INFINITY; + float l = 0.0f; + for (int p = startPos + tid; p < endPos; p += localSize) { + int base = loff + p * kvDim + kvHeadIdx * headSize; + // Packed K·Q: one 32-bit load + one __hfma2 per pair; the two FP16 lane sums + // (over headSize/2 <= 64 terms each) are expanded to FP32 once per row. + Half2 scoreAcc = Half2.fromFloats(0.0f, 0.0f); + for (int d = 0; d < headSize; d += 2) { + scoreAcc = Half2.fma(key_cache.getHalf2(base + d), qPacked[d >> 1], scoreAcc); + } + float score = (Half2.lowFloat(scoreAcc) + Half2.highFloat(scoreAcc)) * invSqrt; + float newM = Math.max(m, score); + float corr = (m == Float.NEGATIVE_INFINITY) ? 0.0f : TornadoMath.exp(m - newM); + float e = TornadoMath.exp(score - newM); + for (int d = 0; d < headSize; d += 2) { + Half2 vPair = value_cache.getHalf2(base + d); + accShared[rowBase + d] = accShared[rowBase + d] * corr + e * Half2.lowFloat(vPair); + accShared[rowBase + d + 1] = accShared[rowBase + d + 1] * corr + e * Half2.highFloat(vPair); + } + l = l * corr + e; + m = newM; + } + mShared[tid] = m; + lShared[tid] = l; + context.localBarrier(); + + // Block max. + if (tid == 0) { + float blockMax = Float.NEGATIVE_INFINITY; + for (int t = 0; t < localSize; t++) { + if (mShared[t] > blockMax) { + blockMax = mShared[t]; + } + } + bcast[0] = blockMax; + } + context.localBarrier(); + float M = bcast[0]; + + corrShared[tid] = (mShared[tid] == Float.NEGATIVE_INFINITY) ? 0.0f : TornadoMath.exp(mShared[tid] - M); + context.localBarrier(); + + // Block sum L = Σ_t l_t · corr_t. + if (tid == 0) { + float blockSum = 0.0f; + for (int t = 0; t < localSize; t++) { + blockSum += lShared[t] * corrShared[t]; + } + bcast[0] = blockSum; + } + context.localBarrier(); + float L = bcast[0]; + + // Write UNNORMALIZED partial numerator (relative to block max M), plus M and L for the combine. + for (int d = tid; d < headSize; d += localSize) { + float acc = 0.0f; + for (int t = 0; t < localSize; t++) { + acc += corrShared[t] * accShared[t * headSize + d]; + } + att.set(outBase + d, acc); + } + if (tid == 0) { + att.set(mBase + s, M); + att.set(lBase + s, L); + } + } + /** * Qwen3-family decode attention, split-KV (flash-decoding) phase 2: combine. * @@ -2078,8 +2210,12 @@ public static void fusedQKVMatmulX( int rowOffset = rowId * dim; float partialSum = 0.0f; - for (int j = localId; j < dim; j += localWorkGroupSize) { - partialSum += wq.get(rowOffset + j).getFloat32() * x.get(j).getFloat32(); + // Packed FP16 pair loads for both the weight row and the FP16 input, multiplied + // with the packed __hmul2 intrinsic; accumulation stays FP32. dim is even. + for (int j = localId * 2; j < dim; j += localWorkGroupSize * 2) { + Half2 product = Half2.mult(wq.getHalf2(rowOffset + j), x.getHalf2(j)); + partialSum += Half2.lowFloat(product); + partialSum += Half2.highFloat(product); } localSum[localId] = partialSum; @@ -2102,8 +2238,12 @@ public static void fusedQKVMatmulX( int rowOffset = kRow * dim; float partialSum = 0.0f; - for (int j = localId; j < dim; j += localWorkGroupSize) { - partialSum += wk.get(rowOffset + j).getFloat32() * x.get(j).getFloat32(); + // Packed FP16 pair loads for both the weight row and the FP16 input, multiplied + // with the packed __hmul2 intrinsic; accumulation stays FP32. dim is even. + for (int j = localId * 2; j < dim; j += localWorkGroupSize * 2) { + Half2 product = Half2.mult(wk.getHalf2(rowOffset + j), x.getHalf2(j)); + partialSum += Half2.lowFloat(product); + partialSum += Half2.highFloat(product); } localSum[localId] = partialSum; @@ -2126,8 +2266,12 @@ public static void fusedQKVMatmulX( int rowOffset = vRow * dim; float partialSum = 0.0f; - for (int j = localId; j < dim; j += localWorkGroupSize) { - partialSum += wv.get(rowOffset + j).getFloat32() * x.get(j).getFloat32(); + // Packed FP16 pair loads for both the weight row and the FP16 input, multiplied + // with the packed __hmul2 intrinsic; accumulation stays FP32. dim is even. + for (int j = localId * 2; j < dim; j += localWorkGroupSize * 2) { + Half2 product = Half2.mult(wv.getHalf2(rowOffset + j), x.getHalf2(j)); + partialSum += Half2.lowFloat(product); + partialSum += Half2.highFloat(product); } localSum[localId] = partialSum; @@ -2557,11 +2701,21 @@ public static float matrixVectorRowMajorOptimized(KernelContext context, int loc int rowOffset = rowId * n; - // Each thread calculates partial dot product + // Each thread accumulates over consecutive FP16 pairs read with a single packed + // 32-bit load; consecutive threads read consecutive pairs, so a warp still issues + // fully coalesced 128-byte transactions with half the memory instructions. + // Rows start at rowId * n with n even for all supported models, so pair indices + // stay even (4-byte aligned) as required by getHalf2. float partialSum = 0.0f; - for (int j = localId; j < n; j += localSize) { + int nEven = n & ~1; + for (int j = localId * 2; j < nEven; j += localSize * 2) { int matrixIdx = rowOffset + j; - partialSum += w.get(matrixIdx).getFloat32() * x.get(j); + Half2 pair = w.getHalf2(matrixIdx); + partialSum += Half2.lowFloat(pair) * x.get(j); + partialSum += Half2.highFloat(pair) * x.get(j + 1); + } + if (nEven != n && localId == 0) { + partialSum += w.get(rowOffset + nEven).getFloat32() * x.get(nEven); } // Store partial sum in local memory @@ -2706,15 +2860,18 @@ public static float matrixVectorRowMajorOptimizedSingle(KernelContext context, i int rowOffset = rowId * n; - HalfFloat partialSum = new HalfFloat(0f); - for (int j = localId; j < n; j += localSize) { + // Both operands are FP16: read them as packed pairs (single 32-bit loads), multiply + // with the packed __hmul2 intrinsic and accumulate in FP32. n is even. + float partialSum = 0.0f; + for (int j = localId * 2; j < n; j += localSize * 2) { int matrixIdx = rowOffset + j; - HalfFloat mul = HalfFloat.mult(w.get(matrixIdx), x.get(j)); - partialSum = HalfFloat.add(partialSum, mul); + Half2 product = Half2.mult(w.getHalf2(matrixIdx), x.getHalf2(j)); + partialSum += Half2.lowFloat(product); + partialSum += Half2.highFloat(product); } // Store partial sum in local memory - localSum[localId] = partialSum.getHalfFloatValue(); + localSum[localId] = partialSum; context.localBarrier(); // Parallel reduction within workgroup diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/LlamaFP16FFNLayers.java b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/LlamaFP16FFNLayers.java index 68e37cda..215b316a 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/LlamaFP16FFNLayers.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/LlamaFP16FFNLayers.java @@ -1,5 +1,6 @@ package org.beehive.gpullama3.tornadovm.layers.type.fp16; +import org.beehive.gpullama3.inference.state.LlamaState; import org.beehive.gpullama3.inference.state.State; import org.beehive.gpullama3.inference.weights.tornado.LlamaTornadoWeights; import org.beehive.gpullama3.model.llama.LlamaConfiguration; @@ -14,9 +15,12 @@ import uk.ac.manchester.tornado.api.enums.DataTransferMode; public class LlamaFP16FFNLayers extends AbstractTransformerLayerTaskGraphs { + private static final boolean SPLIT_KV_ATTENTION = Boolean.getBoolean("llama.attention.splitKv"); + private final LlamaState llamaState; public LlamaFP16FFNLayers(String taskGraph, State state, LlamaTornadoWeights weights, LlamaConfiguration config, SchedulerType schedulerType) { super(taskGraph, state, weights, config, schedulerType); + this.llamaState = (LlamaState) state; setupFFNLayers(); } @@ -30,7 +34,9 @@ public GridScheduler updateGridScheduler(GridScheduler tornadoForwardScheduler) int configHiddenDimRowMajor = config.hiddenDim() * LOCAL_WORK_GROUP_SIZE_ALLOC; WorkerGrid configHiddenDimRowMajorWorker = WorkerGridFactory.genericWorker(configHiddenDimRowMajor, LOCAL_WORK_GROUP_SIZE_ALLOC); - WorkerGrid parallelAttentionWorker = WorkerGridFactory.createAttentionWorker(config.numberOfHeads(), config.headSize()); + int attentionGroups = splitKvAttentionEnabled() ? config.numberOfHeads() * LlamaState.SPLIT_KV : config.numberOfHeads(); + WorkerGrid parallelAttentionWorker = WorkerGridFactory.createAttentionWorker(attentionGroups, config.headSize()); + WorkerGrid attentionCombineWorker = WorkerGridFactory.createAttentionWorker(config.numberOfHeads(), config.headSize()); int fusedQKVRows = config.dim() + 2 * config.kvDim(); int fusedQKVGlobal = fusedQKVRows * LOCAL_WORK_GROUP_SIZE_ALLOC; @@ -45,6 +51,9 @@ public GridScheduler updateGridScheduler(GridScheduler tornadoForwardScheduler) tornadoForwardScheduler.addWorkerGrid("layer_" + i + ".qkv_projection", fusedQKVWorker); tornadoForwardScheduler.addWorkerGrid("layer_" + i + ".rope_and_kv_cache", ropeWithCacheWorker); tornadoForwardScheduler.addWorkerGrid("layer_" + i + ".attention", parallelAttentionWorker); + if (splitKvAttentionEnabled()) { + tornadoForwardScheduler.addWorkerGrid("layer_" + i + ".attention_combine", attentionCombineWorker); + } tornadoForwardScheduler.addWorkerGrid("layer_" + i + ".attn_output_proj", configDimRowMajorGlobalWorker); // === FFN Block === tornadoForwardScheduler.addWorkerGrid("layer_" + i + ".ffn_rms_reduce", rmsNormWorker); @@ -304,6 +313,10 @@ protected String predecessorGraphName(int layerIndex) { return (layerIndex == 0) ? "activationUpdate" : "layer_" + (layerIndex - 1); } + protected boolean splitKvAttentionEnabled() { + return SPLIT_KV_ATTENTION && schedulerType == SchedulerType.NVIDIA; + } + protected TaskGraph configureLayerDataTransfers(TaskGraph unifiedLayer, int layerIndex) { Object keyCache = useFp16KVCache() ? state.wrapKeyCacheFP16 : state.wrapKeyCache; Object valueCache = useFp16KVCache() ? state.wrapValueCacheFP16 : state.wrapValueCache; @@ -324,6 +337,9 @@ protected TaskGraph configureLayerDataTransfers(TaskGraph unifiedLayer, int laye keyCache, valueCache, // Attention & FFN buffers state.wrapAtt, state.wrapHb, state.wrapXbFP16); + if (splitKvAttentionEnabled()) { + unifiedLayer.transferToDevice(DataTransferMode.FIRST_EXECUTION, llamaState.wrapAttSplit); + } } else { // Subsequent layers: consume from the previous layer graph by name. // The no-arg consumeFromDevice form uses the current graph's own name as source key, @@ -342,11 +358,59 @@ protected TaskGraph configureLayerDataTransfers(TaskGraph unifiedLayer, int laye state.wrapAtt, state.wrapHb, // Position & misc state.positionHolder, state.wrapXbFP16); + if (splitKvAttentionEnabled()) { + unifiedLayer.consumeFromDevice(pred, llamaState.wrapAttSplit); + } } return unifiedLayer; } private TaskGraph configureAttention(TaskGraph unifiedLayer, int layerIndex) { + if (splitKvAttentionEnabled()) { + if (useFp16KVCache()) { + unifiedLayer.task("attention", + State.ATTENTION_DEEP_HALF2 + ? TransformerComputeKernelsLayered::processHeadsFlashAttentionSplitKVFP16Packed + : TransformerComputeKernelsLayered::processHeadsFlashAttentionSplitKVFP16, + context, + state.wrapQ, + state.wrapKeyCacheFP16, + state.wrapValueCacheFP16, + llamaState.wrapAttSplit, + config.numberOfHeads(), + config.headSize(), + config.kvDim(), + config.kvMul(), + state.positionHolder, + layerIndex, + config.contextLength(), + LlamaState.SPLIT_KV); + } else { + unifiedLayer.task("attention", + TransformerComputeKernelsLayered::processHeadsFlashAttentionSplitKV, + context, + state.wrapQ, + state.wrapKeyCache, + state.wrapValueCache, + llamaState.wrapAttSplit, + config.numberOfHeads(), + config.headSize(), + config.kvDim(), + config.kvMul(), + state.positionHolder, + layerIndex, + config.contextLength(), + LlamaState.SPLIT_KV); + } + return unifiedLayer.task("attention_combine", + TransformerComputeKernelsLayered::combineSplitKVAttention, + context, + llamaState.wrapAttSplit, + state.wrapXb, + config.numberOfHeads(), + config.headSize(), + LlamaState.SPLIT_KV); + } if (useFp16KVCache()) { // Flash Attention over the half-precision KV cache (FP32 accumulation). // The scalar-read variant is an evaluation aid; the packed variant is the default. diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/Qwen3FP16FFNLayers.java b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/Qwen3FP16FFNLayers.java index ddfadf66..74d27f26 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/Qwen3FP16FFNLayers.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/Qwen3FP16FFNLayers.java @@ -1,6 +1,7 @@ package org.beehive.gpullama3.tornadovm.layers.type.fp16; import org.beehive.gpullama3.inference.state.Qwen3State; +import org.beehive.gpullama3.inference.state.State; import org.beehive.gpullama3.inference.weights.tornado.Qwen3TornadoWeights; import org.beehive.gpullama3.model.qwen3.Qwen3Configuration; import org.beehive.gpullama3.tornadovm.kernels.Qwen3Kernels; @@ -346,7 +347,9 @@ protected TaskGraph createFFNLayerTaskGraph(int layerIndex) { // Phase 1: split each head's KV range across attentionSplits workgroups; partials -> wrapAttSplit. if (useFp16KVCache()) { unifiedLayer.task("attention", - TransformerComputeKernelsLayered::processHeadsFlashAttentionSplitKVFP16, + State.ATTENTION_DEEP_HALF2 + ? TransformerComputeKernelsLayered::processHeadsFlashAttentionSplitKVFP16Packed + : TransformerComputeKernelsLayered::processHeadsFlashAttentionSplitKVFP16, context, qwen3State.wrapQ, // query vectors qwen3State.wrapKeyCacheFP16, // key cache (FP16) 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 index 7d04117c..30a993b8 100644 --- 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 @@ -1,6 +1,7 @@ package org.beehive.gpullama3.tornadovm.layers.type.fp16.prefill; import org.beehive.gpullama3.inference.state.LlamaState; +import org.beehive.gpullama3.inference.state.State; import org.beehive.gpullama3.inference.weights.tornado.LlamaTornadoWeights; import org.beehive.gpullama3.model.llama.LlamaConfiguration; import org.beehive.gpullama3.tornadovm.kernels.TransformerBatchPrefillKernels; @@ -61,6 +62,14 @@ public class LlamaFP16LayersBatchPrefillMMA implements BatchPrefillTransformerLa private final List layerITGs; private String lastLayerTaskGraphID; + /** + * The batched-prefill graphs only run on the CUDA backend (tensor-core gated), + * which is the same backend the FP16 KV cache path targets. + */ + private boolean useFp16KVCache() { + return State.USE_FP16_KV && state.wrapKeyCacheFP16 != null; + } + public LlamaFP16LayersBatchPrefillMMA(LlamaState state, LlamaTornadoWeights weights, LlamaConfiguration config, int batchSize) { this.state = state; @@ -86,6 +95,9 @@ private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) { TaskGraph batchPrefillLayer = new TaskGraph(graphName); + Object keyCache = useFp16KVCache() ? state.wrapKeyCacheFP16 : state.wrapKeyCache; + Object valueCache = useFp16KVCache() ? state.wrapValueCacheFP16 : state.wrapValueCache; + // ── Data Transfers ───────────────────────────────────────────────────── if (layerIndex == 0) { // batchStartPosHolder is set by host before each chunk → EVERY_EXECUTION @@ -96,7 +108,7 @@ private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) { state.attnScaleBatch, state.ffnScaleBatch, state.wrapXbFP16Batch, state.qkvResultBatch, - state.wrapKeyCache, state.wrapValueCache, + keyCache, valueCache, state.normedXFFNFP16, state.gateUpResultBatch, state.attnOutFP16, state.woOut, state.wrapHbFP16Batch, state.w2Out); // wrapXBatch produced by the prefillActivation graph and persists in device memory @@ -113,7 +125,7 @@ private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) { state.attnScaleBatch, state.ffnScaleBatch, state.wrapXbFP16Batch, state.qkvResultBatch, - state.wrapKeyCache, state.wrapValueCache, + keyCache, valueCache, state.normedXFFNFP16, state.gateUpResultBatch, state.attnOutFP16, state.woOut, state.wrapHbFP16Batch, state.w2Out); } @@ -157,22 +169,41 @@ private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) { 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); + if (useFp16KVCache()) { + batchPrefillLayer.task("batch_rope_kv", + TransformerBatchPrefillKernels::batchedRopeWithKVCachePackedFP16, + context, state.batchStartPosHolder, + state.qkvResultBatch, + state.wrapKeyCacheFP16, state.wrapValueCacheFP16, + kvDim, config.headSize(), layerIndex, config.contextLength(), dim); + + batchPrefillLayer.task("batch_attention", + State.ATTENTION_DEEP_HALF2 + ? TransformerBatchPrefillKernels::batchedFlashAttentionFP16OutKVFP16PackedTile + : TransformerBatchPrefillKernels::batchedFlashAttentionFP16OutKVFP16, + context, state.batchStartPosHolder, + state.qkvResultBatch, state.wrapKeyCacheFP16, state.wrapValueCacheFP16, + state.attnOutFP16, + config.numberOfHeads(), config.headSize(), + kvDim, config.kvMul(), layerIndex, config.contextLength(), dim); + } else { + 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, @@ -212,7 +243,7 @@ private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) { // 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); + batchPrefillLayer.persistOnDevice(state.wrapXBatch, keyCache, valueCache); return batchPrefillLayer; } 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 index f12b842d..a62e7a66 100644 --- 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 @@ -1,6 +1,7 @@ package org.beehive.gpullama3.tornadovm.layers.type.fp16.prefill; import org.beehive.gpullama3.inference.state.Qwen3State; +import org.beehive.gpullama3.inference.state.State; import org.beehive.gpullama3.inference.weights.tornado.Qwen3TornadoWeights; import org.beehive.gpullama3.model.qwen3.Qwen3Configuration; import org.beehive.gpullama3.tornadovm.kernels.Qwen3Kernels; @@ -67,6 +68,14 @@ public class Qwen3FP16LayersBatchPrefillMMA implements BatchPrefillTransformerLa private final List layerITGs; private String lastLayerTaskGraphID; + /** + * The batched-prefill graphs only run on the CUDA backend (tensor-core gated), + * which is the same backend the FP16 KV cache path targets. + */ + private boolean useFp16KVCache() { + return State.USE_FP16_KV && state.wrapKeyCacheFP16 != null; + } + public Qwen3FP16LayersBatchPrefillMMA(Qwen3State state, Qwen3TornadoWeights weights, Qwen3Configuration config, int batchSize) { this.state = state; @@ -99,6 +108,9 @@ private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) { int dim = config.dim(); int hidDim = config.hiddenDim(); + Object keyCache = useFp16KVCache() ? state.wrapKeyCacheFP16 : state.wrapKeyCache; + Object valueCache = useFp16KVCache() ? state.wrapValueCacheFP16 : state.wrapValueCache; + // ── Data Transfers ───────────────────────────────────────────────────── if (layerIndex == 0) { batchPrefillLayer.transferToDevice(DataTransferMode.EVERY_EXECUTION, state.batchStartPosHolder); @@ -107,7 +119,7 @@ private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) { state.attnScaleBatch, state.ffnScaleBatch, state.wrapXbFP16Batch, state.qkvResultBatch, - state.wrapKeyCache, state.wrapValueCache, + keyCache, valueCache, state.normedXFFNFP16, state.gateUpResultBatch, state.attnOutFP16, state.woOut, state.wrapHbFP16Batch, state.w2Out); batchPrefillLayer.consumeFromDevice("prefillActivation", state.wrapXBatch); @@ -120,7 +132,7 @@ private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) { state.attnScaleBatch, state.ffnScaleBatch, state.wrapXbFP16Batch, state.qkvResultBatch, - state.wrapKeyCache, state.wrapValueCache, + keyCache, valueCache, state.normedXFFNFP16, state.gateUpResultBatch, state.attnOutFP16, state.woOut, state.wrapHbFP16Batch, state.w2Out); } @@ -169,23 +181,42 @@ private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) { 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); + if (useFp16KVCache()) { + batchPrefillLayer.task("batch_rope_kv", + Qwen3Kernels::batchedRopeWithKVCacheQwen3PackedFP16, + context, state.batchStartPosHolder, + state.qkvResultBatch, + state.wrapKeyCacheFP16, state.wrapValueCacheFP16, + kvDim, nEmbdHead, layerIndex, config.contextLength(), qDim); + + batchPrefillLayer.task("batch_attention", + State.ATTENTION_DEEP_HALF2 + ? TransformerBatchPrefillKernels::batchedFlashAttentionFP16OutKVFP16PackedTile + : TransformerBatchPrefillKernels::batchedFlashAttentionFP16OutKVFP16, + context, state.batchStartPosHolder, + state.qkvResultBatch, state.wrapKeyCacheFP16, state.wrapValueCacheFP16, + state.attnOutFP16, + config.numberOfHeads(), nEmbdHead, + kvDim, gqa, layerIndex, config.contextLength(), qDim); + } else { + batchPrefillLayer.task("batch_rope_kv", + Qwen3Kernels::batchedRopeWithKVCacheQwen3Packed, + context, state.batchStartPosHolder, + state.qkvResultBatch, + state.wrapKeyCache, state.wrapValueCache, + kvDim, nEmbdHead, layerIndex, config.contextLength(), qDim); + + 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, @@ -222,7 +253,7 @@ private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) { .task("w2Resid", TransformerBatchPrefillKernels::batchedResidualAddFP32, context, state.wrapXBatch, state.w2Out); - batchPrefillLayer.persistOnDevice(state.wrapXBatch, state.wrapKeyCache, state.wrapValueCache); + batchPrefillLayer.persistOnDevice(state.wrapXBatch, keyCache, valueCache); return batchPrefillLayer; } From 622704b8580bf9d3613ced7e81cf38276ab84cfa Mon Sep 17 00:00:00 2001 From: Orion Papadakis Date: Wed, 22 Jul 2026 18:34:21 +0300 Subject: [PATCH 4/5] Add llama.bench.ignoreEos flag to fix generation length Optionally keep decoding past the stop token so benchmark runs generate a fixed token count, making before/after throughput directly comparable. --- .../org/beehive/gpullama3/inference/InferenceEngine.java | 7 +++++-- .../inference/InferenceEngineWithBatchPrefillDecode.java | 7 +++++-- .../inference/InferenceEngineWithPrefillDecode.java | 7 +++++-- 3 files changed, 15 insertions(+), 6 deletions(-) diff --git a/src/main/java/org/beehive/gpullama3/inference/InferenceEngine.java b/src/main/java/org/beehive/gpullama3/inference/InferenceEngine.java index d653ba58..2099ca9a 100644 --- a/src/main/java/org/beehive/gpullama3/inference/InferenceEngine.java +++ b/src/main/java/org/beehive/gpullama3/inference/InferenceEngine.java @@ -34,6 +34,9 @@ */ public final class InferenceEngine { + /** Benchmarking aid: keep decoding past the stop token so every run generates the same token count. */ + private static final boolean IGNORE_EOS = Boolean.getBoolean("llama.bench.ignoreEos"); + private InferenceEngine() { //prevent instantiation } @@ -344,7 +347,7 @@ public static List generateTokensGPULlama(Model model, State state, int generatedTokens.add(nextToken); // Check stop condition - if (stopTokens.contains(nextToken)) { + if (!IGNORE_EOS && stopTokens.contains(nextToken)) { break; } } @@ -438,7 +441,7 @@ public static List generateTokensGPUQwen3(Model model, State state, int } // Check for stop condition - if (stopTokens.contains(nextToken)) { + if (!IGNORE_EOS && stopTokens.contains(nextToken)) { break; } diff --git a/src/main/java/org/beehive/gpullama3/inference/InferenceEngineWithBatchPrefillDecode.java b/src/main/java/org/beehive/gpullama3/inference/InferenceEngineWithBatchPrefillDecode.java index 4493340b..45384eff 100644 --- a/src/main/java/org/beehive/gpullama3/inference/InferenceEngineWithBatchPrefillDecode.java +++ b/src/main/java/org/beehive/gpullama3/inference/InferenceEngineWithBatchPrefillDecode.java @@ -36,6 +36,9 @@ */ public final class InferenceEngineWithBatchPrefillDecode { + /** Benchmarking aid: keep decoding past the stop token so every run generates the same token count. */ + private static final boolean IGNORE_EOS = Boolean.getBoolean("llama.bench.ignoreEos"); + private InferenceEngineWithBatchPrefillDecode() { } @@ -128,7 +131,7 @@ public static List generateTokensLlama(Model model, onTokenGenerated.accept(nextToken); } - if (stopTokens.contains(nextToken)) { + if (!IGNORE_EOS && stopTokens.contains(nextToken)) { break; } @@ -233,7 +236,7 @@ public static List generateTokensGPULlama(Model model, onTokenGenerated.accept(nextToken); } - if (stopTokens.contains(nextToken)) { + if (!IGNORE_EOS && stopTokens.contains(nextToken)) { break; } diff --git a/src/main/java/org/beehive/gpullama3/inference/InferenceEngineWithPrefillDecode.java b/src/main/java/org/beehive/gpullama3/inference/InferenceEngineWithPrefillDecode.java index 38450968..eafe9497 100644 --- a/src/main/java/org/beehive/gpullama3/inference/InferenceEngineWithPrefillDecode.java +++ b/src/main/java/org/beehive/gpullama3/inference/InferenceEngineWithPrefillDecode.java @@ -36,6 +36,9 @@ */ public final class InferenceEngineWithPrefillDecode { + /** Benchmarking aid: keep decoding past the stop token so every run generates the same token count. */ + private static final boolean IGNORE_EOS = Boolean.getBoolean("llama.bench.ignoreEos"); + private InferenceEngineWithPrefillDecode() { } @@ -99,7 +102,7 @@ public static List generateTokensLlama(Model model, onTokenGenerated.accept(nextToken); } - if (stopTokens.contains(nextToken)) { + if (!IGNORE_EOS && stopTokens.contains(nextToken)) { break; } @@ -183,7 +186,7 @@ public static List generateTokensGPULlama(Model model, onTokenGenerated.accept(nextToken); } - if (stopTokens.contains(nextToken)) { + if (!IGNORE_EOS && stopTokens.contains(nextToken)) { break; } From 4a6591a65349cc62abca6efe00b3dcade94090a9 Mon Sep 17 00:00:00 2001 From: Orion Papadakis Date: Wed, 22 Jul 2026 18:34:21 +0300 Subject: [PATCH 5/5] Improve Claude Code agents and skills --- .claude/README.md | 45 +++++++++++++++++++ .../gpullama-benchmarking-specialist.md | 29 ++++++++++++ .../gpullama-perf-profiling-specialist.md | 24 ++++++++++ 3 files changed, 98 insertions(+) create mode 100644 .claude/README.md diff --git a/.claude/README.md b/.claude/README.md new file mode 100644 index 00000000..4b4b7b8e --- /dev/null +++ b/.claude/README.md @@ -0,0 +1,45 @@ +# Claude Agents And Skills For GPULlama3 + +This directory contains project-local Claude Code agents and skills for GPULlama3.java work, including benchmarking, profiling, correctness debugging, and TornadoVM integration. + +The scope is broad enough to support future repo work: + +- GPULlama3 benchmarking and profiling +- TornadoVM backend/codegen investigation +- correctness debugging for precision and kernel changes +- accelerator feature validation + +After adding or editing files here, restart Claude Code from the repository root and run `/agents` and `/skills` to verify discovery. + +## Portability Rules + +- Do not hardcode user-specific paths. Discover the repository root from the current working directory and external project roots from environment variables, local config, documentation, or explicit user input. +- Do not bake current experiment settings into durable guidance. Prompt text, model paths, token limits, memory limits, feature flags, and backend choices come from the user request, checked-in scripts, or benchmark manifests. +- Keep current performance conclusions in result artifacts, not in reusable agent policy. Agents should collect evidence before classifying bottlenecks. +- Prefer commands that run from the repository root or use named environment variables such as `TORNADOVM_HOME` and `MODEL_DIR`. +- When local machine details are needed, record them in the benchmark artifact directory rather than in `.claude`. + +## Agent Boundaries + +| Agent | Use For | +|---|---| +| `gpullama-perf-profiling-specialist` | Running GPULlama3 benchmarks, collecting metrics, nsys/ncu profiling, bottleneck classification | +| `tornado-backend-specialist` | TornadoVM backend changes, generated kernel codegen, native accelerator library integration | +| `gpullama-correctness-debug-agent` | Output/accuracy regressions after TornadoVM or GPULlama3 performance changes | + +## Skill Boundaries + +| Skill | Use For | +|---|---| +| `gpullama-benchmarking` | Reproducible GPULlama3 benchmark runs using local scripts | +| `gpullama-nsys-analysis` | Nsight Systems trace collection and system-level analysis | +| `gpullama-ncu-analysis` | Nsight Compute analysis of a specific hot CUDA kernel | +| `tornado-codegen-validation` | TornadoVM API/codegen/correctness validation for generated accelerator kernels | + +## Operating Rules + +- Do not invent performance numbers. Every claim must come from metrics JSON, TornadoVM profiler output, system profilers, kernel profilers, or generated kernel evidence. +- Use `nsys` before `ncu`: identify hot kernels and launch/sync behavior first, then profile a specific kernel. +- Keep artifacts under timestamped directories, preferably `perf-results//` or a clearly named scratch directory. +- For codegen-sensitive work, inspect generated kernel source and add tests that fail if emission regresses. +- Match the backend and profiler to the task. CUDA-specific work should use CUDA evidence; backend comparisons should measure each requested backend under equivalent conditions. diff --git a/.claude/agents/gpullama-benchmarking-specialist.md b/.claude/agents/gpullama-benchmarking-specialist.md index f6074b8a..a0b51d83 100644 --- a/.claude/agents/gpullama-benchmarking-specialist.md +++ b/.claude/agents/gpullama-benchmarking-specialist.md @@ -84,6 +84,35 @@ When testing a feature flag, change one variable at a time and keep the baseline structurally identical to the treatment command — differences beyond the flag under test invalidate the comparison. +## Pre-Flight Checks (do these before trusting ANY A/B) + +1. **Stale jar check.** `llama-tornado` picks the jar by reverse-sorting + `target/gpu-llama3-*.jar`; after a build-suffix change (e.g. `-Djdk.version.suffix`), + a leftover jar from the old suffix can sort ABOVE the fresh one and both "sides" of the + A/B silently run identical old code. Before a flag comparison: `ls -la target/*.jar`, + delete stale jars, and confirm the surviving jar's mtime postdates the last build. +2. **Flag liveness check.** Prove the flag under test actually reaches the JVM and changes + behavior (a cheap probe run where the flag has an observable effect) before spending a + full sweep on it. A flag that is silently ignored produces a perfectly clean null result. +3. **GPU idle check.** Kill leftover `java` processes between runs and confirm + `nvidia-smi` is back to idle memory. Never put the target process name literally inside + a `pkill -f` command line that also does other work — `pkill -f` matches your own shell's + command string and kills it. + +## Flag Ladders (preferred over isolated pairs) + +When several stacked optimizations exist, benchmark them as a cumulative ladder from the +pre-feature baseline (A → A+f1 → A+f1+f2 → ...) and report BOTH the per-step delta and the +cumulative delta vs baseline. This attributes the gain to the right layer; an isolated pair +at the top of the stack can show "no gain" purely because a lower layer already captured it. + +Depth caveat: decode tok/s decays with KV depth, so two runs of different generated length +are NOT comparable. For equal-length runs use `-Dllama.bench.ignoreEos=true` (disables the +stop-token break so generation runs to `--max-tokens`); without it, models stop early at +their end-of-turn token and comparisons are rough estimates only. Flag-gated attention/KV +optimizations also pay off more at depth — a null result at shallow depth (~512) does not +rule out a win at 2048. + ## When Results Point to a Bottleneck State whether the limiting factor looks memory-bound, compute-bound, launch-overhead-bound, diff --git a/.claude/agents/gpullama-perf-profiling-specialist.md b/.claude/agents/gpullama-perf-profiling-specialist.md index 30440f35..da0536fc 100644 --- a/.claude/agents/gpullama-perf-profiling-specialist.md +++ b/.claude/agents/gpullama-perf-profiling-specialist.md @@ -85,6 +85,30 @@ Classify bottlenecks as: Never report a numeric claim without citing the artifact that produced it. +## Two Cheap Cross-Checks (run these before nsys) + +Both need only the profiler dump plus the metrics JSON of a plain (non-profiler) run: + +1. **Kernel-time vs wall-time gap.** Sum profiler kernel time over the run and divide by + generated tokens; compare against wall ms/token from a NON-profiler run of the same + config (profiler runs distort wall time). If kernel-ms/token is well below wall + ms/token (e.g. half), the difference is inter-kernel gap — launch overhead, dispatch, + many small kernels — and the run is launch-overhead-bound no matter what the kernel + shares say. Optimizing the top kernel cannot recover that gap; graphs/fusion can. +2. **Bandwidth roofline sanity.** Decode streams the full weight set once per token: + effective BW = (weights bytes) / (wall s/token). Compare against realistic device + GEMV bandwidth (~80% of peak). Near roofline → memory-bound, only fewer bytes + (quantization) helps; far below roofline with GEMV-dominated kernel shares → the + headroom is in gaps or kernel inefficiency, not in the data volume. + +## Cross-Model Share Table + +When profiling several models, present one table: rows = task type (strip `layer_N.` +prefixes and sum across layers), columns = models, cells = % of total kernel time. Group +related tasks (attention + attention_combine; all RMS variants). This exposes which task +family is the common bottleneck and how it scales with model size far better than five +separate top-10 lists. + ## Bottleneck Method Do not preload a model-specific bottleneck map. Derive it for the current task: