add control for cutlass fp8 blockwise gemm (#3727)

This commit is contained in:
yizhang2077
2025-02-20 16:10:35 +08:00
committed by GitHub
parent 3c7bfd7eab
commit 1eb8eade2b

View File

@@ -1,3 +1,4 @@
import os
from typing import List, Optional, Tuple
import torch
@@ -40,6 +41,8 @@ def normalize_e4m3fn_to_e4m3fnuz(
def cutlass_block_fp8_supported() -> bool:
if os.environ.get("SUPPORT_CUTLASS_BLOCK_FP8") is None:
return False
if _is_cuda:
major, minor = torch.cuda.get_device_capability()
sm_version = major * 10 + minor