chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Fix] add block size logic for sm120 smem size (#14311)

Browse Source
This commit is contained in:
Koushik Dutta
2026-01-29 21:01:57 -08:00
committed by GitHub
parent 046b29be16
commit 632c7afa8c
1 changed files with 9 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
python/sglang/srt/layers/attention/triton_ops/extend_attention.py
View File

@@ -64,7 +64,15 @@ def _get_block_sizes_for_extend_attention(Lq: int, Lv: int):
BLOCK_M, BLOCK_N = (64, 64)
num_warps = 4
else:
if _is_cuda and CUDA_CAPABILITY[0] >= 9:
if _is_cuda and CUDA_CAPABILITY[0] == 12:
# sm120 workstation Blackwell architecture (RTX Pro 6000) has a much smaller shared memory size (100K)
if Lq <= 128:
BLOCK_M, BLOCK_N = (64, 128)
elif Lq <= 256:
BLOCK_M, BLOCK_N = (64, 64)
else:
BLOCK_M, BLOCK_N = (32, 32)
elif _is_cuda and CUDA_CAPABILITY[0] >= 9:
# Hopper architecture (H100, etc.)
if Lq <= 256:
BLOCK_M, BLOCK_N = (128, 64)