[Refactor] Remove Duplicate per_block_cast_to_fp8, Remove Dependencies of DeepGEMM (#21787)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
This commit is contained in:
@@ -222,25 +222,6 @@ def native_per_token_group_quant_int8(x,
|
||||
DEFAULT_BLOCK_SHAPE = [128, 128]
|
||||
|
||||
|
||||
def per_block_cast_to_fp8(
|
||||
x: torch.Tensor,
|
||||
block_shape: list[int] = DEFAULT_BLOCK_SHAPE,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
block_m, block_n = block_shape
|
||||
assert x.dim() == 2
|
||||
m, n = x.shape
|
||||
x_padded = torch.zeros((round_up(m, block_m), round_up(n, block_n)),
|
||||
dtype=x.dtype,
|
||||
device=x.device)
|
||||
x_padded[:m, :n] = x
|
||||
x_view = x_padded.view(-1, block_m, x_padded.size(1) // block_n, block_n)
|
||||
x_amax = x_view.abs().float().amax(dim=(1, 3), keepdim=True).clamp(1e-4)
|
||||
x_scaled = (x_view * (448.0 / x_amax)).to(torch.float8_e4m3fn)
|
||||
x_scaled_sub = x_scaled.view_as(x_padded)[:m, :n].contiguous()
|
||||
scales = (x_amax / 448.0).view(x_view.size(0), x_view.size(2))
|
||||
return x_scaled_sub, scales
|
||||
|
||||
|
||||
def per_block_cast_to_int8(
|
||||
x: torch.Tensor,
|
||||
block_shape: list[int] = DEFAULT_BLOCK_SHAPE,
|
||||
|
||||
Reference in New Issue
Block a user