Skip to content

Commit 3c7c83c

Browse files
xiaoajie738xiongjie
andauthored
[Fix]: kernel meta retrieval for SM7X does not work (#3746)
Co-authored-by: xiongjie <[email protected]>
1 parent a032f2b commit 3c7c83c

File tree

1 file changed

+1
-1
lines changed

1 file changed

+1
-1
lines changed

lmdeploy/pytorch/kernels/cuda/flashattention.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -459,7 +459,7 @@ def grid(args):
459459
num_warps = 4
460460
if _nv_cap[0] < 8:
461461
BLOCK_M, BLOCK_N, num_warps, num_stages = _kernel_meta_sm7x(BLOCK_DK)
462-
if _nv_cap[0] < 9:
462+
elif _nv_cap[0] < 9:
463463
if _nv_cap[1] in [6, 9]:
464464
BLOCK_M, BLOCK_N, num_warps, num_stages = _kernel_meta_sm86(BLOCK_DK, shared_kv)
465465
else:

0 commit comments

Comments
 (0)