Skip to content

Commit 06309c4

Browse files
authored
bugfix: Choose sm90 kernels only for Hopper GPUs. (#719)
Some kernels use instructions specific to Hopper, which might not be compatible with future GPUs. Fallback to non-Hopper kernels for all the other GPUs.
1 parent 9a00cc2 commit 06309c4

File tree

1 file changed

+2
-2
lines changed

1 file changed

+2
-2
lines changed

flashinfer/utils.py

+2-2
Original file line numberDiff line numberDiff line change
@@ -262,7 +262,7 @@ def get_cuda_stream(device: torch.device) -> int:
262262

263263
def determine_gemm_backend(device: torch.device) -> str:
264264
major, _ = get_compute_capability(device)
265-
if major >= 9 and torch.version.cuda >= "12.3":
265+
if major == 9 and torch.version.cuda >= "12.3":
266266
return "sm90"
267267
else:
268268
return "sm80"
@@ -349,7 +349,7 @@ def determine_attention_backend(
349349
major, _ = get_compute_capability(device)
350350

351351
if (
352-
major >= 9
352+
major == 9
353353
and torch.version.cuda >= "12.3"
354354
and is_fa3_backend_supported(
355355
pos_encoding_mode,

0 commit comments

Comments
 (0)