Revert "kernel: use tensor cores for flashinfer gqa kernels" (#1511)

sgl-project · Sep 25, 2024 · f39a019 · f39a019
1 parent 3c93187
commit f39a019
Showing 1 changed file with 3 additions and 11 deletions.
diff --git a/python/sglang/srt/layers/attention_backend.py b/python/sglang/srt/layers/attention_backend.py
@@ -86,17 +86,9 @@ def __init__(self, model_runner: ModelRunner):
         super().__init__()
         self.model_runner = model_runner
 
-        local_num_qo_heads = (
-            model_runner.model_config.num_attention_heads // model_runner.tp_size
-        )
-        local_num_kv_heads = model_runner.model_config.get_num_kv_heads(
-            model_runner.tp_size
-        )
-        if (
-            not _grouped_size_compiled_for_decode_kernels(
-                local_num_qo_heads, local_num_kv_heads
-            )
-            or local_num_qo_heads // local_num_kv_heads > 4
+        if not _grouped_size_compiled_for_decode_kernels(
+            model_runner.model_config.num_attention_heads // model_runner.tp_size,
+            model_runner.model_config.get_num_kv_heads(model_runner.tp_size),
         ):
             self.decode_use_tensor_cores = True
         else: