From 632c7afa8ce20c2a696380ccf5b28337fd9903b7 Mon Sep 17 00:00:00 2001 From: Koushik Dutta Date: Thu, 29 Jan 2026 21:01:57 -0800 Subject: [PATCH] [Fix] add block size logic for sm120 smem size (#14311) --- .../layers/attention/triton_ops/extend_attention.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/layers/attention/triton_ops/extend_attention.py b/python/sglang/srt/layers/attention/triton_ops/extend_attention.py index 62132a340..2fb428fb2 100644 --- a/python/sglang/srt/layers/attention/triton_ops/extend_attention.py +++ b/python/sglang/srt/layers/attention/triton_ops/extend_attention.py @@ -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)