minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Add mxfp8 support for online quantization, Triton dense linear, and CUTLASS MoE (#17449)

Browse Source
This commit is contained in:
Ziang Li
2026-01-29 21:33:57 +08:00
committed by GitHub
parent cfa09d311c
commit 3c9cc44ff5
9 changed files with 725 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/quantization/__init__.py
+1
View File
@@ -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,
python/sglang/srt/layers/quantization/fp8.py
+235 -19
View File
@@ -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)
python/sglang/srt/layers/quantization/fp8_kernel.py
+139
View File
@@ -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,
python/sglang/srt/layers/quantization/fp8_utils.py
+127
View File
@@ -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,