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 b94fe16d..29c5982d 100644 --- a/.claude/agents/gpullama-benchmarking-specialist.md +++ b/.claude/agents/gpullama-benchmarking-specialist.md @@ -90,6 +90,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: diff --git a/pom.xml b/pom.xml index 16936920..b1c582e9 100644 --- a/pom.xml +++ b/pom.xml @@ -39,7 +39,11 @@ 1.0.0 - 5.0.0 + + 5.2.0 -jdk21 ${tornadovm.base.version}${jdk.version.suffix} diff --git a/src/main/java/org/beehive/gpullama3/inference/InferenceEngine.java b/src/main/java/org/beehive/gpullama3/inference/InferenceEngine.java index 049d9b78..f2304372 100644 --- a/src/main/java/org/beehive/gpullama3/inference/InferenceEngine.java +++ b/src/main/java/org/beehive/gpullama3/inference/InferenceEngine.java @@ -35,6 +35,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 } @@ -364,7 +367,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; } } @@ -464,7 +467,9 @@ public static List generateTokensGPUQwen3(Model model, State state, int } // Check for stop condition - if (generatedTokens.size() >= generatedTokenBudget || stopTokens.contains(nextToken)) { + // The budget is a hard limit and is never bypassed; only the stop token is, so that + // an ignore-EOS benchmark run still terminates — on the budget instead of on EOS. + if (generatedTokens.size() >= generatedTokenBudget || (!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 d811c214..8c522eae 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; } @@ -240,7 +243,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; } 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..98372c21 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; @@ -21,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 @@ -70,6 +78,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..9e1be40f 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; @@ -33,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); @@ -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 23055d7e..620fd83c 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,27 @@ */ 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"); + + /** + * 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,) @@ -58,6 +79,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; // On-device greedy sampling: the GPU argmax kernel writes the sampled token id here // (element 0), so only 1 int crosses to the host instead of the full vocab logits row. @@ -138,6 +161,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; @@ -221,6 +246,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 ecceb40d..7ea669db 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 { @@ -360,6 +361,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. @@ -394,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; @@ -418,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; @@ -442,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; @@ -495,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; @@ -520,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; @@ -545,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; @@ -1513,6 +1597,66 @@ 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))); + } + } + /** * Per-slot Qwen3 split-half RoPE + KV-cache write (batched DECODE). * @@ -1635,6 +1779,5 @@ public static void batchedDecodePagedRopeWithKVCacheQwen3Packed( valuePool.set(cacheOff + ic + halfEmbdHead, qkvBatch.get(vHeadBase + ic + halfEmbdHead)); } } - } // @formatter:on diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/kernels/TransformerBatchPrefillKernels.java b/src/main/java/org/beehive/gpullama3/tornadovm/kernels/TransformerBatchPrefillKernels.java index 50e744d0..15359e49 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. @@ -336,121 +337,6 @@ public static void batchedFlashAttention(KernelContext context, } } - /** - * Batched DECODE flash attention: B independent sequences, one query token - * each. Identical online-softmax math to {@link #batchedFlashAttention}, but - * each batch slot has its OWN KV cache region and its OWN position, so slot - * {@code b} attends positions {@code 0..seqPositions[b]} of its own cache — - * the shape produced by batching B concurrent decode requests. - * - *

KV cache layout: one contiguous region of {@code numLayers * - * contextLength * kvDim} per slot, so the base for slot b, layer L is - * {@code b*numLayers*contextLength*kvDim + L*contextLength*kvDim}.

- * - *

One workgroup per (batchIdx, head): {@code groupId = batchIdx*nHeads + h}.

- */ - public static void batchedDecodeAttention(KernelContext context, - IntArray seqPositions, - FloatArray wrapQBatch, - FloatArray wrapKeyCache, - FloatArray wrapValueCache, - FloatArray wrapXbBatch, - int nHeads, int headSize, - int kvDim, int kvMul, - int layerIndex, int numLayers, 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 = seqPositions.get(batchIdx); // per-slot position - int loff = batchIdx * (numLayers * contextLength * kvDim) + layerIndex * contextLength * kvDim; // per-slot KV base - int kvHeadIdx = h / kvMul; - int BLOCK_C = 16; - - 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); - float[] maxHolder = context.allocateFloatLocalArray(1); - - int qOffset = batchIdx * dim + h * headSize; - for (int i = tid; i < headSize; i += localSz) { - qShared[i] = wrapQBatch.get(qOffset + i); - } - context.localBarrier(); - - 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 tileC = 0; tileC <= pos; tileC += BLOCK_C) { - int tileEnd = Math.min(tileC + BLOCK_C - 1, pos); - - for (int t = tileC + tid; t <= tileEnd; t += localSz) { - int tInTile = t - tileC; - int tileMOff = tInTile * headSize; - for (int d = 0; d < headSize; d++) { - int kvOff = loff + t * kvDim + kvHeadIdx * headSize + d; - kTile[tileMOff + d] = wrapKeyCache.get(kvOff); - vTile[tileMOff + d] = wrapValueCache.get(kvOff); - } - } - context.localBarrier(); - - 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 <= tileEnd - tileC; t++) { - if (sTile[t] > tileMax) { - tileMax = sTile[t]; - } - } - if (tid == 0) { - maxHolder[0] = tileMax; - } - context.localBarrier(); - float curTileMax = maxHolder[0]; - - float newMax = Math.max(maxScore, curTileMax); - 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 = 0; t <= tileEnd - tileC; t++) { - float expScore = TornadoMath.exp(sTile[t] - maxScore); - sumExp += expScore; - for (int d = 0; d < headSize; d++) { - output[d] += expScore * vTile[t * headSize + d]; - } - } - context.localBarrier(); - } - - float norm = (sumExp > 0.0f) ? (1.0f / sumExp) : 0.0f; - int xbOffset = batchIdx * dim + h * headSize; - for (int d = tid; d < headSize; d += localSz) { - wrapXbBatch.set(xbOffset + d, output[d] * norm); - } - } - // ── Output / FFN Projections ───────────────────────────────────────────── /** @@ -1633,32 +1519,23 @@ public static void batchedFlashAttentionFP16Out(KernelContext context, } } - // ── Batched DECODE variants (per-slot KV cache + per-slot position) ────── - // - // These two kernels are the only semantic delta between batched PREFILL (B - // tokens of ONE sequence, shared causal KV) and batched DECODE (B independent - // sequences, each with its own KV region and its own position). The math is - // identical to the *Packed / *FP16Out prefill kernels above; only the KV - // addressing changes: - // pos = seqPositions[batchIdx] (per slot) - // base = batchIdx*(numLayers*ctx*kvDim) + layer*ctx*kvDim (per slot) - // The KV cache is therefore sized B*numLayers*contextLength*kvDim. + // ── FP16 KV cache variants (batched prefill) ───────────────────────────── /** - * Per-slot RoPE + KV-cache write over the packed QKV buffer (decode). + * {@link #batchedRopeWithKVCachePacked} writing a half-precision KV cache. * - *

Fork of {@link #batchedRopeWithKVCachePacked}: each batch slot rotates at - * its own position {@code seqPositions[batchIdx]} and writes K/V into its own - * KV region ({@code batchIdx} stride = {@code numLayers*contextLength*kvDim}).

+ *

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 batchedDecodeRopeWithKVCachePacked(KernelContext context, - IntArray seqPositions, - FloatArray qkvBatch, - FloatArray wrapKeyCache, - FloatArray wrapValueCache, - int kvDim, int headSize, - int layerIndex, int numLayers, - int contextLength, int dim) { + 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; @@ -1666,7 +1543,7 @@ public static void batchedDecodeRopeWithKVCachePacked(KernelContext context, int i = pairIdx * 2; int qkvStride = dim + 2 * kvDim; - int pos = seqPositions.get(batchIdx); + int pos = batchStartPosHolder.get(0) + batchIdx; int qOffset = batchIdx * qkvStride; int kOffset = batchIdx * qkvStride + dim; int vOffset = batchIdx * qkvStride + dim + kvDim; @@ -1678,63 +1555,66 @@ public static void batchedDecodeRopeWithKVCachePacked(KernelContext context, 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 slotBase = batchIdx * (numLayers * contextLength * kvDim); - int cacheOff = slotBase + layerIndex * contextLength * kvDim + pos * kvDim; - wrapKeyCache.set(cacheOff + i, rotK0); - wrapKeyCache.set(cacheOff + i + 1, rotK1); - wrapValueCache.set(cacheOff + i, qkvBatch.get(vOffset + i)); - wrapValueCache.set(cacheOff + i + 1, qkvBatch.get(vOffset + i + 1)); + 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))); } } } /** - * Per-slot flash attention over the packed QKV buffer, FP16 output (decode). + * {@link #batchedFlashAttentionFP16Out} reading a half-precision KV cache. * - *

Fork of {@link #batchedFlashAttentionFP16Out}: each batch slot attends - * over {@code 0..seqPositions[batchIdx]} of its OWN KV region. Same - * register-partitioned P·V accumulation and FP16 emission.

+ *

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 batchedDecodeAttentionFP16Out(KernelContext context, - IntArray seqPositions, - FloatArray qkvBatch, - FloatArray wrapKeyCache, - FloatArray wrapValueCache, - HalfFloatArray attnOutFP16, - int nHeads, int headSize, - int kvDim, int kvMul, - int layerIndex, int numLayers, - int contextLength, int dim) { + 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 = seqPositions.get(batchIdx); - int loff = batchIdx * (numLayers * contextLength * kvDim) + layerIndex * contextLength * kvDim; + 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); @@ -1751,15 +1631,22 @@ public static void batchedDecodeAttentionFP16Out(KernelContext context, int tileEnd = Math.min(tileC + BLOCK_C - 1, pos); int tileLen = tileEnd - tileC + 1; - for (int idx = tid; idx < tileLen * headSize; idx += localSz) { - int tInTile = idx / headSize; - int d = idx % headSize; - int kvOff = loff + (tileC + tInTile) * kvDim + kvHeadIdx * headSize + d; - kTile[tInTile * headSize + d] = wrapKeyCache.get(kvOff); - vTile[tInTile * headSize + d] = wrapValueCache.get(kvOff); + // 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; @@ -1805,146 +1692,81 @@ public static void batchedDecodeAttentionFP16Out(KernelContext context, } } - // ── Paged KV variants (block-table indirection) ───────────────────────── - // - // KV lives in a global pool of fixed-size blocks. A block holds `blockSize` - // consecutive positions of ONE sequence across ALL layers: - // pool[ physBlock*(numLayers*blockSize*kvDim) + layer*(blockSize*kvDim) - // + (pos % blockSize)*kvDim + c ] - // The per-slot block table maps a logical block to a physical one: - // physBlock = blockTable[batchIdx*maxBlocksPerSlot + pos/blockSize] - // This removes the fixed per-slot context reservation of the contiguous cache: - // slots draw blocks from a shared pool only for the tokens they actually hold, - // so the pool can be far smaller than B*ctx (and blocks can be shared for - // prefix caching). - /** - * Paged per-slot RoPE + KV write (Llama adjacent-pair). {@code blockCfg} packs - * {@code blockSize | (maxBlocksPerSlot << 16)} to stay within the task arg limit. + * {@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 batchedDecodePagedRopeWithKVCachePacked(KernelContext context, - IntArray seqPositions, - IntArray blockTable, - FloatArray qkvBatch, - FloatArray keyPool, - FloatArray valuePool, - int kvDim, int headSize, - int layerIndex, int numLayers, - int blockCfg, int dim) { - int blockSize = blockCfg & 0xFFFF; - int maxBlocksPerSlot = blockCfg >>> 16; - 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; + 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 pos = seqPositions.get(batchIdx); - int qOffset = batchIdx * qkvStride; - int kOffset = batchIdx * qkvStride + dim; - int vOffset = batchIdx * qkvStride + dim + kvDim; + 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; - 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); + 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); - 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); + // 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(); - 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 physBlock = blockTable.get(batchIdx * maxBlocksPerSlot + pos / blockSize); - int slotInBlock = pos % blockSize; - int cacheOff = physBlock * (numLayers * blockSize * kvDim) - + layerIndex * (blockSize * kvDim) + slotInBlock * kvDim; - keyPool.set(cacheOff + i, rotK0); - keyPool.set(cacheOff + i + 1, rotK1); - valuePool.set(cacheOff + i, qkvBatch.get(vOffset + i)); - valuePool.set(cacheOff + i + 1, qkvBatch.get(vOffset + i + 1)); - } - } - } - - /** - * Paged per-slot flash attention, FP16 output (Llama; {@code dim = nHeads*headSize}). - * {@code blockCfg} packs {@code blockSize | (maxBlocksPerSlot << 16)}. - */ - public static void batchedDecodePagedAttentionFP16Out(KernelContext context, - IntArray seqPositions, - IntArray blockTable, - FloatArray qkvBatch, - FloatArray keyPool, - FloatArray valuePool, - HalfFloatArray attnOutFP16, - int nHeads, int headSize, - int kvDim, int kvMul, - int layerIndex, int numLayers, - int blockCfg) { - int blockSize = blockCfg & 0xFFFF; - int maxBlocksPerSlot = blockCfg >>> 16; - int dim = nHeads * headSize; - int tid = context.localIdx; - int groupId = context.groupIdx; - int localSz = context.localGroupSizeX; - - int batchIdx = groupId / nHeads; - int h = groupId % nHeads; - int pos = seqPositions.get(batchIdx); - int layerOff = layerIndex * (blockSize * kvDim); - int kvHeadIdx = h / kvMul; - int BLOCK_C = 16; - int qkvStride = dim + 2 * kvDim; - int blockStride = numLayers * blockSize * kvDim; - - float[] qShared = context.allocateFloatLocalArray(headSize); - float[] kTile = context.allocateFloatLocalArray(BLOCK_C * headSize); - float[] vTile = context.allocateFloatLocalArray(BLOCK_C * headSize); - float[] sTile = context.allocateFloatLocalArray(BLOCK_C); - - 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; + 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; - for (int idx = tid; idx < tileLen * headSize; idx += localSz) { - int tInTile = idx / headSize; - int d = idx % headSize; - int t = tileC + tInTile; - int physBlock = blockTable.get(batchIdx * maxBlocksPerSlot + t / blockSize); - int kvOff = physBlock * blockStride + layerOff + (t % blockSize) * kvDim + kvHeadIdx * headSize + d; - kTile[tInTile * headSize + d] = keyPool.get(kvOff); - vTile[tInTile * headSize + d] = valuePool.get(kvOff); + // 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; - float score = 0.0f; - for (int d = 0; d < headSize; d++) { - score += qShared[d] * kTile[tInTile * headSize + d]; + 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] = score / TornadoMath.sqrt(headSize); + sTile[tInTile] = (Half2.lowFloat(scoreAcc) + Half2.highFloat(scoreAcc)) / TornadoMath.sqrt(headSize); } context.localBarrier(); @@ -1964,12 +1786,20 @@ public static void batchedDecodePagedAttentionFP16Out(KernelContext context, } 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; - acc0 += p * vTile[t * headSize + tid]; + Half2 v0 = vTile[t * halfHead + pair0]; + acc0 += p * (high0 ? Half2.highFloat(v0) : Half2.lowFloat(v0)); if (d1 < headSize) { - acc1 += p * vTile[t * headSize + d1]; + Half2 v1 = vTile[t * halfHead + pair1]; + acc1 += p * (high1 ? Half2.highFloat(v1) : Half2.lowFloat(v1)); } } context.localBarrier(); @@ -1983,52 +1813,6 @@ public static void batchedDecodePagedAttentionFP16Out(KernelContext context, } } - // ── On-device greedy sampling (argmax) ────────────────────────────────── - - /** - * Per-row argmax over the batched logits: one workgroup per row reduces over the - * whole vocab and writes the winning token id to {@code outTokens[b]}. Keeps the - * full logits tensor on the GPU — only B integers cross to the host, instead of - * the paddedB×vocab (~65–78 MB) D2H copy + a CPU scan every step. - * - * Worker: B workgroups × localSize threads (localSize a power of two, e.g. 256). - */ - public static void batchedArgmaxLogits(KernelContext context, - FloatArray logits, IntArray outTokens, int vocab) { - int b = context.groupIdx; - int tid = context.localIdx; - int localSz = context.localGroupSizeX; - float[] vals = context.allocateFloatLocalArray(256); - int[] idxs = context.allocateIntLocalArray(256); - - int base = b * vocab; - float best = Float.NEGATIVE_INFINITY; - int bestIdx = 0; - for (int i = tid; i < vocab; i += localSz) { - float v = logits.get(base + i); - if (v > best) { - best = v; - bestIdx = i; - } - } - vals[tid] = best; - idxs[tid] = bestIdx; - context.localBarrier(); - - for (int s = localSz / 2; s > 0; s >>= 1) { - if (tid < s) { - if (vals[tid + s] > vals[tid]) { - vals[tid] = vals[tid + s]; - idxs[tid] = idxs[tid + s]; - } - } - context.localBarrier(); - } - if (tid == 0) { - outTokens.set(b, idxs[0]); - } - } - // ── SwiGLU over the packed gate/up buffer, emitting FP16 ───────────────── /** @@ -2502,5 +2286,517 @@ public static void gemmMMAGateUpQ8(KernelContext ctx, ctx.mmaStore(c17, gateUpOut, rBase + 16, cBase + 56, outStride); } + + /** + * Batched DECODE flash attention: B independent sequences, one query token + * each. Identical online-softmax math to {@link #batchedFlashAttention}, but + * each batch slot has its OWN KV cache region and its OWN position, so slot + * {@code b} attends positions {@code 0..seqPositions[b]} of its own cache — + * the shape produced by batching B concurrent decode requests. + * + *

KV cache layout: one contiguous region of {@code numLayers * + * contextLength * kvDim} per slot, so the base for slot b, layer L is + * {@code b*numLayers*contextLength*kvDim + L*contextLength*kvDim}.

+ * + *

One workgroup per (batchIdx, head): {@code groupId = batchIdx*nHeads + h}.

+ */ + public static void batchedDecodeAttention(KernelContext context, + IntArray seqPositions, + FloatArray wrapQBatch, + FloatArray wrapKeyCache, + FloatArray wrapValueCache, + FloatArray wrapXbBatch, + int nHeads, int headSize, + int kvDim, int kvMul, + int layerIndex, int numLayers, 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 = seqPositions.get(batchIdx); // per-slot position + int loff = batchIdx * (numLayers * contextLength * kvDim) + layerIndex * contextLength * kvDim; // per-slot KV base + int kvHeadIdx = h / kvMul; + int BLOCK_C = 16; + + 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); + float[] maxHolder = context.allocateFloatLocalArray(1); + + int qOffset = batchIdx * dim + h * headSize; + for (int i = tid; i < headSize; i += localSz) { + qShared[i] = wrapQBatch.get(qOffset + i); + } + context.localBarrier(); + + 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 tileC = 0; tileC <= pos; tileC += BLOCK_C) { + int tileEnd = Math.min(tileC + BLOCK_C - 1, pos); + + for (int t = tileC + tid; t <= tileEnd; t += localSz) { + int tInTile = t - tileC; + int tileMOff = tInTile * headSize; + for (int d = 0; d < headSize; d++) { + int kvOff = loff + t * kvDim + kvHeadIdx * headSize + d; + kTile[tileMOff + d] = wrapKeyCache.get(kvOff); + vTile[tileMOff + d] = wrapValueCache.get(kvOff); + } + } + context.localBarrier(); + + 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 <= tileEnd - tileC; t++) { + if (sTile[t] > tileMax) { + tileMax = sTile[t]; + } + } + if (tid == 0) { + maxHolder[0] = tileMax; + } + context.localBarrier(); + float curTileMax = maxHolder[0]; + + float newMax = Math.max(maxScore, curTileMax); + 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 = 0; t <= tileEnd - tileC; t++) { + float expScore = TornadoMath.exp(sTile[t] - maxScore); + sumExp += expScore; + for (int d = 0; d < headSize; d++) { + output[d] += expScore * vTile[t * headSize + d]; + } + } + context.localBarrier(); + } + + float norm = (sumExp > 0.0f) ? (1.0f / sumExp) : 0.0f; + int xbOffset = batchIdx * dim + h * headSize; + for (int d = tid; d < headSize; d += localSz) { + wrapXbBatch.set(xbOffset + d, output[d] * norm); + } + } + + // ── Batched DECODE variants (per-slot KV cache + per-slot position) ────── + // + // These two kernels are the only semantic delta between batched PREFILL (B + // tokens of ONE sequence, shared causal KV) and batched DECODE (B independent + // sequences, each with its own KV region and its own position). The math is + // identical to the *Packed / *FP16Out prefill kernels above; only the KV + // addressing changes: + // pos = seqPositions[batchIdx] (per slot) + // base = batchIdx*(numLayers*ctx*kvDim) + layer*ctx*kvDim (per slot) + // The KV cache is therefore sized B*numLayers*contextLength*kvDim. + + /** + * Per-slot RoPE + KV-cache write over the packed QKV buffer (decode). + * + *

Fork of {@link #batchedRopeWithKVCachePacked}: each batch slot rotates at + * its own position {@code seqPositions[batchIdx]} and writes K/V into its own + * KV region ({@code batchIdx} stride = {@code numLayers*contextLength*kvDim}).

+ */ + public static void batchedDecodeRopeWithKVCachePacked(KernelContext context, + IntArray seqPositions, + FloatArray qkvBatch, + FloatArray wrapKeyCache, + FloatArray wrapValueCache, + int kvDim, int headSize, + int layerIndex, int numLayers, + 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 = seqPositions.get(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); + + 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); + + 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 slotBase = batchIdx * (numLayers * contextLength * kvDim); + int cacheOff = slotBase + layerIndex * contextLength * kvDim + pos * kvDim; + wrapKeyCache.set(cacheOff + i, rotK0); + wrapKeyCache.set(cacheOff + i + 1, rotK1); + wrapValueCache.set(cacheOff + i, qkvBatch.get(vOffset + i)); + wrapValueCache.set(cacheOff + i + 1, qkvBatch.get(vOffset + i + 1)); + } + } + } + + /** + * Per-slot flash attention over the packed QKV buffer, FP16 output (decode). + * + *

Fork of {@link #batchedFlashAttentionFP16Out}: each batch slot attends + * over {@code 0..seqPositions[batchIdx]} of its OWN KV region. Same + * register-partitioned P·V accumulation and FP16 emission.

+ * + *

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

+ */ + public static void batchedDecodeAttentionFP16Out(KernelContext context, + IntArray seqPositions, + FloatArray qkvBatch, + FloatArray wrapKeyCache, + FloatArray wrapValueCache, + HalfFloatArray attnOutFP16, + int nHeads, int headSize, + int kvDim, int kvMul, + int layerIndex, int numLayers, + 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 = seqPositions.get(batchIdx); + int loff = batchIdx * (numLayers * contextLength * kvDim) + layerIndex * contextLength * kvDim; + int kvHeadIdx = h / kvMul; + int BLOCK_C = 16; + int qkvStride = dim + 2 * kvDim; + + float[] qShared = context.allocateFloatLocalArray(headSize); + float[] kTile = context.allocateFloatLocalArray(BLOCK_C * headSize); + float[] vTile = context.allocateFloatLocalArray(BLOCK_C * headSize); + float[] sTile = context.allocateFloatLocalArray(BLOCK_C); + + 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; + + for (int idx = tid; idx < tileLen * headSize; idx += localSz) { + int tInTile = idx / headSize; + int d = idx % headSize; + int kvOff = loff + (tileC + tInTile) * kvDim + kvHeadIdx * headSize + d; + kTile[tInTile * headSize + d] = wrapKeyCache.get(kvOff); + vTile[tInTile * headSize + d] = wrapValueCache.get(kvOff); + } + context.localBarrier(); + + 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)); + } + } + + // ── Paged KV variants (block-table indirection) ───────────────────────── + // + // KV lives in a global pool of fixed-size blocks. A block holds `blockSize` + // consecutive positions of ONE sequence across ALL layers: + // pool[ physBlock*(numLayers*blockSize*kvDim) + layer*(blockSize*kvDim) + // + (pos % blockSize)*kvDim + c ] + // The per-slot block table maps a logical block to a physical one: + // physBlock = blockTable[batchIdx*maxBlocksPerSlot + pos/blockSize] + // This removes the fixed per-slot context reservation of the contiguous cache: + // slots draw blocks from a shared pool only for the tokens they actually hold, + // so the pool can be far smaller than B*ctx (and blocks can be shared for + // prefix caching). + + /** + * Paged per-slot RoPE + KV write (Llama adjacent-pair). {@code blockCfg} packs + * {@code blockSize | (maxBlocksPerSlot << 16)} to stay within the task arg limit. + */ + public static void batchedDecodePagedRopeWithKVCachePacked(KernelContext context, + IntArray seqPositions, + IntArray blockTable, + FloatArray qkvBatch, + FloatArray keyPool, + FloatArray valuePool, + int kvDim, int headSize, + int layerIndex, int numLayers, + int blockCfg, int dim) { + int blockSize = blockCfg & 0xFFFF; + int maxBlocksPerSlot = blockCfg >>> 16; + 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 = seqPositions.get(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); + + 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); + + 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 physBlock = blockTable.get(batchIdx * maxBlocksPerSlot + pos / blockSize); + int slotInBlock = pos % blockSize; + int cacheOff = physBlock * (numLayers * blockSize * kvDim) + + layerIndex * (blockSize * kvDim) + slotInBlock * kvDim; + keyPool.set(cacheOff + i, rotK0); + keyPool.set(cacheOff + i + 1, rotK1); + valuePool.set(cacheOff + i, qkvBatch.get(vOffset + i)); + valuePool.set(cacheOff + i + 1, qkvBatch.get(vOffset + i + 1)); + } + } + } + + /** + * Paged per-slot flash attention, FP16 output (Llama; {@code dim = nHeads*headSize}). + * {@code blockCfg} packs {@code blockSize | (maxBlocksPerSlot << 16)}. + */ + public static void batchedDecodePagedAttentionFP16Out(KernelContext context, + IntArray seqPositions, + IntArray blockTable, + FloatArray qkvBatch, + FloatArray keyPool, + FloatArray valuePool, + HalfFloatArray attnOutFP16, + int nHeads, int headSize, + int kvDim, int kvMul, + int layerIndex, int numLayers, + int blockCfg) { + int blockSize = blockCfg & 0xFFFF; + int maxBlocksPerSlot = blockCfg >>> 16; + int dim = nHeads * headSize; + int tid = context.localIdx; + int groupId = context.groupIdx; + int localSz = context.localGroupSizeX; + + int batchIdx = groupId / nHeads; + int h = groupId % nHeads; + int pos = seqPositions.get(batchIdx); + int layerOff = layerIndex * (blockSize * kvDim); + int kvHeadIdx = h / kvMul; + int BLOCK_C = 16; + int qkvStride = dim + 2 * kvDim; + int blockStride = numLayers * blockSize * kvDim; + + float[] qShared = context.allocateFloatLocalArray(headSize); + float[] kTile = context.allocateFloatLocalArray(BLOCK_C * headSize); + float[] vTile = context.allocateFloatLocalArray(BLOCK_C * headSize); + float[] sTile = context.allocateFloatLocalArray(BLOCK_C); + + 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; + + for (int idx = tid; idx < tileLen * headSize; idx += localSz) { + int tInTile = idx / headSize; + int d = idx % headSize; + int t = tileC + tInTile; + int physBlock = blockTable.get(batchIdx * maxBlocksPerSlot + t / blockSize); + int kvOff = physBlock * blockStride + layerOff + (t % blockSize) * kvDim + kvHeadIdx * headSize + d; + kTile[tInTile * headSize + d] = keyPool.get(kvOff); + vTile[tInTile * headSize + d] = valuePool.get(kvOff); + } + context.localBarrier(); + + 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)); + } + } + + // ── On-device greedy sampling (argmax) ────────────────────────────────── + + /** + * Per-row argmax over the batched logits: one workgroup per row reduces over the + * whole vocab and writes the winning token id to {@code outTokens[b]}. Keeps the + * full logits tensor on the GPU — only B integers cross to the host, instead of + * the paddedB×vocab (~65–78 MB) D2H copy + a CPU scan every step. + * + * Worker: B workgroups × localSize threads (localSize a power of two, e.g. 256). + */ + public static void batchedArgmaxLogits(KernelContext context, + FloatArray logits, IntArray outTokens, int vocab) { + int b = context.groupIdx; + int tid = context.localIdx; + int localSz = context.localGroupSizeX; + float[] vals = context.allocateFloatLocalArray(256); + int[] idxs = context.allocateIntLocalArray(256); + + int base = b * vocab; + float best = Float.NEGATIVE_INFINITY; + int bestIdx = 0; + for (int i = tid; i < vocab; i += localSz) { + float v = logits.get(base + i); + if (v > best) { + best = v; + bestIdx = i; + } + } + vals[tid] = best; + idxs[tid] = bestIdx; + context.localBarrier(); + + for (int s = localSz / 2; s > 0; s >>= 1) { + if (tid < s) { + if (vals[tid + s] > vals[tid]) { + vals[tid] = vals[tid + s]; + idxs[tid] = idxs[tid + s]; + } + } + context.localBarrier(); + } + if (tid == 0) { + outTokens.set(b, idxs[0]); + } + } + // @formatter:on } 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..2e027a69 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 { @@ -58,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; @@ -77,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; @@ -541,6 +550,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 +955,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 +1627,250 @@ 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); + } + } + + /** + * 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. * @@ -1695,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; @@ -1719,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; @@ -1743,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; @@ -2174,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 @@ -2323,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/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..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); @@ -203,19 +212,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 +283,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; } @@ -283,7 +313,13 @@ 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; if (layerIndex == 0) { // First layer: Transfer initial data to device (one-time transfer) unifiedLayer.transferToDevice(DataTransferMode.EVERY_EXECUTION, @@ -298,9 +334,12 @@ 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); + 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, @@ -314,17 +353,98 @@ 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 state.positionHolder, state.wrapXbFP16); + if (splitKvAttentionEnabled()) { + unifiedLayer.consumeFromDevice(pred, llamaState.wrapAttSplit); + } } return unifiedLayer; } private TaskGraph configureAttention(TaskGraph unifiedLayer, int layerIndex) { - if (schedulerType == SchedulerType.NVIDIA) { + 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. + 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 e6a205f3..acfed1c2 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; @@ -314,24 +315,43 @@ 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 + } if (isMetalBackend) { - // Metal: single-workgroup-per-head online-softmax attention, writing directly to wrapXb. - // No combine phase needed (TornadoVM fails to JIT the multi-workgroup split-KV kernel here). + // Metal cannot JIT the multi-workgroup split-KV kernel, so use the + // single-workgroup-per-head online-softmax kernel, which writes wrapXb directly and + // needs no combine phase. It reads the FP32 KV cache, so it is not compatible with + // the FP16 KV cache; useFp16KVCache() below is therefore only consulted off Metal. unifiedLayer.task("attention", TransformerComputeKernelsLayered::processHeadsFlashAttention, context, @@ -349,21 +369,41 @@ protected TaskGraph createFFNLayerTaskGraph(int layerIndex) { } else { // Split-KV (flash-decoding) attention. // Phase 1: split each head's KV range across attentionSplits workgroups; partials -> wrapAttSplit. - unifiedLayer.task("attention", - TransformerComputeKernelsLayered::processHeadsFlashAttentionSplitKV, - context, - qwen3State.wrapQ, // query vectors - qwen3State.wrapKeyCache, // key cache - qwen3State.wrapValueCache, // value cache - qwen3State.wrapAttSplit, // scratch: per-head split partials (compact layout) - config.numberOfHeads(), // nHeads - nEmbdHead, // headSize - nEmbdGqa, // kvDim - gqa, // kvMul (nHeads / nHeadKv) - qwen3State.positionHolder, // position - layerIndex, // layer index - config.contextLength(), // context length - attentionSplits); // number of KV splits per head + if (useFp16KVCache()) { + unifiedLayer.task("attention", + State.ATTENTION_DEEP_HALF2 + ? TransformerComputeKernelsLayered::processHeadsFlashAttentionSplitKVFP16Packed + : 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, @@ -475,7 +515,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; } @@ -500,6 +544,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); @@ -508,7 +554,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 { @@ -519,8 +565,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..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 @@ -44,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 @@ -55,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), @@ -66,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/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/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 1e32485e..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 @@ -32,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); @@ -42,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/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); 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; } 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