from dataclasses import dataclass
from enum import Enum
from typing import NamedTuple, Optional
import torch
class FP8Format(NamedTuple):
"""FP8 format specification.
Defines the properties of an FP8 format including range, exponent bits,
mantissa bits, and PyTorch dtype.
Attributes:
max: Maximum representable value
ebits: Number of exponent bits
mbits: Number of mantissa bits
dtype: PyTorch dtype for this format
"""
max: float
ebits: int
mbits: int
dtype: Optional[torch.dtype]
@property
def quant_dtype(self) -> torch.dtype:
"""Get quantization type.
Returns:
PyTorch dtype for quantization
"""
if self.dtype is None:
import torch_npu
return torch_npu.hifloat8
return self.dtype
class FormatEnum(Enum):
"""FP8 format enumeration.
Defines standard FP8 formats with their specifications.
"""
E4M3 = FP8Format(448, 4, 3, torch.float8_e4m3fn)
E5M2 = FP8Format(57344, 5, 2, torch.float8_e5m2)
HIF8 = FP8Format(57344, 5, 2, None)
HIF8_224 = FP8Format(224, 5, 2, None)
HIF8_15 = FP8Format(15, 5, 2, None)
class _FormatConfig(NamedTuple):
"""Format configuration for inputs, weights, and gradients.
Attributes:
fwd: Format for input/wegiht tensors
bwd: Format for grad tensors
"""
fwd: FormatEnum = FormatEnum.E4M3
bwd: FormatEnum = FormatEnum.E4M3
class Format(Enum):
"""FP8 format configuration enumeration.
Defines format configurations for different FP8 recipes.
"""
E4M3 = _FormatConfig()
E5M2 = _FormatConfig(fwd=FormatEnum.E5M2, bwd=FormatEnum.E5M2)
HYBRID = _FormatConfig(bwd=FormatEnum.E5M2)
HIF8 = _FormatConfig(fwd=FormatEnum.HIF8_15, bwd=FormatEnum.HIF8_224)
@classmethod
def from_str(cls, key: str) -> Optional["Format"]:
"""Get format from configuration key.
Args:
key: Configuration key string
Returns:
Format enum value or None
"""
return getattr(cls, key.upper(), None)
class Recipe:
"""
Base recipe class.
"""
fp8_format: Format = Format.E4M3
backward_override: Optional[str] = None
@classmethod
def nvfp4(cls):
"""Whether the given recipe is NVFP4 1D block scaling."""
return False
@classmethod
def mxfp4(cls):
"""Whether the given recipe is MXFP4 block scaling."""
return False
@classmethod
def mxfp8(cls):
"""Whether the given recipe is MXFP8 block scaling."""
return False
@classmethod
def w4a8(cls):
"""Whether the given recipe is W4A8 mixed block scaling."""
return False
@classmethod
def delayed(cls):
"""Whether the given recipe is delayed scaling."""
return False
@classmethod
def float8_current_scaling(cls):
"""Whether the given recipe is (per-tensor) current scaling."""
return False
@classmethod
def float8_per_tensor_scaling(cls):
"""Whether the given recipe is per-tensor scaling."""
return False
@classmethod
def float8_block_scaling(cls):
"""Whether the given recipe is float8 blockwise scaling."""
return False
@classmethod
def custom(cls):
"""Whether the given recipe is custom."""
return False
@dataclass(frozen=True)
class MMParams:
"""Matrix multiplication options.
Parameters
----------
use_split_accumulator : bool, default = True
Use FP8 fast accumulation on Hopper or Ada. For more details,
see CUBLASLT_MATMUL_DESC_FAST_ACCUM option for cublasLtMatmul.
"""
use_split_accumulator: bool = True
@dataclass(frozen=True)
class QParams:
"""Quantization parameters.
power_2_scale: use power of 2 scale parameter
amax_epsilon: optional minimum value of abs max
random_hadamard_transform: whether to use random hadamard transform
stochastic_rounding: whether to use stochastic rounding
"""
power_2_scale: bool = False
amax_epsilon: float = 0.0
random_hadamard_transform: bool = False
stochastic_rounding: bool = False
fp4_2d_quantization: bool = False
def __repr__(self) -> str:
return (
f"Qparams(\npower_2_scale={self.power_2_scale},\n"
f"amax_epsilon={self.amax_epsilon},\n"
f"random_hadamard_transform={self.random_hadamard_transform},\n"
f"stochastic_rounding={self.stochastic_rounding},\n"
f"fp4_2d_quantization={self.fp4_2d_quantization}\n)"
)
