fix: temporary solution for DeepSeek V2 H100 layout conversion issue (#1060)

Co-authored-by: ispobock <ISPObaoke@163.com>

fix: temporary solution for DeepSeek V2 H100 layout conversion issue (#1060)
Co-authored-by: ispobock <ISPObaoke@163.com>
65915f9f · Yineng Zhang · GitHub · 162f3ccb · 65915f9f
Unverified Commit 65915f9f authored Aug 13, 2024 by Yineng Zhang Committed by GitHub Aug 13, 2024
Hide whitespace changes
Inline Side-by-side

Showing with 3 additions and 1 deletion

python/sglang/srt/layers/extend_attention.py python/sglang/srt/layers/extend_attention.py +3 -1

No files found.
--- a/python/sglang/srt/layers/extend_attention.py
+++ b/python/sglang/srt/layers/extend_attention.py
@@ -275,7 +275,9 @@ def extend_attention_fwd(
        BLOCK_DPE = 0
    BLOCK_DV = Lv

-    if CUDA_CAPABILITY[0] >= 8:
+    if CUDA_CAPABILITY[0] >= 9:
+        BLOCK_M, BLOCK_N = (128, 64)
+    elif CUDA_CAPABILITY[0] >= 8:
        BLOCK_M, BLOCK_N = (128, 128) if Lq <= 128 else (64, 64)
    else:
        BLOCK_M, BLOCK_N = (64, 64) if Lq <= 128 else (32, 32)