Enhance GPU token generation by adding token consumer support and removing unused performance metrics

Author: mikepapadim

src/main/java/com/example/LlamaApp.java

@@ -10,6 +10,7 @@
 import com.example.inference.engine.impl.Options;
 import com.example.loader.weights.ModelLoader;
 import com.example.loader.weights.State;
+import com.example.tokenizer.impl.Tokenizer;
 import com.example.tornadovm.FloatArrayUtils;
 import uk.ac.manchester.tornado.api.types.arrays.FloatArray;
 
@@ -160,6 +161,8 @@ static void runInteractive(Llama model, Sampler sampler, Options options) {
                 System.err.println("Ran out of context length...");
                 break;
             }
+            System.out.print("\n");
+
         }
     }
 
@@ -188,16 +191,8 @@ static void runInstructOnce(Llama model, Sampler sampler, Options options) {
         Set<Integer> stopTokens = chatFormat.getStopTokens();
         if (USE_TORNADOVM) {
             // Call generateTokensGPU without the token consumer parameter
-            responseTokens = Llama.generateTokensGPU(model, state, 0, promptTokens, stopTokens, options.maxTokens(), sampler, options.echo());
-            // Handle token output separately if needed
-            // You might need to iterate through responseTokens and process them
-            if (options.stream()) {
-                for (Integer token : responseTokens) {
-                    if (!model.tokenizer().isSpecialToken(token)) {
-                        System.out.print(model.tokenizer().decode(List.of(token)));
-                    }
-                }
-            }
+            responseTokens = Llama.generateTokensGPU(model, state, 0, promptTokens, stopTokens, options.maxTokens(),
+                    sampler, options.echo(), options.stream() ? tokenConsumer : null);
         } else {
             // CPU path still uses the token consumer
             responseTokens = Llama.generateTokens(model, state, 0, promptTokens, stopTokens, options.maxTokens(), sampler, options.echo(), tokenConsumer);

src/main/java/com/example/inference/engine/impl/Llama.java

@@ -192,7 +192,8 @@ public static FloatArray forwardTornadoVM( //
         return tornadoVMMasterPlan.tornadoVMForwardExecuteLayered(position);
     }
 
-    public static List<Integer> generateTokensGPU(Llama model, State state, int startPosition, List<Integer> promptTokens, Set<Integer> stopTokens, int maxTokens, Sampler sampler, boolean echo) {
+    public static List<Integer> generateTokensGPU(Llama model, State state,
+            int startPosition, List<Integer> promptTokens, Set<Integer> stopTokens, int maxTokens, Sampler sampler, boolean echo,   IntConsumer onTokenGenerated) {
         // 1. Pre-allocate the TornadoVM plan just once
         TornadoVMMasterPlan tornadoVMPlan = TornadoVMMasterPlan.initializeTornadoVMPlan(state, model);
 
@@ -247,8 +248,13 @@ public static List<Integer> generateTokensGPU(Llama model, State state, int star
                 // Sample next token - use GPU sampling if available
                 nextToken = sampler.sampleToken(logits);
 
+                // Add token consumer support
+                if (onTokenGenerated != null) {
+                    onTokenGenerated.accept(nextToken);
+                }
+
                 // Output if needed
-                if (echo) {
+                if (echo && onTokenGenerated == null) {
                     System.err.print(Tokenizer.replaceControlCharacters(model.tokenizer().decode(List.of(nextToken))));
                 }
 
@@ -359,23 +365,6 @@ public static List<Integer> generateTokens(Llama model, State state, int startPo
         return generatedTokens;
     }
 
-    /**
-     * Print performance metrics for the generation process
-     */
-    private static void printPerformanceMetrics(long startNanos, long inferenceStartNanos, int promptTokenCount, int generatedTokenCount) {
-        long endNanos = System.nanoTime();
-        long totalNanos = endNanos - startNanos;
-        long inferenceNanos = inferenceStartNanos > 0 ? endNanos - inferenceStartNanos : 0;
-        long promptNanos = inferenceStartNanos - startNanos;
-        int totalTokens = promptTokenCount + generatedTokenCount;
-
-        double totalTokensPerSecond = totalTokens / (totalNanos / 1_000_000_000.0);
-        double promptTokensPerSecond = promptTokenCount > 0 ? promptTokenCount / (promptNanos / 1_000_000_000.0) : 0;
-        double inferenceTokensPerSecond = generatedTokenCount > 0 ? generatedTokenCount / (inferenceNanos / 1_000_000_000.0) : 0;
-
-        System.err.printf("\n%n%.2f tokens/s (%d) [PrEval %.2f tokens/s (%d), TokGen %.2f tokens/s (%d)]%n", totalTokensPerSecond, totalTokens, promptTokensPerSecond, promptTokenCount,
-                inferenceTokensPerSecond, generatedTokenCount);
-    }
 
     public State createNewState() {
         State state = new State(configuration(), -1);