81 lines
2.1 KiB
Python
81 lines
2.1 KiB
Python
"""GPU and backend-op capability detection.
|
|
|
|
Pure functions — no state, no module-level side effects. Cheap to call
|
|
(the underlying ``torch.cuda`` / ``importlib`` probes are fast).
|
|
"""
|
|
|
|
import importlib.util
|
|
import logging
|
|
from typing import List, Tuple
|
|
|
|
import torch
|
|
|
|
|
|
def get_gpu_capability():
|
|
if not torch.cuda.is_available():
|
|
return None, None
|
|
try:
|
|
return torch.cuda.get_device_capability(0)
|
|
except Exception as e:
|
|
logging.warning(f"Failed to get GPU capability: {e}")
|
|
return None, None
|
|
|
|
|
|
def is_fp8_supported_gpu() -> bool:
|
|
major, minor = get_gpu_capability()
|
|
if major is None:
|
|
return False
|
|
return (major == 8 and minor == 9) or (major >= 9)
|
|
|
|
|
|
def is_ada_architecture_gpu() -> bool:
|
|
major, minor = get_gpu_capability()
|
|
if major is None:
|
|
return False
|
|
return major == 8 and minor == 9
|
|
|
|
|
|
def is_module_installed(module_name: str) -> bool:
|
|
try:
|
|
spec = importlib.util.find_spec(module_name)
|
|
return spec is not None
|
|
except ModuleNotFoundError:
|
|
return False
|
|
|
|
|
|
def get_available_ops(op_mapping: dict) -> List[Tuple[str, bool]]:
|
|
return [(op_name, is_module_installed(module_name)) for op_name, module_name in op_mapping.items()]
|
|
|
|
|
|
_QUANT_OP_MAPPING = {
|
|
"sgl": "sgl_kernel",
|
|
"vllm": "vllm",
|
|
"q8f": "q8_kernels",
|
|
"torchao": "torchao",
|
|
}
|
|
|
|
_ATTN_OP_MAPPING = {
|
|
"sage_attn2": "sageattention",
|
|
"sage_attn3": "sageattn3",
|
|
"flash_attn3": "flash_attn_interface",
|
|
"flash_attn2": "flash_attn",
|
|
"torch_sdpa": "torch",
|
|
}
|
|
|
|
|
|
def get_available_quant_ops() -> List[Tuple[str, bool]]:
|
|
available_ops = get_available_ops(_QUANT_OP_MAPPING)
|
|
|
|
# Prefer q8f on Ada (sm_8.9) GPUs — best perf/precision tradeoff there.
|
|
if is_ada_architecture_gpu():
|
|
q8f_available = next((op for op in available_ops if op[0] == "q8f" and op[1]), None)
|
|
if q8f_available:
|
|
available_ops.remove(q8f_available)
|
|
available_ops.insert(0, q8f_available)
|
|
|
|
return available_ops
|
|
|
|
|
|
def get_available_attn_ops() -> List[Tuple[str, bool]]:
|
|
return get_available_ops(_ATTN_OP_MAPPING)
|