Add mxfp8 support for online quantization, Triton dense linear, and CUTLASS MoE (#17449)
This commit is contained in:
Ziang Li
@@ -5,7 +5,7 @@ from typing import Optional, Tuple
|
||||
import torch
|
||||
|
||||
from sglang.srt.layers.moe.cutlass_moe_params import CutlassMoEParams
|
||||
from sglang.srt.utils import is_cuda, is_sm90_supported
|
||||
from sglang.srt.utils import is_cuda, is_sm90_supported, is_sm100_supported
|
||||
|
||||
_is_cuda = is_cuda()
|
||||
if _is_cuda:
|
||||
@@ -13,6 +13,8 @@ if _is_cuda:
|
||||
apply_shuffle_mul_sum,
|
||||
cutlass_fp4_group_mm,
|
||||
es_fp8_blockwise_scaled_grouped_mm,
|
||||
es_sm100_mxfp8_blockscaled_grouped_mm,
|
||||
es_sm100_mxfp8_blockscaled_grouped_quant,
|
||||
fp8_blockwise_scaled_grouped_mm,
|
||||
prepare_moe_input,
|
||||
scaled_fp4_experts_quant,
|
||||
@@ -43,6 +45,7 @@ def cutlass_fused_experts_fp8(
|
||||
problem_sizes1: torch.Tensor,
|
||||
problem_sizes2: torch.Tensor,
|
||||
use_fp8_blockscale: bool = True,
|
||||
use_mxfp8: bool = False,
|
||||
output: Optional[torch.Tensor] = None,
|
||||
enable_es: Tuple[bool, bool] = (False, False),
|
||||
) -> torch.Tensor:
|
||||
@@ -99,6 +102,8 @@ def cutlass_fused_experts_fp8(
|
||||
b_scales_ptrs (torch.Tensor): Pointers container for calculating offsets of the input scales for each expert.
|
||||
use_fp8_blockscale (bool, optional): Flag indicating usage of FP8 with
|
||||
block scaling. Currently, only `True` is supported. Defaults to `True`.
|
||||
use_mxfp8 (bool, optional): Flag indicating usage of MXFP8 (UE8M0 scales)
|
||||
with SM100 expert-specialization kernels. Defaults to `False`.
|
||||
output (torch.Tensor, optional): Output tensor. If not provided, a new tensor will be created.
|
||||
enable_es (tuple(bool, bool)): Flag indicating usage of expert specialization kernel for (up-projection, down-projection)
|
||||
Returns:
|
||||
@@ -137,6 +142,44 @@ def cutlass_fused_experts_fp8(
|
||||
a_map = torch.empty((topk_ids.numel()), dtype=torch.int32, device=device)
|
||||
c_map = torch.empty((topk_ids.numel()), dtype=torch.int32, device=device)
|
||||
|
||||
if use_mxfp8:
|
||||
assert es_up and es_down, "MXFP8 requires expert-specialization for both GEMMs"
|
||||
assert is_sm100_supported(), "MXFP8 requires SM100"
|
||||
assert k % 32 == 0, "MXFP8 requires hidden size to be divisible by 32"
|
||||
assert n % 32 == 0, "MXFP8 requires intermediate size to be divisible by 32"
|
||||
assert w1_scale.dtype == torch.uint8, "MXFP8 w1_scale must be uint8"
|
||||
assert w2_scale.dtype == torch.uint8, "MXFP8 w2_scale must be uint8"
|
||||
expected_w1_scale_shape = (
|
||||
num_experts,
|
||||
w1_q.shape[1] // 32,
|
||||
w1_q.shape[2],
|
||||
)
|
||||
expected_w2_scale_shape = (
|
||||
num_experts,
|
||||
w2_q.shape[1] // 32,
|
||||
w2_q.shape[2],
|
||||
)
|
||||
assert (
|
||||
w1_scale.shape == expected_w1_scale_shape
|
||||
), f"MXFP8 w1_scale must be {expected_w1_scale_shape}, got {w1_scale.shape}"
|
||||
assert (
|
||||
w2_scale.shape == expected_w2_scale_shape
|
||||
), f"MXFP8 w2_scale must be {expected_w2_scale_shape}, got {w2_scale.shape}"
|
||||
|
||||
mxfp8_blockscale_align = 128
|
||||
total_tokens = m * topk
|
||||
nonzero_experts = min(num_experts, total_tokens)
|
||||
max_total = total_tokens + (mxfp8_blockscale_align - 1) * nonzero_experts
|
||||
max_blockscale = (
|
||||
(max_total + mxfp8_blockscale_align - 1) // mxfp8_blockscale_align
|
||||
) * mxfp8_blockscale_align
|
||||
|
||||
blockscale_offsets = None
|
||||
if use_mxfp8 and (es_up or es_down):
|
||||
blockscale_offsets = torch.empty(
|
||||
(num_experts + 1,), dtype=torch.int32, device=device
|
||||
)
|
||||
|
||||
prepare_moe_input(
|
||||
topk_ids,
|
||||
expert_offsets,
|
||||
@@ -147,11 +190,27 @@ def cutlass_fused_experts_fp8(
|
||||
num_experts,
|
||||
n,
|
||||
k,
|
||||
blockscale_offsets,
|
||||
)
|
||||
|
||||
a_q, a1_scale = sglang_per_token_group_quant_fp8(a, 128)
|
||||
rep_a_q = shuffle_rows(a_q, a_map, (m * topk, k))
|
||||
rep_a1_scales = shuffle_rows(a1_scale, a_map, (m * topk, int(k / 128)))
|
||||
if use_mxfp8 and es_up:
|
||||
rep_a = shuffle_rows(a, a_map, (m * topk, k))
|
||||
rep_a_q = torch.empty_like(rep_a, dtype=torch.float8_e4m3fn)
|
||||
rep_a1_scales = torch.empty(
|
||||
(max_blockscale, k // 32), dtype=torch.uint8, device=device
|
||||
)
|
||||
es_sm100_mxfp8_blockscaled_grouped_quant(
|
||||
rep_a,
|
||||
problem_sizes1,
|
||||
expert_offsets[:-1],
|
||||
blockscale_offsets[:-1],
|
||||
rep_a_q,
|
||||
rep_a1_scales,
|
||||
)
|
||||
else:
|
||||
a_q, a1_scale = sglang_per_token_group_quant_fp8(a, 128)
|
||||
rep_a_q = shuffle_rows(a_q, a_map, (m * topk, k))
|
||||
rep_a1_scales = shuffle_rows(a1_scale, a_map, (m * topk, int(k / 128)))
|
||||
|
||||
c1 = torch.empty((m * topk, n * 2), device=device, dtype=out_dtype)
|
||||
c2 = torch.empty((m * topk, k), device=device, dtype=out_dtype)
|
||||
@@ -173,6 +232,17 @@ def cutlass_fused_experts_fp8(
|
||||
expert_offsets[:-1],
|
||||
workspace,
|
||||
)
|
||||
elif use_mxfp8 and es_up:
|
||||
es_sm100_mxfp8_blockscaled_grouped_mm(
|
||||
c1,
|
||||
rep_a_q,
|
||||
w1_q,
|
||||
rep_a1_scales,
|
||||
w1_scale,
|
||||
problem_sizes1,
|
||||
expert_offsets[:-1],
|
||||
blockscale_offsets[:-1],
|
||||
)
|
||||
else:
|
||||
fp8_blockwise_scaled_grouped_mm(
|
||||
c1,
|
||||
@@ -198,7 +268,21 @@ def cutlass_fused_experts_fp8(
|
||||
intermediate = torch.empty((m * topk, n), device=device, dtype=out_dtype)
|
||||
silu_and_mul(c1, intermediate)
|
||||
|
||||
intemediate_q, a2_scale = sglang_per_token_group_quant_fp8(intermediate, 128)
|
||||
if use_mxfp8 and es_down:
|
||||
intemediate_q = torch.empty_like(intermediate, dtype=torch.float8_e4m3fn)
|
||||
a2_scale = torch.empty(
|
||||
(max_blockscale, n // 32), dtype=torch.uint8, device=device
|
||||
)
|
||||
es_sm100_mxfp8_blockscaled_grouped_quant(
|
||||
intermediate,
|
||||
problem_sizes2,
|
||||
expert_offsets[:-1],
|
||||
blockscale_offsets[:-1],
|
||||
intemediate_q,
|
||||
a2_scale,
|
||||
)
|
||||
else:
|
||||
intemediate_q, a2_scale = sglang_per_token_group_quant_fp8(intermediate, 128)
|
||||
|
||||
if is_sm90_supported() and es_down:
|
||||
es_fp8_blockwise_scaled_grouped_mm(
|
||||
@@ -214,6 +298,17 @@ def cutlass_fused_experts_fp8(
|
||||
expert_offsets[:-1],
|
||||
workspace,
|
||||
)
|
||||
elif use_mxfp8 and es_down:
|
||||
es_sm100_mxfp8_blockscaled_grouped_mm(
|
||||
c2,
|
||||
intemediate_q,
|
||||
w2_q,
|
||||
a2_scale,
|
||||
w2_scale,
|
||||
problem_sizes2,
|
||||
expert_offsets[:-1],
|
||||
blockscale_offsets[:-1],
|
||||
)
|
||||
else:
|
||||
fp8_blockwise_scaled_grouped_mm(
|
||||
c2,
|
||||
|
||||
@@ -52,6 +52,7 @@ if TYPE_CHECKING:
|
||||
# Base quantization methods
|
||||
BASE_QUANTIZATION_METHODS: Dict[str, Type[QuantizationConfig]] = {
|
||||
"fp8": Fp8Config,
|
||||
"mxfp8": Fp8Config,
|
||||
"blockwise_int8": BlockInt8Config,
|
||||
"modelopt": ModelOptFp8Config, # Auto-detect, defaults to FP8
|
||||
"modelopt_fp8": ModelOptFp8Config,
|
||||
|
||||
@@ -47,8 +47,10 @@ from sglang.srt.layers.quantization.fp8_utils import (
|
||||
cutlass_fp8_supported,
|
||||
dispatch_w8a8_block_fp8_linear,
|
||||
input_to_float8,
|
||||
mxfp8_group_quantize,
|
||||
normalize_e4m3fn_to_e4m3fnuz,
|
||||
requant_weight_ue8m0_inplace,
|
||||
triton_mxfp8_blockscaled_linear,
|
||||
)
|
||||
from sglang.srt.layers.quantization.kv_cache import BaseKVCacheMethod
|
||||
from sglang.srt.layers.quantization.marlin_utils_fp8 import (
|
||||
@@ -112,6 +114,7 @@ class Fp8Config(QuantizationConfig):
|
||||
activation_scheme: str = "dynamic",
|
||||
ignored_layers: Optional[List[str]] = None,
|
||||
weight_block_size: List[int] = None,
|
||||
use_mxfp8: bool = False,
|
||||
) -> None:
|
||||
self.is_checkpoint_fp8_serialized = is_checkpoint_fp8_serialized
|
||||
if is_checkpoint_fp8_serialized:
|
||||
@@ -120,6 +123,7 @@ class Fp8Config(QuantizationConfig):
|
||||
raise ValueError(f"Unsupported activation scheme {activation_scheme}")
|
||||
self.activation_scheme = activation_scheme
|
||||
self.ignored_layers = ignored_layers or []
|
||||
self.use_mxfp8 = use_mxfp8
|
||||
if weight_block_size is not None:
|
||||
if not is_checkpoint_fp8_serialized:
|
||||
raise ValueError(
|
||||
@@ -133,19 +137,22 @@ class Fp8Config(QuantizationConfig):
|
||||
raise ValueError(
|
||||
f"The block-wise quantization only supports dynamic activation scheme for now, but got {activation_scheme} activation scheme."
|
||||
)
|
||||
if self.use_mxfp8:
|
||||
if weight_block_size is None:
|
||||
weight_block_size = [1, 32]
|
||||
elif weight_block_size != [1, 32]:
|
||||
raise ValueError("MXFP8 requires weight_block_size=[1, 32].")
|
||||
self.weight_block_size = weight_block_size
|
||||
|
||||
@classmethod
|
||||
def get_name(cls) -> str:
|
||||
return "fp8"
|
||||
def get_name(self) -> str:
|
||||
return "mxfp8" if self.use_mxfp8 else "fp8"
|
||||
|
||||
@classmethod
|
||||
def get_supported_act_dtypes(cls) -> List[torch.dtype]:
|
||||
return [torch.bfloat16, torch.half]
|
||||
|
||||
@classmethod
|
||||
def get_min_capability(cls) -> int:
|
||||
return 80
|
||||
def get_min_capability(self) -> int:
|
||||
return 100 if self.use_mxfp8 else 80
|
||||
|
||||
@classmethod
|
||||
def get_config_filenames(cls) -> List[str]:
|
||||
@@ -154,7 +161,8 @@ class Fp8Config(QuantizationConfig):
|
||||
@classmethod
|
||||
def from_config(cls, config: Dict[str, Any]) -> Fp8Config:
|
||||
quant_method = cls.get_from_keys(config, ["quant_method"])
|
||||
is_checkpoint_fp8_serialized = "fp8" in quant_method
|
||||
use_mxfp8 = "mxfp8" in quant_method
|
||||
is_checkpoint_fp8_serialized = ("fp8" in quant_method) or use_mxfp8
|
||||
activation_scheme = cls.get_from_keys(config, ["activation_scheme"])
|
||||
ignored_layers = cls.get_from_keys_or(
|
||||
config, ["ignored_layers", "modules_to_not_convert"], None
|
||||
@@ -163,11 +171,17 @@ class Fp8Config(QuantizationConfig):
|
||||
# hack for ministral
|
||||
ignored_layers = [layer.replace("model.", "") for layer in ignored_layers]
|
||||
weight_block_size = cls.get_from_keys_or(config, ["weight_block_size"], None)
|
||||
if use_mxfp8 and weight_block_size is not None:
|
||||
logger.warning(
|
||||
"MXFP8 ignoring incoming weight_block_size in config.json; it is fixed to [1, 32]."
|
||||
)
|
||||
weight_block_size = [1, 32]
|
||||
return cls(
|
||||
is_checkpoint_fp8_serialized=is_checkpoint_fp8_serialized,
|
||||
activation_scheme=activation_scheme,
|
||||
ignored_layers=ignored_layers,
|
||||
weight_block_size=weight_block_size,
|
||||
use_mxfp8=use_mxfp8,
|
||||
)
|
||||
|
||||
def get_quant_method(
|
||||
@@ -223,7 +237,10 @@ class Fp8LinearMethod(LinearMethodBase):
|
||||
auto_enable = can_auto_enable_marlin_fp8()
|
||||
self.use_marlin = force_marlin or auto_enable
|
||||
|
||||
self.block_quant = self.quant_config.weight_block_size is not None
|
||||
self.use_mxfp8 = getattr(self.quant_config, "use_mxfp8", False)
|
||||
self.block_quant = (
|
||||
self.use_mxfp8 or self.quant_config.weight_block_size is not None
|
||||
)
|
||||
self.w8a8_block_fp8_linear = dispatch_w8a8_block_fp8_linear()
|
||||
self.is_checkpoint_fp8_serialized = (
|
||||
self.quant_config.is_checkpoint_fp8_serialized
|
||||
@@ -324,18 +341,25 @@ class Fp8LinearMethod(LinearMethodBase):
|
||||
assert self.quant_config.activation_scheme == "dynamic"
|
||||
elif hasattr(self.quant_config, "linear_activation_scheme"):
|
||||
assert self.quant_config.linear_activation_scheme == "dynamic"
|
||||
if self.use_mxfp8 and not self.is_checkpoint_fp8_serialized:
|
||||
raise ValueError(
|
||||
"MXFP8 requires fp8-serialized checkpoint for linear layers."
|
||||
)
|
||||
scale_dtype = torch.uint8 if self.use_mxfp8 else torch.float32
|
||||
scale_init = torch.zeros if scale_dtype == torch.uint8 else torch.empty
|
||||
scale = BlockQuantScaleParameter(
|
||||
data=torch.empty(
|
||||
data=scale_init(
|
||||
(output_size_per_partition + block_n - 1) // block_n,
|
||||
(input_size_per_partition + block_k - 1) // block_k,
|
||||
dtype=torch.float32,
|
||||
dtype=scale_dtype,
|
||||
),
|
||||
input_dim=1,
|
||||
output_dim=0,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
scale.format_ue8m0 = False
|
||||
scale[:] = torch.finfo(torch.float32).min
|
||||
scale.format_ue8m0 = self.use_mxfp8
|
||||
if scale_dtype != torch.uint8:
|
||||
scale[:] = torch.finfo(torch.float32).min
|
||||
layer.register_parameter("weight_scale_inv", scale)
|
||||
else:
|
||||
scale = PerTensorScaleParameter(
|
||||
@@ -382,6 +406,15 @@ class Fp8LinearMethod(LinearMethodBase):
|
||||
layer.weight_scale_inv.data, requires_grad=False
|
||||
)
|
||||
return
|
||||
elif self.use_mxfp8:
|
||||
if not self.is_checkpoint_fp8_serialized:
|
||||
self._quantize_mxfp8_weights(layer)
|
||||
return
|
||||
# MXFP8 scales are stored as UE8M0 uint8; no requantization here.
|
||||
# Keep parameter object to preserve weight_loader attrs for hot reload.
|
||||
layer.weight_scale_inv.requires_grad_(False)
|
||||
layer.weight_scale_inv.format_ue8m0 = True
|
||||
return
|
||||
else:
|
||||
# For fp8 linear weights run with deepgemm, the weights and scales need be requantized to ue8m0
|
||||
from sglang.srt.layers.quantization.fp8_utils import (
|
||||
@@ -414,6 +447,23 @@ class Fp8LinearMethod(LinearMethodBase):
|
||||
layer.weight.data = weight.data
|
||||
layer.weight_scale_inv.data = weight_scale.data
|
||||
|
||||
def _quantize_mxfp8_weights(self, layer: Module) -> None:
|
||||
weight = layer.weight.data
|
||||
qweight, weight_scale = mxfp8_group_quantize(weight)
|
||||
# Keep parameter objects to preserve weight_loader attrs for hot reload.
|
||||
layer.weight.data = qweight
|
||||
layer.weight.requires_grad_(False)
|
||||
if hasattr(layer, "weight_scale_inv") and layer.weight_scale_inv is not None:
|
||||
layer.weight_scale_inv.data = weight_scale
|
||||
layer.weight_scale_inv.requires_grad_(False)
|
||||
else:
|
||||
# First-time online MXFP8 quantization (no serialized scales).
|
||||
layer.register_parameter(
|
||||
"weight_scale_inv", Parameter(weight_scale, requires_grad=False)
|
||||
)
|
||||
layer.weight_scale_inv.format_ue8m0 = True
|
||||
layer.input_scale = None
|
||||
|
||||
def process_weights_after_loading(self, layer: Module) -> None:
|
||||
if self.block_quant:
|
||||
self.process_weights_after_loading_block_quant(layer)
|
||||
@@ -543,6 +593,23 @@ class Fp8LinearMethod(LinearMethodBase):
|
||||
bias=bias,
|
||||
)
|
||||
|
||||
if self.use_mxfp8:
|
||||
if isinstance(x, tuple):
|
||||
return triton_mxfp8_blockscaled_linear(
|
||||
input=x[0],
|
||||
weight=layer.weight,
|
||||
weight_scale=layer.weight_scale_inv,
|
||||
input_scale=x[1],
|
||||
bias=bias,
|
||||
)
|
||||
return triton_mxfp8_blockscaled_linear(
|
||||
input=x,
|
||||
weight=layer.weight,
|
||||
weight_scale=layer.weight_scale_inv,
|
||||
input_scale=None,
|
||||
bias=bias,
|
||||
)
|
||||
|
||||
if self.block_quant:
|
||||
if use_intel_amx_backend(layer):
|
||||
return torch.ops.sgl_kernel.fp8_scaled_mm_cpu(
|
||||
@@ -600,7 +667,10 @@ class Fp8MoEMethod(FusedMoEMethodBase):
|
||||
|
||||
def __init__(self, quant_config: Fp8Config):
|
||||
self.quant_config = quant_config
|
||||
self.block_quant = self.quant_config.weight_block_size is not None
|
||||
self.use_mxfp8 = getattr(self.quant_config, "use_mxfp8", False)
|
||||
self.block_quant = (
|
||||
self.use_mxfp8 or self.quant_config.weight_block_size is not None
|
||||
)
|
||||
if get_moe_runner_backend().is_cutlass():
|
||||
assert (
|
||||
cutlass_fp8_supported()
|
||||
@@ -708,27 +778,29 @@ class Fp8MoEMethod(FusedMoEMethodBase):
|
||||
|
||||
# WEIGHT_SCALES
|
||||
if self.block_quant:
|
||||
scale_dtype = torch.uint8 if self.use_mxfp8 else torch.float32
|
||||
scale_init = torch.zeros if scale_dtype == torch.uint8 else torch.ones
|
||||
w13_weight_scale = torch.nn.Parameter(
|
||||
torch.ones(
|
||||
scale_init(
|
||||
num_experts,
|
||||
2 * ((intermediate_size_per_partition + block_n - 1) // block_n),
|
||||
(hidden_size + block_k - 1) // block_k,
|
||||
dtype=torch.float32,
|
||||
dtype=scale_dtype,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
w2_weight_scale = torch.nn.Parameter(
|
||||
torch.ones(
|
||||
scale_init(
|
||||
num_experts,
|
||||
(hidden_size + block_n - 1) // block_n,
|
||||
(intermediate_size_per_partition + block_k - 1) // block_k,
|
||||
dtype=torch.float32,
|
||||
dtype=scale_dtype,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
# w13_weight and w2_weight are always requanted together
|
||||
w13_weight_scale.format_ue8m0 = False
|
||||
w2_weight_scale.format_ue8m0 = False
|
||||
w13_weight_scale.format_ue8m0 = self.use_mxfp8
|
||||
w2_weight_scale.format_ue8m0 = self.use_mxfp8
|
||||
layer.register_parameter("w13_weight_scale_inv", w13_weight_scale)
|
||||
layer.register_parameter("w2_weight_scale_inv", w2_weight_scale)
|
||||
assert self.quant_config.activation_scheme == "dynamic"
|
||||
@@ -856,6 +928,10 @@ class Fp8MoEMethod(FusedMoEMethodBase):
|
||||
_is_cpu_amx_available
|
||||
), "Fp8MoEMethod on CPU requires that CPU has AMX support"
|
||||
_amx_process_weight_after_loading(layer, ["w13_weight", "w2_weight"])
|
||||
elif self.use_mxfp8:
|
||||
self._process_mxfp8_moe_weights(
|
||||
layer, quantize=not self.quant_config.is_checkpoint_fp8_serialized
|
||||
)
|
||||
else:
|
||||
# For fp8 moe run with deepgemm, the expert weights and scales need be requantized to ue8m0
|
||||
from sglang.srt.layers.moe.ep_moe.layer import DeepEPMoE
|
||||
@@ -888,6 +964,143 @@ class Fp8MoEMethod(FusedMoEMethodBase):
|
||||
layer.w13_weight_scale_inv.format_ue8m0 = True
|
||||
layer.w2_weight_scale_inv.format_ue8m0 = True
|
||||
|
||||
def _process_mxfp8_moe_weights(self, layer: Module, quantize: bool = True) -> None:
|
||||
|
||||
if not (_is_cuda and is_sm100_supported()):
|
||||
raise RuntimeError("MXFP8 MoE quantization requires SM100.")
|
||||
|
||||
def _quantize_and_swizzle_with_cutlass_es_kernel(weight: torch.Tensor):
|
||||
from sgl_kernel import es_sm100_mxfp8_blockscaled_grouped_quant
|
||||
|
||||
weight = weight.contiguous()
|
||||
num_experts, m, k = weight.shape
|
||||
assert k % 32 == 0, f"{k=} must be divisible by 32 for MXFP8"
|
||||
|
||||
weight_flat = weight.view(-1, k).contiguous()
|
||||
problem_sizes = torch.empty(
|
||||
(num_experts, 3), dtype=torch.int32, device=weight.device
|
||||
)
|
||||
problem_sizes[:, 0] = m
|
||||
problem_sizes[:, 1] = 0
|
||||
problem_sizes[:, 2] = k
|
||||
expert_offsets = torch.arange(
|
||||
0, num_experts * m, m, dtype=torch.int32, device=weight.device
|
||||
)
|
||||
aligned_m = ((m + 127) // 128) * 128
|
||||
blockscale_offsets = torch.arange(
|
||||
0,
|
||||
num_experts * aligned_m,
|
||||
aligned_m,
|
||||
dtype=torch.int32,
|
||||
device=weight.device,
|
||||
)
|
||||
qweight = torch.empty_like(weight_flat, dtype=torch.float8_e4m3fn)
|
||||
scale = torch.empty(
|
||||
(num_experts * aligned_m, k // 32),
|
||||
dtype=torch.uint8,
|
||||
device=weight.device,
|
||||
)
|
||||
es_sm100_mxfp8_blockscaled_grouped_quant(
|
||||
weight_flat,
|
||||
problem_sizes,
|
||||
expert_offsets,
|
||||
blockscale_offsets,
|
||||
qweight,
|
||||
scale,
|
||||
)
|
||||
qweight = qweight.view_as(weight)
|
||||
scale = scale.view(num_experts, aligned_m, k // 32)
|
||||
if aligned_m != m:
|
||||
scale = scale[:, :m, :]
|
||||
return qweight, scale
|
||||
|
||||
def _swizzle_mxfp8_sf(scale, num_warps):
|
||||
from triton_kernels.tensor import convert_layout, wrap_torch_tensor
|
||||
from triton_kernels.tensor_details import layout
|
||||
|
||||
scale_layout, scale_layout_opts = (
|
||||
layout.make_default_matmul_mxfp4_w_scale_layout(
|
||||
mx_axis=1, num_warps=num_warps
|
||||
)
|
||||
)
|
||||
scale = scale.transpose(-2, -1)
|
||||
scale = convert_layout(
|
||||
wrap_torch_tensor(scale), scale_layout, **scale_layout_opts
|
||||
)
|
||||
return scale
|
||||
|
||||
def _swizzle_with_triton_kernel(
|
||||
weight_shape: tuple[int, int, int], scale: torch.Tensor
|
||||
):
|
||||
num_experts, m, k = weight_shape
|
||||
aligned_m = ((m + 127) // 128) * 128
|
||||
scale = scale.view(num_experts, aligned_m, k // 32)
|
||||
num_warps = 8
|
||||
scale = _swizzle_mxfp8_sf(scale, num_warps)
|
||||
scale = scale.data.view(num_experts, aligned_m, k // 32)
|
||||
return scale
|
||||
|
||||
def _quantize_and_swizzle_with_triton_kernel(weight: torch.Tensor):
|
||||
|
||||
weight = weight.contiguous()
|
||||
_, _, k = weight.shape
|
||||
assert k % 32 == 0, f"{k=} must be divisible by 32 for MXFP8"
|
||||
|
||||
weight_flat = weight.view(-1, k).contiguous()
|
||||
qweight, scale = mxfp8_group_quantize(weight_flat)
|
||||
qweight = qweight.view_as(weight)
|
||||
scale = _swizzle_with_triton_kernel(weight.shape, scale)
|
||||
return qweight, scale
|
||||
|
||||
if quantize:
|
||||
if get_moe_runner_backend().is_cutlass():
|
||||
w13_q, w13_s = _quantize_and_swizzle_with_cutlass_es_kernel(
|
||||
layer.w13_weight.data
|
||||
)
|
||||
w2_q, w2_s = _quantize_and_swizzle_with_cutlass_es_kernel(
|
||||
layer.w2_weight.data
|
||||
)
|
||||
else:
|
||||
w13_q, w13_s = _quantize_and_swizzle_with_triton_kernel(
|
||||
layer.w13_weight.data
|
||||
)
|
||||
w2_q, w2_s = _quantize_and_swizzle_with_triton_kernel(
|
||||
layer.w2_weight.data
|
||||
)
|
||||
else:
|
||||
w13_q = layer.w13_weight.data
|
||||
w2_q = layer.w2_weight.data
|
||||
w13_s = _swizzle_with_triton_kernel(
|
||||
layer.w13_weight.data.shape, layer.w13_weight_scale_inv.data
|
||||
)
|
||||
w2_s = _swizzle_with_triton_kernel(
|
||||
layer.w2_weight.data.shape, layer.w2_weight_scale_inv.data
|
||||
)
|
||||
|
||||
# Keep parameter objects to preserve weight_loader attrs for hot reload.
|
||||
# Prefer in-place copy; rebind only when shape/dtype changes (online quantize).
|
||||
def _copy_or_rebind(param: Parameter, new_value: torch.Tensor) -> None:
|
||||
if (
|
||||
param.data.shape == new_value.shape
|
||||
and param.data.dtype == new_value.dtype
|
||||
):
|
||||
param.data.copy_(new_value)
|
||||
else:
|
||||
param.data = new_value
|
||||
|
||||
_copy_or_rebind(layer.w13_weight, w13_q)
|
||||
_copy_or_rebind(layer.w2_weight, w2_q)
|
||||
_copy_or_rebind(layer.w13_weight_scale_inv, w13_s)
|
||||
_copy_or_rebind(layer.w2_weight_scale_inv, w2_s)
|
||||
layer.w13_weight.requires_grad_(False)
|
||||
layer.w2_weight.requires_grad_(False)
|
||||
layer.w13_weight_scale_inv.requires_grad_(False)
|
||||
layer.w2_weight_scale_inv.requires_grad_(False)
|
||||
layer.w13_weight_scale_inv.format_ue8m0 = True
|
||||
layer.w2_weight_scale_inv.format_ue8m0 = True
|
||||
layer.w13_input_scale = None
|
||||
layer.w2_input_scale = None
|
||||
|
||||
def process_weights_after_loading(self, layer: Module) -> None:
|
||||
if _is_hip and _use_hip_int4:
|
||||
self.process_weights_hip_int4(layer)
|
||||
@@ -1173,6 +1386,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
|
||||
symm_output = torch.empty_like(x)
|
||||
|
||||
topk_weights, topk_ids, _ = dispatch_output.topk_output
|
||||
use_mxfp8 = getattr(self.quant_config, "use_mxfp8", False)
|
||||
output = cutlass_fused_experts_fp8(
|
||||
x,
|
||||
layer.w13_weight.transpose(1, 2),
|
||||
@@ -1195,7 +1409,9 @@ class Fp8MoEMethod(FusedMoEMethodBase):
|
||||
self.problem_sizes1,
|
||||
self.problem_sizes2,
|
||||
use_fp8_blockscale=True,
|
||||
use_mxfp8=use_mxfp8,
|
||||
output=symm_output,
|
||||
enable_es=(use_mxfp8, use_mxfp8),
|
||||
)
|
||||
return StandardCombineInput(hidden_states=output)
|
||||
|
||||
|
||||
@@ -23,6 +23,11 @@ import torch
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
try:
|
||||
from triton.tools.tensor_descriptor import TensorDescriptor
|
||||
except:
|
||||
pass
|
||||
|
||||
from sglang.srt.layers import deep_gemm_wrapper
|
||||
from sglang.srt.utils import (
|
||||
ceil_align,
|
||||
@@ -1175,6 +1180,140 @@ def w8a8_block_fp8_matmul(
|
||||
)
|
||||
|
||||
|
||||
# Copied and adapted from https://github.com/triton-lang/triton/blob/main/python/tutorials/10-block-scaled-matmul.py
|
||||
@triton.jit
|
||||
def _mxfp8_block_scaled_matmul_kernel( #
|
||||
a_desc, #
|
||||
a_scale_desc, #
|
||||
b_desc, #
|
||||
b_scale_desc, #
|
||||
c_desc, #
|
||||
M: tl.constexpr, #
|
||||
N: tl.constexpr, #
|
||||
K: tl.constexpr, #
|
||||
output_type: tl.constexpr, #
|
||||
BLOCK_M: tl.constexpr, #
|
||||
BLOCK_N: tl.constexpr, #
|
||||
BLOCK_K: tl.constexpr, #
|
||||
rep_m: tl.constexpr, #
|
||||
rep_n: tl.constexpr, #
|
||||
rep_k: tl.constexpr, #
|
||||
NUM_STAGES: tl.constexpr, #
|
||||
): #
|
||||
if output_type == 0:
|
||||
output_dtype = tl.float32
|
||||
elif output_type == 1:
|
||||
output_dtype = tl.float16
|
||||
elif output_type == 2:
|
||||
output_dtype = tl.bfloat16
|
||||
|
||||
pid = tl.program_id(axis=0)
|
||||
num_pid_m = tl.cdiv(M, BLOCK_M)
|
||||
pid_m = pid % num_pid_m
|
||||
pid_n = pid // num_pid_m
|
||||
offs_am = pid_m * BLOCK_M
|
||||
offs_bn = pid_n * BLOCK_N
|
||||
offs_k_a = 0
|
||||
offs_k_b = 0
|
||||
offs_scale_m = pid_m * rep_m
|
||||
offs_scale_n = pid_n * rep_n
|
||||
offs_scale_k = 0
|
||||
|
||||
VEC_SIZE: tl.constexpr = 32
|
||||
|
||||
accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
|
||||
for k in tl.range(0, tl.cdiv(K, BLOCK_K), num_stages=NUM_STAGES):
|
||||
a = a_desc.load([offs_am, offs_k_a])
|
||||
b = b_desc.load([offs_bn, offs_k_b])
|
||||
scale_a = a_scale_desc.load([0, offs_scale_m, offs_scale_k, 0, 0])
|
||||
scale_b = b_scale_desc.load([0, offs_scale_n, offs_scale_k, 0, 0])
|
||||
|
||||
scale_a = (
|
||||
scale_a.reshape(rep_m, rep_k, 32, 4, 4)
|
||||
.trans(0, 3, 2, 1, 4)
|
||||
.reshape(BLOCK_M, BLOCK_K // VEC_SIZE)
|
||||
)
|
||||
scale_b = (
|
||||
scale_b.reshape(rep_n, rep_k, 32, 4, 4)
|
||||
.trans(0, 3, 2, 1, 4)
|
||||
.reshape(BLOCK_N, BLOCK_K // VEC_SIZE)
|
||||
)
|
||||
|
||||
accumulator = tl.dot_scaled(
|
||||
a, scale_a, "e4m3", b.T, scale_b, "e4m3", accumulator
|
||||
)
|
||||
|
||||
offs_k_a += BLOCK_K
|
||||
offs_k_b += BLOCK_K
|
||||
offs_scale_k += rep_k
|
||||
|
||||
c_desc.store([offs_am, offs_bn], accumulator.to(output_dtype))
|
||||
|
||||
|
||||
# Copied and adapted from https://github.com/triton-lang/triton/blob/main/python/tutorials/10-block-scaled-matmul.py
|
||||
def mxfp8_block_scaled_matmul_triton(
|
||||
a: torch.Tensor,
|
||||
a_scale: torch.Tensor,
|
||||
b: torch.Tensor,
|
||||
b_scale: torch.Tensor,
|
||||
output_dtype: torch.dtype,
|
||||
*,
|
||||
block_m: int = 128,
|
||||
block_n: int = 256,
|
||||
block_k: int = 128,
|
||||
num_stages: int = 4,
|
||||
) -> torch.Tensor:
|
||||
"""Block-scaled matmul for MXFP8 using Triton dot_scaled."""
|
||||
M, K = a.shape
|
||||
N, K_b = b.shape
|
||||
assert K == K_b
|
||||
|
||||
if output_dtype == torch.float32:
|
||||
output_type = 0
|
||||
elif output_dtype == torch.float16:
|
||||
output_type = 1
|
||||
elif output_dtype == torch.bfloat16:
|
||||
output_type = 2
|
||||
else:
|
||||
raise ValueError(f"Unsupported output dtype: {output_dtype}")
|
||||
|
||||
rep_m = block_m // 128
|
||||
rep_n = block_n // 128
|
||||
rep_k = block_k // 32 // 4
|
||||
|
||||
a_desc = TensorDescriptor.from_tensor(a, [block_m, block_k])
|
||||
b_desc = TensorDescriptor.from_tensor(b, [block_n, block_k])
|
||||
|
||||
scale_block_shape = [1, rep_m, rep_k, 2, 256]
|
||||
a_scale_desc = TensorDescriptor.from_tensor(a_scale, block_shape=scale_block_shape)
|
||||
scale_block_shape = [1, rep_n, rep_k, 2, 256]
|
||||
b_scale_desc = TensorDescriptor.from_tensor(b_scale, block_shape=scale_block_shape)
|
||||
|
||||
output = torch.empty((M, N), dtype=output_dtype, device=a.device)
|
||||
c_desc = TensorDescriptor.from_tensor(output, [block_m, block_n])
|
||||
|
||||
grid = (triton.cdiv(M, block_m) * triton.cdiv(N, block_n), 1)
|
||||
_mxfp8_block_scaled_matmul_kernel[grid](
|
||||
a_desc,
|
||||
a_scale_desc,
|
||||
b_desc,
|
||||
b_scale_desc,
|
||||
c_desc,
|
||||
M,
|
||||
N,
|
||||
K,
|
||||
output_type,
|
||||
block_m,
|
||||
block_n,
|
||||
block_k,
|
||||
rep_m,
|
||||
rep_n,
|
||||
rep_k,
|
||||
num_stages,
|
||||
)
|
||||
return output
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _per_tensor_quant_mla_fp8_stage1(
|
||||
x_ptr,
|
||||
|
||||
@@ -19,6 +19,7 @@ from sglang.srt.layers.quantization.fp8_kernel import (
|
||||
fp8_dtype,
|
||||
fp8_max,
|
||||
is_fp8_fnuz,
|
||||
mxfp8_block_scaled_matmul_triton,
|
||||
per_token_group_quant_fp8,
|
||||
scaled_fp8_quant,
|
||||
sglang_per_token_quant_fp8,
|
||||
@@ -38,6 +39,7 @@ from sglang.srt.utils import (
|
||||
is_flashinfer_available,
|
||||
is_hip,
|
||||
is_sm90_supported,
|
||||
is_sm100_supported,
|
||||
offloader,
|
||||
)
|
||||
|
||||
@@ -536,6 +538,131 @@ def triton_w8a8_block_fp8_linear(
|
||||
return output.to(dtype=input_2d.dtype).view(*output_shape)
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def _get_triton_mxfp8_downcast():
|
||||
try:
|
||||
from triton_kernels.numerics_details.mxfp import downcast_to_mxfp
|
||||
except Exception as err:
|
||||
raise RuntimeError(
|
||||
"MXFP8 quantization requires triton_kernels with MXFP8 support."
|
||||
) from err
|
||||
return downcast_to_mxfp
|
||||
|
||||
|
||||
def mxfp8_group_quantize(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Quantize a 2D contiguous tensor to MXFP8 with UE8M0 scales per group (32)."""
|
||||
assert x.dim() == 2, f"Expected 2D input, got {x.dim()}D"
|
||||
assert x.is_contiguous(), "MXFP8 quantization requires a contiguous 2D tensor."
|
||||
_, k = x.shape
|
||||
assert k % 32 == 0, f"{k=} must be divisible by 32"
|
||||
downcast_to_mxfp = _get_triton_mxfp8_downcast()
|
||||
q_input, scale_u8 = downcast_to_mxfp(x, torch.float8_e4m3fn, axis=1)
|
||||
return q_input.contiguous(), scale_u8.contiguous()
|
||||
|
||||
|
||||
def _pack_mxfp8_scales(scale_u8: torch.Tensor) -> torch.Tensor:
|
||||
# Pack (M, K//32) UE8M0 scales into the layout expected by tl.dot_scaled.
|
||||
assert scale_u8.dim() == 2, f"Expected 2D scale tensor, got {scale_u8.dim()}D"
|
||||
scale_u8 = scale_u8.contiguous()
|
||||
m, k_groups = scale_u8.shape
|
||||
assert (
|
||||
k_groups % 4 == 0
|
||||
), f"{k_groups=} must be divisible by 4 (K must be multiple of 128)"
|
||||
|
||||
scale_m = ceil_div(m, 128)
|
||||
if m % 128 != 0:
|
||||
pad_rows = scale_m * 128 - m
|
||||
pad = torch.full(
|
||||
(pad_rows, k_groups),
|
||||
127,
|
||||
dtype=scale_u8.dtype,
|
||||
device=scale_u8.device,
|
||||
)
|
||||
scale_u8 = torch.cat([scale_u8, pad], dim=0)
|
||||
|
||||
scale_k = k_groups // 4
|
||||
scale_u8 = scale_u8.view(scale_m, 128, scale_k, 4)
|
||||
scale_u8 = scale_u8.view(scale_m, 4, 32, scale_k, 4)
|
||||
packed = scale_u8.permute(0, 3, 2, 1, 4).contiguous()
|
||||
return packed.view(1, scale_m, scale_k, 2, 256)
|
||||
|
||||
|
||||
def triton_mxfp8_blockscaled_linear(
|
||||
input: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
weight_scale: torch.Tensor,
|
||||
input_scale: Optional[torch.Tensor] = None,
|
||||
bias: Optional[torch.Tensor] = None,
|
||||
output_dtype: Optional[torch.dtype] = None,
|
||||
) -> torch.Tensor:
|
||||
if not (_is_cuda and is_sm100_supported()):
|
||||
raise RuntimeError("MXFP8 dense linear requires Blackwell GPUs (SM100+).")
|
||||
|
||||
input_2d = input.view(-1, input.shape[-1]).contiguous()
|
||||
output_shape = [*input.shape[:-1], weight.shape[0]]
|
||||
|
||||
block_m = 128
|
||||
block_n = 256 if weight.shape[0] % 256 == 0 else 128
|
||||
block_k = 128
|
||||
|
||||
m, k = input_2d.shape
|
||||
n, k_w = weight.shape
|
||||
assert k == k_w, f"{k=} does not match {k_w=}"
|
||||
assert k % 128 == 0, f"{k=} must be divisible by 128 for MXFP8"
|
||||
assert n % block_n == 0, f"{n=} must be divisible by {block_n}"
|
||||
assert weight.dtype == torch.float8_e4m3fn, "MXFP8 weight must be FP8 E4M3."
|
||||
assert weight_scale.dtype == torch.uint8, "MXFP8 weight_scale must be UE8M0 uint8."
|
||||
|
||||
if input_scale is None:
|
||||
q_input, x_scale_u8 = mxfp8_group_quantize(input_2d)
|
||||
else:
|
||||
q_input = input_2d
|
||||
x_scale_u8 = input_scale
|
||||
assert x_scale_u8.dtype == torch.uint8, "MXFP8 input_scale must be UE8M0 uint8."
|
||||
assert x_scale_u8.shape == (m, k // 32)
|
||||
|
||||
if output_dtype is None:
|
||||
if input_2d.dtype in (torch.float16, torch.bfloat16, torch.float32):
|
||||
output_dtype = input_2d.dtype
|
||||
else:
|
||||
output_dtype = torch.bfloat16
|
||||
|
||||
if m % block_m != 0:
|
||||
pad_rows = ceil_div(m, block_m) * block_m - m
|
||||
q_input = torch.cat(
|
||||
[
|
||||
q_input,
|
||||
torch.zeros((pad_rows, k), device=q_input.device, dtype=q_input.dtype),
|
||||
],
|
||||
dim=0,
|
||||
)
|
||||
pad_scale = torch.full(
|
||||
(pad_rows, k // 32),
|
||||
127,
|
||||
device=x_scale_u8.device,
|
||||
dtype=x_scale_u8.dtype,
|
||||
)
|
||||
x_scale_u8 = torch.cat([x_scale_u8, pad_scale], dim=0)
|
||||
|
||||
a_scale_packed = _pack_mxfp8_scales(x_scale_u8)
|
||||
b_scale_packed = _pack_mxfp8_scales(weight_scale)
|
||||
|
||||
output = mxfp8_block_scaled_matmul_triton(
|
||||
q_input,
|
||||
a_scale_packed,
|
||||
weight.contiguous(),
|
||||
b_scale_packed,
|
||||
output_dtype=output_dtype,
|
||||
block_m=block_m,
|
||||
block_n=block_n,
|
||||
block_k=block_k,
|
||||
)
|
||||
output = output[:m, :]
|
||||
if bias is not None:
|
||||
output += bias
|
||||
return output.to(dtype=output_dtype).view(*output_shape)
|
||||
|
||||
|
||||
def dequant_mxfp4(
|
||||
w_block: torch.Tensor,
|
||||
w_scale: torch.Tensor,
|
||||
|
||||
@@ -36,6 +36,7 @@ from sglang.srt.configs.model_config import ModelConfig
|
||||
from sglang.srt.distributed import get_tensor_model_parallel_rank
|
||||
from sglang.srt.layers.dp_attention import get_attention_tp_rank
|
||||
from sglang.srt.layers.quantization import QuantizationConfig, get_quantization_config
|
||||
from sglang.srt.layers.quantization.fp8 import Fp8Config
|
||||
from sglang.srt.layers.quantization.modelopt_quant import (
|
||||
ModelOptFp4Config,
|
||||
ModelOptFp8Config,
|
||||
@@ -227,6 +228,8 @@ def get_quant_config(
|
||||
|
||||
# If the quantization config is not found, use the default config.
|
||||
if not possible_config_filenames:
|
||||
if model_config.quantization == "mxfp8":
|
||||
return Fp8Config(use_mxfp8=True, is_checkpoint_fp8_serialized=False)
|
||||
return quant_cls()
|
||||
|
||||
config_files = glob.glob(os.path.join(hf_folder, "*.json"))
|
||||
|
||||
@@ -93,6 +93,7 @@ LOAD_FORMAT_CHOICES = [
|
||||
QUANTIZATION_CHOICES = [
|
||||
"awq",
|
||||
"fp8",
|
||||
"mxfp8",
|
||||
"gptq",
|
||||
"marlin",
|
||||
"gptq_marlin",
|
||||
@@ -2013,6 +2014,14 @@ class ServerArgs:
|
||||
), "Please enable dp attention when setting enable_dp_lm_head. "
|
||||
|
||||
def _handle_moe_kernel_config(self):
|
||||
if self.quantization == "mxfp8":
|
||||
if self.moe_runner_backend not in ["auto", "cutlass"]:
|
||||
logger.warning(
|
||||
"mxfp8 quantization forces --moe-runner-backend=cutlass. "
|
||||
f"Overriding {self.moe_runner_backend!r}."
|
||||
)
|
||||
self.moe_runner_backend = "cutlass"
|
||||
|
||||
if self.moe_runner_backend == "flashinfer_cutlass":
|
||||
assert self.quantization in [
|
||||
"modelopt_fp4",
|
||||
@@ -2041,14 +2050,18 @@ class ServerArgs:
|
||||
logger.warning(
|
||||
"SGLANG_CUTLASS_MOE is deprecated, use --moe-runner-backend=cutlass and/or --speculative-moe-runner-backend=cutlass instead"
|
||||
)
|
||||
assert (
|
||||
self.quantization == "fp8"
|
||||
), "cutlass MoE is only supported with fp8 quantization"
|
||||
assert self.quantization in [
|
||||
"fp8",
|
||||
"mxfp8",
|
||||
], "cutlass MoE is only supported with fp8/mxfp8 quantization"
|
||||
self.moe_runner_backend = "cutlass"
|
||||
if self.moe_runner_backend == "cutlass" and self.quantization == "fp8":
|
||||
if self.moe_runner_backend == "cutlass" and self.quantization in [
|
||||
"fp8",
|
||||
"mxfp8",
|
||||
]:
|
||||
assert (
|
||||
self.ep_size == 1
|
||||
), "FP8 Cutlass MoE is only supported with ep_size == 1"
|
||||
), "FP8/MXFP8 Cutlass MoE is only supported with ep_size == 1"
|
||||
|
||||
def _handle_a2a_moe(self):
|
||||
if self.moe_a2a_backend == "deepep":
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
import itertools
|
||||
import unittest
|
||||
from functools import lru_cache
|
||||
|
||||
import torch
|
||||
|
||||
@@ -13,12 +14,29 @@ from sglang.srt.layers.quantization.fp8_kernel import (
|
||||
static_quant_fp8,
|
||||
w8a8_block_fp8_matmul,
|
||||
)
|
||||
from sglang.srt.layers.quantization.fp8_utils import input_to_float8
|
||||
from sglang.srt.layers.quantization.fp8_utils import (
|
||||
input_to_float8,
|
||||
mxfp8_group_quantize,
|
||||
triton_mxfp8_blockscaled_linear,
|
||||
)
|
||||
from sglang.srt.utils import is_sm100_supported
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
_is_cuda = torch.cuda.is_available() and torch.version.cuda
|
||||
|
||||
|
||||
# For test
|
||||
@lru_cache(maxsize=1)
|
||||
def _get_triton_mxfp8_upcast():
|
||||
try:
|
||||
from triton_kernels.numerics_details.mxfp import upcast_from_mxfp_torch
|
||||
except Exception as err:
|
||||
raise RuntimeError(
|
||||
"MXFP8 dequantization requires triton_kernels with MXFP8 support."
|
||||
) from err
|
||||
return upcast_from_mxfp_torch
|
||||
|
||||
|
||||
# For test
|
||||
def native_per_token_group_quant_fp8(
|
||||
x, group_size, eps=1e-10, dtype=torch.float8_e4m3fn
|
||||
@@ -414,6 +432,88 @@ class TestW8A8BlockFP8Matmul(CustomTestCase):
|
||||
self._w8a8_block_fp8_matmul(*params)
|
||||
|
||||
|
||||
def _mxfp8_group_dequant(q: torch.Tensor, scale_u8: torch.Tensor) -> torch.Tensor:
|
||||
upcast_from_mxfp_torch = _get_triton_mxfp8_upcast()
|
||||
return upcast_from_mxfp_torch(q, scale_u8, torch.float32, axis=1)
|
||||
|
||||
|
||||
class TestMXFP8DenseLinear(CustomTestCase):
|
||||
DTYPES = [torch.bfloat16]
|
||||
M = [1, 127, 128, 129, 255, 256]
|
||||
NKs = [
|
||||
(256, 512),
|
||||
(384, 1024),
|
||||
(512, 2048),
|
||||
(768, 1024),
|
||||
]
|
||||
SEEDS = [0]
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
if not torch.cuda.is_available():
|
||||
raise unittest.SkipTest("CUDA is not available")
|
||||
if not is_sm100_supported():
|
||||
raise unittest.SkipTest("MXFP8 requires Blackwell (SM100+)")
|
||||
torch.set_default_device("cuda")
|
||||
|
||||
def _mxfp8_dense_linear(self, M, NK, dtype, seed):
|
||||
N, K = NK
|
||||
torch.manual_seed(seed)
|
||||
|
||||
input_fp32 = torch.randn((M, K), dtype=torch.float32) / 4
|
||||
input_fp16 = input_fp32.to(dtype)
|
||||
|
||||
weight_fp32 = torch.randn((N, K), dtype=torch.float32) / 4
|
||||
weight_q, weight_scale_u8 = mxfp8_group_quantize(weight_fp32)
|
||||
|
||||
with torch.inference_mode():
|
||||
q_input, input_scale_u8 = mxfp8_group_quantize(input_fp16.to(torch.float32))
|
||||
a_dq = _mxfp8_group_dequant(q_input, input_scale_u8)
|
||||
b_dq = _mxfp8_group_dequant(weight_q, weight_scale_u8)
|
||||
ref_out = torch.matmul(a_dq, b_dq.t()).to(dtype)
|
||||
|
||||
out = triton_mxfp8_blockscaled_linear(
|
||||
input=input_fp16,
|
||||
weight=weight_q,
|
||||
weight_scale=weight_scale_u8,
|
||||
)
|
||||
out_prequant = triton_mxfp8_blockscaled_linear(
|
||||
input=q_input,
|
||||
weight=weight_q,
|
||||
weight_scale=weight_scale_u8,
|
||||
input_scale=input_scale_u8,
|
||||
output_dtype=dtype,
|
||||
)
|
||||
|
||||
self.assertTrue(
|
||||
torch.mean(torch.abs(out.to(torch.float32) - ref_out.to(torch.float32)))
|
||||
/ torch.mean(torch.abs(ref_out.to(torch.float32)))
|
||||
< 0.02
|
||||
)
|
||||
self.assertTrue(
|
||||
torch.mean(
|
||||
torch.abs(out_prequant.to(torch.float32) - ref_out.to(torch.float32))
|
||||
)
|
||||
/ torch.mean(torch.abs(ref_out.to(torch.float32)))
|
||||
< 0.02
|
||||
)
|
||||
|
||||
def test_mxfp8_dense_linear(self):
|
||||
for params in itertools.product(
|
||||
self.M,
|
||||
self.NKs,
|
||||
self.DTYPES,
|
||||
self.SEEDS,
|
||||
):
|
||||
with self.subTest(
|
||||
M=params[0],
|
||||
NKs=params[1],
|
||||
dtype=params[2],
|
||||
seed=params[3],
|
||||
):
|
||||
self._mxfp8_dense_linear(*params)
|
||||
|
||||
|
||||
# For test
|
||||
def torch_w8a8_block_fp8_moe(a, w1, w2, w1_s, w2_s, score, topk, block_shape):
|
||||
"""This function performs fused moe with block-wise quantization using native torch."""
|
||||
|
||||
Reference in New Issue
Block a user