[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:
Yuan Luo
2026-02-20 21:01:05 +08:00
committed by GitHub
parent 38a69652e6
commit 7d953440ec
6 changed files with 811 additions and 14 deletions

View File

@@ -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",
}