fix.

Author: b4rtaz

src/app.cpp

@@ -128,6 +128,7 @@ AppCliArgs::~AppCliArgs() {
 static NnDevice *createDevice(AppCliArgs *args, NnNetConfig *netConfig, NnNodeConfig *nodeConfig, NnNetExecution *netExecution) {
     if (args->gpuIndex >= 0) {
 #if defined(DLLAMA_VULKAN)
+        args->nBatches = 1; // TODO: this should be fixed
         return new NnVulkanDevice(args->gpuIndex, netConfig, nodeConfig, netExecution);
 #else
         throw std::runtime_error("This build does not support GPU");

src/nn/nn-vulkan.cpp

@@ -456,12 +456,12 @@ static std::vector<NnVulkanBatchInfo> buildBatchInfo(NnOpConfig *opConfig, NnVul
     return offset;
 }
 
-static void resolveShaderGroups(const NnOpCode opCode, const NnUint batchSize, NnUint *groupCount) {
+static void resolveShaderGroups(const NnOpConfig *opConfig, const NnUint batchSize, NnUint *groupCount) {
     groupCount[0] = 1;
     groupCount[1] = batchSize;
     groupCount[2] = 1;
 
-    if (opCode == OP_MATMUL)
+    if (opConfig->code == OP_MATMUL)
         groupCount[2] = 32;
 }
 
@@ -732,11 +732,11 @@ void NnVulkanDeviceSegment::forward(NnUint opIndex, NnUint nThreads, NnUint thre
 
     if (lastBatchSize != batchSize) {
         lastBatchSize = batchSize;
-        commandBuffer.begin({ vk::CommandBufferUsageFlags{ vk::CommandBufferUsageFlagBits::eSimultaneousUse } });
+        commandBuffer.begin({ vk::CommandBufferUsageFlags{} });
 
         NnUint opGroupCount[3];
         for (NnUint opIndex = 0; opIndex < segmentConfig->nOps; opIndex++) {
-            resolveShaderGroups(segmentConfig->ops[opIndex].code, batchSize, opGroupCount);
+            resolveShaderGroups(&segmentConfig->ops[opIndex], batchSize, opGroupCount);
 
             if (opIndex > 0) {
                 vk::MemoryBarrier memoryBarrier(

src/nn/vulkan/multi-head-att-forward-f32-f32.comp

@@ -1,6 +1,6 @@
 #version 450
 
-#define N_THREADS 32
+#define N_THREADS 256
 
 layout(local_size_x = N_THREADS, local_size_y = 1, local_size_z = 1) in;