EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 Projects Releases Wiki Activity

Use Tensor Core Decode when gqa group size >= 4 (#8624)

Browse Source 
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
This commit is contained in:
Wenxuan Tan
2025-08-22 10:25:15 -05:00
committed by GitHub
parent 6078d5fcc0
commit 0f587e80d3
1 changed files with 3 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
python/sglang/srt/layers/attention/flashinfer_backend.py
View File

@@ -1263,11 +1263,12 @@ def should_use_tensor_core(
# Calculate GQA group size # Calculate GQA group size
gqa_group_size = num_attention_heads // num_kv_heads gqa_group_size = num_attention_heads // num_kv_heads
# Determine based on dtype and GQA group size # For Flashinfer, a GQA group size of at least 4 is needed to efficiently
# use Tensor Cores, as it fuses the head group with the token dimension in MMA.
if kv_cache_dtype in (torch.float8_e4m3fn, torch.float8_e5m2): if kv_cache_dtype in (torch.float8_e4m3fn, torch.float8_e5m2):
return True return True
elif kv_cache_dtype in (torch.float16, torch.half, torch.bfloat16): elif kv_cache_dtype in (torch.float16, torch.half, torch.bfloat16):
return gqa_group_size > 4 return gqa_group_size >= 4
else: else:
return False return False