fix

flexflow · Oct 12, 2023 · 5ee8587 · 5ee8587
1 parent d71cf29
commit 5ee8587
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/src/ops/spec_inc_multihead_self_attention.cu b/src/ops/spec_inc_multihead_self_attention.cu
@@ -212,7 +212,7 @@ void update_kv_cache_kernel(SpecIncMultiHeadSelfAttentionMeta const *m,
                                     BatchConfig::max_sequence_length(),
                                     BeamSearchBatchConfig::MAX_BEAM_WIDTH,
                                     /*root*/ curr_depth == 0,
-                                    m->qSize);
+                                    m->num_kv_heads * m->kProjSize);
   }
 }
 

diff --git a/src/ops/tree_inc_multihead_self_attention.cu b/src/ops/tree_inc_multihead_self_attention.cu
@@ -119,7 +119,7 @@ void commit_tokens(TreeIncMultiHeadSelfAttentionMeta const *m,
         m->num_q_heads,
         m->num_kv_heads,
         BatchConfig::max_sequence_length(),
-        m->qSize);
+        m->num_kv_heads * m->kProjSize);
   }
 }
 
@@ -270,7 +270,7 @@ void compute_attention_kernel(TreeIncMultiHeadSelfAttentionMeta const *m,
             m->num_q_heads,
             m->num_kv_heads,
             BatchConfig::max_sequence_length(),
-            m->qSize);
+            m->num_kv_heads * m->kProjSize);
       }
 
       // bc->token_last_available_idx[i] + 1;