[Iluvatar] Add vdot_heur_block_size

FlagOpen · Feb 11, 2025 · c988738 · c988738
1 parent 5aa3ab7
commit c988738
Showing 1 changed file with 13 additions and 0 deletions.
diff --git a/src/flag_gems/runtime/backend/_iluvatar/heuristics_config_utils.py b/src/flag_gems/runtime/backend/_iluvatar/heuristics_config_utils.py
@@ -214,6 +214,16 @@ def batch_norm_heur_block_n(args):
     return min(BLOCK_N, max(1, 2**14 // BLOCK_M))
 
 
+def vdot_heur_block_size(args):
+    n = args["n_elements"]
+    if n < 1024:
+        return 32
+    elif n < 8192:
+        return 256
+    else:
+        return 1024
+
+
 HEURISTICS_CONFIGS = {
     "argmax": {
         "BLOCK_M": argmax_heur_block_m,
@@ -289,4 +299,7 @@ def batch_norm_heur_block_n(args):
         "BLOCK_M": batch_norm_heur_block_m,
         "BLOCK_N": batch_norm_heur_block_n,
     },
+    "vdot": {
+        "BLOCK_SIZE": vdot_heur_block_size,
+    },
 }