wehub-resource-sync
282 lines
8.7 KiB
Python
282 lines
8.7 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
import torch
|
|
|
|
from vllm.utils.torch_utils import direct_register_custom_op
|
|
|
|
# MXFP8 constants
|
|
MXFP8_VALUE_DTYPE = torch.float8_e4m3fn
|
|
MXFP8_SCALE_DTYPE = torch.uint8
|
|
MXFP8_BLOCK_SIZE = 32
|
|
|
|
|
|
def swizzle_mxfp8_scale(sf: torch.Tensor, M: int, K: int) -> torch.Tensor:
|
|
"""Swizzle MXFP8 scales from row-major 2D to F8_128x4 layout."""
|
|
scaling_vector_size = MXFP8_BLOCK_SIZE # 32 for MXFP8
|
|
factor = scaling_vector_size * 4 # 128
|
|
|
|
num_m_tiles = (M + 127) // 128
|
|
num_k_tiles = (K + factor - 1) // factor
|
|
|
|
m_padded = num_m_tiles * 128
|
|
k_scale_padded = num_k_tiles * 4
|
|
|
|
scale_cols = K // scaling_vector_size
|
|
sf_padded = torch.zeros(
|
|
(m_padded, k_scale_padded), dtype=sf.dtype, device=sf.device
|
|
)
|
|
sf_padded[:M, :scale_cols] = sf
|
|
|
|
sf_reshaped = sf_padded.view(num_m_tiles, 4, 32, num_k_tiles, 4)
|
|
|
|
sf_swizzled = sf_reshaped.transpose(1, 3)
|
|
|
|
return sf_swizzled.contiguous().view(-1)
|
|
|
|
|
|
def _mxfp8_e4m3_quantize_torch(
|
|
x: torch.Tensor,
|
|
is_sf_swizzled_layout: bool = False,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
"""Naive MXFP8 quantization.
|
|
For each block of 32 elements along the last dimension, compute a
|
|
shared e8m0 scale (the biased exponent of the block-wise amax)
|
|
and quantize each element to float8_e4m3fn.
|
|
|
|
Returns (quantized_values [same shape, fp8], scales uint8).
|
|
Scale shape depends on is_sf_swizzled_layout:
|
|
False -> [..., K//32] (row-major 2D)
|
|
True -> [flat swizzled 1D]
|
|
"""
|
|
assert x.shape[-1] % MXFP8_BLOCK_SIZE == 0
|
|
orig_shape = x.shape
|
|
num_blocks = x.shape[-1] // MXFP8_BLOCK_SIZE
|
|
|
|
x_fp32 = x.to(torch.float32)
|
|
x_blocked = x_fp32.view(*orig_shape[:-1], num_blocks, MXFP8_BLOCK_SIZE)
|
|
|
|
amax = x_blocked.abs().amax(dim=-1)
|
|
amax = amax.clamp(min=torch.finfo(torch.float32).tiny)
|
|
scale_biased = torch.floor(torch.log2(amax)) + 127.0
|
|
scale_biased = scale_biased.clamp(0, 254)
|
|
scales_uint8 = scale_biased.to(torch.uint8)
|
|
|
|
descale = torch.exp2(scale_biased - 127.0)
|
|
x_scaled = x_blocked / descale.unsqueeze(-1)
|
|
|
|
x_fp8 = x_scaled.view(orig_shape).to(MXFP8_VALUE_DTYPE)
|
|
|
|
if x.ndim == 2:
|
|
M, K = x.shape
|
|
scales_uint8 = scales_uint8.view(M, -1)
|
|
if is_sf_swizzled_layout:
|
|
scales_uint8 = swizzle_mxfp8_scale(scales_uint8, M=M, K=K)
|
|
elif x.ndim == 3:
|
|
B, M, K = x.shape
|
|
scales_uint8 = scales_uint8.view(B, M, -1)
|
|
if is_sf_swizzled_layout:
|
|
swizzled = []
|
|
for i in range(B):
|
|
swizzled.append(swizzle_mxfp8_scale(scales_uint8[i], M=M, K=K))
|
|
scales_uint8 = torch.cat(swizzled)
|
|
|
|
return x_fp8, scales_uint8
|
|
|
|
|
|
def _mxfp8_quant_triton_kernel():
|
|
"""Lazily-built Triton kernel: per-32-block E8M0 scale + FP8-E4M3 quant.
|
|
|
|
Fuses what ``_mxfp8_e4m3_quantize_torch`` does in several elementwise passes
|
|
into one launch. Each program handles ``[BLOCK_M, 32]`` (one MX block).
|
|
"""
|
|
from vllm.triton_utils import tl, triton
|
|
|
|
@triton.jit
|
|
def _kernel(
|
|
x_ptr,
|
|
xq_ptr,
|
|
s_ptr,
|
|
M,
|
|
K,
|
|
sxm,
|
|
sxk,
|
|
sqm,
|
|
sqk,
|
|
ssm,
|
|
ssk,
|
|
BLOCK_M: tl.constexpr,
|
|
):
|
|
pid_m = tl.program_id(0)
|
|
pid_b = tl.program_id(1) # which 32-element block along K
|
|
offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
|
|
offs_k = pid_b * 32 + tl.arange(0, 32)
|
|
m_mask = offs_m < M
|
|
x = tl.load(
|
|
x_ptr + offs_m[:, None] * sxm + offs_k[None, :] * sxk,
|
|
mask=m_mask[:, None],
|
|
other=0.0,
|
|
).to(tl.float32)
|
|
amax = tl.maximum(tl.max(tl.abs(x), axis=1), 1e-30) # [BLOCK_M]
|
|
sb = tl.floor(tl.log2(amax)) + 127.0
|
|
sb = tl.minimum(tl.maximum(sb, 0.0), 254.0)
|
|
descale = tl.exp2(sb - 127.0)
|
|
xq = (x / descale[:, None]).to(xq_ptr.dtype.element_ty)
|
|
tl.store(
|
|
xq_ptr + offs_m[:, None] * sqm + offs_k[None, :] * sqk,
|
|
xq,
|
|
mask=m_mask[:, None],
|
|
)
|
|
tl.store(s_ptr + offs_m * ssm + pid_b * ssk, sb.to(tl.uint8), mask=m_mask)
|
|
|
|
return _kernel
|
|
|
|
|
|
_MXFP8_QUANT_KERNEL = None
|
|
|
|
|
|
def _mxfp8_e4m3_quantize_triton(
|
|
x: torch.Tensor,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
"""Fused 2D MXFP8 quant (non-swizzled, row-major [M, K//32] scales)."""
|
|
from vllm.triton_utils import triton
|
|
|
|
global _MXFP8_QUANT_KERNEL
|
|
if _MXFP8_QUANT_KERNEL is None:
|
|
_MXFP8_QUANT_KERNEL = _mxfp8_quant_triton_kernel()
|
|
|
|
M, K = x.shape
|
|
x = x.contiguous()
|
|
xq = torch.empty((M, K), dtype=MXFP8_VALUE_DTYPE, device=x.device)
|
|
scales = torch.empty(
|
|
(M, K // MXFP8_BLOCK_SIZE), dtype=MXFP8_SCALE_DTYPE, device=x.device
|
|
)
|
|
BLOCK_M = 64
|
|
grid = (triton.cdiv(M, BLOCK_M), K // MXFP8_BLOCK_SIZE)
|
|
_MXFP8_QUANT_KERNEL[grid](
|
|
x,
|
|
xq,
|
|
scales,
|
|
M,
|
|
K,
|
|
x.stride(0),
|
|
x.stride(1),
|
|
xq.stride(0),
|
|
xq.stride(1),
|
|
scales.stride(0),
|
|
scales.stride(1),
|
|
BLOCK_M=BLOCK_M,
|
|
)
|
|
return xq, scales
|
|
|
|
|
|
def _mxfp8_e4m3_quantize_impl(
|
|
x: torch.Tensor,
|
|
is_sf_swizzled_layout: bool = False,
|
|
alignment: int = 0,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
from vllm.platforms import current_platform
|
|
|
|
if current_platform.has_device_capability(100):
|
|
from flashinfer import mxfp8_quantize as flashinfer_mxfp8_quantize
|
|
|
|
x_q, x_scales = flashinfer_mxfp8_quantize(
|
|
x,
|
|
is_sf_swizzled_layout=is_sf_swizzled_layout,
|
|
alignment=alignment if alignment > 0 else 32,
|
|
backend="cute-dsl",
|
|
)
|
|
if x_scales.ndim == 1 and x.ndim == 2 and not is_sf_swizzled_layout:
|
|
x_scales = x_scales.view(x.size(0), -1)
|
|
return x_q, x_scales
|
|
|
|
# ROCm: a single fused Triton kernel beats the multi-pass torch path for the
|
|
# common 2D, non-swizzled activation-quant case (used by the native MX
|
|
# linear/MoE). Falls back to torch otherwise (3D weights, swizzled layout).
|
|
if (
|
|
current_platform.is_rocm()
|
|
and not is_sf_swizzled_layout
|
|
and x.ndim == 2
|
|
and x.shape[-1] % MXFP8_BLOCK_SIZE == 0
|
|
):
|
|
return _mxfp8_e4m3_quantize_triton(x)
|
|
|
|
return _mxfp8_e4m3_quantize_torch(x, is_sf_swizzled_layout)
|
|
|
|
|
|
def mxfp8_e4m3_quantize(
|
|
x: torch.Tensor,
|
|
is_sf_swizzled_layout: bool = False,
|
|
alignment: int = 0,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
return torch.ops.vllm.mxfp8_quantize(x, is_sf_swizzled_layout, alignment)
|
|
|
|
|
|
def dequant_mxfp8_to_bf16(x: torch.Tensor, scales: torch.Tensor) -> torch.Tensor:
|
|
"""Dequantize MXFP8 tensor to BF16."""
|
|
x_float = x.to(torch.float32)
|
|
|
|
num_blocks = x.shape[-1] // MXFP8_BLOCK_SIZE
|
|
x_blocked = x_float.view(*x.shape[:-1], num_blocks, MXFP8_BLOCK_SIZE)
|
|
|
|
descale = torch.exp2(scales.to(torch.float32) - 127.0)
|
|
|
|
dequantized = x_blocked * descale.unsqueeze(-1)
|
|
|
|
dequantized = dequantized.view(*x.shape)
|
|
|
|
return dequantized.to(torch.bfloat16)
|
|
|
|
|
|
def mxfp8_e4m3_quantize_fake(
|
|
x: torch.Tensor,
|
|
is_sf_swizzled_layout: bool = False,
|
|
alignment: int = 0,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
"""Fake implementation for torch.compile tracing."""
|
|
fp_data = torch.empty_like(x, dtype=MXFP8_VALUE_DTYPE)
|
|
|
|
block_size = MXFP8_BLOCK_SIZE
|
|
|
|
if x.ndim == 2:
|
|
M, N = x.shape
|
|
K = (N + block_size - 1) // block_size
|
|
if is_sf_swizzled_layout:
|
|
M_padded = ((M + 127) // 128) * 128
|
|
K_padded = ((K + 3) // 4) * 4
|
|
scales = torch.empty(
|
|
M_padded * K_padded, dtype=MXFP8_SCALE_DTYPE, device=x.device
|
|
)
|
|
else:
|
|
scales = torch.empty((M, K), dtype=MXFP8_SCALE_DTYPE, device=x.device)
|
|
elif x.ndim == 3:
|
|
B, M, N = x.shape
|
|
K = (N + block_size - 1) // block_size
|
|
if is_sf_swizzled_layout:
|
|
M_padded = ((M + 127) // 128) * 128
|
|
K_padded = ((K + 3) // 4) * 4
|
|
scales = torch.empty(
|
|
B * M_padded * K_padded, dtype=MXFP8_SCALE_DTYPE, device=x.device
|
|
)
|
|
else:
|
|
scales = torch.empty((B, M, K), dtype=MXFP8_SCALE_DTYPE, device=x.device)
|
|
else:
|
|
scale_shape = list(x.shape)
|
|
scale_shape[-1] = (x.shape[-1] + block_size - 1) // block_size
|
|
scales = torch.empty(scale_shape, dtype=MXFP8_SCALE_DTYPE, device=x.device)
|
|
|
|
return fp_data, scales
|
|
|
|
|
|
direct_register_custom_op(
|
|
op_name="mxfp8_quantize",
|
|
op_func=_mxfp8_e4m3_quantize_impl,
|
|
fake_impl=mxfp8_e4m3_quantize_fake,
|
|
)
|
|
|
|
|
|
def xpu_mxfp8_quantize(
|
|
x: torch.Tensor, dtype: torch.dtype | None = None
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
return torch.ops.vllm.xpu_mxfp8_quantize(x, dtype)
|