251 lines
7.8 KiB
Python
251 lines
7.8 KiB
Python
import logging
|
|
from typing import Any, Dict, Optional
|
|
|
|
import torch
|
|
|
|
_PETIT_INSTALL_ERROR = (
|
|
"Petit is not installed. Please install it with `pip install petit-kernel`."
|
|
)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
try:
|
|
from petit_kernel import (
|
|
mul_mxfp4_a16,
|
|
mul_nvfp4_a16,
|
|
process_mxfp4_scales,
|
|
process_nvfp4_scales,
|
|
repack_mxfp4,
|
|
repack_nvfp4,
|
|
)
|
|
except ImportError:
|
|
mul_mxfp4_a16 = None
|
|
mul_nvfp4_a16 = None
|
|
process_mxfp4_scales = None
|
|
process_nvfp4_scales = None
|
|
repack_mxfp4 = None
|
|
repack_nvfp4 = None
|
|
|
|
|
|
def _require_petit_kernel() -> None:
|
|
if mul_nvfp4_a16 is None:
|
|
raise ValueError(_PETIT_INSTALL_ERROR)
|
|
|
|
|
|
def _check_petit_nvfp4_supported(
|
|
quant_method: str, group_size: Optional[int]
|
|
) -> tuple[bool, Optional[str]]:
|
|
if quant_method.upper() != "NVFP4":
|
|
return (
|
|
False,
|
|
"Petit currently only supports: NVFP4 "
|
|
"quantizations in sglang. Please check the "
|
|
"`hf_quant_config.json` file for your model's "
|
|
"quant configuration.",
|
|
)
|
|
if group_size is not None and group_size != 16:
|
|
return (
|
|
False,
|
|
"Petit currently only supports: group_size=16 quantizations.",
|
|
)
|
|
return (True, None)
|
|
|
|
|
|
def verify_petit_nvfp4_supported(quant_method: str, group_size: Optional[int]) -> None:
|
|
supported, error_msg = _check_petit_nvfp4_supported(quant_method, group_size)
|
|
if not supported:
|
|
raise ValueError(error_msg)
|
|
|
|
|
|
def _check_petit_mxfp4_supported(
|
|
quant_method: str,
|
|
group_size: Optional[int],
|
|
quant_config: Optional[Dict[str, Any]] = None,
|
|
) -> tuple[bool, Optional[str]]:
|
|
quant_method_lower = quant_method.lower()
|
|
is_mxfp4_method = "mxfp4" in quant_method_lower
|
|
is_quark_method = quant_method_lower == "quark"
|
|
|
|
if not is_mxfp4_method and not is_quark_method:
|
|
return (
|
|
False,
|
|
"Petit MXFP4 currently only supports MXFP4 or Quark-MXFP4 quantizations "
|
|
"in sglang. Please check the model quantization config.",
|
|
)
|
|
|
|
if is_quark_method and quant_config is not None:
|
|
if not is_quark_mxfp4_compatible_config(quant_config):
|
|
return (
|
|
False,
|
|
"Detected quant_method=quark, but the quark quantization config "
|
|
"does not look like MXFP4 weights "
|
|
"(fp4/per_group/group_size=32/e8m0).",
|
|
)
|
|
|
|
if group_size is not None and group_size != 32:
|
|
return (
|
|
False,
|
|
"Petit MXFP4 currently only supports: group_size=32 quantizations.",
|
|
)
|
|
return (True, None)
|
|
|
|
|
|
def verify_petit_mxfp4_supported(
|
|
quant_method: str,
|
|
group_size: Optional[int],
|
|
quant_config: Optional[Dict[str, Any]] = None,
|
|
) -> None:
|
|
supported, error_msg = _check_petit_mxfp4_supported(
|
|
quant_method, group_size, quant_config
|
|
)
|
|
if not supported:
|
|
raise ValueError(error_msg)
|
|
|
|
|
|
def _is_quark_mxfp4_layer_quant_config(config: Dict[str, Any]) -> bool:
|
|
weight_quant = config.get("weight")
|
|
input_quant = config.get("input_tensors")
|
|
if not isinstance(weight_quant, dict):
|
|
return False
|
|
|
|
if isinstance(input_quant, dict):
|
|
if hasattr(logger, "warning_once"):
|
|
logger.warning_once(
|
|
"Quark input_tensors quant config is ignored for petit_mxfp4 "
|
|
"(kernel path is w4a16). Only weight quant config is validated."
|
|
)
|
|
else:
|
|
logger.warning(
|
|
"Quark input_tensors quant config is ignored for petit_mxfp4 "
|
|
"(kernel path is w4a16). Only weight quant config is validated."
|
|
)
|
|
|
|
return (
|
|
weight_quant.get("dtype") == "fp4"
|
|
and weight_quant.get("qscheme") == "per_group"
|
|
and weight_quant.get("group_size") == 32
|
|
and weight_quant.get("is_dynamic") is False
|
|
and weight_quant.get("scale_format") == "e8m0"
|
|
)
|
|
|
|
|
|
def is_quark_mxfp4_compatible_config(quant_config: Dict[str, Any]) -> bool:
|
|
"""Best-effort detection for Quark MXFP4 dense configs.
|
|
|
|
Some checkpoints only expose `quant_method=quark` in config.json without
|
|
full layer quant metadata. In that case we return True and defer validation
|
|
to weight-loading/runtime checks when user explicitly selects petit_mxfp4.
|
|
"""
|
|
candidates: list[Dict[str, Any]] = []
|
|
|
|
global_quant = quant_config.get("global_quant_config")
|
|
if isinstance(global_quant, dict):
|
|
candidates.append(global_quant)
|
|
|
|
layer_quant = quant_config.get("layer_quant_config")
|
|
if isinstance(layer_quant, dict):
|
|
candidates.extend(v for v in layer_quant.values() if isinstance(v, dict))
|
|
|
|
layer_type_quant = quant_config.get("layer_type_quant_config")
|
|
if isinstance(layer_type_quant, dict):
|
|
candidates.extend(v for v in layer_type_quant.values() if isinstance(v, dict))
|
|
|
|
if not candidates:
|
|
return True
|
|
|
|
return all(_is_quark_mxfp4_layer_quant_config(cfg) for cfg in candidates)
|
|
|
|
|
|
def prepare_nvfp4_layer_for_petit(layer: torch.nn.Module) -> None:
|
|
_require_petit_kernel()
|
|
|
|
# Repack weights to petit format
|
|
part_size_n = layer.output_size_per_partition
|
|
part_size_k = layer.input_size_per_partition
|
|
qweight = layer.weight.view(torch.int32).contiguous()
|
|
petit_qweight = repack_nvfp4(qweight, size_n=part_size_n, size_k=part_size_k)
|
|
layer.weight = torch.nn.Parameter(petit_qweight, requires_grad=False)
|
|
|
|
# Permute scales
|
|
weight_scale = process_nvfp4_scales(
|
|
scales=layer.weight_scale, size_k=part_size_k, size_n=part_size_n
|
|
)
|
|
layer.weight_scale = torch.nn.Parameter(weight_scale, requires_grad=False)
|
|
|
|
|
|
def prepare_mxfp4_layer_for_petit(layer: torch.nn.Module) -> None:
|
|
_require_petit_kernel()
|
|
|
|
part_size_n = layer.output_size_per_partition
|
|
part_size_k = layer.input_size_per_partition
|
|
qweight = layer.weight.view(torch.int32).contiguous()
|
|
petit_qweight = repack_mxfp4(qweight, size_n=part_size_n, size_k=part_size_k)
|
|
layer.weight = torch.nn.Parameter(petit_qweight, requires_grad=False)
|
|
|
|
weight_scale = process_mxfp4_scales(
|
|
scales=layer.weight_scale, size_k=part_size_k, size_n=part_size_n
|
|
)
|
|
layer.weight_scale = torch.nn.Parameter(weight_scale, requires_grad=False)
|
|
|
|
|
|
def apply_petit_nvfp4_linear(
|
|
input: torch.Tensor,
|
|
weight: torch.Tensor,
|
|
weight_scale: torch.Tensor,
|
|
weight_scale_2: torch.Tensor,
|
|
size_n: int,
|
|
size_k: int,
|
|
bias: Optional[torch.Tensor] = None,
|
|
) -> torch.Tensor:
|
|
_require_petit_kernel()
|
|
|
|
reshaped_x = input.reshape(-1, input.shape[-1])
|
|
out_shape = input.shape[:-1] + (size_n,)
|
|
|
|
# TODO: Use auto-tuning to find the performant solution_id
|
|
output = mul_nvfp4_a16(
|
|
a=reshaped_x,
|
|
b=weight,
|
|
s=weight_scale,
|
|
global_scale=weight_scale_2,
|
|
size_m=reshaped_x.size(0),
|
|
size_n=size_n,
|
|
size_k=size_k,
|
|
solution_id=-1,
|
|
)
|
|
if bias is not None:
|
|
output.add_(bias) # In-place add
|
|
|
|
return output.reshape(out_shape)
|
|
|
|
|
|
def apply_petit_mxfp4_linear(
|
|
input: torch.Tensor,
|
|
weight: torch.Tensor,
|
|
weight_scale: torch.Tensor,
|
|
size_n: int,
|
|
size_k: int,
|
|
bias: Optional[torch.Tensor] = None,
|
|
global_scale: Optional[torch.Tensor] = None,
|
|
) -> torch.Tensor:
|
|
_require_petit_kernel()
|
|
|
|
reshaped_x = input.reshape(-1, input.shape[-1])
|
|
out_shape = input.shape[:-1] + (size_n,)
|
|
if global_scale is None:
|
|
global_scale = torch.ones(1, device=reshaped_x.device, dtype=torch.float32)
|
|
|
|
output = mul_mxfp4_a16(
|
|
a=reshaped_x,
|
|
b=weight,
|
|
s=weight_scale,
|
|
global_scale=global_scale,
|
|
size_m=reshaped_x.size(0),
|
|
size_n=size_n,
|
|
size_k=size_k,
|
|
solution_id=-1,
|
|
)
|
|
if bias is not None:
|
|
output.add_(bias)
|
|
|
|
return output.reshape(out_shape)
|