fix: rope gpu config memory alignment. (#245)

Author: b4rtaz

src/llm.cpp

@@ -313,15 +313,15 @@ LlmNet buildLlmNet(LlmHeader *h, NnUint nNodes, NnUint nBatches) {
                 pointerBatchConfig(SRC_BUFFER, qBufferIndex),
                 pointerBatchConfig(SRC_BUFFER, qBufferIndex),
                 size0(),
-                NnRopeOpConfig{n.header->ropeType, true, n.positionPipeIndex, ropeCacheBufferIndex, 
+                NnRopeOpConfig{n.header->ropeType, 1, n.positionPipeIndex, ropeCacheBufferIndex, 
                     h->ropeScalingFactor, h->ropeScalingLowFreqFactor, h->ropeScalingHighFreqFactory, h->ropeScalingOrigMaxSeqLen,
                     ropeSlice});
             att.addOp(
                 OP_ROPE, "block_rope_k", layerIndex,
                 pointerBatchConfig(SRC_BUFFER, kTempBufferIndex),
                 pointerBatchConfig(SRC_BUFFER, kTempBufferIndex),
                 size0(),
-                NnRopeOpConfig{n.header->ropeType, false, n.positionPipeIndex, ropeCacheBufferIndex, 
+                NnRopeOpConfig{n.header->ropeType, 0, n.positionPipeIndex, ropeCacheBufferIndex, 
                     h->ropeScalingFactor, h->ropeScalingLowFreqFactor, h->ropeScalingHighFreqFactory, h->ropeScalingOrigMaxSeqLen,
                     ropeSlice});
             att.addOp(

src/nn/nn-core.hpp

@@ -205,7 +205,7 @@ typedef struct {
 
 typedef struct {
     NnRopeType type;
-    bool isQ;
+    NnUint isQ; // Cannot use `bool` here due to GPU memory alignment
     NnUint positionPipeIndex;
     NnUint ropeCacheBufferIndex;
     float ropeScalingFactor;

src/nn/nn-cpu-ops.cpp

@@ -1156,14 +1156,15 @@ static void ropeForward_F32_F32(NnUint nThreads, NnUint threadIndex, NnUint batc
     const NnRopeSlice *slice = &config->slice;
     const float *positions = (float *)context->pipes[config->positionPipeIndex];
     const float *cache = (float *)context->buffers[config->ropeCacheBufferIndex];
+    const bool isQ = config->isQ == 1;
 
     for (NnUint batchIndex = 0; batchIndex < batchSize; batchIndex++) {
         float *x = (float *)context->input[batchIndex];
         const NnUint pos = (NnUint)positions[batchIndex];
         if (config->type == ROPE_LLAMA || config->type == ROPE_LLAMA3_1)
-            ropeLlama_F32(x, cache, config->isQ, pos, slice, nThreads, threadIndex);
+            ropeLlama_F32(x, cache, isQ, pos, slice, nThreads, threadIndex);
         else if (config->type == ROPE_FALCON)
-            ropeFalcon_F32(x, cache, config->isQ, pos, slice, nThreads, threadIndex);
+            ropeFalcon_F32(x, cache, isQ, pos, slice, nThreads, threadIndex);
         else
             throw std::runtime_error("Unsupported rope type");
     }

src/nn/nn-vulkan-test.cpp

@@ -440,7 +440,7 @@ void testRope_F32_F32() {
             NnUint xPipeIndex = netBuilder->addPipe("X", size2D(F_32, N_BATCHES, ROPE_DIM));
             NnUint posPipeIndex = netBuilder->addPipe("POS", size2D(F_32, N_BATCHES, 1));
             NnUint ropeCacheBufferIndex = nodeBuilder->addBuffer("ropeCache", slice.cacheSize);
-            bool isQ = true;
+            NnUint isQ = 1;
 
             segmentBuilder->addOp(
                 OP_ROPE, "rope_llama", 0,

src/nn/vulkan/rope-forward-f32-f32.comp

@@ -32,7 +32,7 @@ layout(binding = 1) writeonly buffer outputBuffer { float y[]; };
 layout(binding = 2) readonly buffer batchInfosBuffer { BatchInfo infos[]; };
 layout(binding = 3) readonly uniform configBuffer {
     uint ropeType;
-    bool isQ;
+    uint isQ;
     uint positionPipeIndex;
     uint ropeCacheBufferIndex;
     float ropeScalingFactor;
@@ -56,7 +56,7 @@ void main() {
 
         if (ropeType == 0 || ropeType == 2 /* Llama */) {
             sharedOffset = position * slice.sliceDim;
-            if (isQ) {
+            if (isQ == 1) {
                 sharedOffset += slice.qShift;
             }
         } else if (ropeType == 1 /* Falcon */) {
@@ -70,7 +70,7 @@ void main() {
 
     const uint xOffset = sharedInfo.inputOffset;
     const uint yOffset = sharedInfo.outputOffset;
-    const uint dim0 = isQ ? slice.qDim0 : slice.kvDim0;
+    const uint dim0 = isQ == 1 ? slice.qDim0 : slice.kvDim0;
 
     if (ropeType == 0 || ropeType == 2 /* Llama */) {
         const uint dim0Half = dim0 / 2;