[prf/dec][refactor] Standardize task graph and grid scheduler naming for prefill and decode paths in TornadoVM

Author: orionpapadakis

src/main/java/org/beehive/gpullama3/tornadovm/TornadoVMMasterPlanWithBatchPrefillDecode.java

@@ -110,10 +110,10 @@ public class TornadoVMMasterPlanWithBatchPrefillDecode implements TornadoVMMaste
 
     /** Graph 0: B×dim FP16 embeddings → FP32 wrapXBatch. */
     private TaskGraph buildBatchPrefillActivationGraph(KernelContext ctx) {
-        return new TaskGraph("batchActivation")
+        return new TaskGraph("prefillActivation")
                 .transferToDevice(DataTransferMode.FIRST_EXECUTION, ctx, state.wrapXBatch)
                 .transferToDevice(DataTransferMode.EVERY_EXECUTION, state.embeddingXBatch)
-                .task("batchUpdateX", TransformerComputeKernels::convertFP16toFP32,
+                .task("updateX", TransformerComputeKernels::convertFP16toFP32,
                         ctx, state.embeddingXBatch, state.wrapXBatch)
                 .persistOnDevice(state.wrapXBatch);
     }
@@ -128,7 +128,7 @@ private TaskGraph buildBatchPrefillActivationGraph(KernelContext ctx) {
      * not forwarded in interpreter (non-CUDA-graph) mode.</p>
      */
     private TaskGraph buildDecodeActivationGraph(KernelContext ctx, String lastBatchLayerID) {
-        return new TaskGraph("decodeActivationUpdate")
+        return new TaskGraph("decodeActivation")
                 .consumeFromDevice(lastBatchLayerID, state.wrapKeyCache, state.wrapValueCache)   // KV pass-through
                 .transferToDevice(DataTransferMode.EVERY_EXECUTION, state.embeddingX)
                 .task("updateX",
@@ -153,7 +153,7 @@ public TornadoExecutionPlan createExecutionPlan() {
         // [0] Batch prefill activation ────────────────────────────────────────────────
         KernelContext batchActCtx = new KernelContext();
         all.add(buildBatchPrefillActivationGraph(batchActCtx).snapshot());
-        gridScheduler.addWorkerGrid("batchActivation.batchUpdateX",
+        gridScheduler.addWorkerGrid("prefillActivation.updateX",
                 WorkerGridFactory.genericWorker(batchSize * config.dim(), 128));
 
         // [1..N] Batch prefill layer graphs ───────────────────────────────────────────
@@ -165,14 +165,14 @@ public TornadoExecutionPlan createExecutionPlan() {
         // [N+1] Decode activation (with KV-cache pass-through) ────────────────
         KernelContext decodeActCtx = new KernelContext();
         all.add(buildDecodeActivationGraph(decodeActCtx, batchLayers.getLastLayerTaskGraphID()).snapshot());
-        gridScheduler.addWorkerGrid("decodeActivationUpdate.updateX",
+        gridScheduler.addWorkerGrid("decodeActivation.updateX",
                 WorkerGridFactory.genericWorker(config.dim(), 128));
 
         // [N+2..2N+1] Decode layer graphs  ────────────────────────────────────
         // Layer 0 uses consumeFromDevice for KV cache (no FIRST_EXECUTION upload).
         LlamaFP16FFNLayersDecode decodeLayers =
                 new LlamaFP16FFNLayersDecode(
-                        "llamaFFNDecode", state, weights, config, schedulerType);
+                        "decode", state, weights, config, schedulerType);
         all.addAll(decodeLayers.getFFNLayerImmutableTaskGraphs());
         decodeLayers.updateGridScheduler(gridScheduler);
 

src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/decode/LlamaFP16FFNLayersDecode.java

@@ -43,7 +43,7 @@ public LlamaFP16FFNLayersDecode(String taskGraph, LlamaState state,
      */
     @Override
     protected String predecessorGraphName(int layerIndex) {
-        return (layerIndex == 0) ? "decodeActivationUpdate" : "layer_" + (layerIndex - 1);
+        return (layerIndex == 0) ? "decodeActivation" : "layer_" + (layerIndex - 1);
     }
 
     @Override
@@ -60,7 +60,7 @@ protected TaskGraph configureLayerDataTransfers(TaskGraph layer, int layerIndex)
                     state.wrapQ, state.wrapK, state.wrapV,
                     state.wrapAtt, state.wrapHb, state.wrapXbFP16);
             // Explicit source — must match the TaskGraph name in buildDecodeActivationGraph().
-            layer.consumeFromDevice("decodeActivationUpdate", state.wrapKeyCache, state.wrapValueCache);
+            layer.consumeFromDevice("decodeActivation", state.wrapKeyCache, state.wrapValueCache);
         } 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),

src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/prefill/LlamaFP16LayersBatchPrefill.java

@@ -51,37 +51,32 @@ public LlamaFP16LayersBatchPrefill(LlamaState state, LlamaTornadoWeights weights
 
     // @formatter:off
     private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) {
-        String graphName = "batchLayer_" + layerIndex;
+        String graphName = "batchPrefillLayer_" + layerIndex;
         if (layerIndex == config.numberOfLayers() - 1) lastLayerTaskGraphID = graphName;
 
-        TaskGraph layer = new TaskGraph(graphName);
+        TaskGraph batchPrefillLayer = new TaskGraph(graphName);
 
         // ── Data Transfers ─────────────────────────────────────────────────────
         if (layerIndex == 0) {
             // batchStartPosHolder is set by host before each chunk → EVERY_EXECUTION
-            layer.transferToDevice(DataTransferMode.EVERY_EXECUTION, state.batchStartPosHolder);
+            batchPrefillLayer.transferToDevice(DataTransferMode.EVERY_EXECUTION, state.batchStartPosHolder);
             // Allocate persistent GPU-side intermediates once
-            layer.transferToDevice(DataTransferMode.FIRST_EXECUTION,
+            batchPrefillLayer.transferToDevice(DataTransferMode.FIRST_EXECUTION,
                     context,
                     state.attnScaleBatch, state.ffnScaleBatch,
                     state.wrapXbFP16Batch,
                     state.wrapQBatch, state.wrapKBatch, state.wrapVBatch,
                     state.wrapXbBatch,
                     state.wrapHbBatch,
                     state.wrapKeyCache, state.wrapValueCache);
-            // wrapXBatch produced by the batch activation graph.
-            // Explicit source name required: the no-arg form uses the current graph's own
-            // name ("batchLayer_0") which never matches "batchActivation" in interpreter mode,
-            // causing wrapXBatch to be re-uploaded from host (zeros) instead of using the
-            // FP32 embeddings computed by the activation graph's convertFP16toFP32 kernel.
-            layer.consumeFromDevice("batchActivation", state.wrapXBatch);
+            // wrapXBatch produced by the prefillActivation graph and persists in device memory
+            // to consume it from there we should use the explicit uniqueTaskGraph name
+            // the no-arg form would use current graph name, which causes NPE without CUDA Graphs
+            batchPrefillLayer.consumeFromDevice("prefillActivation", state.wrapXBatch);
         } else {
-            // Explicit predecessor name for all objects.
-            // The no-arg form would use "batchLayer_k" as the source key, which never matches
-            // "batchLayer_{k-1}" in interpreter mode — every object would be re-uploaded from
-            // host (zeros or stale), corrupting the KV cache written by the previous layer.
-            String pred = "batchLayer_" + (layerIndex - 1);
-            layer.consumeFromDevice(pred,
+            // for the same reasons as above, we should use the explicit uniqueTaskGraph name to consume
+            String pred = "batchPrefillLayer_" + (layerIndex - 1);
+            batchPrefillLayer.consumeFromDevice(pred,
                     context,
                     state.wrapXBatch,
                     state.wrapXbFP16Batch,
@@ -94,7 +89,7 @@ private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) {
         }
 
         // Per-layer weights: upload once
-        layer.transferToDevice(DataTransferMode.FIRST_EXECUTION,
+        batchPrefillLayer.transferToDevice(DataTransferMode.FIRST_EXECUTION,
                 weights.rms_att_weightLayered[layerIndex].asFloatArray(),
                 weights.wqLayered[layerIndex].asHalfFloatArray(),
                 weights.wkLayered[layerIndex].asHalfFloatArray(),
@@ -110,18 +105,18 @@ private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) {
         int hidDim   = config.hiddenDim();
 
         // ── Attention Block ────────────────────────────────────────────────────
-        layer.task("batch_attn_rms",
+        batchPrefillLayer.task("batch_attn_rms",
                 TransformerBatchPrefillKernels::batchedRmsReduce,
                 context, state.wrapXBatch, state.attnScaleBatch,
                 dim, config.rmsNormEps());
 
-        layer.task("batch_attn_rms_apply",
+        batchPrefillLayer.task("batch_attn_rms_apply",
                 TransformerBatchPrefillKernels::batchedRmsApplyFP16,
                 context, state.wrapXbFP16Batch, state.wrapXBatch,
                 weights.rms_att_weightLayered[layerIndex].asFloatArray(),
                 state.attnScaleBatch, dim);
 
-        layer.task("batch_qkv",
+        batchPrefillLayer.task("batch_qkv",
                 TransformerBatchPrefillKernels::batchedFusedQKVMatmul,
                 context,
                 state.wrapXbFP16Batch,
@@ -131,34 +126,34 @@ private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) {
                 weights.wvLayered[layerIndex].asHalfFloatArray(),
                 dim, kvDim, LOCAL_WORK_GROUP_SIZE);
 
-        layer.task("batch_rope_kv",
+        batchPrefillLayer.task("batch_rope_kv",
                 TransformerBatchPrefillKernels::batchedRopeWithKVCache,
                 context, state.batchStartPosHolder,
                 state.wrapQBatch, state.wrapKBatch, state.wrapVBatch,
                 state.wrapKeyCache, state.wrapValueCache,
                 kvDim, config.headSize(), layerIndex, config.contextLength(), dim);
 
-        layer.task("batch_attention",
+        batchPrefillLayer.task("batch_attention",
                 TransformerBatchPrefillKernels::batchedFlashAttention,
                 context, state.batchStartPosHolder,
                 state.wrapQBatch, state.wrapKeyCache, state.wrapValueCache,
                 state.wrapXbBatch,
                 config.numberOfHeads(), config.headSize(),
                 kvDim, config.kvMul(), layerIndex, config.contextLength(), dim);
 
-        layer.task("batch_attn_out",
+        batchPrefillLayer.task("batch_attn_out",
                 TransformerBatchPrefillKernels::batchedMatVecWithResidual,
                 context, state.wrapXbBatch, state.wrapXBatch,
                 weights.woLayered[layerIndex].asHalfFloatArray(),
                 dim, dim, LOCAL_WORK_GROUP_SIZE);
 
         // ── FFN Block ──────────────────────────────────────────────────────────
-        layer.task("batch_ffn_rms",
+        batchPrefillLayer.task("batch_ffn_rms",
                 TransformerBatchPrefillKernels::batchedFFNRmsReduce,
                 context, state.wrapXBatch, state.ffnScaleBatch,
                 dim, config.rmsNormEps());
 
-        layer.task("batch_ffn_gate_up",
+        batchPrefillLayer.task("batch_ffn_gate_up",
                 TransformerBatchPrefillKernels::batchedFusedRmsNormFFNGateUp,
                 context, state.wrapXBatch, state.wrapHbBatch,
                 weights.rms_ffn_weightLayered[layerIndex].asFloatArray(),
@@ -167,17 +162,17 @@ private TaskGraph createBatchPrefillLayerTaskGraph(int layerIndex) {
                 weights.w3Layered[layerIndex].asHalfFloatArray(),
                 dim, hidDim, LOCAL_WORK_GROUP_SIZE);
 
-        layer.task("batch_ffn_down",
+        batchPrefillLayer.task("batch_ffn_down",
                 TransformerBatchPrefillKernels::batchedMatVecWithResidual,
                 context, state.wrapHbBatch, state.wrapXBatch,
                 weights.w2Layered[layerIndex].asHalfFloatArray(),
                 hidDim, dim, LOCAL_WORK_GROUP_SIZE);
 
         // Persist wrapXBatch for the next layer, and KV cache so the decode
         // layers can consume it via the activation graph pass-through.
-        layer.persistOnDevice(state.wrapXBatch, state.wrapKeyCache, state.wrapValueCache);
+        batchPrefillLayer.persistOnDevice(state.wrapXBatch, state.wrapKeyCache, state.wrapValueCache);
 
-        return layer;
+        return batchPrefillLayer;
     }
     // @formatter:on
 
@@ -218,7 +213,7 @@ public void updateGridScheduler(GridScheduler scheduler) {
                 batchSize * hidDim * LOCAL_WORK_GROUP_SIZE, LOCAL_WORK_GROUP_SIZE);
 
         for (int i = 0; i < config.numberOfLayers(); i++) {
-            String p = "batchLayer_" + i + ".";
+            String p = "batchPrefillLayer_" + i + ".";
             scheduler.addWorkerGrid(p + "batch_attn_rms",     rmsWorker);
             scheduler.addWorkerGrid(p + "batch_attn_rms_apply", rmsApplyWorker);
             scheduler.addWorkerGrid(p + "batch_qkv",          qkvWorker);