fix wrong dispatch

InternLM · Nov 5, 2024 · b7f394a · b7f394a
1 parent 80e1775
commit b7f394a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/turbomind/kernels/attention/kv_cache_utils_v2.cu b/src/turbomind/kernels/attention/kv_cache_utils_v2.cu
@@ -543,7 +543,7 @@ void invokeFlattenKV_v2(T*           k,
             return invoke(tkv, std::integral_constant<int, 64>{});
         }
         else if (head_dim == 128) {
-            return invoke(tkv, std::integral_constant<int, 64>{});
+            return invoke(tkv, std::integral_constant<int, 128>{});
         }
         FT_CHECK(0);
     };