FP16 KV cache with packed half2 split-KV attention

.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/<timestamp>/` 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.

.claude/agents/gpullama-benchmarking-specialist.md

@@ -84,6 +84,35 @@ When testing a feature flag, change one variable at a time and keep the baseline
 structurally identical to the treatment command — differences beyond the flag under test
 invalidate the comparison.
 
+## Pre-Flight Checks (do these before trusting ANY A/B)
+
+1. **Stale jar check.** `llama-tornado` picks the jar by reverse-sorting
+   `target/gpu-llama3-*.jar`; after a build-suffix change (e.g. `-Djdk.version.suffix`),
+   a leftover jar from the old suffix can sort ABOVE the fresh one and both "sides" of the
+   A/B silently run identical old code. Before a flag comparison: `ls -la target/*.jar`,
+   delete stale jars, and confirm the surviving jar's mtime postdates the last build.
+2. **Flag liveness check.** Prove the flag under test actually reaches the JVM and changes
+   behavior (a cheap probe run where the flag has an observable effect) before spending a
+   full sweep on it. A flag that is silently ignored produces a perfectly clean null result.
+3. **GPU idle check.** Kill leftover `java` processes between runs and confirm
+   `nvidia-smi` is back to idle memory. Never put the target process name literally inside
+   a `pkill -f` command line that also does other work — `pkill -f` matches your own shell's
+   command string and kills it.
+
+## Flag Ladders (preferred over isolated pairs)
+
+When several stacked optimizations exist, benchmark them as a cumulative ladder from the
+pre-feature baseline (A → A+f1 → A+f1+f2 → ...) and report BOTH the per-step delta and the
+cumulative delta vs baseline. This attributes the gain to the right layer; an isolated pair
+at the top of the stack can show "no gain" purely because a lower layer already captured it.
+
+Depth caveat: decode tok/s decays with KV depth, so two runs of different generated length
+are NOT comparable. For equal-length runs use `-Dllama.bench.ignoreEos=true` (disables the
+stop-token break so generation runs to `--max-tokens`); without it, models stop early at
+their end-of-turn token and comparisons are rough estimates only. Flag-gated attention/KV
+optimizations also pay off more at depth — a null result at shallow depth (~512) does not
+rule out a win at 2048.
+
 ## When Results Point to a Bottleneck
 
 State whether the limiting factor looks memory-bound, compute-bound, launch-overhead-bound,

.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:

pom.xml

@@ -39,8 +39,8 @@
         <properties>
             <!-- CI-friendly version: resolved by flatten-maven-plugin at build time -->
             <revision>0.5.0</revision>
-            <tornadovm.base.version>5.0.0</tornadovm.base.version>
-            <jdk.version.suffix>-jdk21</jdk.version.suffix>
+            <tornadovm.base.version>5.1.1</tornadovm.base.version>
+            <jdk.version.suffix>-jdk21-dev</jdk.version.suffix>
             <!-- CUDA backend is only available after 5.0.0 TornadoVM version -->
             <tornadovm.version>${tornadovm.base.version}${jdk.version.suffix}</tornadovm.version>
             <!-- Compiler defaults (overridden by JDK profiles below) -->
@@ -147,7 +147,7 @@
                 <properties>
                     <maven.compiler.source>21</maven.compiler.source>
                     <maven.compiler.target>21</maven.compiler.target>
-                    <jdk.version.suffix>-jdk21</jdk.version.suffix>
+                    <jdk.version.suffix>-jdk21-dev</jdk.version.suffix>
                     <tornadovm.version>${tornadovm.base.version}${jdk.version.suffix}</tornadovm.version>
                 </properties>
                 <build>

src/main/java/org/beehive/gpullama3/inference/InferenceEngine.java

@@ -34,6 +34,9 @@
  */
 public final class InferenceEngine {
 
+    /** Benchmarking aid: keep decoding past the stop token so every run generates the same token count. */
+    private static final boolean IGNORE_EOS = Boolean.getBoolean("llama.bench.ignoreEos");
+
     private InferenceEngine() {
         //prevent instantiation
     }
@@ -344,7 +347,7 @@ public static List<Integer> generateTokensGPULlama(Model model, State state, int
                 generatedTokens.add(nextToken);
 
                 // Check stop condition
-                if (stopTokens.contains(nextToken)) {
+                if (!IGNORE_EOS && stopTokens.contains(nextToken)) {
                     break;
                 }
             }
@@ -438,7 +441,7 @@ public static List<Integer> generateTokensGPUQwen3(Model model, State state, int
             }
 
             // Check for stop condition
-            if (stopTokens.contains(nextToken)) {
+            if (!IGNORE_EOS && stopTokens.contains(nextToken)) {
                 break;
             }
 

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<Integer> generateTokensLlama(Model model,
                 onTokenGenerated.accept(nextToken);
             }
 
-            if (stopTokens.contains(nextToken)) {
+            if (!IGNORE_EOS && stopTokens.contains(nextToken)) {
                 break;
             }
 
@@ -233,7 +236,7 @@ public static List<Integer> generateTokensGPULlama(Model model,
                 onTokenGenerated.accept(nextToken);
             }
 
-            if (stopTokens.contains(nextToken)) {
+            if (!IGNORE_EOS && stopTokens.contains(nextToken)) {
                 break;
             }
 

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<Integer> generateTokensLlama(Model model,
                 onTokenGenerated.accept(nextToken);
             }
 
-            if (stopTokens.contains(nextToken)) {
+            if (!IGNORE_EOS && stopTokens.contains(nextToken)) {
                 break;
             }
 
@@ -183,7 +186,7 @@ public static List<Integer> generateTokensGPULlama(Model model,
                 onTokenGenerated.accept(nextToken);
             }
 
-            if (stopTokens.contains(nextToken)) {
+            if (!IGNORE_EOS && stopTokens.contains(nextToken)) {
                 break;
             }
 

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);
 

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);
 

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;
 
     public TornadoNativeArray embeddingX;
@@ -135,6 +158,8 @@ protected State(Configuration config, int batchsize) {
         // dim vs kvdim
         this.wrapKeyCache = fields.wrapKeyCache;
         this.wrapValueCache = fields.wrapValueCache;
+        this.wrapKeyCacheFP16 = fields.wrapKeyCacheFP16;
+        this.wrapValueCacheFP16 = fields.wrapValueCacheFP16;
         this.wrapAtt = fields.wrapAtt;
         this.positionHolder = fields.positionHolder;
 
@@ -218,6 +243,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;