feat: add should_use_tensor_core (#2179)

python/sglang/srt/layers/attention/flashinfer_backend.py

@@ -18,7 +18,11 @@ import triton.language as tl
 from sglang.global_config import global_config
 from sglang.srt.layers.attention import AttentionBackend
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
-from sglang.srt.utils import get_bool_env_var, is_flashinfer_available
+from sglang.srt.utils import (
+    get_bool_env_var,
+    is_flashinfer_available,
+    should_use_tensor_core,
+)
 
 if TYPE_CHECKING:
     from sglang.srt.layers.radix_attention import RadixAttention
@@ -31,7 +35,6 @@ if is_flashinfer_available():
         BatchPrefillWithRaggedKVCacheWrapper,
     )
     from flashinfer.cascade import merge_state
-    from flashinfer.decode import _grouped_size_compiled_for_decode_kernels
 
 
 class WrapperDispatch(Enum):
@@ -45,19 +48,14 @@ class FlashInferAttnBackend(AttentionBackend):
     def __init__(self, model_runner: ModelRunner):
         super().__init__()
 
-        # Parse constants
-        if "SGLANG_FLASHINFER_USE_TENSOR_CORE" in os.environ:
-            self.decode_use_tensor_cores = get_bool_env_var(
-                "SGLANG_FLASHINFER_USE_TENSOR_CORE"
-            )
-        else:
-            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:
-                self.decode_use_tensor_cores = False
+        self.decode_use_tensor_cores = should_use_tensor_core(
+            kv_cache_dtype=model_runner.kv_cache_dtype,
+            num_attention_heads=model_runner.model_config.num_attention_heads
+            // model_runner.tp_size,
+            num_kv_heads=model_runner.model_config.get_num_kv_heads(
+                model_runner.tp_size
+            ),
+        )
 
         self.max_context_len = model_runner.model_config.context_len