forked from ccf-ai-infra/TileOPs-Metax
126 lines
5.1 KiB
Python
126 lines
5.1 KiB
Python
# 2026 - Modified by MetaX Integrated Circuits (Shanghai) Co., Ltd. All Rights Reserved.
|
|
|
|
import functools
|
|
import itertools
|
|
from typing import Optional, Tuple
|
|
|
|
import tilelang
|
|
import tilelang.language as T
|
|
import torch
|
|
|
|
from tileops.kernels.kernel_base import Kernel
|
|
|
|
__all__ = ["FP8QuantKernel"]
|
|
|
|
|
|
@functools.lru_cache(maxsize=32)
|
|
def _fp8_quant_kernel(batch, seq_len_kv, kv_group, index_dim, in_dtype: str):
|
|
|
|
@tilelang.jit(out_idx=[1, 2])
|
|
def _fp8_quant_fwd_func(num_stages, block_m):
|
|
out_dtype = T.float8_e4m3fn
|
|
scale_dtype = T.float32
|
|
fp8_min = -448.0
|
|
fp8_max = 448.0
|
|
fp8_max_inv = 1 / fp8_max
|
|
|
|
@T.prim_func
|
|
def _fp8_quant_fwd_main(input_tensor: T.Tensor[(batch, seq_len_kv, kv_group, index_dim),
|
|
in_dtype],
|
|
scale_tensor: T.Tensor[(batch, seq_len_kv, kv_group), scale_dtype],
|
|
output_tensor: T.Tensor[(batch, seq_len_kv, kv_group, index_dim),
|
|
out_dtype]):
|
|
with T.Kernel(
|
|
batch, T.ceildiv(seq_len_kv, block_m), kv_group, threads=128) as (bx, pid_m, g):
|
|
input_local = T.alloc_fragment((block_m, index_dim), in_dtype)
|
|
amax_local = T.alloc_fragment((block_m,), scale_dtype)
|
|
scale_local = T.alloc_fragment((block_m,), scale_dtype)
|
|
output_local = T.alloc_fragment((block_m, index_dim), out_dtype)
|
|
|
|
# Load a (block_m, index_dim) tile explicitly to avoid stride bugs with extra dims
|
|
for i, j in T.Parallel(block_m, index_dim):
|
|
input_local[i, j] = input_tensor[bx, pid_m * block_m + i, g, j]
|
|
|
|
# Reduce over index_dim to get amax per sequence position
|
|
T.reduce_absmax(input_local, amax_local, dim=1)
|
|
for i in T.Parallel(block_m):
|
|
amax_local[i] = T.max(amax_local[i], 1e-4)
|
|
scale_local[i] = amax_local[i] * fp8_max_inv
|
|
|
|
# Quantize: q = clamp(input / scale, [-448, 448])
|
|
for i, j in T.Parallel(block_m, index_dim):
|
|
output_local[i, j] = T.clamp(input_local[i, j] / scale_local[i], fp8_min,
|
|
fp8_max)
|
|
|
|
# Write back scale and output
|
|
for i in T.Parallel(block_m):
|
|
scale_tensor[bx, pid_m * block_m + i, g] = scale_local[i]
|
|
for i, j in T.Parallel(block_m, index_dim):
|
|
output_tensor[bx, pid_m * block_m + i, g, j] = output_local[i, j]
|
|
|
|
return _fp8_quant_fwd_main
|
|
|
|
return _fp8_quant_fwd_func
|
|
|
|
|
|
@torch.library.custom_op("top::fp8_quant_wrapped_kernel", mutates_args=())
|
|
def _fp8_quant_wrapped_kernel(batch: int, seq_len_kv: int, kv_group: int, index_dim: int,
|
|
in_dtype: str, num_stages: int, block_m: int,
|
|
input_tensor: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
|
|
return _fp8_quant_kernel(batch, seq_len_kv, kv_group, index_dim, in_dtype)(num_stages, block_m)(
|
|
input_tensor)
|
|
|
|
|
|
@_fp8_quant_wrapped_kernel.register_fake
|
|
def _(batch, seq_len_kv, kv_group, index_dim, in_dtype, num_stages, block_m, *inputs):
|
|
return torch.empty((batch, seq_len_kv, kv_group), dtype=torch.float32,
|
|
device=inputs[0].device), torch.empty(
|
|
(batch, seq_len_kv, kv_group, index_dim),
|
|
dtype=torch.float8_e4m3fn,
|
|
device=inputs[0].device)
|
|
|
|
|
|
class FP8QuantKernel(Kernel):
|
|
|
|
supported_archs: list[int] = [80, 89]
|
|
|
|
def __init__(self,
|
|
batch: int,
|
|
seq_len_kv: int,
|
|
kv_group: int,
|
|
index_dim: int,
|
|
in_dtype: torch.dtype,
|
|
config: Optional[dict] = None,
|
|
tune: bool = False):
|
|
super().__init__()
|
|
self.batch = batch
|
|
self.seq_len_kv = seq_len_kv
|
|
self.kv_group = kv_group
|
|
self.index_dim = index_dim
|
|
self.dtype = in_dtype
|
|
self.config = config or {}
|
|
self.kernel = _fp8_quant_kernel(self.batch, self.seq_len_kv, self.kv_group, self.index_dim,
|
|
self.dtype_str)
|
|
self.init_config(config, tune)
|
|
|
|
@property
|
|
def dtype_str(self) -> str:
|
|
return str(self.dtype).replace("torch.", "")
|
|
|
|
@property
|
|
def default_config(self) -> dict:
|
|
return {"num_stages": 0, "block_m": 32}
|
|
|
|
@property
|
|
def autotune_configs(self) -> list[dict]:
|
|
num_stages = [0, 2]
|
|
block_m = [32, 64]
|
|
_configs = list(itertools.product(num_stages, block_m))
|
|
|
|
return [{'num_stages': c[0], 'block_m': c[1]} for c in _configs]
|
|
|
|
def forward(self, input_tensor: torch.Tensor):
|
|
return _fp8_quant_wrapped_kernel(self.batch, self.seq_len_kv, self.kv_group, self.index_dim,
|
|
self.dtype_str, self.config["num_stages"],
|
|
self.config["block_m"], input_tensor)
|