Bug fix for DPU kernel weight offset and layer id for multiple layers.

Yinan · Yinan · commit 9cd419028af8 · 2025-01-15T17:05:11.000+08:00
diff --git a/dpu/dpu_main.c b/dpu/dpu_main.c
@@ -107,8 +107,9 @@ int main() {
     ptable_f32_f16 = (__mram_ptr float *)DPU_MRAM_HEAP_POINTER;
     uint32_t table_f32_f16_len = (1 << 16)*sizeof(float);
     uint32_t offset = table_f32_f16_len;
-    int input_row_size,input_cols;
-    float *psumf;
+    int input_row_size = 0;
+    int input_cols = 0;
+    float *psumf = NULL;
 
 #if PRINT
     printf("table_f32_f16_len=%d\n",table_f32_f16_len);
@@ -124,7 +125,7 @@ int main() {
     mram_read((__mram_ptr void const*) (weightmetadatabase), cache_meta, sizeof(struct pim_meta));
 
 #if PRINT
-    printf("layer_num: %d, weight_type=%d,rows_per_dpu=%d,rest_rows=%d,input_offset=%d",
+    printf("layer_num: %d, weight_type=%d, rows_per_dpu=%d, rest_rows=%d, input_offset=%d",
         cache_meta->layer_num,cache_meta->weight_type,cache_meta->rows_per_dpu,cache_meta->rest_rows,cache_meta->input_offset);
 #endif
 
@@ -142,14 +143,15 @@ int main() {
     //input metadata
     offset += (cache_meta->layer_len * cache_meta->layer_num);
 #if PRINT
-    printf("layer_len=%d,offset=%d\n",cache_meta->layer_len,offset);
+    printf("layer_len=%d, input metadata offset=%d\n",cache_meta->layer_len,offset);
 #endif
     uint32_t inputmetadatabase = weightmetadatabase + sizeof(struct pim_meta) + cache_meta->layer_len * cache_meta->layer_num;   
     pim_matrix_des *pinputcache = (pim_matrix_des *) mem_alloc(sizeof(pim_matrix_des));
     mram_read((__mram_ptr void const*) (inputmetadatabase), pinputcache, sizeof(pim_matrix_des));
     input_cols = pinputcache->ne[1];
+    assert(input_cols == 1 && "Only support vector as input.");
 #if PRINT
-    printf("input_type=%d,layerID=%d\n",pinputcache->type,pinputcache->layerid);
+    printf("input_type=%d, layerID=%d\n",pinputcache->type,pinputcache->layerid);
     for(int nn=0;nn<GGML_MAX_DIMS;nn++) {
         printf("ne[%d]=%lld\n",nn,pinputcache->ne[nn]);
     }
@@ -165,19 +167,19 @@ int main() {
         int nb = pinputcache->ne[0]/QK8_0;
         int qk = QK8_0;
         input_row_size = nb*sizeof(block_q8_0);
-        __mram_ptr block_q4_0 *pweight_base = (__mram_ptr  block_q4_0 *)(weightmetadatabase + sizeof(struct pim_meta));
-        __mram_ptr block_q8_0 *pinput_base = (__mram_ptr block_q8_0 *)(DPU_MRAM_HEAP_POINTER + cache_meta->input_offset + sizeof(pim_matrix_des));
+        __mram_ptr void *pweight_base = (__mram_ptr void *)(weightmetadatabase + sizeof(struct pim_meta));
+        __mram_ptr void *pinput_base = DPU_MRAM_HEAP_POINTER + cache_meta->input_offset + sizeof(pim_matrix_des);
         psumf = (float *)mem_alloc(sizeof(float)*input_cols*weight_rows_cur_thread);
         memset(psumf, 0 ,sizeof(float)*input_cols*weight_rows_cur_thread);
 #if PRINT
-        printf("input_cols=%d,rows_cur_thread=%d,nb=%d,input_row_size=%d\n",input_cols,weight_rows_cur_thread,nb,input_row_size);
+        printf("input_cols=%d, rows_cur_thread=%d, nb=%d, input_row_size=%d\n",input_cols,weight_rows_cur_thread,nb,input_row_size);
 #endif
         block_q4_0 *pweight_cache = (block_q4_0 *) mem_alloc(sizeof(block_q4_0)*nb);
         block_q8_0 *pinput_cache = (block_q8_0 *) mem_alloc(sizeof(block_q8_0)*nb);          
 
         // weight_rows_cur_thread = 16;
         for(int l = 0;l < input_cols;l++) {
-            __mram_ptr block_q8_0 *pinput = pinput_base + l*nb;
+          __mram_ptr block_q8_0 *pinput = pinput_base + l*nb*sizeof(block_q8_0);
             mram2wram(pinput, pinput_cache, sizeof(block_q8_0)*nb);
 #if PRINT
             printf("input:\n");
@@ -191,8 +193,7 @@ int main() {
             printf("pweight_base: %p\n", pweight_base);
 #endif
             for(int k = 0;k < weight_rows_cur_thread;k++) {
-                //block_q4_0 *pqlayer0weight = (block_q4_0 *)(weightmetadatabase + sizeof(struct pim_meta) + cache_meta->layer_len*k);
-                __mram_ptr block_q4_0 *pweight = pweight_base + pinputcache->layerid*cache_meta->layer_len + k*nb;
+              __mram_ptr block_q4_0 *pweight = pweight_base + pinputcache->layerid*cache_meta->layer_len + k*nb*sizeof(block_q4_0);
                 mram2wram(pweight, pweight_cache, sizeof(block_q4_0)*nb);
 #if PRINT
                 if (k % 64 == 0) {
@@ -207,11 +208,10 @@ int main() {
 #endif
 
                 for (int i = 0; i < nb; i++) {
-                    //printf("input_col:%d,weight_row:%d\n",l,k);
+                    //printf("input_col:%d, current inner weight row idx:%d\n",l,k);
 
                     int sumi = 0;
                     for (int j = 0; j < qk/2; ++j) {
-                        //printf("nb:%d,qk=%d,qs=%d\n",i,j,pweight_cache[i].qs[j]);
                         const int v0 = (pweight_cache[i].qs[j] & 0x0F) - 8;
                         const int v1 = (pweight_cache[i].qs[j] >>   4) - 8;
 
@@ -230,9 +230,10 @@ int main() {
         printf("psumf[%d]=%f\n",iii,psumf[iii]);
     }
 
-    printf("offset=%d\n",offset);
+    printf("output offset=%d\n",offset);
 #endif
     // Write C Matrix to current MRAM block
-    wram2mram((__mram_ptr void *) (DPU_MRAM_HEAP_POINTER + offset),psumf,sizeof(float)*input_cols*weight_rows_cur_thread);
+    // Note: with input_cols > 1, the results should be rearranged on host
+    wram2mram((__mram_ptr void *) (DPU_MRAM_HEAP_POINTER + offset), psumf, sizeof(float)*input_cols*weight_rows_cur_thread);
     return 0;
 }
diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c
@@ -12541,7 +12541,8 @@ static void ggml_compute_forward_mul_mat(
     //   compute by src0 rows
 
     // export the first gemv's tensor
-    if (type == GGML_TYPE_Q4_0 && src1->type == GGML_TYPE_F32 &&
+    if (dst->flags & GGML_TENSOR_FLAG_PIM &&
+        type == GGML_TYPE_Q4_0 && src1->type == GGML_TYPE_F32 &&
         ne00 == 4096 && ne01 == 4096 &&
         ne02 == 1 && ne03 == 1 &&
         ne10 == 4096 && ne11 == 1 &&
@@ -12613,7 +12614,7 @@ UseGgmlGemm1:;
         }
 
 	if ((dst->flags & GGML_TENSOR_FLAG_PIM)) {
-          dpu_launch_gemv_async(src1, wdata, src0, dst, 0);
+          dpu_launch_gemv_async(src1, wdata, src0, dst, dst->layerid);
           dpu_kernel_barrier(*(dst->dpu_set));
 
           pim_res->type = dst->type;
@@ -17405,6 +17406,9 @@ static __inline__ int dpu_get_gemv_res(struct ggml_tensor *input, struct ggml_te
     dpu_get_nr_dpus(dpu_set, &nr_dpus);
     int rows_per_dpu = w->ne[1] / nr_dpus;
 
+    // Only support vector as input
+    GGML_ASSERT(input->ne[1] == 1);
+
     uint32_t i;
     DPU_FOREACH(dpu_set, dpu, i) {
         DPU_ASSERT(dpu_prepare_xfer(dpu, mul_max_res + i * rows_per_dpu*input->ne[1]));
diff --git a/include/llama.h b/include/llama.h
@@ -424,6 +424,9 @@ extern "C" {
             struct llama_context_params   params);
 
 #ifdef PIM_KERNEL
+#define NR_DPUS 64
+#define NR_LAYER 2
+#define DPU_BINARY "./dpu/gemv_dpu"
   enum WeightId {
     WQ,
     WCNT
diff --git a/src/llama.cpp b/src/llama.cpp
@@ -9348,7 +9348,7 @@ int load_weight2dpu(enum WeightId w_id, struct dpu_set_t dpu_set, struct llama_m
   uint32_t nr_dpus;
   dpu_get_nr_dpus(dpu_set, &nr_dpus);
   ggml_tensor *w = NULL;
-  for (uint32_t layeridx = 0; layeridx < 1; layeridx++) {
+  for (uint32_t layeridx = 0; layeridx < pim_metadata->layer_num; layeridx++) {
     switch (w_id) {
     case WQ:
       w = model->layers[layeridx].wq;
@@ -9366,7 +9366,7 @@ int load_weight2dpu(enum WeightId w_id, struct dpu_set_t dpu_set, struct llama_m
     uint32_t layer_len = pim_metadata->layer_len;
     uint32_t i;
 
-    printf("%s: size_per_row: %d, rows_per_dpu: %d, offset_base: %d, layer_len: %d\n", __FUNCTION__, size_per_row, pim_metadata->rows_per_dpu, offset_base, layer_len);
+    printf("%s: size_per_row: %d, rows_per_dpu: %d, offset_base (the 1st weight base): %d, layer_len: %d, layer_id: %d/%d\n", __FUNCTION__, size_per_row, pim_metadata->rows_per_dpu, offset_base, layer_len, layeridx, pim_metadata->layer_num);
 
     // row is send to dpu
     DPU_FOREACH(dpu_set, dpu, i) {
@@ -9382,8 +9382,6 @@ int load_weight2dpu(enum WeightId w_id, struct dpu_set_t dpu_set, struct llama_m
 }
 
 int llama_load2dpu(struct llama_context *ctx, struct llama_model *model) {
-    #define NR_DPUS 64
-    #define DPU_BINARY "./dpu/gemv_dpu"
     uint32_t nr_of_dpus;
     uint32_t pim_offset = 0;
     int i;
@@ -9409,7 +9407,7 @@ int llama_load2dpu(struct llama_context *ctx, struct llama_model *model) {
     // WQ metadata is loaded to dpu WRAM, make WQ's param in every layer is same
 
     //uint32_t n_layer = model->layers.size();
-    uint32_t n_layer = 1;
+    uint32_t n_layer = NR_LAYER;
     uint32_t il = 0;
     dpu_get_nr_dpus(pqcontext->dpu_set, &nr_of_dpus);
     pqcontext->pim_metadata.layer_num = n_layer;
@@ -10755,7 +10753,7 @@ struct llm_build_context {
                 // compute Q and K and RoPE them
                 struct ggml_tensor * Qcur = llm_build_lora_mm(lctx, ctx0, model.layers[il].wq, cur);
 
-                if (il == 0 && Qcur->op == GGML_OP_MUL_MAT && n_tokens == 1) {
+                if (il < NR_LAYER && Qcur->op == GGML_OP_MUL_MAT && n_tokens == 1) {
                   Qcur->flags |= GGML_TENSOR_FLAG_PIM;
                   Qcur->dpu_set = &(lctx.pim_context_map[WQ]->dpu_set);
                   Qcur->inout_offset = (lctx.pim_context_map[WQ]->pim_metadata).input_offset;
