Add missing pieces for Qwen2.5 & Deepseek-r1-distill-qwen with tornado

Author: orionpapadakis

src/main/java/org/beehive/gpullama3/inference/state/Qwen2State.java

@@ -4,17 +4,16 @@
 import org.beehive.gpullama3.core.model.tensor.FloatTensor;
 import org.beehive.gpullama3.model.Configuration;
 import org.beehive.gpullama3.model.qwen2.Qwen2Configuration;
+import uk.ac.manchester.tornado.api.types.arrays.FloatArray;
+import uk.ac.manchester.tornado.api.types.arrays.IntArray;
 
 import java.util.stream.Stream;
 
 public class Qwen2State extends State {
 
-    //Qwen2 specific fields TODO
-
     public Qwen2State(Configuration config, int batchsize) {
         super(config, batchsize);
-        // Initialize Qwen2-specific fields TODO
-        Qwen2Configuration qwen2Config = (Qwen2Configuration) config;
+        this.localSize = 32;
     }
     @Override
     protected StateFields createStateFields(Configuration configuration) {
@@ -40,6 +39,30 @@ protected StateFields createStateFields(Configuration configuration) {
         fields.keyCache = Stream.generate(() -> ArrayFloatTensor.allocate(config.contextLength(), nEmbdGqa)).limit(config.numberOfLayers()).toArray(FloatTensor[]::new);
         fields.valueCache = Stream.generate(() -> ArrayFloatTensor.allocate(config.contextLength(), nEmbdGqa)).limit(config.numberOfLayers()).toArray(FloatTensor[]::new);
 
+        // TornadoVM wrappers with Qwen2 dimensions
+        fields.wrapX = new FloatArray(config.dim());
+        fields.wrapXb = new FloatArray(config.dim());
+        fields.wrapXb2 = new FloatArray(config.dim());
+        fields.wrapHb = new FloatArray(config.hiddenDim());
+        fields.wrapHb2 = new FloatArray(config.hiddenDim());
+
+        fields.wrapLogits = new FloatArray(config.vocabularySize());
+        fields.wrapQ = new FloatArray(config.dim());
+        fields.wrapK = new FloatArray(config.kvDim());
+        fields.wrapV = new FloatArray(config.kvDim());
+
+        fields.wrapKeyCache = new FloatArray(config.contextLength() * nEmbdGqa * config.numberOfLayers());
+        fields.wrapValueCache = new FloatArray(config.contextLength() * nEmbdGqa * config.numberOfLayers());
+        fields.wrapValueCache.init(0.f);
+        fields.wrapKeyCache.init(0.f);
+        fields.wrapAtt = new FloatArray(config.numberOfHeads() * config.contextLength());
+        fields.positionHolder = new IntArray(1);
+
+        // Temporary arrays
+        fields.temp = new FloatArray(1 + ((config.dim() + localSize - 1) / localSize));
+        fields.tempFFN = new FloatArray(1 + ((config.dim() + localSize - 1) / localSize));
+        fields.tempLogits = new FloatArray(1 + ((config.dim() + localSize - 1) / localSize));
+
         return fields;
 
     }

src/main/java/org/beehive/gpullama3/inference/weights/tornado/Qwen2TornadoWeights.java

@@ -7,9 +7,9 @@
 public class Qwen2TornadoWeights extends TornadoWeights {
 
     // Qwen2-specific tornado weights
-    FloatArray[] q_biasLayered;
-    FloatArray[] k_biasLayered;
-    FloatArray[] v_biasLayered;
+    public FloatArray[] q_biasLayered;
+    public FloatArray[] k_biasLayered;
+    public FloatArray[] v_biasLayered;
 
     public Qwen2TornadoWeights(FloatArray tokenEmbeddingTable, FloatArray[] rms_att_weightLayered, HalfFloatArray[] wqLayered, HalfFloatArray[] wkLayered, HalfFloatArray[] wvLayered,
             FloatArray[] wqBiasLayered,

src/main/java/org/beehive/gpullama3/model/loader/Qwen2ModelLoader.java

@@ -150,9 +150,9 @@ public Weights createTornadoVMWeights(Map<String, GGMLTensorEntry> tensorEntries
                 loadArrayAsHalfFloatArray(config.numberOfLayers(), i -> tensorEntries.get("blk." + i + ".attn_k.weight")),
                 loadArrayAsHalfFloatArray(config.numberOfLayers(), i -> tensorEntries.get("blk." + i + ".attn_v.weight")),
                 // Qwen2-specific: qkv bias
-                loadArrayAsFloatArray(config.numberOfLayers(), i -> tensorEntries.get("blk." + i + ".attn_q.bias")),
-                loadArrayAsFloatArray(config.numberOfLayers(), i -> tensorEntries.get("blk." + i + ".attn_k.bias")),
-                loadArrayAsFloatArray(config.numberOfLayers(), i -> tensorEntries.get("blk." + i + ".attn_v.bias")),
+                loadArrayAsFloatArrayFromBuffer(config.numberOfLayers(), i -> tensorEntries.get("blk." + i + ".attn_q.bias")),
+                loadArrayAsFloatArrayFromBuffer(config.numberOfLayers(), i -> tensorEntries.get("blk." + i + ".attn_k.bias")),
+                loadArrayAsFloatArrayFromBuffer(config.numberOfLayers(), i -> tensorEntries.get("blk." + i + ".attn_v.bias")),
 
                 loadArrayAsHalfFloatArray(config.numberOfLayers(), i -> tensorEntries.get("blk." + i + ".attn_output.weight")),
                 loadArrayAsFloatArrayFromBuffer(config.numberOfLayers(), i -> tensorEntries.get("blk." + i + ".ffn_norm.weight")),

src/main/java/org/beehive/gpullama3/model/qwen2/Qwen2Configuration.java

@@ -27,7 +27,7 @@ public int kvDim() {
 
     @Override
     public int kvMul() {
-        throw new UnsupportedOperationException("Not supported for Qwen2.");
+        return numberOfHeads / numberOfKeyValueHeads;
     }
 
     @Override