[Bugfix] Fix batch-invariant fp32 matmul OOR on SM89 for N=1 (#52960)
Signed-off-by: Juqi Li <2223621784@qq.com>
vhagor committed
Aug 21, 2026 at 11:46 UTC
ba07e4a48fc951300d97eb506217dd530583dea3
1 file changed
+19
-6
vllm/model_executor/layers/batch_invariant.py
+19
-6
@@ -174,10 +174,10 @@ def matmul_persistent(
174
},
175
torch.float32: {
176
"BLOCK_SIZE_M": 128,
177
- "BLOCK_SIZE_N": 128,
177
+ "BLOCK_SIZE_N": _fp32_block_size_n if N == 1 else 128,
178
"BLOCK_SIZE_K": 32,
179
"GROUP_SIZE_M": 8,
180
- "num_stages": 3,
180
+ "num_stages": _fp32_num_stages if N == 1 else 3,
181
"num_warps": 8,
182
},
183
}
@@ -687,9 +687,9 @@ def bmm_batch_invariant(a, b, *, out=None):
687
},
688
torch.float32: {
689
"BLOCK_SIZE_M": 128,
690
- "BLOCK_SIZE_N": 128,
690
+ "BLOCK_SIZE_N": _fp32_block_size_n if N == 1 else 128,
691
"BLOCK_SIZE_K": 32,
692
- "num_stages": 3,
692
+ "num_stages": _fp32_num_stages if N == 1 else 3,
693
"num_warps": 8,
694
},
695
}
@@ -897,11 +897,13 @@ def linear_batch_invariant(input, weight, bias=None):
897
_batch_invariant_MODE = False
898
_batch_invariant_LIB = None
899
_fp16_block_size_n = 256
900
+_fp32_block_size_n = 128
901
+_fp32_num_stages = 3
902
903
904
def enable_batch_invariant_mode():
905
global _batch_invariant_MODE, _batch_invariant_LIB
904
- global _fp16_block_size_n
906
+ global _fp16_block_size_n, _fp32_block_size_n, _fp32_num_stages
907
908
if _batch_invariant_MODE:
909
return
@@ -926,7 +928,18 @@ def enable_batch_invariant_mode():
928
os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":16:8"
929
os.environ["CUBLASLT_WORKSPACE_SIZE"] = "1"
930
929
- _fp16_block_size_n = 256 if get_max_shared_memory_bytes() > 106496 else 128
931
+ # Query the shared memory size and set block size
932
+ # accordingly to avoid triton OutOfResources.
933
+ if get_max_shared_memory_bytes() > 106496:
934
+ _fp16_block_size_n = 256
935
+ _fp32_block_size_n = 128
936
+ _fp32_num_stages = 3
937
+ else:
938
+ _fp16_block_size_n = 128
939
+ # SM89 N=1 fp32 128/stages=3 compiles to 131072B (>101376).
940
+ # Only N=1 OORs; keep 128/stages=3 for every other N.
941
+ _fp32_block_size_n = 32
942
+ _fp32_num_stages = 2
943
elif current_platform.is_xpu():
944
_batch_invariant_LIB.impl("aten::mm", mm_batch_invariant, key)
945
_batch_invariant_LIB.impl("aten::addmm", addmm_batch_invariant, key)