[ROCm] Add AMD GPU support on Deepseek v3.2 and SparseMLA (#26670)
Signed-off-by: ganyi <ygan@amd.com>
This commit is contained in:
@@ -325,6 +325,7 @@ DEFAULT_BLOCK_SIZE = [128, 128]
|
||||
def per_block_cast_to_fp8(
|
||||
x: torch.Tensor, block_size: list[int] = DEFAULT_BLOCK_SIZE, use_ue8m0: bool = False
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
fp8_dtype = current_platform.fp8_dtype()
|
||||
assert x.dim() == 2
|
||||
m, n = x.shape
|
||||
block_m, block_n = block_size
|
||||
@@ -334,9 +335,9 @@ def per_block_cast_to_fp8(
|
||||
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)
|
||||
sf = x_amax / 448.0
|
||||
sf = x_amax / 224.0 if current_platform.is_fp8_fnuz() else x_amax / 448.0
|
||||
sf = _ceil_to_ue8m0(sf) if use_ue8m0 else sf
|
||||
x_scaled = (x_view * (1.0 / sf)).to(torch.float8_e4m3fn)
|
||||
x_scaled = (x_view * (1.0 / sf)).to(fp8_dtype)
|
||||
return x_scaled.view_as(x_padded)[:m, :n].contiguous(), sf.view(
|
||||
x_view.size(0), x_view.size(2)
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user