007c645f87
Separate source and destination FP4 scale packing in requant_fp4_to_gran_k so group16-to-group32 conversion always recomputes UE8M0 runtime scales by default. Make prepare_fp4_weights_for_mega_moe accept raw grouped FP4 weights and scales, then perform optional requantization, DeepGEMM scale layout transform, and MegaMoE UTCCP weight transform internally. Update the MegaMoE synthetic benchmark so baseline grouped GEMM uses runtime-layout weights while fused MegaMoE uses transformed weights from the same raw source tensors. Tested: PYTHONPYCACHEPREFIX=/private/tmp/deepgemm_pycache python3 -m py_compile deep_gemm/__init__.py deep_gemm/mega/__init__.py deep_gemm/utils/math.py tests/test_layout.py tests/test_mega_moe.py Tested: git diff --check Not-tested: CUDA build, SM100/B300 runtime, and GLM-5.2 accuracy validation are not available locally.
143 lines
6.7 KiB
Python
143 lines
6.7 KiB
Python
import torch
|
|
import random
|
|
from deep_gemm.testing import bench_kineto, count_bytes, get_arch_major
|
|
from deep_gemm.utils import (
|
|
align, ceil_div,
|
|
per_token_cast_to_fp8, per_channel_cast_to_fp8,
|
|
per_token_cast_to_fp4, cast_back_from_fp4, requant_fp4_to_gran_k,
|
|
get_tma_aligned_size,
|
|
get_mn_major_tma_aligned_tensor,
|
|
get_mn_major_tma_aligned_packed_ue8m0_tensor,
|
|
get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor
|
|
)
|
|
|
|
from generators import (
|
|
enumerate_sf_layout,
|
|
enumerate_k_grouped_sf_layout
|
|
)
|
|
|
|
|
|
def get_mn_major_tma_aligned_packed_ue8m0_tensor_torch_impl(x: torch.Tensor) -> torch.Tensor:
|
|
assert x.dtype == torch.float and x.dim() in (2, 3)
|
|
|
|
# First, convert into UE8M0 `uint8_t`
|
|
ue8m0_tensor = (x.view(torch.int) >> 23).to(torch.uint8)
|
|
|
|
# Second, make padded packed tensors
|
|
mn, k = x.shape[-2], x.shape[-1]
|
|
remove_dim = False
|
|
if x.dim() == 2:
|
|
x, remove_dim = x.unsqueeze(0), True
|
|
b = x.shape[0]
|
|
aligned_mn = get_tma_aligned_size(mn, 4)
|
|
aligned_k = align(k, 4)
|
|
padded = torch.zeros((b, aligned_mn, aligned_k), device=x.device, dtype=torch.uint8)
|
|
padded[:, :mn, :k] = ue8m0_tensor
|
|
padded = padded.view(-1).view(dtype=torch.int).view(b, aligned_mn, aligned_k // 4)
|
|
|
|
# Finally, transpose
|
|
transposed = torch.zeros((b, aligned_k // 4, aligned_mn), device=x.device, dtype=torch.int).mT
|
|
transposed[:, :, :] = padded
|
|
aligned_x = transposed[:, :mn, :]
|
|
return aligned_x.squeeze(0) if remove_dim else aligned_x
|
|
|
|
|
|
def test_sf_layout_kernels() -> None:
|
|
print('Testing SF layout kernels:')
|
|
for mn, k, with_transpose, use_ue8m0, num_groups, gran_k in enumerate_sf_layout():
|
|
x = torch.randn((num_groups * mn, k), dtype=torch.bfloat16, device='cuda')
|
|
x, fp32_sf = per_token_cast_to_fp8(x, use_ue8m0=use_ue8m0, gran_k=gran_k)
|
|
fp32_sf = fp32_sf if num_groups == 1 else fp32_sf.view(num_groups, mn, -1)
|
|
fp32_sf = fp32_sf if with_transpose else fp32_sf.transpose(-1, -2).contiguous().transpose(-1, -2)
|
|
|
|
# Correctness
|
|
if use_ue8m0:
|
|
impl, name = get_mn_major_tma_aligned_packed_ue8m0_tensor, 'pack_fp32_into_ue8m0'
|
|
packed_sf = get_mn_major_tma_aligned_packed_ue8m0_tensor(fp32_sf)
|
|
ref_packed_sf = get_mn_major_tma_aligned_packed_ue8m0_tensor_torch_impl(fp32_sf)
|
|
assert torch.equal(packed_sf, ref_packed_sf), f'{mn=}, {k=}, {with_transpose=}, {num_groups=}'
|
|
assert packed_sf.shape == ref_packed_sf.shape
|
|
assert all([packed_sf.stride(i) == ref_packed_sf.stride(i) for i in range(packed_sf.dim())])
|
|
else:
|
|
impl, name = get_mn_major_tma_aligned_tensor, 'transpose'
|
|
transposed_sf = get_mn_major_tma_aligned_tensor(fp32_sf)
|
|
tma_aligned_mn, sf_k = get_tma_aligned_size(mn, fp32_sf.element_size()), ceil_div(k, gran_k)
|
|
if num_groups > 1:
|
|
assert transposed_sf.size(0) == num_groups
|
|
assert transposed_sf.stride(0) == tma_aligned_mn * sf_k
|
|
assert transposed_sf.shape[-2:] == (mn, sf_k)
|
|
assert transposed_sf.stride()[-2:] == (1, tma_aligned_mn)
|
|
assert torch.equal(fp32_sf, transposed_sf)
|
|
|
|
# Performance
|
|
try:
|
|
t = bench_kineto(lambda: impl(fp32_sf), name)
|
|
except AssertionError as e:
|
|
# Some cases may fallback to PyTorch impl
|
|
t = 0
|
|
print(f' > Perf ({num_groups=:2}, {mn=:5}, {k=:5}, transpose={int(with_transpose)}, use_ue8m0={int(use_ue8m0)}, gran_k={gran_k:3}): '
|
|
f'{t * 1e6:4.0f} us | {count_bytes(fp32_sf, impl(fp32_sf)) / 1e9 / t if t else 0:4.0f} GB/s')
|
|
print()
|
|
|
|
|
|
def test_k_grouped_sf_layout_kernels() -> None:
|
|
print('Testing k-grouped SF layout kernels:')
|
|
for mn, ks, num_groups, gran_k in enumerate_k_grouped_sf_layout():
|
|
sf_ks = [k // gran_k for k in ks]
|
|
packed_sf_ks = [ceil_div(k, gran_k * 4) for k in ks]
|
|
ks_tensor = torch.tensor(ks, dtype=torch.int, device='cuda')
|
|
x = torch.randn((sum(ks), mn), dtype=torch.bfloat16, device='cuda')
|
|
x, fp32_sf = per_channel_cast_to_fp8(x, use_ue8m0=True, gran_k=gran_k)
|
|
|
|
# Correctness
|
|
packed_sf = get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor(fp32_sf, ks_tensor, ks, gran_k)
|
|
split_packed_sf = packed_sf.split(packed_sf_ks)
|
|
split_fp32_sf = fp32_sf.split(sf_ks)
|
|
for i in range(num_groups):
|
|
ref_packed_sf = get_mn_major_tma_aligned_packed_ue8m0_tensor_torch_impl(split_fp32_sf[i].T).T
|
|
assert torch.equal(split_packed_sf[i], ref_packed_sf), f'{i=}'
|
|
|
|
# Performance
|
|
t = bench_kineto(lambda: get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor(fp32_sf, ks_tensor, ks, gran_k), 'pack_fp32_into_ue8m0')
|
|
print(f' > Perf ({num_groups=:3}, {mn=:5}, sum_k={sum(ks):5}, gran_k={gran_k:3}):'
|
|
f'{t * 1e6:4.0f} us | '
|
|
f'{count_bytes(fp32_sf, packed_sf, ks_tensor) / 1e9 / t:4.0f} GB/s')
|
|
print()
|
|
|
|
|
|
def test_fp4_requant_granularity() -> None:
|
|
print('Testing FP4 requant granularity:')
|
|
for m, n in ((128, 7168), (257, 3072)):
|
|
x = torch.randn((m, n), dtype=torch.bfloat16, device='cuda')
|
|
fp4_g16 = per_token_cast_to_fp4(x, use_ue8m0=True, gran_k=16)
|
|
repacked, rescaled = requant_fp4_to_gran_k(fp4_g16[0], fp4_g16[1], 16, 32)
|
|
restored = cast_back_from_fp4(fp4_g16[0], fp4_g16[1], gran_k=16)
|
|
ref_repacked, ref_rescaled = per_token_cast_to_fp4(restored, use_ue8m0=True, gran_k=32)
|
|
fp4_g16_packed = per_token_cast_to_fp4(x, use_ue8m0=True, gran_k=16, use_packed_ue8m0=True)
|
|
repacked_from_packed, rescaled_from_packed = requant_fp4_to_gran_k(
|
|
fp4_g16_packed[0], fp4_g16_packed[1], 16, 32, src_scale_packed_ue8m0=True)
|
|
ref_repacked_packed, ref_rescaled_packed = per_token_cast_to_fp4(
|
|
restored, use_ue8m0=True, gran_k=32, use_packed_ue8m0=True)
|
|
repacked_to_packed, rescaled_to_packed = requant_fp4_to_gran_k(
|
|
fp4_g16[0], fp4_g16[1], 16, 32, dst_scale_packed_ue8m0=True)
|
|
|
|
assert repacked.shape == fp4_g16[0].shape
|
|
assert rescaled.shape[-1] == ceil_div(n, 32)
|
|
assert torch.equal(repacked, ref_repacked)
|
|
assert torch.equal(rescaled, ref_rescaled)
|
|
assert torch.equal(repacked_from_packed, ref_repacked)
|
|
assert torch.equal(rescaled_from_packed, ref_rescaled)
|
|
assert torch.equal(repacked_to_packed, ref_repacked_packed)
|
|
assert torch.equal(rescaled_to_packed, ref_rescaled_packed)
|
|
print(f' > Requant ({m=}, {n=}): group16 -> group32')
|
|
print()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
torch.manual_seed(1)
|
|
random.seed(1)
|
|
|
|
test_sf_layout_kernels()
|
|
test_k_grouped_sf_layout_kernels()
|
|
test_fp4_requant_granularity()
|