[jit kernel] Support per_token_group_quant_8bit jit kernel (#18905)
Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>
This commit is contained in:
@@ -10,7 +10,6 @@ import torch
|
||||
if TYPE_CHECKING:
|
||||
from tvm_ffi import Module
|
||||
|
||||
|
||||
F = TypeVar("F", bound=Callable[..., Any])
|
||||
|
||||
|
||||
@@ -73,7 +72,9 @@ class CPPArgList(list[str]):
|
||||
CPP_DTYPE_MAP = {
|
||||
torch.float: "fp32_t",
|
||||
torch.float16: "fp16_t",
|
||||
torch.float8_e4m3fn: "fp8_e4m3_t",
|
||||
torch.bfloat16: "bf16_t",
|
||||
torch.int8: "int8_t",
|
||||
}
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user