Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-18 16:54:33 +02:00
Code Issues Packages Projects Releases Wiki Activity

Move SDNQ to upstream SDNQ repo

Browse Source
This commit is contained in:
Dity0
2026-08-10 12:13:10 +03:00
parent 40fb45a546
commit 712a13f1a0
42 changed files with 79 additions and 6027 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cli/sdnq-attention-benchmark.py
+28 -28
View File
@@ -2,7 +2,7 @@
"""
Benchmark and validate SDNQ attention and weight dequantization on the local GPU.
The attention section runs the kernel from modules/sdnq/kernels/triton_atten.py directly and
The attention section runs the kernel from sdnq/kernels/triton_atten.py directly and
compares speed and numerical error against torch scaled_dot_product_attention and
sageattention when installed. Verifies mask, causal, GQA, cross-attention and padding code
paths, probes float8 hardware support and the torch.compile input prep, and prints
@@ -358,7 +358,7 @@ def load_sdnext():
with capture_console_output() as startup_log:
from modules import shared as shared_module
from modules import devices as devices_module
from modules.sdnq.kernels.triton_atten import sdnq_triton_atten as atten
from sdnq.kernels.triton_atten import sdnq_triton_atten as atten
except BaseException as e: # pylint: disable=broad-exception-caught # SystemExit is not an Exception: the bootstrap exits on a failed torch or library import
if isinstance(e, KeyboardInterrupt):
raise
@@ -727,7 +727,7 @@ def fp8_compile_gate_flag():
# False on gpus where sdnq upcasts e4m3 storage to the scale dtype before the compiled
# dequant, because triton cannot convert e4m3 there; absent on builds without the gate
try:
from modules.sdnq import kernel_wrappers as sdnq_kernel_wrappers
from sdnq import kernel_wrappers as sdnq_kernel_wrappers
return getattr(sdnq_kernel_wrappers, "is_fp8_compile_supported", None)
except Exception:
return None
@@ -839,13 +839,13 @@ def probe_fp8():
# toggling torch._dynamo.config.disable: torch 2.13+ raises "found no compiled frames"
# for fullgraph-compiled functions called inside a disable window
try:
from modules.sdnq import kernel_wrappers as sdnq_kernel_wrappers
from sdnq import kernel_wrappers as sdnq_kernel_wrappers
is_fp8_mm_supported = getattr(sdnq_kernel_wrappers, "is_fp8_mm_supported", True)
except Exception:
is_fp8_mm_supported = True
if not is_fp8_mm_supported:
return dict(qk=(False, "FP8 matmul is not supported in this architecture"), pv=(False, "FP8 matmul is not supported in this architecture"))
from modules.sdnq.kernels import triton_atten as atten_module
from sdnq.kernels import triton_atten as atten_module
q, k, v = make_qkv(1, 2, 256, 64, structured=False)
result = {}
compiled_prep = atten_module.get_attn_inputs
@@ -884,7 +884,7 @@ def probe_compiled_prep():
torch._dynamo.reset() # pylint: disable=protected-access # drop failed compile state
detail = f"{type(e).__name__}: {error_summary(e, 120)}"
# check whether the dynamic=false workaround holds
from modules.sdnq.kernels import triton_atten as atten_module
from sdnq.kernels import triton_atten as atten_module
compiled_prep = atten_module.get_attn_inputs
inner = getattr(compiled_prep, "_torchdynamo_orig_callable", None)
if inner is None:
@@ -948,8 +948,8 @@ def make_block_master():
def build_bench_block(master_sd, weights_cfg, use_mm, attention_spec):
from modules.sdnq import SDNQConfig
from modules.sdnq.quantizer import apply_sdnq_to_module
from sdnq import SDNQConfig
from sdnq.quantizer import apply_sdnq_to_module
hidden, heads, mlp_dim = block_geometry["hidden"], block_geometry["heads"], block_geometry["mlp_dim"]
block = BenchBlock(hidden, heads, mlp_dim, device=torch_device, dtype=bench_dtype)
block.load_state_dict(master_sd)
@@ -990,8 +990,8 @@ def make_quantized_linear(weight, weights_dtype, group_size=0, use_quantized_mat
# quantize through the same entry point model loading uses, so forward benches measure
# the production wrapper classes and dequantizer configuration; returns the quantize wall
# time, a one-shot measurement of what on-the-fly quantization pays per layer at load
from modules.sdnq import SDNQConfig
from modules.sdnq.quantizer import sdnq_quantize_layer
from sdnq import SDNQConfig
from sdnq.quantizer import sdnq_quantize_layer
out_features, in_features = weight.shape
linear = torch.nn.Linear(in_features, out_features, bias=False, device=device, dtype=bench_dtype)
with torch.no_grad():
@@ -1058,7 +1058,7 @@ def get_compiled_dequantize_weight():
for limit_name in ("recompile_limit", "cache_size_limit", "accumulated_recompile_limit", "accumulated_cache_size_limit"):
if hasattr(torch._dynamo.config, limit_name): # pylint: disable=protected-access
setattr(torch._dynamo.config, limit_name, max(8192, getattr(torch._dynamo.config, limit_name) or 0)) # pylint: disable=protected-access
from modules.sdnq.dequantizer import dequantize_weight
from sdnq.dequantizer import dequantize_weight
compiled_dequantize_weight = torch.compile(dequantize_weight, fullgraph=True, dynamic=False)
return compiled_dequantize_weight
@@ -1134,7 +1134,7 @@ def run_correctness():
# do not toggle torch._dynamo.config.disable for this: newer torch raises "found no
# compiled frames" when a fullgraph-compiled function is called inside a disable window,
# failing every check and poisoning the first compiled call afterwards
from modules.sdnq.kernels import triton_atten as atten_module
from sdnq.kernels import triton_atten as atten_module
compiled_prep = atten_module.get_attn_inputs
inner_prep = getattr(compiled_prep, "_torchdynamo_orig_callable", None)
if inner_prep is not None:
@@ -1194,9 +1194,9 @@ def run_correctness():
def make_prep_fn(q, k, v, attn_mask, kwargs, is_causal=False, enable_gqa=False):
# mirror sdnq_triton_atten's prep call so the prep column measures the same code path
from modules.sdnq.kernels import triton_atten as atten_module
from modules.sdnq.quant_utils import get_hadamard, get_hadamard_group_size
from modules.sdnq.utils import next_power_of_2
from sdnq.kernels import triton_atten as atten_module
from sdnq.quant_utils import get_hadamard, get_hadamard_group_size
from sdnq.utils import next_power_of_2
matmul_dtype = kwargs.get("matmul_dtype", "int8")
do_quantize = kwargs.get("do_quantize", True)
hadamard_group_size = kwargs.get("hadamard_group_size", 256)
@@ -1377,7 +1377,7 @@ def bench_shape(preset, iters, warmup, position=None, config_timeout=300, fp8_re
def bench_dequant_shape(shape_label, out_features, in_features, iters, warmup, position=None, config_timeout=300, selected_dtypes=None):
from modules.sdnq.common import check_torch_compile
from sdnq.common import check_torch_compile
compile_on = check_torch_compile()
dtype_configs = [(dtype_id, label, cfg) for dtype_id, label, cfg in dequant_dtype_configs if selected_dtypes is None or dtype_id in selected_dtypes]
@@ -1490,7 +1490,7 @@ def bench_dequant_shape(shape_label, out_features, in_features, iters, warmup, p
# bench matches the webui with Dequantize using torch.compile off. do not toggle
# torch._dynamo.config.disable instead: code objects called during a disable window
# keep their skip marking and never compile again in this process
from modules.sdnq import dequantizer as dequantizer_module
from sdnq import dequantizer as dequantizer_module
saved_compiled_fn = dequantizer_module.dequantize_weight_compiled
try:
phase("timing linear forward, eager dequant")
@@ -1754,10 +1754,10 @@ def bench_float_mm_alternatives(shape_label, out_features, in_features, plain_re
# never defined and the row needs SDNQ_USE_TRITON_MM=0.
mm_swap_targets = [
("modules.sdnq.layers.linear.linear_int8", "int_scaled_mm_func"),
("modules.sdnq.layers.linear.linear_uint8", "int_scaled_mm_func"),
("modules.sdnq.layers.linear.linear_fp16", "fp_scaled_mm_func"),
("modules.sdnq.layers.linear.linear_fp8", "fp8_scaled_mm_func"),
("sdnq.layers.linear.linear_int8", "int_scaled_mm_func"),
("sdnq.layers.linear.linear_uint8", "int_scaled_mm_func"),
("sdnq.layers.linear.linear_fp16", "fp_scaled_mm_func"),
("sdnq.layers.linear.linear_fp8", "fp8_scaled_mm_func"),
]
@@ -1774,7 +1774,7 @@ def mm_backend_bindings():
if func is not None:
bound[(module_path, attr)] = func
try:
from modules.sdnq.kernels.triton_scaled_mm import sdnq_scaled_mm
from sdnq.kernels.triton_scaled_mm import sdnq_scaled_mm
except Exception as e:
return {}, {"triton": f"triton scaled mm unavailable: {error_summary(e, 120)}"}
@@ -2284,8 +2284,8 @@ def fp32_conv_reference(x, weight_fp32, padding):
def make_quantized_conv(weight, weights_dtype, use_quantized_matmul=False):
from modules.sdnq import SDNQConfig
from modules.sdnq.quantizer import sdnq_quantize_layer
from sdnq import SDNQConfig
from sdnq.quantizer import sdnq_quantize_layer
out_channels, in_channels, kh, kw = weight.shape
conv = torch.nn.Conv2d(in_channels, out_channels, (kh, kw), padding=(kh // 2, kw // 2), bias=False, device=torch_device, dtype=bench_dtype)
with torch.no_grad():
@@ -3403,11 +3403,11 @@ def main():
if free_vram_gb() < 2.0:
emit(f"[yellow]skipping dequant benchmarks: needs about 2 gb free vram, {free_vram_gb():.1f} gb available[/yellow]")
else:
from modules.sdnq.common import check_torch_compile
from sdnq.common import check_torch_compile
if not check_torch_compile():
# the module-level compiled dequant is a passthrough with the option off; swap in
# a real compiled variant so the compiled fwd rows measure what enabling it gives
from modules.sdnq import dequantizer as dequantizer_module
from sdnq import dequantizer as dequantizer_module
dequantizer_module.dequantize_weight_compiled = get_compiled_dequantize_weight()
emit("[yellow]Dequantize using torch.compile is off in the current config: compiled fwd rows are measured with a tool-compiled dequant, matching the webui after enabling it[/yellow]")
for index, (shape_label, out_features, in_features) in enumerate(dequant_shapes, start=1):
@@ -3446,13 +3446,13 @@ def main():
if "attention" in sections:
# bench the prep mode the advice points to: compiled, static workaround, or eager
if prep_status == "failing_dynamic":
from modules.sdnq.kernels import triton_atten as atten_module
from sdnq.kernels import triton_atten as atten_module
inner = getattr(atten_module.get_attn_inputs, "_torchdynamo_orig_callable", None)
atten_module.get_attn_inputs = torch.compile(inner, fullgraph=True, dynamic=False)
emit("[yellow]dynamic-shape compile is broken here: benchmarking with the dynamic=false workaround applied, numbers match the webui after setting SDNQ_COMPILE_KWARGS='{\"dynamic\": false}'[/yellow]")
elif prep_status == "failing":
emit("[yellow]torch compile is broken here: benchmarking with eager input prep, numbers match the webui after disabling Dequantize using torch.compile[/yellow]")
from modules.sdnq.kernels import triton_atten as atten_module
from sdnq.kernels import triton_atten as atten_module
inner = getattr(atten_module.get_attn_inputs, "_torchdynamo_orig_callable", None)
if inner is not None: # swap in the eager prep; a disable toggle raises on torch 2.13+
atten_module.get_attn_inputs = inner
installer.py
+25 -1
View File
@@ -67,6 +67,7 @@ args = Dot({
'uv': False,
})
git_commit = "unknown"
sdnq_commit = "unknown"
diffusers_commit = "unknown"
transformers_commit = "unknown"
restart_required = False
@@ -549,6 +550,30 @@ def check_python(supported_minors=None, experimental_minors=None, reason=None):
ts('python', t_start)
# check sdnq version
def check_sdnq():
t_start = time.time()
if args.skip_all:
return
target_commit = "27a54aed657ec583620b76b42c65d907d13290ad"
pkg = package_spec('sdnq')
parts = pkg.version.split('.') if pkg is not None else []
minor = int(parts[1]) if len(parts) > 1 else -1
current = package_commit(pkg) if minor > -1 else ''
if (minor == -1) or ((current != target_commit) and (not args.experimental)):
if minor == -1:
log.info(f'Install: package="sdnq" commit={target_commit}')
else:
log.info(f'Update: package="sdnq" current={pkg.version} commit={current} target={target_commit}')
pip('uninstall --yes sdnq', ignore=True, quiet=True, uv=False)
if args.skip_git:
log.warning('Git: marked as not available but required for sdnq installation')
pip(f'install git+https://github.com/Disty0/sdnq@{target_commit}', ignore=False, quiet=True, uv=False)
global sdnq_commit # pylint: disable=global-statement
sdnq_commit = target_commit
ts('sdnq', t_start)
# check diffusers version
def check_diffusers():
t_start = time.time()
@@ -1378,7 +1403,6 @@ def install_requirements():
# set environment variables controlling the behavior of various libraries
def set_environment():
log.debug('Setting environment tuning')
os.environ.setdefault('SDNQ_REGISTER_DIFFUSERS', '1')
os.environ.setdefault('ACCELERATE', 'True')
os.environ.setdefault('ATTN_PRECISION', 'fp16')
os.environ.setdefault('ClDeviceGlobalMemSizeAvailablePercent', '100')
launch.py
+1
View File
@@ -261,6 +261,7 @@ def main():
installer.install_gradio()
installer.check_torch()
installer.check_onnx()
installer.check_sdnq()
installer.check_transformers()
installer.check_diffusers()
installer.check_modified_files()
modules/attention.py
+1 -1
View File
@@ -20,7 +20,7 @@ def set_dynamic_attention():
def set_sdnq_attention():
try:
from modules import shared
from modules.sdnq.kernels.triton_atten import sdnq_triton_atten
from sdnq.kernels.triton_atten import sdnq_triton_atten
sdpa_pre_sdnq_atten = torch.nn.functional.scaled_dot_product_attention
@wraps(sdpa_pre_sdnq_atten)
def sdpa_sdnq_atten(query: torch.FloatTensor, key: torch.FloatTensor, value: torch.FloatTensor, attn_mask: torch.Tensor | None = None, dropout_p: float = 0.0, is_causal: bool = False, scale: float | None = None, enable_gqa: bool = False, **kwargs) -> torch.FloatTensor:
modules/lora/lora_apply.py
+1 -1
View File
@@ -167,7 +167,7 @@ def network_add_weights(self: torch.nn.Conv2d | torch.nn.Linear | torch.nn.Group
weight, new_weight = None, None
if not bias and hasattr(self, "sdnq_dequantizer"):
try:
from modules.sdnq import SDNQConfig, sdnq_quantize_layer
from sdnq import SDNQConfig, sdnq_quantize_layer
if hasattr(self, "sdnq_dequantizer_backup"):
use_svd = bool(self.sdnq_svd_up_backup is not None)
dequantize_fp32 = bool(self.sdnq_scale_backup.dtype == torch.float32)
modules/model_quant.py
+2 -2
View File
@@ -104,7 +104,7 @@ def create_sdnq_config(kwargs = None,
):
from modules import shared
if allow and (shared.opts.sdnq_quantize_mode in {'pre', 'auto'}) and (module == 'any' or module in shared.opts.sdnq_quantize_weights):
from modules.sdnq import SDNQConfig
from sdnq import SDNQConfig
if weights_dtype is None:
if module in {"TE", "LLM"} and shared.opts.sdnq_quantize_weights_mode_te not in {"Same as model", "default"}:
@@ -363,7 +363,7 @@ def apply_layerwise(sd_model, quiet:bool=False):
def sdnq_quantize_model(model, op=None, sd_model=None, do_gc: bool = True, weights_dtype: str | None = None, quantized_matmul_dtype: str | None = None, modules_to_not_convert: list | None = None, modules_dtype_dict: dict | None = None):
global quant_last_model_name, quant_last_model_device # pylint: disable=global-statement
from modules import devices, shared, timer
from modules.sdnq import sdnq_post_load_quant
from sdnq import sdnq_post_load_quant
if (
hasattr(model, "quantization_config")
modules/sd_models.py
+5 -5
View File
@@ -202,7 +202,7 @@ def set_diffuser_options(sd_model, vae=None, op:str='model', offload:bool=True,
else:
sdnq_use_quantized_matmul = shared.opts.sdnq_quantize_matmul_mode != "disabled"
if module.quantization_config.use_quantized_matmul != sdnq_use_quantized_matmul:
from modules.sdnq.loader import apply_sdnq_options_to_model
from sdnq.loader import apply_sdnq_options_to_model
# log.debug(f'Setting {op} {module_name}: sdnq_use_quantized_matmul={sdnq_use_quantized_matmul}')
module = apply_sdnq_options_to_model(module, use_quantized_matmul=sdnq_use_quantized_matmul)
setattr(sd_model, module_name, module)
@@ -359,7 +359,7 @@ def hf_prefetch_configs(checkpoint_info: CheckpointInfo | str, diffusers_load_co
def load_diffuser_force(detected_model_type: str, checkpoint_info: CheckpointInfo, diffusers_load_config: dict, op='model'):
from modules import sdnq # pylint: disable=unused-import
import sdnq # pylint: disable=unused-import
sd_model = None
global allow_post_quant # pylint: disable=global-statement
unload_model_weights(op=op)
@@ -773,7 +773,7 @@ def load_sdnq_module(fn: str, module_name: str, load_method: str):
return None, module_name, 0
model_name = os.path.join(fn, module_name)
try:
from modules import sdnq
import sdnq
module = sdnq.load_sdnq_model(
model_path=model_name,
quantization_config=quantization_config,
@@ -1573,7 +1573,7 @@ def unload_model_weights(op='model'):
disable_offload(model_data.sd_model)
move_model(model_data.sd_model, 'meta')
model_data.sd_model = None
from modules.sdnq.common import reset_compile_caches
from sdnq.common import reset_compile_caches
reset_compile_caches() # dead compiled-dequant graphs and their lifetime recompile counters otherwise accumulate across switches
devices.torch_gc(force=True, reason='unload')
log.debug(f'Unload {op}: {memory_stats()} fn={fn}')
@@ -1620,7 +1620,7 @@ def save_model(name: str, path: str | None = None, shard: str = "5GB", overwrite
if not shared.sd_loaded:
log.error('Save model: model not loaded')
return 'Model not loaded'
from modules.sdnq import save_sdnq_model
from sdnq import save_sdnq_model
if path is None:
path = shared.opts.diffusers_dir
model_name = os.path.join(path.strip(), name.strip())
modules/sdnq/__init__.py
-16
View File
@@ -1,16 +0,0 @@
from .quantizer import QuantizationMethod, SDNQConfig, SDNQQuantizer, apply_sdnq_to_module, sdnq_post_load_quant, sdnq_quantize_layer
from .loader import save_sdnq_model, load_sdnq_model
from .common import sdnq_version
__version__ = sdnq_version
__all__ = [
"QuantizationMethod",
"SDNQConfig",
"SDNQQuantizer",
"apply_sdnq_to_module",
"load_sdnq_model",
"save_sdnq_model",
"sdnq_post_load_quant",
"sdnq_quantize_layer",
]
modules/sdnq/common.py
-524
View File
@@ -1,524 +0,0 @@
# pylint: disable=redefined-builtin,no-member,protected-access
import os
import json
import torch
from modules import shared
sdnq_version = "0.2.5"
sdnq_keys = {"weight", "scale", "zero_point", "svd_up", "svd_down"}
torch_version = torch.__version__[:4]
if torch_version[-1] not in {"0", "1", "2", "3", "4", "5", "6", "7", "8", "9"}:
torch_version = torch_version[:-1]
torch_version = torch_version.split(".")
torch_version[0], torch_version[1] = int(torch_version[0]), int(torch_version[1])
dtype_dict = {
### Integers
"int32": {"min": -2147483648, "max": 2147483647, "num_bits": 32, "sign": 1, "exponent": 0, "mantissa": 31, "target_dtype": torch.int32, "torch_dtype": torch.int32, "storage_dtype": torch.int32, "is_unsigned": False, "is_integer": True, "is_packed": False},
"int16": {"min": -32768, "max": 32767, "num_bits": 16, "sign": 1, "exponent": 0, "mantissa": 15, "target_dtype": torch.int16, "torch_dtype": torch.int16, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": True, "is_packed": False},
"int8": {"min": -128, "max": 127, "num_bits": 8, "sign": 1, "exponent": 0, "mantissa": 7, "target_dtype": torch.int8, "torch_dtype": torch.int8, "storage_dtype": torch.int8, "is_unsigned": False, "is_integer": True, "is_packed": False},
### Custom Integers
"int15": {"min": -16384, "max": 16383, "num_bits": 15, "sign": 1, "exponent": 0, "mantissa": 14, "target_dtype": "int15", "torch_dtype": torch.int16, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": True, "is_packed": True},
"int14": {"min": -8192, "max": 8191, "num_bits": 14, "sign": 1, "exponent": 0, "mantissa": 13, "target_dtype": "int14", "torch_dtype": torch.int16, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": True, "is_packed": True},
"int13": {"min": -4096, "max": 4095, "num_bits": 13, "sign": 1, "exponent": 0, "mantissa": 12, "target_dtype": "int13", "torch_dtype": torch.int16, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": True, "is_packed": True},
"int12": {"min": -2048, "max": 2047, "num_bits": 12, "sign": 1, "exponent": 0, "mantissa": 11, "target_dtype": "int12", "torch_dtype": torch.int16, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": True, "is_packed": True},
"int11": {"min": -1024, "max": 1023, "num_bits": 11, "sign": 1, "exponent": 0, "mantissa": 10, "target_dtype": "int11", "torch_dtype": torch.int16, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": True, "is_packed": True},
"int10": {"min": -512, "max": 511, "num_bits": 10, "sign": 1, "exponent": 0, "mantissa": 9, "target_dtype": "int10", "torch_dtype": torch.int16, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": True, "is_packed": True},
"int9": {"min": -256, "max": 255, "num_bits": 9, "sign": 1, "exponent": 0, "mantissa": 8, "target_dtype": "int9", "torch_dtype": torch.int16, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": True, "is_packed": True},
#
"int7": {"min": -64, "max": 63, "num_bits": 7, "sign": 1, "exponent": 0, "mantissa": 6, "target_dtype": "int7", "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True, "is_packed": True},
"int6": {"min": -32, "max": 31, "num_bits": 6, "sign": 1, "exponent": 0, "mantissa": 5, "target_dtype": "int6", "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True, "is_packed": True},
"int5": {"min": -16, "max": 15, "num_bits": 5, "sign": 1, "exponent": 0, "mantissa": 4, "target_dtype": "int5", "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True, "is_packed": True},
"int4": {"min": -8, "max": 7, "num_bits": 4, "sign": 1, "exponent": 0, "mantissa": 3, "target_dtype": "int4", "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True, "is_packed": True},
"int3": {"min": -4, "max": 3, "num_bits": 3, "sign": 1, "exponent": 0, "mantissa": 2, "target_dtype": "int3", "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True, "is_packed": True},
"int2": {"min": -2, "max": 1, "num_bits": 2, "sign": 1, "exponent": 0, "mantissa": 1, "target_dtype": "int2", "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True, "is_packed": True},
### Unsigned Integers
"uint32": {"min": 0, "max": 4294967295, "num_bits": 32, "sign": 0, "exponent": 0, "mantissa": 32, "target_dtype": torch.uint32, "torch_dtype": torch.uint32, "storage_dtype": torch.uint32, "is_unsigned": True, "is_integer": True, "is_packed": False},
"uint16": {"min": 0, "max": 65535, "num_bits": 16, "sign": 0, "exponent": 0, "mantissa": 16, "target_dtype": torch.uint16, "torch_dtype": torch.uint16, "storage_dtype": torch.uint16, "is_unsigned": True, "is_integer": True, "is_packed": False},
"uint8": {"min": 0, "max": 255, "num_bits": 8, "sign": 0, "exponent": 0, "mantissa": 8, "target_dtype": torch.uint8, "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True, "is_packed": False},
### Custom Unsigned Integers
"uint15": {"min": 0, "max": 32768, "num_bits": 15, "sign": 0, "exponent": 0, "mantissa": 15, "target_dtype": "uint15", "torch_dtype": torch.int16, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": True, "is_packed": True},
"uint14": {"min": 0, "max": 16384, "num_bits": 14, "sign": 0, "exponent": 0, "mantissa": 14, "target_dtype": "uint14", "torch_dtype": torch.int16, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": True, "is_packed": True},
"uint13": {"min": 0, "max": 8192, "num_bits": 13, "sign": 0, "exponent": 0, "mantissa": 13, "target_dtype": "uint13", "torch_dtype": torch.int16, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": True, "is_packed": True},
"uint12": {"min": 0, "max": 4096, "num_bits": 12, "sign": 0, "exponent": 0, "mantissa": 12, "target_dtype": "uint12", "torch_dtype": torch.int16, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": True, "is_packed": True},
"uint11": {"min": 0, "max": 2048, "num_bits": 11, "sign": 0, "exponent": 0, "mantissa": 11, "target_dtype": "uint11", "torch_dtype": torch.int16, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": True, "is_packed": True},
"uint10": {"min": 0, "max": 1024, "num_bits": 10, "sign": 0, "exponent": 0, "mantissa": 10, "target_dtype": "uint10", "torch_dtype": torch.int16, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": True, "is_packed": True},
"uint9": {"min": 0, "max": 512, "num_bits": 9, "sign": 0, "exponent": 0, "mantissa": 9, "target_dtype": "uint9", "torch_dtype": torch.int16, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": True, "is_packed": True},
#
"uint7": {"min": 0, "max": 127, "num_bits": 7, "sign": 0, "exponent": 0, "mantissa": 7, "target_dtype": "uint7", "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True, "is_packed": True},
"uint6": {"min": 0, "max": 63, "num_bits": 6, "sign": 0, "exponent": 0, "mantissa": 6, "target_dtype": "uint6", "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True, "is_packed": True},
"uint5": {"min": 0, "max": 31, "num_bits": 5, "sign": 0, "exponent": 0, "mantissa": 5, "target_dtype": "uint5", "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True, "is_packed": True},
"uint4": {"min": 0, "max": 15, "num_bits": 4, "sign": 0, "exponent": 0, "mantissa": 4, "target_dtype": "uint4", "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True, "is_packed": True},
"uint3": {"min": 0, "max": 7, "num_bits": 3, "sign": 0, "exponent": 0, "mantissa": 3, "target_dtype": "uint3", "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True, "is_packed": True},
"uint2": {"min": 0, "max": 3, "num_bits": 2, "sign": 0, "exponent": 0, "mantissa": 2, "target_dtype": "uint2", "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True, "is_packed": True},
"uint1": {"min": 0, "max": 1, "num_bits": 1, "sign": 0, "exponent": 0, "mantissa": 1, "target_dtype": torch.bool, "torch_dtype": torch.bool, "storage_dtype": torch.bool, "is_unsigned": True, "is_integer": True, "is_packed": True},
### Floats
"float32": {"min": -3.40282e+38, "max": 3.40282e+38, "num_bits": 32, "sign": 1, "exponent": 8, "mantissa": 23, "target_dtype": torch.float32, "torch_dtype": torch.float32, "storage_dtype": torch.float32, "is_unsigned": False, "is_integer": False, "is_packed": False},
"bfloat16": {"min": -3.38953e+38, "max": 3.38953e+38, "num_bits": 16, "sign": 1, "exponent": 8, "mantissa": 7, "target_dtype": torch.bfloat16, "torch_dtype": torch.bfloat16, "storage_dtype": torch.bfloat16, "is_unsigned": False, "is_integer": False, "is_packed": False},
"float16": {"min": -65504.0, "max": 65504.0, "num_bits": 16, "sign": 1, "exponent": 5, "mantissa": 10, "target_dtype": torch.float16, "torch_dtype": torch.float16, "storage_dtype": torch.float16, "is_unsigned": False, "is_integer": False, "is_packed": False},
"float8_e4m3fn": {"min": -448.0, "max": 448.0, "num_bits": 8, "sign": 1, "exponent": 4, "mantissa": 3, "target_dtype": torch.float8_e4m3fn, "torch_dtype": torch.float8_e4m3fn, "storage_dtype": torch.float8_e4m3fn, "is_unsigned": False, "is_integer": False, "is_packed": False},
"float8_e5m2": {"min": -57344.0, "max": 57344.0, "num_bits": 8, "sign": 1, "exponent": 5, "mantissa": 2, "target_dtype": torch.float8_e5m2, "torch_dtype": torch.float8_e5m2, "storage_dtype": torch.float8_e5m2, "is_unsigned": False, "is_integer": False, "is_packed": False},
### Custom Floats
"float16_e1m14fn": {"min": -3.9998779296875, "max": 3.9998779296875, "num_bits": 16, "sign": 1, "exponent": 1, "mantissa": 14, "target_dtype": "fp16", "torch_dtype": torch.float32, "storage_dtype": torch.uint16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float16_e2m13fn": {"min": -7.99951171875, "max": 7.99951171875, "num_bits": 16, "sign": 1, "exponent": 2, "mantissa": 13, "target_dtype": "fp16", "torch_dtype": torch.float32, "storage_dtype": torch.uint16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float16_e3m12fn": {"min": -31.99609375, "max": 31.99609375, "num_bits": 16, "sign": 1, "exponent": 3, "mantissa": 12, "target_dtype": "fp16", "torch_dtype": torch.float32, "storage_dtype": torch.uint16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float16_e4m11fn": {"min": -511.875, "max": 511.875, "num_bits": 16, "sign": 1, "exponent": 4, "mantissa": 11, "target_dtype": "fp16", "torch_dtype": torch.float32, "storage_dtype": torch.uint16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float16_e5m10fn": {"min": -131008.0, "max": 131008.0, "num_bits": 16, "sign": 1, "exponent": 5, "mantissa": 10, "target_dtype": "fp16", "torch_dtype": torch.float32, "storage_dtype": torch.uint16, "is_unsigned": False, "is_integer": False, "is_packed": True},
#
"float15_e1m13fn": {"min": -3.999755859375, "max": 3.999755859375, "num_bits": 15, "sign": 1, "exponent": 1, "mantissa": 13, "target_dtype": "fp15", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float15_e2m12fn": {"min": -7.9990234375, "max": 7.9990234375, "num_bits": 15, "sign": 1, "exponent": 2, "mantissa": 12, "target_dtype": "fp15", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float15_e3m11fn": {"min": -31.9921875, "max": 31.9921875, "num_bits": 15, "sign": 1, "exponent": 3, "mantissa": 11, "target_dtype": "fp15", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float15_e4m10fn": {"min": -511.75, "max": 511.75, "num_bits": 15, "sign": 1, "exponent": 4, "mantissa": 10, "target_dtype": "fp15", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float15_e5m9fn": {"min": -130944.0, "max": 130944.0, "num_bits": 15, "sign": 1, "exponent": 5, "mantissa": 9, "target_dtype": "fp15", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
#
"float14_e1m12fn": {"min": -3.99951171875, "max": 3.99951171875, "num_bits": 14, "sign": 1, "exponent": 1, "mantissa": 12, "target_dtype": "fp14", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float14_e2m11fn": {"min": -7.998046875, "max": 7.998046875, "num_bits": 14, "sign": 1, "exponent": 2, "mantissa": 11, "target_dtype": "fp14", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float14_e3m10fn": {"min": -31.984375, "max": 31.984375, "num_bits": 14, "sign": 1, "exponent": 3, "mantissa": 10, "target_dtype": "fp14", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float14_e4m9fn": {"min": -511.5, "max": 511.5, "num_bits": 14, "sign": 1, "exponent": 4, "mantissa": 9, "target_dtype": "fp14", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float14_e5m8fn": {"min": -130816.0, "max": 130816.0, "num_bits": 14, "sign": 1, "exponent": 5, "mantissa": 8, "target_dtype": "fp14", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
#
"float13_e1m11fn": {"min": -3.9990234375, "max": 3.9990234375, "num_bits": 13, "sign": 1, "exponent": 1, "mantissa": 11, "target_dtype": "fp13", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float13_e2m10fn": {"min": -7.99609375, "max": 7.99609375, "num_bits": 13, "sign": 1, "exponent": 2, "mantissa": 10, "target_dtype": "fp13", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float13_e3m9fn": {"min": -31.96875, "max": 31.96875, "num_bits": 13, "sign": 1, "exponent": 3, "mantissa": 9, "target_dtype": "fp13", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float13_e4m8fn": {"min": -511.0, "max": 511.0, "num_bits": 13, "sign": 1, "exponent": 4, "mantissa": 8, "target_dtype": "fp13", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float13_e5m7fn": {"min": -130560.0, "max": 130560.0, "num_bits": 13, "sign": 1, "exponent": 5, "mantissa": 7, "target_dtype": "fp13", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
#
"float12_e1m10fn": {"min": -3.998046875, "max": 3.998046875, "num_bits": 12, "sign": 1, "exponent": 1, "mantissa": 10, "target_dtype": "fp12", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float12_e2m9fn": {"min": -7.9921875, "max": 7.9921875, "num_bits": 12, "sign": 1, "exponent": 2, "mantissa": 9, "target_dtype": "fp12", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float12_e3m8fn": {"min": -31.9375, "max": 31.9375, "num_bits": 12, "sign": 1, "exponent": 3, "mantissa": 8, "target_dtype": "fp12", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float12_e4m7fn": {"min": -510.0, "max": 510.0, "num_bits": 12, "sign": 1, "exponent": 4, "mantissa": 7, "target_dtype": "fp12", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float12_e5m6fn": {"min": -130048.0, "max": 130048.0, "num_bits": 12, "sign": 1, "exponent": 5, "mantissa": 6, "target_dtype": "fp12", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
#
"float11_e1m9fn": {"min": -3.99609375, "max": 3.99609375, "num_bits": 11, "sign": 1, "exponent": 1, "mantissa": 9, "target_dtype": "fp11", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float11_e2m8fn": {"min": -7.984375, "max": 7.984375, "num_bits": 11, "sign": 1, "exponent": 2, "mantissa": 8, "target_dtype": "fp11", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float11_e3m7fn": {"min": -31.875, "max": 31.875, "num_bits": 11, "sign": 1, "exponent": 3, "mantissa": 7, "target_dtype": "fp11", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float11_e4m6fn": {"min": -508.0, "max": 508.0, "num_bits": 11, "sign": 1, "exponent": 4, "mantissa": 6, "target_dtype": "fp11", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float11_e5m5fn": {"min": -129024.0, "max": 129024.0, "num_bits": 11, "sign": 1, "exponent": 5, "mantissa": 5, "target_dtype": "fp11", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
#
"float10_e1m8fn": {"min": -3.9921875, "max": 3.9921875, "num_bits": 10, "sign": 1, "exponent": 1, "mantissa": 8, "target_dtype": "fp10", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float10_e2m7fn": {"min": -7.96875, "max": 7.96875, "num_bits": 10, "sign": 1, "exponent": 2, "mantissa": 7, "target_dtype": "fp10", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float10_e3m6fn": {"min": -31.75, "max": 31.75, "num_bits": 10, "sign": 1, "exponent": 3, "mantissa": 6, "target_dtype": "fp10", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float10_e4m5fn": {"min": -504.0, "max": 504.0, "num_bits": 10, "sign": 1, "exponent": 4, "mantissa": 5, "target_dtype": "fp10", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float10_e5m4fn": {"min": -126976.0, "max": 126976.0, "num_bits": 10, "sign": 1, "exponent": 5, "mantissa": 4, "target_dtype": "fp10", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
#
"float9_e1m7fn": {"min": -3.984375, "max": 3.984375, "num_bits": 9, "sign": 1, "exponent": 1, "mantissa": 7, "target_dtype": "fp9", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float9_e2m6fn": {"min": -7.9375, "max": 7.9375, "num_bits": 9, "sign": 1, "exponent": 2, "mantissa": 6, "target_dtype": "fp9", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float9_e3m5fn": {"min": -31.5, "max": 31.5, "num_bits": 9, "sign": 1, "exponent": 3, "mantissa": 5, "target_dtype": "fp9", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float9_e4m4fn": {"min": -496.0, "max": 496.0, "num_bits": 9, "sign": 1, "exponent": 4, "mantissa": 4, "target_dtype": "fp9", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float9_e5m3fn": {"min": -122880.0, "max": 122880.0, "num_bits": 9, "sign": 1, "exponent": 5, "mantissa": 3, "target_dtype": "fp9", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": False, "is_integer": False, "is_packed": True},
#
"float8_e1m6fn": {"min": -3.96875, "max": 3.96875, "num_bits": 8, "sign": 1, "exponent": 1, "mantissa": 6, "target_dtype": "fp8", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float8_e2m5fn": {"min": -7.875, "max": 7.875, "num_bits": 8, "sign": 1, "exponent": 2, "mantissa": 5, "target_dtype": "fp8", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float8_e3m4fn": {"min": -31.0, "max": 31.0, "num_bits": 8, "sign": 1, "exponent": 3, "mantissa": 4, "target_dtype": "fp8", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float8_e4m3fn_sdnq": {"min": -480.0, "max": 480.0, "num_bits": 8, "sign": 1, "exponent": 4, "mantissa": 3, "target_dtype": "fp8", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float8_e5m2fn": {"min": -114688.0, "max": 114688.0, "num_bits": 8, "sign": 1, "exponent": 5, "mantissa": 2, "target_dtype": "fp8", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
#
"float7_e1m5fn": {"min": -3.9375, "max": 3.9375, "num_bits": 7, "sign": 1, "exponent": 1, "mantissa": 5, "target_dtype": "fp7", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float7_e2m4fn": {"min": -7.75, "max": 7.75, "num_bits": 7, "sign": 1, "exponent": 2, "mantissa": 4, "target_dtype": "fp7", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float7_e3m3fn": {"min": -30.0, "max": 30.0, "num_bits": 7, "sign": 1, "exponent": 3, "mantissa": 3, "target_dtype": "fp7", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float7_e4m2fn": {"min": -448.0, "max": 448.0, "num_bits": 7, "sign": 1, "exponent": 4, "mantissa": 2, "target_dtype": "fp7", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float7_e5m1fn": {"min": -98304.0, "max": 98304.0, "num_bits": 7, "sign": 1, "exponent": 5, "mantissa": 1, "target_dtype": "fp7", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
#
"float6_e1m4fn": {"min": -3.875, "max": 3.875, "num_bits": 6, "sign": 1, "exponent": 1, "mantissa": 4, "target_dtype": "fp6", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float6_e2m3fn": {"min": -7.5, "max": 7.5, "num_bits": 6, "sign": 1, "exponent": 2, "mantissa": 3, "target_dtype": "fp6", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float6_e3m2fn": {"min": -28.0, "max": 28.0, "num_bits": 6, "sign": 1, "exponent": 3, "mantissa": 2, "target_dtype": "fp6", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float6_e4m1fn": {"min": -384.0, "max": 384.0, "num_bits": 6, "sign": 1, "exponent": 4, "mantissa": 1, "target_dtype": "fp6", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float6_e5m0fn": {"min": -65536.0, "max": 65536.0, "num_bits": 6, "sign": 1, "exponent": 5, "mantissa": 0, "target_dtype": "fp6", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
#
"float5_e1m3fn": {"min": -3.75, "max": 3.75, "num_bits": 5, "sign": 1, "exponent": 1, "mantissa": 3, "target_dtype": "fp5", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float5_e2m2fn": {"min": -7.0, "max": 7.0, "num_bits": 5, "sign": 1, "exponent": 2, "mantissa": 2, "target_dtype": "fp5", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float5_e3m1fn": {"min": -24.0, "max": 24.0, "num_bits": 5, "sign": 1, "exponent": 3, "mantissa": 1, "target_dtype": "fp5", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float5_e4m0fn": {"min": -256.0, "max": 256.0, "num_bits": 5, "sign": 1, "exponent": 4, "mantissa": 0, "target_dtype": "fp5", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
#
"float4_e1m2fn": {"min": -3.5, "max": 3.5, "num_bits": 4, "sign": 1, "exponent": 1, "mantissa": 2, "target_dtype": "fp4", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float4_e2m1fn": {"min": -6.0, "max": 6.0, "num_bits": 4, "sign": 1, "exponent": 2, "mantissa": 1, "target_dtype": "fp4", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float4_e3m0fn": {"min": -16.0, "max": 16.0, "num_bits": 4, "sign": 1, "exponent": 3, "mantissa": 0, "target_dtype": "fp4", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
#
"float3_e1m1fn": {"min": -3.0, "max": 3.0, "num_bits": 3, "sign": 1, "exponent": 1, "mantissa": 1, "target_dtype": "fp3", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
"float3_e2m0fn": {"min": -4.0, "max": 4.0, "num_bits": 3, "sign": 1, "exponent": 2, "mantissa": 0, "target_dtype": "fp3", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
#
"float2_e1m0fn": {"min": -2.0, "max": 2.0, "num_bits": 2, "sign": 1, "exponent": 1, "mantissa": 0, "target_dtype": "fp2", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": False, "is_packed": True},
### Custom Unsigned Floats
"float16_e1m15fnu": {"min": 0, "max": 3.99993896484375, "num_bits": 16, "sign": 0, "exponent": 1, "mantissa": 15, "target_dtype": "fp16", "torch_dtype": torch.float32, "storage_dtype": torch.uint16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float16_e2m14fnu": {"min": 0, "max": 7.999755859375, "num_bits": 16, "sign": 0, "exponent": 2, "mantissa": 14, "target_dtype": "fp16", "torch_dtype": torch.float32, "storage_dtype": torch.uint16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float16_e3m13fnu": {"min": 0, "max": 31.998046875, "num_bits": 16, "sign": 0, "exponent": 3, "mantissa": 13, "target_dtype": "fp16", "torch_dtype": torch.float32, "storage_dtype": torch.uint16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float16_e4m12fnu": {"min": 0, "max": 511.9375, "num_bits": 16, "sign": 0, "exponent": 4, "mantissa": 12, "target_dtype": "fp16", "torch_dtype": torch.float32, "storage_dtype": torch.uint16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float16_e5m11fnu": {"min": 0, "max": 131040.0, "num_bits": 16, "sign": 0, "exponent": 5, "mantissa": 11, "target_dtype": "fp16", "torch_dtype": torch.float32, "storage_dtype": torch.uint16, "is_unsigned": True, "is_integer": False, "is_packed": True},
#
"float15_e1m14fnu": {"min": 0, "max": 3.9998779296875, "num_bits": 15, "sign": 0, "exponent": 1, "mantissa": 14, "target_dtype": "fp15", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float15_e2m13fnu": {"min": 0, "max": 7.99951171875, "num_bits": 15, "sign": 0, "exponent": 2, "mantissa": 13, "target_dtype": "fp15", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float15_e3m12fnu": {"min": 0, "max": 31.99609375, "num_bits": 15, "sign": 0, "exponent": 3, "mantissa": 12, "target_dtype": "fp15", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float15_e4m11fnu": {"min": 0, "max": 511.875, "num_bits": 15, "sign": 0, "exponent": 4, "mantissa": 11, "target_dtype": "fp15", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float15_e5m10fnu": {"min": 0, "max": 131008.0, "num_bits": 15, "sign": 0, "exponent": 5, "mantissa": 10, "target_dtype": "fp15", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
#
"float14_e1m13fnu": {"min": 0, "max": 3.999755859375, "num_bits": 14, "sign": 0, "exponent": 1, "mantissa": 13, "target_dtype": "fp14", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float14_e2m12fnu": {"min": 0, "max": 7.9990234375, "num_bits": 14, "sign": 0, "exponent": 2, "mantissa": 12, "target_dtype": "fp14", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float14_e3m11fnu": {"min": 0, "max": 31.9921875, "num_bits": 14, "sign": 0, "exponent": 3, "mantissa": 11, "target_dtype": "fp14", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float14_e4m10fnu": {"min": 0, "max": 511.75, "num_bits": 14, "sign": 0, "exponent": 4, "mantissa": 10, "target_dtype": "fp14", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float14_e5m9fnu": {"min": 0, "max": 130944.0, "num_bits": 14, "sign": 0, "exponent": 5, "mantissa": 9, "target_dtype": "fp14", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
#
"float13_e1m12fnu": {"min": 0, "max": 3.99951171875, "num_bits": 13, "sign": 0, "exponent": 1, "mantissa": 12, "target_dtype": "fp13", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float13_e2m11fnu": {"min": 0, "max": 7.998046875, "num_bits": 13, "sign": 0, "exponent": 2, "mantissa": 11, "target_dtype": "fp13", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float13_e3m10fnu": {"min": 0, "max": 31.984375, "num_bits": 13, "sign": 0, "exponent": 3, "mantissa": 10, "target_dtype": "fp13", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float13_e4m9fnu": {"min": 0, "max": 511.5, "num_bits": 13, "sign": 0, "exponent": 4, "mantissa": 9, "target_dtype": "fp13", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float13_e5m8fnu": {"min": 0, "max": 130816.0, "num_bits": 13, "sign": 0, "exponent": 5, "mantissa": 8, "target_dtype": "fp13", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
#
"float12_e1m11fnu": {"min": 0, "max": 3.9990234375, "num_bits": 12, "sign": 0, "exponent": 1, "mantissa": 11, "target_dtype": "fp12", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float12_e2m10fnu": {"min": 0, "max": 7.99609375, "num_bits": 12, "sign": 0, "exponent": 2, "mantissa": 10, "target_dtype": "fp12", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float12_e3m9fnu": {"min": 0, "max": 31.96875, "num_bits": 12, "sign": 0, "exponent": 3, "mantissa": 9, "target_dtype": "fp12", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float12_e4m8fnu": {"min": 0, "max": 511.0, "num_bits": 12, "sign": 0, "exponent": 4, "mantissa": 8, "target_dtype": "fp12", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float12_e5m7fnu": {"min": 0, "max": 130560.0, "num_bits": 12, "sign": 0, "exponent": 5, "mantissa": 7, "target_dtype": "fp12", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
#
"float11_e1m10fnu": {"min": 0, "max": 3.998046875, "num_bits": 11, "sign": 0, "exponent": 1, "mantissa": 10, "target_dtype": "fp11", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float11_e2m9fnu": {"min": 0, "max": 7.9921875, "num_bits": 11, "sign": 0, "exponent": 2, "mantissa": 9, "target_dtype": "fp11", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float11_e3m8fnu": {"min": 0, "max": 31.9375, "num_bits": 11, "sign": 0, "exponent": 3, "mantissa": 8, "target_dtype": "fp11", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float11_e4m7fnu": {"min": 0, "max": 510.0, "num_bits": 11, "sign": 0, "exponent": 4, "mantissa": 7, "target_dtype": "fp11", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float11_e5m6fnu": {"min": 0, "max": 130048.0, "num_bits": 11, "sign": 0, "exponent": 5, "mantissa": 6, "target_dtype": "fp11", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
#
"float10_e1m9fnu": {"min": 0, "max": 3.99609375, "num_bits": 10, "sign": 0, "exponent": 1, "mantissa": 9, "target_dtype": "fp10", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float10_e2m8fnu": {"min": 0, "max": 7.984375, "num_bits": 10, "sign": 0, "exponent": 2, "mantissa": 8, "target_dtype": "fp10", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float10_e3m7fnu": {"min": 0, "max": 31.875, "num_bits": 10, "sign": 0, "exponent": 3, "mantissa": 7, "target_dtype": "fp10", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float10_e4m6fnu": {"min": 0, "max": 508.0, "num_bits": 10, "sign": 0, "exponent": 4, "mantissa": 6, "target_dtype": "fp10", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float10_e5m5fnu": {"min": 0, "max": 129024.0, "num_bits": 10, "sign": 0, "exponent": 5, "mantissa": 5, "target_dtype": "fp10", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
#
"float9_e1m8fnu": {"min": 0, "max": 3.9921875, "num_bits": 9, "sign": 0, "exponent": 1, "mantissa": 8, "target_dtype": "fp9", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float9_e2m7fnu": {"min": 0, "max": 7.96875, "num_bits": 9, "sign": 0, "exponent": 2, "mantissa": 7, "target_dtype": "fp9", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float9_e3m6fnu": {"min": 0, "max": 31.75, "num_bits": 9, "sign": 0, "exponent": 3, "mantissa": 6, "target_dtype": "fp9", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float9_e4m5fnu": {"min": 0, "max": 504.0, "num_bits": 9, "sign": 0, "exponent": 4, "mantissa": 5, "target_dtype": "fp9", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float9_e5m4fnu": {"min": 0, "max": 126976.0, "num_bits": 9, "sign": 0, "exponent": 5, "mantissa": 4, "target_dtype": "fp9", "torch_dtype": torch.float32, "storage_dtype": torch.int16, "is_unsigned": True, "is_integer": False, "is_packed": True},
#
"float8_e1m7fnu": {"min": 0, "max": 3.984375, "num_bits": 8, "sign": 0, "exponent": 1, "mantissa": 7, "target_dtype": "fp8", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float8_e2m6fnu": {"min": 0, "max": 7.9375, "num_bits": 8, "sign": 0, "exponent": 2, "mantissa": 6, "target_dtype": "fp8", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float8_e3m5fnu": {"min": 0, "max": 31.5, "num_bits": 8, "sign": 0, "exponent": 3, "mantissa": 5, "target_dtype": "fp8", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float8_e4m4fnu": {"min": 0, "max": 496.0, "num_bits": 8, "sign": 0, "exponent": 4, "mantissa": 4, "target_dtype": "fp8", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float8_e5m3fnu": {"min": 0, "max": 122880.0, "num_bits": 8, "sign": 0, "exponent": 5, "mantissa": 3, "target_dtype": "fp8", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
#
"float7_e1m6fnu": {"min": 0, "max": 3.96875, "num_bits": 7, "sign": 0, "exponent": 1, "mantissa": 6, "target_dtype": "fp7", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float7_e2m5fnu": {"min": 0, "max": 7.875, "num_bits": 7, "sign": 0, "exponent": 2, "mantissa": 5, "target_dtype": "fp7", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float7_e3m4fnu": {"min": 0, "max": 31.0, "num_bits": 7, "sign": 0, "exponent": 3, "mantissa": 4, "target_dtype": "fp7", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float7_e4m3fnu": {"min": 0, "max": 480.0, "num_bits": 7, "sign": 0, "exponent": 4, "mantissa": 3, "target_dtype": "fp7", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float7_e5m2fnu": {"min": 0, "max": 114688.0, "num_bits": 7, "sign": 0, "exponent": 5, "mantissa": 2, "target_dtype": "fp7", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
#
"float6_e1m5fnu": {"min": 0, "max": 3.9375, "num_bits": 6, "sign": 0, "exponent": 1, "mantissa": 5, "target_dtype": "fp6", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float6_e2m4fnu": {"min": 0, "max": 7.75, "num_bits": 6, "sign": 0, "exponent": 2, "mantissa": 4, "target_dtype": "fp6", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float6_e3m3fnu": {"min": 0, "max": 30.0, "num_bits": 6, "sign": 0, "exponent": 3, "mantissa": 3, "target_dtype": "fp6", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float6_e4m2fnu": {"min": 0, "max": 448.0, "num_bits": 6, "sign": 0, "exponent": 4, "mantissa": 2, "target_dtype": "fp6", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float6_e5m1fnu": {"min": 0, "max": 98304.0, "num_bits": 6, "sign": 0, "exponent": 5, "mantissa": 1, "target_dtype": "fp6", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
#
"float5_e1m4fnu": {"min": 0, "max": 3.875, "num_bits": 5, "sign": 0, "exponent": 1, "mantissa": 4, "target_dtype": "fp5", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float5_e2m3fnu": {"min": 0, "max": 7.5, "num_bits": 5, "sign": 0, "exponent": 2, "mantissa": 3, "target_dtype": "fp5", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float5_e3m2fnu": {"min": 0, "max": 28.0, "num_bits": 5, "sign": 0, "exponent": 3, "mantissa": 2, "target_dtype": "fp5", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float5_e4m1fnu": {"min": 0, "max": 384.0, "num_bits": 5, "sign": 0, "exponent": 4, "mantissa": 1, "target_dtype": "fp5", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float5_e5m0fnu": {"min": 0, "max": 65536.0, "num_bits": 5, "sign": 0, "exponent": 5, "mantissa": 0, "target_dtype": "fp5", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
#
"float4_e1m3fnu": {"min": 0, "max": 3.75, "num_bits": 4, "sign": 0, "exponent": 1, "mantissa": 3, "target_dtype": "fp4", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float4_e2m2fnu": {"min": 0, "max": 7.0, "num_bits": 4, "sign": 0, "exponent": 2, "mantissa": 2, "target_dtype": "fp4", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float4_e3m1fnu": {"min": 0, "max": 24.0, "num_bits": 4, "sign": 0, "exponent": 3, "mantissa": 1, "target_dtype": "fp4", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float4_e4m0fnu": {"min": 0, "max": 256.0, "num_bits": 4, "sign": 0, "exponent": 4, "mantissa": 0, "target_dtype": "fp4", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
#
"float3_e1m2fnu": {"min": 0, "max": 3.5, "num_bits": 3, "sign": 0, "exponent": 1, "mantissa": 2, "target_dtype": "fp3", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float3_e2m1fnu": {"min": 0, "max": 6.0, "num_bits": 3, "sign": 0, "exponent": 2, "mantissa": 1, "target_dtype": "fp3", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float3_e3m0fnu": {"min": 0, "max": 16.0, "num_bits": 3, "sign": 0, "exponent": 3, "mantissa": 0, "target_dtype": "fp3", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
#
"float2_e1m1fnu": {"min": 0, "max": 3.0, "num_bits": 2, "sign": 0, "exponent": 1, "mantissa": 1, "target_dtype": "fp2", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
"float2_e2m0fnu": {"min": 0, "max": 4.0, "num_bits": 2, "sign": 0, "exponent": 2, "mantissa": 0, "target_dtype": "fp2", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
#
"float1_e1m0fnu": {"min": 0, "max": 2.0, "num_bits": 1, "sign": 0, "exponent": 1, "mantissa": 0, "target_dtype": "fp1", "torch_dtype": torch.float32, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": False, "is_packed": True},
}
dtype_dict["fp32"] = dtype_dict["float32"]
dtype_dict["bf16"] = dtype_dict["bfloat16"]
dtype_dict["fp16"] = dtype_dict["float16"]
dtype_dict["fp15"] = dtype_dict["float15_e5m9fn"]
dtype_dict["fp14"] = dtype_dict["float14_e5m8fn"]
dtype_dict["fp13"] = dtype_dict["float13_e5m7fn"]
dtype_dict["fp12"] = dtype_dict["float12_e5m6fn"]
dtype_dict["fp11"] = dtype_dict["float11_e5m5fn"]
dtype_dict["fp10"] = dtype_dict["float10_e5m4fn"]
dtype_dict["fp9"] = dtype_dict["float9_e4m4fn"]
dtype_dict["fp8"] = dtype_dict["float8_e4m3fn"]
dtype_dict["fp7"] = dtype_dict["float7_e3m3fn"]
dtype_dict["fp6"] = dtype_dict["float6_e3m2fn"]
dtype_dict["fp5"] = dtype_dict["float5_e2m2fn"]
dtype_dict["fp4"] = dtype_dict["float4_e2m1fn"]
dtype_dict["fp3"] = dtype_dict["float3_e1m1fn"]
dtype_dict["fp2"] = dtype_dict["float2_e1m0fn"]
dtype_dict["ufp16"] = dtype_dict["float16_e5m11fnu"]
dtype_dict["ufp15"] = dtype_dict["float15_e5m10fnu"]
dtype_dict["ufp14"] = dtype_dict["float14_e5m9fnu"]
dtype_dict["ufp13"] = dtype_dict["float13_e5m8fnu"]
dtype_dict["ufp12"] = dtype_dict["float12_e5m7fnu"]
dtype_dict["ufp11"] = dtype_dict["float11_e5m6fnu"]
dtype_dict["ufp10"] = dtype_dict["float10_e5m5fnu"]
dtype_dict["ufp9"] = dtype_dict["float9_e4m5fnu"]
dtype_dict["ufp8"] = dtype_dict["float8_e4m4fnu"]
dtype_dict["ufp7"] = dtype_dict["float7_e3m4fnu"]
dtype_dict["ufp6"] = dtype_dict["float6_e3m3fnu"]
dtype_dict["ufp5"] = dtype_dict["float5_e2m3fnu"]
dtype_dict["ufp4"] = dtype_dict["float4_e2m2fnu"]
dtype_dict["ufp3"] = dtype_dict["float3_e1m2fnu"]
dtype_dict["ufp2"] = dtype_dict["float2_e1m1fnu"]
dtype_dict["ufp1"] = dtype_dict["float1_e1m0fnu"]
dtype_dict["fp1"] = dtype_dict["ufp1"]
dtype_dict["int1"] = dtype_dict["uint1"]
dtype_dict["bool"] = dtype_dict["uint1"]
torch_dtype_dict = {
torch.int32: "int32",
torch.int16: "int16",
torch.int8: "int8",
torch.uint32: "uint32",
torch.uint16: "uint16",
torch.uint8: "uint8",
torch.float32: "float32",
torch.bfloat16: "bfloat16",
torch.float16: "float16",
torch.float8_e4m3fn: "float8_e4m3fn",
torch.float8_e5m2: "float8_e5m2",
}
if hasattr(torch, "float8_e8m0fnu"):
dtype_dict["float8_e8m0fnu"] = {"min": -1.70141e+38, "max": 1.70141e+38, "num_bits": 8, "sign": 1, "exponent": 8, "mantissa": 0, "target_dtype": "fp8", "torch_dtype": torch.float8_e8m0fnu, "storage_dtype": torch.float8_e8m0fnu, "is_unsigned": False, "is_integer": False, "is_packed": False}
torch_dtype_dict[torch.float8_e8m0fnu] = "float8_e8m0fnu"
if hasattr(torch, "float8_e4m3fnuz"):
dtype_dict["float8_e4m3fnuz"] = {"min": -240.0, "max": 240.0, "num_bits": 8, "sign": 1, "exponent": 4, "mantissa": 3, "target_dtype": "fp8", "torch_dtype": torch.float8_e4m3fnuz, "storage_dtype": torch.float8_e4m3fnuz, "is_unsigned": False, "is_integer": False, "is_packed": False}
torch_dtype_dict[torch.float8_e4m3fnuz] = "float8_e4m3fnuz"
if hasattr(torch, "float8_e5m2fnuz"):
dtype_dict["float8_e5m2fnuz"] = {"min": -57344.0, "max": 57344.0, "num_bits": 8, "sign": 1, "exponent": 5, "mantissa": 2, "target_dtype": "fp8", "torch_dtype": torch.float8_e5m2fnuz, "storage_dtype": torch.float8_e5m2fnuz, "is_unsigned": False, "is_integer": False, "is_packed": False}
torch_dtype_dict[torch.float8_e5m2fnuz] = "float8_e5m2fnuz"
linear_types = {"Linear", "SDNQLinear"}
embedding_types = {"Embedding", "SDNQEmbedding", "Gemma4TextScaledWordEmbedding"}
conv_types = {"Conv1d", "Conv2d", "Conv3d", "SDNQConv1d", "SDNQConv2d", "SDNQConv3d"}
conv_transpose_types = {"ConvTranspose1d", "ConvTranspose2d", "ConvTranspose3d", "SDNQConvTranspose1d", "SDNQConvTranspose2d", "SDNQConvTranspose3d"}
allowed_types = set.union(linear_types, embedding_types, conv_types, conv_transpose_types)
accepted_weight_dtypes = set(dtype_dict.keys())
accepted_matmul_dtypes = {"int8", "uint8", "fp8", "fp16", "float8_e4m3fn", "float16"}
weights_dtype_order = [
"uint1", "float1_e1m0fnu",
"int2", "float2_e1m0fn",
"uint2", "float2_e1m1fnu", "float2_e2m0fnu",
"int3", "float3_e1m1fn", "float3_e2m0fn",
"uint3", "float3_e1m2fnu", "float3_e2m1fnu", "float3_e3m0fnu",
"int4", "float4_e1m2fn", "float4_e2m1fn", "float4_e3m0fn",
"uint4", "float4_e1m3fnu", "float4_e2m2fnu", "float4_e3m1fnu", "float4_e4m0fnu",
"int5", "float5_e1m3fn", "float5_e2m2fn", "float5_e3m1fn", "float5_e4m0fn",
"uint5", "float5_e1m4fnu", "float5_e2m3fnu", "float5_e3m2fnu", "float5_e4m1fnu", "float5_e5m0fnu",
"int6", "float6_e1m4fn", "float6_e2m3fn", "float6_e3m2fn", "float6_e4m1fn", "float6_e5m0fn",
"uint6", "float6_e1m5fnu", "float6_e2m4fnu", "float6_e3m3fnu", "float6_e4m2fnu", "float6_e5m1fnu",
"int7", "float7_e1m5fn", "float7_e2m4fn", "float7_e3m3fn", "float7_e4m2fn", "float7_e5m1fn",
"uint7", "float7_e1m6fnu", "float7_e2m5fnu", "float7_e3m4fnu", "float7_e4m3fnu", "float7_e5m2fnu",
"int8", "float8_e4m3fn", "float8_e5m2", "float8_e1m6fn", "float8_e2m5fn", "float8_e3m4fn", "float8_e4m3fn_sdnq", "float8_e5m2fn",
"uint8", "float8_e1m7fnu", "float8_e2m6fnu", "float8_e3m5fnu", "float8_e4m4fnu", "float8_e5m3fnu",
"int9", "float9_e1m7fn", "float9_e2m6fn", "float9_e3m5fn", "float9_e4m4fn", "float9_e5m3fn",
"uint9", "float9_e1m8fnu", "float9_e2m7fnu", "float9_e3m6fnu", "float9_e4m5fnu", "float9_e5m4fnu",
"int10", "float10_e1m8fn", "float10_e2m7fn", "float10_e3m6fn", "float10_e4m5fn", "float10_e5m4fn",
"uint10", "float10_e1m9fnu", "float10_e2m8fnu", "float10_e3m7fnu", "float10_e4m6fnu", "float10_e5m5fnu",
"int11", "float11_e1m9fn", "float11_e2m8fn", "float11_e3m7fn", "float11_e4m6fn", "float11_e5m5fn",
"uint11", "float11_e1m10fnu", "float11_e2m9fnu", "float11_e3m8fnu", "float11_e4m7fnu", "float11_e5m6fnu",
"int12", "float12_e1m10fn", "float12_e2m9fn", "float12_e3m8fn", "float12_e4m7fn", "float12_e5m6fn",
"uint12", "float12_e1m11fnu", "float12_e2m10fnu", "float12_e3m9fnu", "float12_e4m8fnu", "float12_e5m7fnu",
"int13", "float13_e1m11fn", "float13_e2m10fn", "float13_e3m9fn", "float13_e4m8fn", "float13_e5m7fn",
"uint13", "float13_e1m12fnu", "float13_e2m11fnu", "float13_e3m10fnu", "float13_e4m9fnu", "float13_e5m8fnu",
"int14", "float14_e1m12fn", "float14_e2m11fn", "float14_e3m10fn", "float14_e4m9fn", "float14_e5m8fn",
"uint14", "float14_e1m13fnu", "float14_e2m12fnu", "float14_e3m11fnu", "float14_e4m10fnu", "float14_e5m9fnu",
"int15", "float15_e1m13fn", "float15_e2m12fn", "float15_e3m11fn", "float15_e4m10fn", "float15_e5m9fn",
"uint15", "float15_e1m14fnu", "float15_e2m13fnu", "float15_e3m12fnu", "float15_e4m11fnu", "float15_e5m10fnu",
"int16", "float16", "float16_e1m14fn", "float16_e2m13fn", "float16_e3m12fn", "float16_e4m11fn", "float16_e5m10fn",
"uint16", "float16_e1m15fnu", "float16_e2m14fnu", "float16_e3m13fnu", "float16_e4m12fnu", "float16_e5m11fnu",
]
use_torch_compile = shared.opts.sdnq_dequantize_compile # this setting requires a full restart of the webui to apply
def check_torch_compile() -> bool: # dynamo can be disabled after startup
return use_torch_compile and not torch._dynamo.config.disable # pylint: disable=protected-access
if use_torch_compile:
if hasattr(torch._dynamo.config, "recompile_limit"):
torch._dynamo.config.recompile_limit = max(8192, getattr(torch._dynamo.config, "recompile_limit", 0))
if hasattr(torch._dynamo.config, "cache_size_limit"):
torch._dynamo.config.cache_size_limit = max(8192, getattr(torch._dynamo.config, "cache_size_limit", 0))
if hasattr(torch._dynamo.config, "accumulated_recompile_limit"):
torch._dynamo.config.accumulated_recompile_limit = max(8192, getattr(torch._dynamo.config, "accumulated_recompile_limit", 0))
if hasattr(torch._dynamo.config, "accumulated_cache_size_limit"):
torch._dynamo.config.accumulated_cache_size_limit = max(8192, getattr(torch._dynamo.config, "accumulated_cache_size_limit", 0))
def compile_func(fn, **kwargs):
if kwargs.get("fullgraph", None) is None:
kwargs["fullgraph"] = True
if kwargs.get("dynamic", None) is None:
kwargs["dynamic"] = False
if (torch_version[0] > 2 or (torch_version[0] == 2 and torch_version[1] >= 12)) and kwargs.get("recompile_limit", None) is None:
kwargs["recompile_limit"] = max(8192, getattr(torch._dynamo.config, "recompile_limit", 0))
if os.environ.get("SDNQ_COMPILE_KWARGS", None) is not None:
for key, value in json.loads(os.environ.get("SDNQ_COMPILE_KWARGS")).items():
kwargs[key] = value
return torch.compile(fn, **kwargs)
else:
def compile_func(fn, **kwargs): # pylint: disable=unused-argument
return fn
def reset_compile_caches():
if check_torch_compile():
shared.log.debug('SDNQ compile: dynamo reset')
torch._dynamo.reset()
from .kernel_wrappers import use_openvino_mm
if use_openvino_mm:
shared.log.debug('SDNQ compile: openvino reset')
from .kernels.openvino_mm import OV_COMPILED_CACHE
OV_COMPILED_CACHE.clear()
common_skip_keys = (
".time_embed",
".context_embedder",
".condition_embedder",
".x_embedder",
".t_embedder",
".y_embedder",
".emb_in",
".txt_in",
".img_in",
".vid_in",
".proj_out",
".norm_out",
".emb_out",
".txt_out",
".img_out",
".vid_out",
".final_layer",
"multi_modal_projector",
"time_text_embed",
"patch_embedding",
"patch_embed",
"patch_emb",
"lm_head",
"wte",
)
# modules_to_not_convert: ["x_embedder", "y_embedder"]
# modules_dtype_dict: {"minimum_6bit": ["x_embedder", "y_embedder"]}
# modules_to_not_use_matmul: {"int8": ["x_embedder", "y_embedder"], "float8_e4m3fn": ["x_embedder", "y_embedder"]}
module_skip_keys_dict = {
"FluxTransformer2DModel": [
["single_transformer_blocks.0.norm.linear.weight", "time_text_embed", "time_embed", "context_embedder", "x_embedder", ".proj_out", "norm_out"],
{},
{},
],
"Flux2Transformer2DModel": [
["double_stream_modulation_img", "double_stream_modulation_txt", "single_stream_modulation", "time_guidance_embed", "context_embedder", "x_embedder", ".proj_out", "norm_out"],
{},
{},
],
"ChromaTransformer2DModel": [
["distilled_guidance_layer", "time_text_embed", "context_embedder", "x_embedder", ".proj_out", "norm_out"],
{},
{},
],
"QwenImageTransformer2DModel": [
["transformer_blocks.0.img_mod.1.weight", "time_text_embed", "txt_in", "img_in", "proj_out", "norm_out"],
{},
{},
],
"WanTransformer3DModel": [
["scale_shift_table", "patch_embedding", "condition_embedder", "proj_out", "norm_out"],
{},
{},
],
"LongCatVideoTransformer3DModel": [
["blocks.0.adaLN_modulation.1.weight", "x_embedder", "t_embedder", "y_embedder", "final_layer"],
{},
{},
],
"LTX2VideoTransformer3DModel": [
[
"audio_time_embed", "time_embed", "audio_caption_projection", "caption_projection", "proj_in", "audio_proj_in", "proj_out", "audio_proj_out",
"av_cross_attn_audio_scale_shift", "av_cross_attn_audio_v2a_gate", "av_cross_attn_video_a2v_gate", "av_cross_attn_video_scale_shift",
],
{},
{},
],
"Lumina2Transformer2DModel": [
["layers.0.norm1.linear.weight", "time_caption_embed", "x_embedder", "norm_out"],
{},
{},
],
"ZImageTransformer2DModel": [
["layers.0.adaLN_modulation.0.weight", "t_embedder", "cap_embedder", "siglip_embedder", "all_x_embedder", "all_final_layer"],
{},
{},
],
"Ideogram4Transformer2DModel": [
["layers.0.adaln_modulation.weight", "input_proj", "llm_cond_proj", "llm_cond_norm", "final_layer", "t_embedding", "adaln_proj", "embed_image_indicator"],
{},
{},
],
"CosmosTransformer3DModel": [
["transformer_blocks.0.norm*", "patch_embed", "time_embed", "norm_out", "proj_out", "crossattn_proj"],
{},
{},
],
"GlmImageTransformer2DModel": [
["transformer_blocks.0.norm1.linear.weight", "image_projector", "glyph_projector", "prior_projector", "time_condition_embed", "norm_out", "proj_out"],
{},
{},
],
"GlmImageForConditionalGeneration": [
["lm_head", "patch_embed", "embeddings", "embed_tokens", "vqmodel"],
{},
{},
],
"HunyuanImage3ForCausalMM": [
["lm_head", "patch_embed", "time_embed", "time_embed_2", "final_layer", "wte", "ln_f", "timestep_emb", "vae", "vision_aligner", "head", "post_layernorm", "embeddings"],
{},
{},
],
"Emu3ForCausalLM": [
["lm_head", "vq_model", "tokenizer"],
{},
{},
],
"Gemma3nForCausalLM": [
["lm_head", "correction_coefs", "prediction_coefs", "embedding_projection"],
{},
{},
],
"Gemma4ForConditionalGeneration": [
["lm_head", "embed_audio", "embed_vision", "patch_embedder", "embed_tokens", "subsample_conv_projection", "output_proj"],
{},
{},
],
"MoondreamModel": [
["lm_head", "region", "wte", "post_ln", "proj_mlp", "patch_emb", "pos_emb"],
{},
{},
],
"NaDiT": [
[".emb_in", ".txt_in", ".vid_in", ".emb_scale", ".vid_out", ".vid_out_norm", ".vid_out_ada"],
{},
{},
],
"HiDreamO1Qwen3VLTransformer": [
["lm_head", "embed_tokens", "x_embedder", "t_embedder1", "final_layer2", "patch_embed", "pos_embed"],
{},
{},
],
}
module_skip_keys_dict["LongCatImageTransformer2DModel"] = module_skip_keys_dict["FluxTransformer2DModel"]
module_skip_keys_dict["ChronoEditTransformer3DModel"] = module_skip_keys_dict["WanTransformer3DModel"]
module_skip_keys_dict["Gemma3nForConditionalGeneration"] = module_skip_keys_dict["Gemma3nForCausalLM"]
module_skip_keys_dict["HfMoondream"] = module_skip_keys_dict["MoondreamModel"]
module_skip_keys_dict["NaDiTUpscaler"] = module_skip_keys_dict["NaDiT"]
modules/sdnq/dequantizer.py
-364
View File
@@ -1,364 +0,0 @@
# pylint: disable=redefined-builtin,no-member,protected-access
from dataclasses import dataclass
import torch
from modules import devices
from .common import dtype_dict, compile_func
from .kernel_wrappers import use_contiguous_int8_mm, use_contiguous_fp16_mm, use_tensorwise_fp8_matmul, is_fp8_compile_supported
from .quant_utils import quantize_int_mm, quantize_uint_mm, quantize_fp_mm, rotate_hadamard, get_hadamard
from .packed_int import unpack_int
from .packed_float import unpack_float
from .layers import SDNQLayer
@devices.inference_context()
def dequantize_asymmetric(
weight: torch.Tensor,
scale: torch.FloatTensor,
zero_point: torch.FloatTensor,
svd_up: torch.FloatTensor | None = None,
svd_down: torch.FloatTensor | None = None,
hadamard: torch.FloatTensor | None = None,
dtype: torch.dtype | None = None,
result_shape: torch.Size | None = None,
skip_quantized_matmul: bool = False,
re_quantize_for_matmul: bool = False,
) -> torch.FloatTensor:
result = torch.addcmul(zero_point, weight.to(dtype=scale.dtype), scale)
if skip_quantized_matmul and not re_quantize_for_matmul:
result.t_()
if result_shape is not None:
result = result.view(result_shape)
is_conv = bool(result.ndim > 2 and weight.ndim > 2)
if svd_up is not None:
if skip_quantized_matmul:
svd_up = svd_up.t().contiguous()
if use_contiguous_fp16_mm:
svd_down = svd_down.t().contiguous()
else:
svd_down = svd_down.contiguous().t()
if is_conv:
result = result.add_(torch.mm(svd_up, svd_down).unflatten(-1, (*result.shape[1:],)))
else:
result = result.to(dtype=svd_up.dtype).addmm_(svd_up, svd_down)
if dtype is not None:
result = result.to(dtype=dtype)
if hadamard is not None:
result = rotate_hadamard(result, hadamard=hadamard, is_conv=is_conv)
return result
@devices.inference_context()
def dequantize_symmetric(
weight: torch.Tensor,
scale: torch.FloatTensor,
svd_up: torch.FloatTensor | None = None,
svd_down: torch.FloatTensor | None = None,
hadamard: torch.FloatTensor | None = None,
dtype: torch.dtype | None = None,
result_shape: torch.Size | None = None,
skip_quantized_matmul: bool = False,
re_quantize_for_matmul: bool = False,
) -> torch.FloatTensor:
result = weight.to(dtype=scale.dtype).mul_(scale)
if skip_quantized_matmul and not re_quantize_for_matmul:
result.t_()
if result_shape is not None:
result = result.view(result_shape)
is_conv = bool(result.ndim > 2 and weight.ndim > 2)
if svd_up is not None:
if skip_quantized_matmul:
svd_up = svd_up.t().contiguous()
if use_contiguous_fp16_mm:
svd_down = svd_down.t().contiguous()
else:
svd_down = svd_down.contiguous().t()
if is_conv:
result = result.add_(torch.mm(svd_up, svd_down).unflatten(-1, (*result.shape[1:],)))
else:
result = result.to(dtype=svd_up.dtype).addmm_(svd_up, svd_down)
if dtype is not None:
result = result.to(dtype=dtype)
if hadamard is not None:
result = rotate_hadamard(result, hadamard=hadamard, is_conv=is_conv)
return result
def dequantize_weight(
weights_dtype: str,
weight: torch.Tensor,
scale: torch.FloatTensor,
zero_point: torch.FloatTensor | None = None,
svd_up: torch.FloatTensor | None = None,
svd_down: torch.FloatTensor | None = None,
hadamard: torch.FloatTensor | None = None,
dtype: torch.dtype | None = None,
result_shape: torch.Size | None = None,
quantized_weight_shape: torch.Size | None = None,
skip_quantized_matmul: bool = False,
re_quantize_for_matmul: bool = False,
) -> torch.FloatTensor:
if dtype_dict[weights_dtype]["is_packed"]:
if dtype_dict[weights_dtype]["is_integer"]:
weight = unpack_int(weight, weights_dtype, quantized_weight_shape, dtype=scale.dtype)
else:
weight = unpack_float(weight, weights_dtype, quantized_weight_shape)
if dtype_dict[weights_dtype]["is_unsigned"]:
return dequantize_asymmetric(weight, scale, zero_point, svd_up=svd_up, svd_down=svd_down, hadamard=hadamard, dtype=dtype, result_shape=result_shape, skip_quantized_matmul=skip_quantized_matmul, re_quantize_for_matmul=re_quantize_for_matmul)
else:
return dequantize_symmetric(weight, scale, svd_up=svd_up, svd_down=svd_down, hadamard=hadamard, dtype=dtype, result_shape=result_shape, skip_quantized_matmul=skip_quantized_matmul, re_quantize_for_matmul=re_quantize_for_matmul)
@devices.inference_context()
def re_quantize_int_mm(weight: torch.FloatTensor, matmul_dtype: str = "int8") -> tuple[torch.Tensor, torch.FloatTensor]:
if weight.ndim > 2: # convs
weight = weight.flatten(1,-1)
if use_contiguous_int8_mm:
weight, scale = quantize_int_mm(weight.t().contiguous(), dim=0, matmul_dtype=matmul_dtype)
else:
weight, scale = quantize_int_mm(weight.contiguous(), dim=-1, matmul_dtype=matmul_dtype)
weight, scale = weight.t_(), scale.t_().contiguous()
return weight, scale
@devices.inference_context()
def re_quantize_uint_mm(weight: torch.FloatTensor, matmul_dtype: str = "uint8") -> tuple[torch.Tensor, torch.FloatTensor, torch.FloatTensor]:
if weight.ndim > 2: # convs
weight = weight.flatten(1,-1)
if use_contiguous_int8_mm:
weight, scale, zero_point = quantize_uint_mm(weight.t().contiguous(), dim=0, matmul_dtype=matmul_dtype)
else:
weight, scale, zero_point = quantize_uint_mm(weight.contiguous(), dim=-1, matmul_dtype=matmul_dtype)
weight, scale, zero_point = weight.t_(), scale.t_().contiguous(), zero_point.t_().contiguous()
return weight, scale, zero_point
@devices.inference_context()
def re_quantize_fp_mm(weight: torch.FloatTensor, matmul_dtype: str = "float8_e4m3fn") -> tuple[torch.Tensor, torch.FloatTensor]:
if weight.ndim > 2: # convs
weight = weight.flatten(1,-1)
if use_contiguous_fp16_mm and matmul_dtype in {"fp16", "float16"}:
weight, scale = quantize_fp_mm(weight.t().contiguous(), dim=0, matmul_dtype=matmul_dtype)
else:
weight, scale = quantize_fp_mm(weight.contiguous(), dim=-1, matmul_dtype=matmul_dtype)
weight, scale = weight.t_(), scale.t_().contiguous()
if not use_tensorwise_fp8_matmul and dtype_dict[matmul_dtype]["num_bits"] == 8:
scale = scale.to(dtype=torch.float32)
return weight, scale
def re_quantize_matmul(
weights_dtype: str,
weight: torch.Tensor,
scale: torch.FloatTensor,
zero_point: torch.FloatTensor | None = None,
svd_up: torch.FloatTensor | None = None,
svd_down: torch.FloatTensor | None = None,
hadamard: torch.FloatTensor | None = None,
matmul_dtype: str = "int8",
result_shape: torch.Size | None = None,
quantized_weight_shape: torch.Size | None = None,
) -> tuple[torch.Tensor, torch.FloatTensor] | tuple[torch.Tensor, torch.FloatTensor, torch.FloatTensor]:
if dtype_dict[weights_dtype]["is_packed"]:
if dtype_dict[weights_dtype]["is_integer"]:
weight = unpack_int(weight, weights_dtype, quantized_weight_shape, dtype=scale.dtype)
else:
weight = unpack_float(weight, weights_dtype, quantized_weight_shape)
if dtype_dict[weights_dtype]["is_unsigned"]:
weight = dequantize_asymmetric(weight, scale, zero_point, svd_up=svd_up, svd_down=svd_down, hadamard=hadamard, dtype=scale.dtype, result_shape=result_shape)
else:
weight = dequantize_symmetric(weight, scale, svd_up=svd_up, svd_down=svd_down, hadamard=hadamard, dtype=scale.dtype, result_shape=result_shape)
if dtype_dict[matmul_dtype]["is_integer"]:
if dtype_dict[matmul_dtype]["is_unsigned"]:
return re_quantize_uint_mm(weight, matmul_dtype=matmul_dtype)
else:
return re_quantize_int_mm(weight, matmul_dtype=matmul_dtype)
else:
return re_quantize_fp_mm(weight, matmul_dtype=matmul_dtype)
@devices.inference_context()
def dequantize_sdnq_module(model: torch.nn.Module) -> torch.nn.Module:
if isinstance(model, SDNQLayer):
model = model.dequantize()
has_children = list(model.children())
if not has_children:
return model
for module_name, module in model.named_children():
if isinstance(module, SDNQLayer):
setattr(model, module_name, module.dequantize())
else:
setattr(model, module_name, dequantize_sdnq_model(module))
return model
@devices.inference_context()
def dequantize_sdnq_model(model: torch.nn.Module) -> torch.nn.Module:
model = dequantize_sdnq_module(model)
if hasattr(model, "quantization_method"):
del model.quantization_method
if hasattr(model, "quantization_config"):
del model.quantization_config
if hasattr(model, "config"):
try:
if hasattr(model.config, "quantization_config"):
del model.config.quantization_config
except Exception:
pass
try:
if hasattr(model.config, "pop"):
model.config.pop("quantization_config", None)
except Exception:
pass
return model
# SDNQDequantizer has to be a dataclass for torch.compile
@dataclass
class SDNQDequantizer:
result_dtype: torch.dtype
result_shape: torch.Size
original_shape: torch.Size
original_stride: list[int]
quantized_weight_shape: torch.Size
weights_dtype: str
quantized_matmul_dtype: str
hadamard_group_size: int
group_size: int
svd_rank: int
svd_steps: int
use_quantized_matmul: bool
re_quantize_for_matmul: bool
use_stochastic_rounding: bool
layer_class_name: str
is_packed: bool
is_unsigned: bool
is_integer: bool
is_integer_matmul: bool
def __init__(
self,
result_dtype: torch.dtype,
result_shape: torch.Size,
original_shape: torch.Size,
original_stride: list[int],
quantized_weight_shape: torch.Size,
weights_dtype: str,
quantized_matmul_dtype: str,
hadamard_group_size: int,
group_size: int,
svd_rank: int,
svd_steps: int,
use_quantized_matmul: bool,
re_quantize_for_matmul: bool,
use_stochastic_rounding: bool,
use_hadamard: bool,
layer_class_name: str,
):
self.result_dtype = result_dtype
self.result_shape = result_shape
self.original_shape = original_shape
self.original_stride = original_stride
self.quantized_weight_shape = quantized_weight_shape
self.weights_dtype = weights_dtype
self.quantized_matmul_dtype = quantized_matmul_dtype
self.hadamard_group_size = hadamard_group_size
self.group_size = group_size
self.svd_rank = svd_rank
self.svd_steps = svd_steps
self.use_quantized_matmul = use_quantized_matmul
self.re_quantize_for_matmul = re_quantize_for_matmul
self.use_stochastic_rounding = use_stochastic_rounding
self.use_hadamard = use_hadamard
self.layer_class_name = layer_class_name
self.num_bits = dtype_dict[weights_dtype]["num_bits"]
self.is_packed = dtype_dict[weights_dtype]["is_packed"]
self.is_integer = dtype_dict[weights_dtype]["is_integer"]
self.is_unsigned = dtype_dict[weights_dtype]["is_unsigned"]
self.num_bits_matmul = dtype_dict[quantized_matmul_dtype]["num_bits"]
self.is_packed_matmul = dtype_dict[quantized_matmul_dtype]["is_packed"]
self.is_integer_matmul = dtype_dict[quantized_matmul_dtype]["is_integer"]
self.is_unsigned_matmul = dtype_dict[quantized_matmul_dtype]["is_unsigned"]
@devices.inference_context()
def re_quantize_matmul(
self,
weight: torch.Tensor,
scale: torch.FloatTensor,
zero_point: torch.FloatTensor | None = None,
svd_up: torch.FloatTensor | None = None,
svd_down: torch.FloatTensor | None = None,
hadamard: torch.FloatTensor | None = None,
non_hadamard: bool = True,
skip_compile: bool = False,
) -> tuple[torch.Tensor, torch.FloatTensor]: # pylint: disable=unused-argument
if hadamard is None and self.use_hadamard and not non_hadamard:
hadamard = get_hadamard(self.hadamard_group_size, dtype=self.result_dtype, device=weight.device)
if skip_compile:
re_quantize_matmul_func = re_quantize_matmul
else:
re_quantize_matmul_func = re_quantize_matmul_compiled
if not is_fp8_compile_supported and weight.dtype == torch.float8_e4m3fn:
weight = weight.to(dtype=scale.dtype)
return re_quantize_matmul_func(
self.weights_dtype,
weight,
scale,
zero_point=zero_point,
svd_up=svd_up,
svd_down=svd_down,
hadamard=hadamard,
matmul_dtype=self.quantized_matmul_dtype,
result_shape=self.result_shape,
quantized_weight_shape=self.quantized_weight_shape,
)
@devices.inference_context()
def __call__(
self,
weight: torch.Tensor,
scale: torch.FloatTensor,
zero_point: torch.FloatTensor | None = None,
svd_up: torch.FloatTensor | None = None,
svd_down: torch.FloatTensor | None = None,
hadamard: torch.FloatTensor | None = None,
skip_quantized_matmul: bool = False,
non_hadamard: bool = False,
skip_compile: bool = False,
dtype: torch.dtype | None = None,
) -> torch.FloatTensor: # pylint: disable=unused-argument
if dtype is None:
dtype = self.result_dtype
if hadamard is None and self.use_hadamard and not non_hadamard:
hadamard = get_hadamard(self.hadamard_group_size, dtype=dtype, device=weight.device)
re_quantize_for_matmul = self.re_quantize_for_matmul or self.is_packed
if skip_compile:
dequantize_weight_func = dequantize_weight
else:
dequantize_weight_func = dequantize_weight_compiled
if not is_fp8_compile_supported and weight.dtype == torch.float8_e4m3fn:
weight = weight.to(dtype=scale.dtype)
return dequantize_weight_func(
self.weights_dtype,
weight,
scale,
zero_point=zero_point,
svd_up=svd_up,
svd_down=svd_down,
hadamard=hadamard,
dtype=dtype,
result_shape=self.result_shape,
quantized_weight_shape=self.quantized_weight_shape,
skip_quantized_matmul=skip_quantized_matmul,
re_quantize_for_matmul=re_quantize_for_matmul,
)
dequantize_asymmetric_compiled = compile_func(dequantize_asymmetric)
dequantize_symmetric_compiled = compile_func(dequantize_symmetric)
dequantize_weight_compiled = compile_func(dequantize_weight)
re_quantize_matmul_compiled = compile_func(re_quantize_matmul)
torch.serialization.add_safe_globals([SDNQDequantizer])
modules/sdnq/file_loader.py
-64
View File
@@ -1,64 +0,0 @@
import re
import concurrent.futures
import torch
def map_keys(key: str, key_mapping: dict | None) -> str:
new_key = key
if key_mapping:
for pattern, replacement in key_mapping.items():
new_key, n_replace = re.subn(pattern, replacement, new_key)
if n_replace > 0:
break
return new_key
def load_safetensors(files: list[str], state_dict: dict | None = None, key_mapping: dict | None = None, device: torch.device = "cpu") -> None:
from safetensors.torch import safe_open
if state_dict is None:
state_dict = {}
for fn in files:
with safe_open(fn, framework="pt", device=str(device)) as f:
for key in f.keys(): # safe_open exposes keys() but is not iterable
state_dict[map_keys(key, key_mapping)] = f.get_tensor(key)
def load_threaded(files: list[str], state_dict: dict | None = None, key_mapping: dict | None = None, device: torch.device = "cpu") -> None:
future_items = {}
if state_dict is None:
state_dict = {}
with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor:
for fn in files:
future_items[executor.submit(load_safetensors, [fn], key_mapping=key_mapping, device=device, state_dict=state_dict)] = fn
for future in concurrent.futures.as_completed(future_items):
future.result()
def load_streamer(files: list[str], state_dict: dict | None = None, key_mapping: dict | None = None, device: torch.device = "cpu") -> None:
# requires pip install runai_model_streamer
from runai_model_streamer import SafetensorsStreamer
if state_dict is None:
state_dict = {}
with SafetensorsStreamer() as streamer:
streamer.stream_files(files)
for key, tensor in streamer.get_tensors():
state_dict[map_keys(key, key_mapping)] = tensor.to(device)
def load_files(files: list[str], state_dict: dict | None = None, key_mapping: dict | None = None, device: torch.device = "cpu", method: str | None = None) -> dict:
# note: files is list-of-files within a module for chunked loading, not across model
if isinstance(files, str):
files = [files]
if method is None:
method = "safetensors"
if state_dict is None:
state_dict = {}
if method == "safetensors":
load_safetensors(files, state_dict=state_dict, key_mapping=key_mapping, device=device)
elif method == "threaded":
load_threaded(files, state_dict=state_dict, key_mapping=key_mapping, device=device)
elif method == "streamer":
load_streamer(files, state_dict=state_dict, key_mapping=key_mapping, device=device)
else:
raise ValueError(f"Unsupported loading method: {method}")
return state_dict
modules/sdnq/forward.py
-59
View File
@@ -1,59 +0,0 @@
# pylint: disable=protected-access
from collections.abc import Callable
from .common import dtype_dict, embedding_types, conv_types, conv_transpose_types
def get_forward_func(layer_class_name: str, quantized_matmul_dtype: str, use_quantized_matmul: bool) -> Callable: # pylint: disable=inconsistent-return-statements
if layer_class_name in embedding_types:
from .layers.embedding.forward import quantized_embedding_forward
return quantized_embedding_forward
elif layer_class_name in conv_types:
if use_quantized_matmul:
if dtype_dict[quantized_matmul_dtype]["is_integer"]:
if dtype_dict[quantized_matmul_dtype]["is_unsigned"]:
from .layers.conv.conv_uint8 import quantized_conv_forward_uint8_matmul
return quantized_conv_forward_uint8_matmul
else:
from .layers.conv.conv_int8 import quantized_conv_forward_int8_matmul
return quantized_conv_forward_int8_matmul
else:
if dtype_dict[quantized_matmul_dtype]["num_bits"] == 8:
from .layers.conv.conv_fp8 import quantized_conv_forward_fp8_matmul
return quantized_conv_forward_fp8_matmul
else:
from .layers.conv.conv_fp16 import quantized_conv_forward_fp16_matmul
return quantized_conv_forward_fp16_matmul
else:
from .layers.conv.forward import quantized_conv_forward
return quantized_conv_forward
elif layer_class_name in conv_transpose_types:
if layer_class_name.endswith("1d"):
from .layers.conv.forward import quantized_conv_transpose_1d_forward
return quantized_conv_transpose_1d_forward
elif layer_class_name.endswith("2d"):
from .layers.conv.forward import quantized_conv_transpose_2d_forward
return quantized_conv_transpose_2d_forward
elif layer_class_name.endswith("3d"):
from .layers.conv.forward import quantized_conv_transpose_3d_forward
return quantized_conv_transpose_3d_forward
else:
if use_quantized_matmul:
if dtype_dict[quantized_matmul_dtype]["is_integer"]:
if dtype_dict[quantized_matmul_dtype]["is_unsigned"]:
from .layers.linear.linear_uint8 import quantized_linear_forward_uint8_matmul
return quantized_linear_forward_uint8_matmul
else:
from .layers.linear.linear_int8 import quantized_linear_forward_int8_matmul
return quantized_linear_forward_int8_matmul
else:
if dtype_dict[quantized_matmul_dtype]["num_bits"] == 8:
from .layers.linear.linear_fp8 import quantized_linear_forward_fp8_matmul
return quantized_linear_forward_fp8_matmul
else:
from .layers.linear.linear_fp16 import quantized_linear_forward_fp16_matmul
return quantized_linear_forward_fp16_matmul
else:
from .layers.linear.forward import quantized_linear_forward
return quantized_linear_forward
modules/sdnq/kernel_wrappers.py
-215
View File
@@ -1,215 +0,0 @@
# pylint: disable=protected-access
import os
import sys
import torch
from modules import devices, shared
from .common import compile_func
if os.environ.get("SDNQ_ALLOW_FP8_MM", None) is None:
if devices.backend == "cuda":
is_fp8_mm_supported = bool(torch.cuda.get_device_capability(devices.device) >= (8,9))
elif devices.backend == "rocm":
gfx_version = devices.get_hip_agent().gfx_version
is_fp8_mm_supported = bool(gfx_version >= 0x1200 or (gfx_version >= 0x940 and gfx_version < 0x1000))
else:
is_fp8_mm_supported = False
else:
is_fp8_mm_supported = os.environ.get("SDNQ_ALLOW_FP8_MM", "0").lower() not in {"0", "false", "no"}
if os.environ.get("SDNQ_ALLOW_FP8_COMPILE", None) is None:
if devices.backend == "cuda" and "linux" in sys.platform:
is_fp8_compile_supported = bool(torch.cuda.get_device_capability(devices.device) >= (8,9)) # triton has no e4m3 conversions before sm_89
else:
is_fp8_compile_supported = True
else:
is_fp8_compile_supported = bool(os.environ.get("SDNQ_ALLOW_FP8_COMPILE", "0").lower() not in {"0", "false", "no"})
if devices.backend == "rocm":
gfx_version = devices.get_hip_agent().gfx_version
is_rdna2_and_older = bool(gfx_version < 0x940 or (gfx_version < 0x1100 and gfx_version >= 0x1000))
else:
is_rdna2_and_older = False
if devices.backend in {"ipex", "xpu"}:
is_alchemist_or_igpu = bool(not torch.xpu.get_device_capability(devices.device).get("has_subgroup_2d_block_io", False))
else:
is_alchemist_or_igpu = False
if os.environ.get("SDNQ_USE_TRITON_MM", None) is None:
use_triton_mm = bool(not is_alchemist_or_igpu and (devices.backend in {"cuda", "rocm", "ipex", "xpu", "zluda"}))
else:
use_triton_mm = bool(os.environ.get("SDNQ_USE_TRITON_MM", "0").lower() not in {"0", "false", "no"})
if os.environ.get("SDNQ_USE_TENSORWISE_FP8_MM", None) is None:
# row-wise FP8 only exist on H100 hardware, sdnq will use software row-wise with tensorwise hardware with this setting
use_tensorwise_fp8_matmul = bool(devices.backend != "cuda" or (devices.backend == "cuda" and torch.cuda.get_device_capability(devices.device) < (9,0)))
else:
use_tensorwise_fp8_matmul = bool(os.environ.get("SDNQ_USE_TENSORWISE_FP8_MM", "0").lower() not in {"0", "false", "no"})
use_openvino_mm = bool(os.environ.get("SDNQ_USE_OPENVINO_MM", "1").lower() not in {"0", "false", "no"})
use_triton_scaled_mm = bool(use_triton_mm and os.environ.get("SDNQ_USE_TRITON_SCALED_MM", "1").lower() not in {"0", "false", "no"})
if use_openvino_mm:
try:
from .kernels.openvino_mm import openvino_int_mm, openvino_fp_mm
except Exception as e:
use_openvino_mm = False
openvino_int_mm = None
openvino_fp_mm = None
shared.log.warning(f"SDNQ: OpenVINO MM kernels are not available! Falling back to PyTorch Eager kernels for CPU device. Error message: {e}")
else:
openvino_int_mm = None
openvino_fp_mm = None
if use_triton_mm:
try:
from .kernels.triton_mm import sdnq_triton_mm
if is_fp8_mm_supported:
use_tensorwise_fp8_matmul = True
except Exception as e:
use_triton_mm = False
sdnq_triton_mm = None
shared.log.warning(f"SDNQ: Triton MM kernels are not available! Falling back to PyTorch Eager kernels. Error message: {e}")
else:
sdnq_triton_mm = None
if use_triton_scaled_mm:
try:
from .kernels.triton_scaled_mm import sdnq_scaled_mm
except Exception as e:
use_triton_scaled_mm = False
sdnq_scaled_mm = None
shared.log.warning(f"SDNQ: Triton Scaled MM kernels are not available! Falling back to PyTorch Eager kernels. Error message: {e}")
else:
sdnq_scaled_mm = None
if os.environ.get("SDNQ_INCLUDE_MM_KERNEL_IN_COMPILE", None) is None:
include_mm_kernel_in_compile = bool(not use_triton_scaled_mm)
else:
include_mm_kernel_in_compile = bool(os.environ.get("SDNQ_INCLUDE_MM_KERNEL_IN_COMPILE", "0").lower() not in {"0", "false", "no"})
if os.environ.get("SDNQ_USE_CONTIGUOUS_MM", None) is None:
use_contiguous_int8_mm = bool(is_rdna2_and_older or devices.backend in {"ipex", "xpu", "cpu", "mps", "openvino", "zluda"})
use_contiguous_fp16_mm = bool(use_contiguous_int8_mm or devices.backend == "rocm")
use_contiguous_fp8_mm = use_contiguous_fp16_mm and (is_fp8_mm_supported and use_triton_mm)
else:
use_contiguous_int8_mm = bool(os.environ.get("SDNQ_USE_CONTIGUOUS_MM", "0").lower() not in {"0", "false", "no"})
use_contiguous_fp16_mm = use_contiguous_int8_mm
use_contiguous_fp8_mm = use_contiguous_fp16_mm and (is_fp8_mm_supported and use_triton_mm)
def int_mm_torch(a: torch.Tensor, b: torch.Tensor, out_dtype: torch.dtype = torch.int32) -> torch.FloatTensor:
return torch._int_mm(a,b).to(dtype=out_dtype)
def fp8_mm_torch(a: torch.Tensor, b: torch.Tensor, out_dtype: torch.dtype = torch.float32) -> torch.FloatTensor:
dummy_input_scale = torch.ones(1, device=a.device, dtype=torch.float32)
return torch._scaled_mm(a, b, scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=out_dtype)
def fp_mm_torch(a: torch.Tensor, b: torch.Tensor, out_dtype: torch.dtype = torch.float32) -> torch.FloatTensor:
if b.dtype == torch.float8_e4m3fn:
fp16_scale = 4 * b.shape[-2]
else:
fp16_scale = 65536 * b.shape[-2]
in_scale = fp16_scale**0.5
a = a.to(dtype=torch.float32).div_(in_scale).to(dtype=torch.float16)
b = b.to(dtype=torch.float32).div_(in_scale).to(dtype=torch.float16)
return torch.mm(a,b).to(dtype=torch.float32).mul_(fp16_scale).to(dtype=out_dtype)
def int_scaled_mm_torch(a: torch.Tensor, b: torch.Tensor, scale_a: torch.Tensor, scale_b: torch.Tensor, bias: torch.FloatTensor | None = None, out_dtype: torch.dtype = torch.float32) -> torch.FloatTensor:
if bias is None:
return int_mm_func(a,b, out_dtype=scale_a.dtype).mul_(scale_a).mul_(scale_b).to(dtype=out_dtype)
else:
return torch.addcmul(bias, int_mm_func(a,b, out_dtype=scale_a.dtype).mul_(scale_a), scale_b).to(dtype=out_dtype)
if use_tensorwise_fp8_matmul or not is_fp8_mm_supported:
def fp8_scaled_mm_torch(a: torch.Tensor, b: torch.Tensor, scale_a: torch.Tensor, scale_b: torch.Tensor, bias: torch.FloatTensor | None = None, out_dtype: torch.dtype = torch.float32) -> torch.FloatTensor:
if bias is None:
return fp8_mm_func(a,b, out_dtype=scale_a.dtype).mul_(scale_a).mul_(scale_b).to(dtype=out_dtype)
else:
return torch.addcmul(bias, fp8_mm_func(a,b, out_dtype=scale_a.dtype).mul_(scale_a), scale_b).to(dtype=out_dtype)
else:
def fp8_scaled_mm_torch(a: torch.Tensor, b: torch.Tensor, scale_a: torch.Tensor, scale_b: torch.Tensor, bias: torch.FloatTensor | None = None, out_dtype: torch.dtype = torch.float32) -> torch.FloatTensor:
if bias is not None and bias.ndim != 1:
return torch._scaled_mm(a, b, scale_a=scale_a, scale_b=scale_b, bias=None, out_dtype=out_dtype).add_(bias)
else:
return torch._scaled_mm(a, b, scale_a=scale_a, scale_b=scale_b, bias=bias.to(dtype=out_dtype) if bias is not None else None, out_dtype=out_dtype)
def fp_scaled_mm_torch(a: torch.Tensor, b: torch.Tensor, scale_a: torch.Tensor, scale_b: torch.Tensor, bias: torch.FloatTensor | None = None, out_dtype: torch.dtype = torch.float32) -> torch.FloatTensor:
if bias is None:
return fp_mm_func(a,b, out_dtype=scale_a.dtype).mul_(scale_a).mul_(scale_b).to(dtype=out_dtype)
else:
return torch.addcmul(bias, fp_mm_func(a,b, out_dtype=scale_a.dtype).mul_(scale_a), scale_b).to(dtype=out_dtype)
def int_mm_func(a: torch.Tensor, b: torch.Tensor, out_dtype: torch.dtype = torch.int32) -> torch.FloatTensor:
if sdnq_triton_mm is not None and a.device.type in {"cuda", "xpu"}:
return sdnq_triton_mm(a, b, out_dtype=out_dtype)
elif openvino_int_mm is not None and a.device.type == "cpu":
return openvino_int_mm(a, b, out_dtype=out_dtype)
else:
return int_mm_torch(a, b, out_dtype=out_dtype)
def fp8_mm_func(a: torch.Tensor, b: torch.Tensor, out_dtype: torch.dtype = torch.float32) -> torch.FloatTensor:
if is_fp8_mm_supported:
if sdnq_triton_mm is not None and a.device.type in {"cuda", "xpu"}:
return sdnq_triton_mm(a, b, out_dtype=out_dtype)
elif openvino_fp_mm is not None and a.device.type == "cpu":
return openvino_fp_mm(a, b, out_dtype=out_dtype)
else:
return fp8_mm_torch(a, b, out_dtype=out_dtype)
else:
if openvino_fp_mm is not None and a.device.type == "cpu":
return openvino_fp_mm(a, b, out_dtype=out_dtype)
else:
return fp_mm_torch(a, b, out_dtype=out_dtype)
def fp_mm_func(a: torch.Tensor, b: torch.Tensor, out_dtype: torch.dtype = torch.float32) -> torch.FloatTensor:
if sdnq_triton_mm is not None and a.device.type in {"cuda", "xpu"}:
return sdnq_triton_mm(a, b, out_dtype=out_dtype)
elif openvino_fp_mm is not None and a.device.type == "cpu":
return openvino_fp_mm(a, b, out_dtype=out_dtype)
else:
return fp_mm_torch(a, b, out_dtype=out_dtype)
def int_scaled_mm_func(a: torch.Tensor, b: torch.Tensor, scale_a: torch.Tensor, scale_b: torch.Tensor, bias: torch.FloatTensor | None = None, out_dtype: torch.dtype = torch.float32) -> torch.FloatTensor:
if sdnq_scaled_mm is not None and a.device.type in {"cuda", "xpu"}:
return sdnq_scaled_mm(a, b, scale_a, scale_b, bias=bias, out_dtype=out_dtype)
else:
return int_scaled_mm_torch(a, b, scale_a, scale_b, bias=bias, out_dtype=out_dtype)
def fp8_scaled_mm_func(a: torch.Tensor, b: torch.Tensor, scale_a: torch.Tensor, scale_b: torch.Tensor, bias: torch.FloatTensor | None = None, out_dtype: torch.dtype = torch.float32) -> torch.FloatTensor:
if is_fp8_mm_supported and sdnq_scaled_mm is not None and a.device.type in {"cuda", "xpu"}:
return sdnq_scaled_mm(a, b, scale_a, scale_b, bias=bias, out_dtype=out_dtype)
else:
return fp8_scaled_mm_torch(a, b, scale_a, scale_b, bias=bias, out_dtype=out_dtype)
def fp_scaled_mm_func(a: torch.Tensor, b: torch.Tensor, scale_a: torch.Tensor, scale_b: torch.Tensor, bias: torch.FloatTensor | None = None, out_dtype: torch.dtype = torch.float32) -> torch.FloatTensor:
if sdnq_scaled_mm is not None and a.device.type in {"cuda", "xpu"}:
return sdnq_scaled_mm(a, b, scale_a, scale_b, bias=bias, out_dtype=out_dtype)
else:
return fp_scaled_mm_torch(a, b, scale_a, scale_b, bias=bias, out_dtype=out_dtype)
int_mm_torch = compile_func(int_mm_torch)
fp8_mm_torch = compile_func(fp8_mm_torch)
fp_mm_torch = compile_func(fp_mm_torch)
int_scaled_mm_torch = compile_func(int_scaled_mm_torch)
fp8_scaled_mm_torch = compile_func(fp8_scaled_mm_torch)
fp_scaled_mm_torch = compile_func(fp_scaled_mm_torch)
modules/sdnq/kernels/openvino_mm.py
-155
View File
@@ -1,155 +0,0 @@
import os
import torch
import openvino as ov
from openvino import opset16 as ov_ops
OV_CORE = None
OV_DEVICE: str = None
OV_COMPILED_CACHE: dict[tuple[str, tuple[int,int] | None, str, tuple[int,int] | None], tuple[ov.InferRequest, str]] = {}
def get_ov_core():
global OV_CORE, OV_DEVICE # pylint: disable=global-statement
if OV_CORE is not None:
return OV_CORE, OV_DEVICE
from openvino.properties import hint as ov_hints
OV_CORE = ov.Core()
OV_DEVICE = os.environ.get("SDNQ_OPENVINO_DEVICE", "HETERO:NPU,CPU" if "NPU" in OV_CORE.get_available_devices() else "CPU")
if OV_DEVICE == "NPU":
OV_DEVICE = "HETERO:NPU,CPU"
for ov_device in OV_DEVICE.removeprefix("HETERO:").split(","):
if ov_device != "NPU":
OV_CORE.set_property(ov_device, {ov_hints.execution_mode: ov_hints.ExecutionMode.ACCURACY})
return OV_CORE, OV_DEVICE
def ov_mm(infer_request: ov.InferRequest, out_name: str, A: torch.Tensor, B: torch.Tensor, out_dtype: torch.dtype = torch.float32) -> torch.Tensor:
C = torch.empty((A.shape[0], B.shape[-1]), device="cpu", dtype=torch.float32)
A, B = A.contiguous(), B.contiguous()
infer_request.set_tensor("A", ov.Tensor(A.detach().to("cpu").numpy(), shared_memory=True))
infer_request.set_tensor("B", ov.Tensor(B.detach().to("cpu").numpy(), shared_memory=True))
infer_request.set_tensor(out_name, ov.Tensor(C.numpy(), shared_memory=True))
infer_request.infer()
C = C.to(A.device, dtype=out_dtype)
return C
@torch.library.custom_op("sdnq::openvino_int_mm", mutates_args=())
def openvino_int_mm(Tensor_A: torch.Tensor, Tensor_B: torch.Tensor, out_dtype: torch.dtype = torch.float32) -> torch.Tensor:
ov_core, ov_device = get_ov_core()
if "GPU" not in ov_device:
cache_key = (ov_device, "int8", Tensor_A.shape, Tensor_B.shape)
else:
cache_key = (ov_device, "int8", None, None)
infer_request, out_name = OV_COMPILED_CACHE.get(cache_key, (None, None))
if infer_request is not None:
return ov_mm(infer_request, out_name, Tensor_A, Tensor_B, out_dtype=out_dtype)
if "GPU" not in ov_device:
shape_a = ov.Shape(Tensor_A.shape)
shape_b = ov.Shape(Tensor_B.shape)
else:
shape_a = ov.PartialShape([-1,-1])
shape_b = ov.PartialShape([-1,-1])
input_a = ov_ops.parameter(shape_a, ov.Type.i8, name="A")
input_b = ov_ops.parameter(shape_b, ov.Type.i8, name="B")
a = ov_ops.convert(input_a, ov.Type.f32)
b = ov_ops.convert(input_b, ov.Type.f32)
low = ov_ops.constant(-128.0, dtype=ov.Type.f32)
high = ov_ops.constant(127.0, dtype=ov.Type.f32)
a = ov_ops.fake_quantize(a, low, high, low, high, 256)
b = ov_ops.fake_quantize(b, low, high, low, high, 256)
# NPU uses FP16 x INT8 -> FP16 instead of INT8 x INT8 -> INT32 and FP16 output overflows
if "NPU" in ov_device:
fp16_scale = 0.25012213 * Tensor_B.shape[-2]
in_scale = ov_ops.constant(fp16_scale ** 0.5, dtype=ov.Type.f32)
out_scale = ov_ops.constant(fp16_scale, dtype=ov.Type.f32, name="out_scale_const")
a = ov_ops.divide(a, in_scale)
b = ov_ops.divide(b, in_scale)
out = ov_ops.matmul(a, b, False, False)
out = ov_ops.multiply(out, out_scale, name="out_scale")
else:
out = ov_ops.matmul(a, b, False, False)
ov_model = ov.Model([out], [input_a, input_b], "ov_int8_mm")
if "NPU" in ov_device: # NPU can't use FP32 for regular multiplications
for node in ov_model.get_ops():
if node.get_friendly_name() in {"out_scale", "out_scale_const"}:
node.get_rt_info()["affinity"] = "CPU"
else:
node.get_rt_info()["affinity"] = "NPU"
ov_model = ov_core.compile_model(ov_model, ov_device)
infer_request = ov_model.create_infer_request()
out_name = ov_model.outputs[0]
OV_COMPILED_CACHE[cache_key] = (infer_request, out_name)
return ov_mm(infer_request, out_name, Tensor_A, Tensor_B, out_dtype=out_dtype)
@openvino_int_mm.register_fake
def openvino_int_mm_fake(A: torch.Tensor, B: torch.Tensor, out_dtype: torch.dtype = torch.float32) -> torch.Tensor:
return torch.mm(A.to(dtype=torch.float32), B.to(dtype=torch.float32)).to(dtype=out_dtype)
@torch.library.custom_op("sdnq::openvino_fp_mm", mutates_args=())
def openvino_fp_mm(Tensor_A: torch.Tensor, Tensor_B: torch.Tensor, out_dtype: torch.dtype = torch.float32) -> torch.Tensor:
ov_core, ov_device = get_ov_core()
mm_dtype = "fp16" if Tensor_B.dtype == torch.float16 else "fp8"
if mm_dtype == "fp8":
Tensor_A = Tensor_A.to(dtype=torch.float16)
Tensor_B = Tensor_B.to(dtype=torch.float16)
if "GPU" not in ov_device:
cache_key = (ov_device, mm_dtype, Tensor_A.shape, Tensor_B.shape)
else:
cache_key = (ov_device, mm_dtype, None, None)
infer_request, out_name = OV_COMPILED_CACHE.get(cache_key, (None, None))
if infer_request is not None:
return ov_mm(infer_request, out_name, Tensor_A, Tensor_B, out_dtype=out_dtype)
if "GPU" not in ov_device:
shape_a = ov.Shape(Tensor_A.shape)
shape_b = ov.Shape(Tensor_B.shape)
else:
shape_a = ov.PartialShape([-1,-1])
shape_b = ov.PartialShape([-1,-1])
input_a = ov_ops.parameter(shape_a, ov.Type.f16, name="A")
input_b = ov_ops.parameter(shape_b, ov.Type.f16, name="B")
a = ov_ops.convert(input_a, ov.Type.f32)
b = ov_ops.convert(input_b, ov.Type.f32)
if mm_dtype == "fp8":
low = ov_ops.constant(-448.0, dtype=ov.Type.f32)
high = ov_ops.constant(448.0, dtype=ov.Type.f32)
a = ov_ops.fake_quantize(a, low, high, low, high, 256)
b = ov_ops.fake_quantize(b, low, high, low, high, 256)
fp16_scale = 4 * Tensor_B.shape[-2]
else:
fp16_scale = 65536 * Tensor_B.shape[-2]
in_scale = ov_ops.constant(fp16_scale**0.5, dtype=ov.Type.f32)
out_scale = ov_ops.constant(fp16_scale, dtype=ov.Type.f32, name="out_scale_const")
a = ov_ops.convert(ov_ops.divide(a, in_scale), ov.Type.f16)
b = ov_ops.convert(ov_ops.divide(b, in_scale), ov.Type.f16)
out = ov_ops.matmul(a, b, False, False, name="fp_mm")
out = ov_ops.multiply(ov_ops.convert(out, ov.Type.f32), out_scale, name="out_scale")
ov_model = ov.Model([out], [input_a, input_b], "ov_fp_mm")
if "NPU" in ov_device: # NPU can't use FP32 for regular multiplications
for node in ov_model.get_ops():
if node.get_friendly_name() in {"out_scale", "out_scale_const"}:
node.get_rt_info()["affinity"] = "CPU"
else:
node.get_rt_info()["affinity"] = "NPU"
ov_model = ov_core.compile_model(ov_model, ov_device)
infer_request = ov_model.create_infer_request()
out_name = ov_model.outputs[0]
OV_COMPILED_CACHE[cache_key] = (infer_request, out_name)
return ov_mm(infer_request, out_name, Tensor_A, Tensor_B, out_dtype=out_dtype)
@openvino_fp_mm.register_fake
def openvino_fp_mm_fake(A: torch.Tensor, B: torch.Tensor, out_dtype: torch.dtype = torch.float32) -> torch.Tensor:
return torch.mm(A.to(dtype=torch.float32), B.to(dtype=torch.float32)).to(dtype=out_dtype)
modules/sdnq/kernels/triton_atten.py
-355
View File
@@ -1,355 +0,0 @@
import os
import math
import torch
import triton
import triton.language as tl
from ..common import compile_func # pylint: disable=relative-beyond-top-level
from ..quant_utils import quantize_int_mm, quantize_fp_mm, get_hadamard, get_hadamard_group_size, apply_hadamard # pylint: disable=relative-beyond-top-level
from ..utils import is_pow2, next_power_of_2 # pylint: disable=relative-beyond-top-level
min_block_size = int(os.environ.get("SDNQ_TRITON_ATTEN_MIN_BLOCK_SIZE", "256"))
matmul_configs = [
triton.Config({"BLOCK_SIZE_M": BM, "BLOCK_SIZE_N": BN}, num_warps=w, num_stages=s)
for BM in [int(BM) for BM in os.environ.get("SDNQ_TRITON_ATTEN_BLOCK_SIZE_M_LIST", "64,128").replace(" ","").split(",")]
for BN in [int(BN) for BN in os.environ.get("SDNQ_TRITON_ATTEN_BLOCK_SIZE_N_LIST", "32,64").replace(" ","").split(",")]
for w in [int(w) for w in os.environ.get("SDNQ_TRITON_ATTEN_NUM_WARPS_LIST", "8,16" if torch.xpu.is_available() else "4,8").replace(" ","").split(",")]
for s in [int(s) for s in os.environ.get("SDNQ_TRITON_ATTEN_NUM_STAGES_LIST", "1" if (torch.cuda.is_available() and torch.version.hip) else "2").replace(" ","").split(",")]
]
@triton.autotune(
configs=matmul_configs,
key=[
"is_causal", "do_mask",
"QZ", "QH", "QN_AT", "QHD",
"KZ", "KH", "KN_AT", "KHD",
"VZ", "VH", "VN_AT", "VHD",
"qk_is_quantized",
"pv_is_quantized",
"q_dtype", "v_dtype",
"out_dtype", "mask_dtype",
],
cache_results=True,
)
@triton.jit
def sdnq_attn_kernel(
q_ptr, k_ptr, v_ptr,
q_scale_ptr, k_scale_ptr, v_scale_ptr,
out_ptr, mask_ptr,
is_causal: tl.constexpr,
do_mask: tl.constexpr,
QZ: tl.constexpr, QH: tl.constexpr, QN: tl.constexpr, QHD: tl.constexpr,
KZ: tl.constexpr, KH: tl.constexpr, KN: tl.constexpr, KHD: tl.constexpr,
VZ: tl.constexpr, VH: tl.constexpr, VN: tl.constexpr, VHD: tl.constexpr,
OZ: tl.constexpr, OH: tl.constexpr, ON: tl.constexpr, OHD: tl.constexpr,
MZ: tl.constexpr, MH: tl.constexpr, MQN: tl.constexpr, MKN: tl.constexpr,
QN_AT: tl.constexpr, # pylint: disable=unused-argument
KN_AT: tl.constexpr, # pylint: disable=unused-argument
VN_AT: tl.constexpr, # pylint: disable=unused-argument
qk_is_quantized: tl.constexpr,
pv_is_quantized: tl.constexpr,
q_dtype: tl.constexpr, # pylint: disable=unused-argument
v_dtype: tl.constexpr, # pylint: disable=unused-argument
out_dtype: tl.constexpr, # pylint: disable=unused-argument
mask_dtype: tl.constexpr, # pylint: disable=unused-argument
BLOCK_SIZE_M: tl.constexpr,
BLOCK_SIZE_N: tl.constexpr,
) -> None:
start_m = tl.program_id(0)
off_h = tl.program_id(1)
off_z = tl.program_id(2)
tl.assume(QZ > 0)
tl.assume(QH > 0)
tl.assume(QN > 0)
tl.assume(QHD > 0)
tl.assume(KZ > 0)
tl.assume(KH > 0)
tl.assume(KN > 0)
tl.assume(KHD > 0)
tl.assume(VZ > 0)
tl.assume(VH > 0)
tl.assume(VN > 0)
tl.assume(VHD > 0)
tl.assume(OZ > 0)
tl.assume(OH > 0)
tl.assume(ON > 0)
tl.assume(OHD > 0)
tl.assume(MZ >= 0)
tl.assume(MH >= 0)
tl.assume(MQN >= 0)
tl.assume(MKN >= 0)
tl.assume(off_h >= 0)
tl.assume(off_z >= 0)
tl.assume(start_m >= 0)
tl.assume(BLOCK_SIZE_M > 0)
tl.assume(BLOCK_SIZE_N > 0)
tl.assume(do_mask == 0 or do_mask == 1) # pylint: disable=consider-using-in
tl.assume(is_causal == 0 or is_causal == 1) # pylint: disable=consider-using-in
tl.assume(qk_is_quantized == 0 or qk_is_quantized == 1) # pylint: disable=consider-using-in
tl.assume(pv_is_quantized == 0 or pv_is_quantized == 1) # pylint: disable=consider-using-in
do_k_mask = KN % BLOCK_SIZE_N != 0
start_m_block = start_m * BLOCK_SIZE_M
offs_m = start_m_block + tl.arange(0, BLOCK_SIZE_M)
offs_n = tl.arange(0, BLOCK_SIZE_N)
offset_y = off_z * (QN * QH) + off_h * QN
offset_y_k = off_z * (KN * KH) + ((off_h * KH) // QH) * KN
offset_y_v = off_z * (VN * VH) + ((off_h * VH) // QH) * VN
q_desc = tl.make_tensor_descriptor(q_ptr + offset_y * QHD, shape=[QN, QHD], strides=[QHD, 1], block_shape=[BLOCK_SIZE_M, QHD])
k_desc = tl.make_tensor_descriptor(k_ptr + offset_y_k * KHD, shape=[KN, KHD], strides=[KHD, 1], block_shape=[BLOCK_SIZE_N, KHD])
v_desc = tl.make_tensor_descriptor(v_ptr + offset_y_v * VHD, shape=[VN, VHD], strides=[VHD, 1], block_shape=[BLOCK_SIZE_N, VHD])
if qk_is_quantized:
q_scale_desc = tl.make_tensor_descriptor(q_scale_ptr + offset_y, shape=[QN], strides=[1,], block_shape=[BLOCK_SIZE_M])
k_scale_desc = tl.make_tensor_descriptor(k_scale_ptr + offset_y_k, shape=[KN], strides=[1,], block_shape=[BLOCK_SIZE_N])
q_scale = q_scale_desc.load([start_m_block])[:, None]
if pv_is_quantized:
v_scale_desc = tl.make_tensor_descriptor(v_scale_ptr + offset_y_v, shape=[VN], strides=[1,], block_shape=[BLOCK_SIZE_N])
if do_mask:
mask_desc = tl.make_tensor_descriptor(mask_ptr + offset_y * MKN, shape=[MQN, MKN], strides=[MKN, 1], block_shape=[BLOCK_SIZE_M, BLOCK_SIZE_N])
m_i = tl.full([BLOCK_SIZE_M], float("-inf"), dtype=tl.float32)
l_i = tl.full([BLOCK_SIZE_M], 1.0, dtype=tl.float32)
acc = tl.zeros([BLOCK_SIZE_M, VHD], dtype=tl.float32)
q = q_desc.load([start_m_block, 0])
for start_n_idx in tl.range(0, tl.cdiv(KN, BLOCK_SIZE_N)):
start_n = start_n_idx * BLOCK_SIZE_N
skip = False
if is_causal and ((start_m_block + BLOCK_SIZE_M) <= start_n):
skip = True
if do_mask:
mask = mask_desc.load([start_m_block, start_n])
mask_max = tl.max(mask)
if mask.dtype == tl.int8:
mask = mask.to(tl.int1)
else:
mask = mask.to(tl.float32)
if mask.dtype == tl.int1:
if mask_max == 0:
skip = True
elif mask_max == float("-inf"):
skip = True
if not skip:
k = k_desc.load([start_n, 0]).T
if qk_is_quantized:
k_scale = k_scale_desc.load([start_n])[None, :]
if q.dtype == tl.int8:
qk = tl.mul(tl.mul(tl.dot(q, k, out_dtype=tl.int32).to(tl.float32), q_scale), k_scale)
else:
qk = tl.mul(tl.mul(tl.dot(q, k, out_dtype=tl.float32), q_scale), k_scale)
else:
qk = tl.dot(q, k, out_dtype=tl.float32)
if is_causal:
qk = tl.where((offs_m[:, None] >= (start_n + offs_n[None, :])) & (offs_m[:, None] < QN), qk, float("-inf"))
if do_mask:
if mask.dtype == tl.int1:
qk = tl.where(mask, qk, float("-inf"))
else:
qk += mask
if do_k_mask and (start_n + BLOCK_SIZE_N) > KN:
qk = tl.where(offs_n[None, :] < (KN - start_n), qk, float("-inf"))
m_ij = tl.maximum(m_i, tl.max(qk, 1))
if do_mask:
alpha = tl.exp2(tl.where((m_i == float("-inf")) & (m_ij == float("-inf")), 0.0, (m_i - m_ij)))
qk -= tl.where(m_ij == float("-inf"), 0.0, m_ij)[:, None]
else:
alpha = tl.exp2(m_i - m_ij)
qk = qk - m_ij[:, None]
p = tl.exp2(qk)
l_i = tl.fma(l_i, alpha, tl.sum(p, 1))
acc *= alpha[:, None]
v = v_desc.load([start_n, 0])
if pv_is_quantized:
v_scale = v_scale_desc.load([start_n])[None, :]
p *= v_scale
if v.dtype == tl.int8:
p_scale = tl.mul(tl.max(p, 1)[:, None], (1 / 127.0))
p_scale = tl.where(p_scale <= 2e-38, 1.0, p_scale)
p = tl.floor(tl.fma(p, (1 / p_scale), 0.5)).to(tl.int8)
acc = tl.fma(tl.dot(p, v, out_dtype=tl.int32).to(tl.float32), p_scale, acc)
else:
p_scale = tl.mul(tl.max(p, 1)[:, None], (1 / (65504.0 if v.dtype == tl.float16 else 448.0)))
p_scale = tl.where(p_scale <= 2e-38, 1.0, p_scale)
p = tl.mul(p, (1 / p_scale)).to(v.dtype)
acc = tl.fma(tl.dot(p, v, out_dtype=tl.float32), p_scale, acc)
else:
p = p.to(v.dtype)
acc = tl.dot(p, v, acc, out_dtype=tl.float32)
m_i = m_ij
l_i = 1 / l_i[:, None]
acc *= l_i
acc = acc.to(out_ptr.type.element_ty)
out_desc = tl.make_tensor_descriptor(out_ptr + offset_y * OHD, shape=[ON, OHD], strides=[OHD, 1], block_shape=[BLOCK_SIZE_M, OHD])
out_desc.store([start_m_block, 0], acc)
def quantize_attn(
q, k, v,
scale: float | None = None,
smooth_k: bool = False,
hadamard: torch.FloatTensor | None = None,
hadamard_group_size: int = 256,
matmul_dtype: str = "int8",
pv_matmul_dtype: str | None = None,
) -> tuple[torch.Tensor]:
if matmul_dtype in {"auto", "enabled", "uint8"}:
matmul_dtype = "int8"
if pv_matmul_dtype in {"enabled", "uint8"}:
pv_matmul_dtype = "int8"
if scale is None:
scale = q.shape[-1] ** -0.5
if smooth_k:
if k.dtype != torch.float32:
k = k.to(dtype=torch.float32)
k = k.sub_(k.mean(dim=2, keepdim=True))
else:
k = k.sub(k.mean(dim=2, keepdim=True))
if matmul_dtype not in {None, "none", "no", "disabled"}:
if hadamard is not None:
q, use_hadamard, hadamard_group_size = apply_hadamard(q, group_size=hadamard_group_size, hadamard=hadamard, layer_class_name="Linear")
if use_hadamard:
k = apply_hadamard(k.to(dtype=hadamard.dtype), group_size=hadamard_group_size, hadamard=hadamard, layer_class_name="Linear")[0]
quantize_mm_func = quantize_int_mm if matmul_dtype.startswith("int") else quantize_fp_mm
q_q, q_scale = quantize_mm_func(q.contiguous().to(dtype=torch.float32), dim=-1, matmul_dtype=matmul_dtype)
k_q, k_scale = quantize_mm_func(k.contiguous().to(dtype=torch.float32), dim=-1, matmul_dtype=matmul_dtype)
q_scale = q_scale.squeeze(-1).mul_(scale * 1.4426950408889634)
k_scale = k_scale.squeeze(-1)
else:
q_q = q.contiguous().mul(scale * 1.4426950408889634)
k_q = k.contiguous().to(dtype=q.dtype)
q_scale = None
k_scale = None
if pv_matmul_dtype not in {None, "auto", "none", "no", "disabled"}:
quantize_mm_func_pv = quantize_int_mm if pv_matmul_dtype.startswith("int") else quantize_fp_mm
v_q, v_scale = quantize_mm_func_pv(v.contiguous().to(dtype=torch.float32), dim=-1, matmul_dtype=pv_matmul_dtype)
v_scale = v_scale.squeeze(-1)
else:
v_q = v.contiguous()
v_scale = None
return q_q, q_scale, k_q, k_scale, v_q, v_scale
def get_attn_inputs(
query: torch.FloatTensor,
key: torch.FloatTensor,
value: torch.FloatTensor,
hadamard: torch.FloatTensor | None = None,
attn_mask: torch.Tensor | None = None,
dropout_p: float = 0.0, # pylint: disable=unused-argument
is_causal: bool = False, # pylint: disable=unused-argument
scale: float | None = None,
enable_gqa: bool = False, # pylint: disable=unused-argument
smooth_k: bool = False,
hadamard_group_size: int = 256,
matmul_dtype: str = "int8",
pv_matmul_dtype: str | None = None,
do_quantize: bool = True,
out_dtype: torch.dtype | None = None,
) -> tuple[torch.Tensor, float, torch.dtype]:
QZ, QH, QN, QHD = query.shape
_, _, KN, KHD = key.shape
_, _, _, VHD = value.shape
if out_dtype is None:
out_dtype = query.dtype
if scale is None:
scale = QHD ** -0.5
if not is_pow2(QHD):
query = torch.nn.functional.pad(query, (0, next_power_of_2(QHD) - QHD))
key = torch.nn.functional.pad(key, (0, next_power_of_2(KHD) - KHD))
value = torch.nn.functional.pad(value, (0, next_power_of_2(VHD) - VHD))
if attn_mask is not None:
attn_mask = attn_mask.expand((QZ, QH, QN, KN))
if not is_pow2(KN):
pad_value = float("-inf") if torch.is_floating_point(attn_mask) else 0
attn_mask = torch.nn.functional.pad(attn_mask, (0, next_power_of_2(KN) - KN), value=pad_value)
if attn_mask.dtype == torch.bool:
attn_mask = attn_mask.to(dtype=torch.int8)
attn_mask = attn_mask.contiguous()
query, query_scale, key, key_scale, value, value_scale = quantize_attn(
query, key, value,
scale=scale,
smooth_k=smooth_k,
hadamard=hadamard,
hadamard_group_size=hadamard_group_size,
matmul_dtype=matmul_dtype if do_quantize else "disabled",
pv_matmul_dtype=pv_matmul_dtype if do_quantize else "disabled",
)
return query, query_scale, key, key_scale, value, value_scale, attn_mask, scale, out_dtype
def sdnq_triton_atten(
query: torch.FloatTensor,
key: torch.FloatTensor,
value: torch.FloatTensor,
attn_mask: torch.Tensor | None = None,
dropout_p: float = 0.0, # pylint: disable=unused-argument
is_causal: bool = False,
scale: float | None = None,
enable_gqa: bool = False, # pylint: disable=unused-argument
smooth_k: bool = False,
use_hadamard: bool = False,
hadamard_group_size: int = 256,
matmul_dtype: str = "int8",
pv_matmul_dtype: str | None = None,
do_quantize: bool = True,
out_dtype: torch.dtype | None = None,
) -> torch.FloatTensor:
QZ, QH, QN, QHD = query.shape
_, _, KN, KHD = key.shape
_, _, VN, VHD = value.shape
hadamard = None
if use_hadamard and do_quantize and matmul_dtype not in {None, "none", "no", "disabled"}:
hadamard_channel_size = next_power_of_2(min(QHD, KHD))
hadamard_group_size = min(hadamard_group_size, hadamard_channel_size)
use_hadamard, hadamard_group_size = get_hadamard_group_size(hadamard_channel_size, hadamard_group_size)
if use_hadamard:
hadamard = get_hadamard(hadamard_group_size, dtype=query.dtype, device=query.device)
(
query, query_scale,
key, key_scale,
value, value_scale,
attn_mask, scale, out_dtype,
) = get_attn_inputs(
query=query, key=key, value=value,
hadamard=hadamard, attn_mask=attn_mask,
dropout_p=dropout_p, is_causal=is_causal,
scale=scale, enable_gqa=enable_gqa,
smooth_k=smooth_k, hadamard_group_size=hadamard_group_size,
matmul_dtype=matmul_dtype, pv_matmul_dtype=pv_matmul_dtype,
do_quantize=do_quantize, out_dtype=out_dtype,
)
def grid(META):
return (triton.cdiv(QN, META["BLOCK_SIZE_M"]), QH, QZ)
out = torch.empty((QZ, QH, QN, value.shape[-1]), dtype=out_dtype, device=query.device)
sdnq_attn_kernel[grid](
query, key, value,
query_scale, key_scale, value_scale,
out, attn_mask,
(1 if is_causal else 0),
(1 if attn_mask is not None else 0),
*query.shape, *key.shape, *value.shape, *out.shape,
*(attn_mask.shape if attn_mask is not None else (0, 0, 0, 0)),
math.ceil(QN / min_block_size),
math.ceil(KN / min_block_size),
math.ceil(VN / min_block_size),
(1 if query_scale is not None else 0),
(1 if value_scale is not None else 0),
str(query.dtype), str(value.dtype), str(out.dtype),
str(attn_mask.dtype if attn_mask is not None else None),
)
return out[..., :VHD]
get_attn_inputs = compile_func(get_attn_inputs)
modules/sdnq/kernels/triton_mm.py
-131
View File
@@ -1,131 +0,0 @@
import os
import math
import torch
import triton
import triton.language as tl
min_block_size = int(os.environ.get("SDNQ_TRITON_MM_MIN_BLOCK_SIZE", "256"))
matmul_configs = [
triton.Config({"BLOCK_SIZE_M": BM, "BLOCK_SIZE_N": BN, "BLOCK_SIZE_K": BK, "GROUP_SIZE_M": GM}, num_warps=w, num_stages=s)
for BM in [int(BM) for BM in os.environ.get("SDNQ_TRITON_MM_BLOCK_SIZE_M_LIST", "64,128").replace(" ","").split(",")]
for BN in [int(BN) for BN in os.environ.get("SDNQ_TRITON_MM_BLOCK_SIZE_N_LIST", "64,128,256").replace(" ","").split(",")]
for BK in [int(BK) for BK in os.environ.get("SDNQ_TRITON_MM_BLOCK_SIZE_K_LIST", "32,64,128").replace(" ","").split(",")]
for GM in [int(GM) for GM in os.environ.get("SDNQ_TRITON_MM_GROUP_SIZE_M_LIST", "8").replace(" ","").split(",")]
for w in [int(w) for w in os.environ.get("SDNQ_TRITON_MM_NUM_WARPS_LIST", "16" if torch.xpu.is_available() else "4").replace(" ","").split(",")]
for s in [int(s) for s in os.environ.get("SDNQ_TRITON_MM_NUM_STAGES_LIST", "1" if (torch.cuda.is_available() and torch.version.hip) else "2").replace(" ","").split(",")]
]
@triton.autotune(configs=matmul_configs, key=["b_is_contiguous", "bias_ndim", "M_AT", "N_AT", "K_AT", "a_dtype", "out_dtype"], cache_results=True)
@triton.jit
def sdnq_triton_mm_kernel(
a_ptr, b_ptr, c_ptr, bias_ptr,
M: tl.constexpr,
N: tl.constexpr,
K: tl.constexpr,
b_is_contiguous: tl.constexpr,
bias_ndim: tl.constexpr,
M_AT: tl.constexpr, # pylint: disable=unused-argument
N_AT: tl.constexpr, # pylint: disable=unused-argument
K_AT: tl.constexpr, # pylint: disable=unused-argument
a_dtype: tl.constexpr, # pylint: disable=unused-argument
out_dtype: tl.constexpr, # pylint: disable=unused-argument
BLOCK_SIZE_M: tl.constexpr,
BLOCK_SIZE_N: tl.constexpr,
BLOCK_SIZE_K: tl.constexpr,
GROUP_SIZE_M: tl.constexpr,
) -> None:
pid = tl.program_id(axis=0)
num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
num_pid_in_group = GROUP_SIZE_M * num_pid_n
group_id = pid // num_pid_in_group
first_pid_m = group_id * GROUP_SIZE_M
group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
pid_n = (pid % num_pid_in_group) // group_size_m
off_m = pid_m * BLOCK_SIZE_M
off_n = pid_n * BLOCK_SIZE_N
tl.assume(M > 0)
tl.assume(N > 0)
tl.assume(K > 0)
tl.assume(pid_m >= 0)
tl.assume(pid_n >= 0)
tl.assume(off_m >= 0)
tl.assume(off_n >= 0)
tl.assume(BLOCK_SIZE_M > 0)
tl.assume(BLOCK_SIZE_N > 0)
tl.assume(BLOCK_SIZE_K > 0)
tl.assume(GROUP_SIZE_M > 0)
tl.assume(b_is_contiguous == 0 or b_is_contiguous == 1) # pylint: disable=consider-using-in
tl.assume(bias_ndim >= 0 and bias_ndim <= 2) # pylint: disable=consider-using-in
a_desc = tl.make_tensor_descriptor(base=a_ptr, shape=(M, K), strides=(K, 1), block_shape=(BLOCK_SIZE_M, BLOCK_SIZE_K))
if b_is_contiguous:
b_desc = tl.make_tensor_descriptor(base=b_ptr, shape=(K, N), strides=(N, 1), block_shape=(BLOCK_SIZE_K, BLOCK_SIZE_N))
else:
offs_k = tl.arange(0, BLOCK_SIZE_K)
offs_bn = (off_n + tl.arange(0, BLOCK_SIZE_N)) % N
b_ptrs = b_ptr + (offs_k[:, None] + offs_bn[None, :] * K)
off_k = 0
accumulator_dtype = tl.int32 if a_ptr.type.element_ty == tl.int8 else tl.float32
accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=accumulator_dtype)
for _ in tl.range(0, tl.cdiv(K, BLOCK_SIZE_K)):
a = a_desc.load([off_m, off_k])
if b_is_contiguous:
b = b_desc.load([off_k, off_n])
else:
b = tl.load(b_ptrs, mask=offs_k[:, None] < K - off_k, other=0.0)
b_ptrs += BLOCK_SIZE_K
accumulator = tl.dot(a, b, accumulator, out_dtype=accumulator_dtype)
off_k += BLOCK_SIZE_K
if bias_ndim == 1:
accumulator = accumulator.to(tl.float32)
bias_desc = tl.make_tensor_descriptor(base=bias_ptr, shape=(N,), strides=(1,), block_shape=(BLOCK_SIZE_N,))
bias = bias_desc.load([off_n])[None, :].to(tl.float32)
accumulator += bias
elif bias_ndim == 2:
accumulator = accumulator.to(tl.float32)
bias_desc = tl.make_tensor_descriptor(base=bias_ptr, shape=(M, N), strides=(N, 1), block_shape=(BLOCK_SIZE_M, BLOCK_SIZE_N))
bias = bias_desc.load([off_m, off_n]).to(tl.float32)
accumulator += bias
accumulator = accumulator.to(c_ptr.type.element_ty)
c_desc = tl.make_tensor_descriptor(base=c_ptr, shape=(M, N), strides=(N, 1), block_shape=(BLOCK_SIZE_M, BLOCK_SIZE_N))
c_desc.store([off_m, off_n], accumulator)
def sdnq_triton_mm(
a: torch.Tensor,
b: torch.Tensor,
bias: torch.FloatTensor | None = None,
out_dtype: torch.dtype | None = None,
) -> torch.Tensor:
assert a.shape[1] == b.shape[0], "Incompatible dimensions"
assert a.is_contiguous(), "Matrix A must be contiguous"
if bias is not None:
assert bias.is_contiguous(), "Bias must be contiguous"
assert bias.ndim in {1, 2}, "Bias must be 1D or 2D"
M, K = a.shape
K, N = b.shape
if out_dtype is None:
out_dtype = torch.int32 if a.dtype == torch.int8 else torch.float32
c = torch.empty((M, N), device=a.device, dtype=out_dtype)
def grid(META):
return (triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]), )
sdnq_triton_mm_kernel[grid](
a, b, c, bias,
M, N, K,
(1 if b.is_contiguous() else 0),
(0 if bias is None else bias.ndim),
math.ceil(M / min_block_size),
math.ceil(N / min_block_size),
math.ceil(K / min_block_size),
str(a.dtype), str(c.dtype),
)
return c
modules/sdnq/kernels/triton_scaled_mm.py
-142
View File
@@ -1,142 +0,0 @@
import os
import math
import torch
import triton
import triton.language as tl
min_block_size = int(os.environ.get("SDNQ_TRITON_MM_MIN_BLOCK_SIZE", "256"))
matmul_configs = [
triton.Config({"BLOCK_SIZE_M": BM, "BLOCK_SIZE_N": BN, "BLOCK_SIZE_K": BK, "GROUP_SIZE_M": GM}, num_warps=w, num_stages=s)
for BM in [int(BM) for BM in os.environ.get("SDNQ_TRITON_MM_BLOCK_SIZE_M_LIST", "64,128").replace(" ","").split(",")]
for BN in [int(BN) for BN in os.environ.get("SDNQ_TRITON_MM_BLOCK_SIZE_N_LIST", "64,128,256").replace(" ","").split(",")]
for BK in [int(BK) for BK in os.environ.get("SDNQ_TRITON_MM_BLOCK_SIZE_K_LIST", "32,64,128").replace(" ","").split(",")]
for GM in [int(GM) for GM in os.environ.get("SDNQ_TRITON_MM_GROUP_SIZE_M_LIST", "8").replace(" ","").split(",")]
for w in [int(w) for w in os.environ.get("SDNQ_TRITON_MM_NUM_WARPS_LIST", "16" if torch.xpu.is_available() else "4").replace(" ","").split(",")]
for s in [int(s) for s in os.environ.get("SDNQ_TRITON_MM_NUM_STAGES_LIST", "1" if (torch.cuda.is_available() and torch.version.hip) else "2").replace(" ","").split(",")]
]
@triton.autotune(configs=matmul_configs, key=["b_is_contiguous", "bias_ndim", "M_AT", "N_AT", "K_AT", "a_dtype", "out_dtype"], cache_results=True)
@triton.jit
def sdnq_scaled_mm_kernel(
a_ptr, b_ptr, c_ptr, bias_ptr,
scale_a_ptr, scale_b_ptr,
M: tl.constexpr,
N: tl.constexpr,
K: tl.constexpr,
b_is_contiguous: tl.constexpr,
bias_ndim: tl.constexpr,
M_AT: tl.constexpr, # pylint: disable=unused-argument
N_AT: tl.constexpr, # pylint: disable=unused-argument
K_AT: tl.constexpr, # pylint: disable=unused-argument
a_dtype: tl.constexpr, # pylint: disable=unused-argument
out_dtype: tl.constexpr, # pylint: disable=unused-argument
BLOCK_SIZE_M: tl.constexpr,
BLOCK_SIZE_N: tl.constexpr,
BLOCK_SIZE_K: tl.constexpr,
GROUP_SIZE_M: tl.constexpr,
) -> None:
pid = tl.program_id(axis=0)
num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
num_pid_in_group = GROUP_SIZE_M * num_pid_n
group_id = pid // num_pid_in_group
first_pid_m = group_id * GROUP_SIZE_M
group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
pid_n = (pid % num_pid_in_group) // group_size_m
off_m = pid_m * BLOCK_SIZE_M
off_n = pid_n * BLOCK_SIZE_N
tl.assume(M > 0)
tl.assume(N > 0)
tl.assume(K > 0)
tl.assume(pid_m >= 0)
tl.assume(pid_n >= 0)
tl.assume(off_m >= 0)
tl.assume(off_n >= 0)
tl.assume(BLOCK_SIZE_M > 0)
tl.assume(BLOCK_SIZE_N > 0)
tl.assume(BLOCK_SIZE_K > 0)
tl.assume(GROUP_SIZE_M > 0)
tl.assume(b_is_contiguous == 0 or b_is_contiguous == 1) # pylint: disable=consider-using-in
tl.assume(bias_ndim >= 0 and bias_ndim <= 2) # pylint: disable=consider-using-in
a_desc = tl.make_tensor_descriptor(base=a_ptr, shape=(M, K), strides=(K, 1), block_shape=(BLOCK_SIZE_M, BLOCK_SIZE_K))
if b_is_contiguous:
b_desc = tl.make_tensor_descriptor(base=b_ptr, shape=(K, N), strides=(N, 1), block_shape=(BLOCK_SIZE_K, BLOCK_SIZE_N))
else:
offs_k = tl.arange(0, BLOCK_SIZE_K)
offs_bn = (off_n + tl.arange(0, BLOCK_SIZE_N)) % N
b_ptrs = b_ptr + (offs_k[:, None] + offs_bn[None, :] * K)
off_k = 0
accumulator_dtype = tl.int32 if a_ptr.type.element_ty == tl.int8 else tl.float32
accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=accumulator_dtype)
for _ in tl.range(0, tl.cdiv(K, BLOCK_SIZE_K)):
a = a_desc.load([off_m, off_k])
if b_is_contiguous:
b = b_desc.load([off_k, off_n])
else:
b = tl.load(b_ptrs, mask=offs_k[:, None] < K - off_k, other=0.0)
b_ptrs += BLOCK_SIZE_K
accumulator = tl.dot(a, b, accumulator, out_dtype=accumulator_dtype)
off_k += BLOCK_SIZE_K
scale_a_desc = tl.make_tensor_descriptor(base=scale_a_ptr, shape=(M,), strides=(1,), block_shape=(BLOCK_SIZE_M,))
scale_b_desc = tl.make_tensor_descriptor(base=scale_b_ptr, shape=(N,), strides=(1,), block_shape=(BLOCK_SIZE_N,))
scale_a = scale_a_desc.load([off_m])[:, None].to(tl.float32)
scale_b = scale_b_desc.load([off_n])[None, :].to(tl.float32)
if bias_ndim == 1:
accumulator = tl.mul(accumulator.to(tl.float32), scale_a)
bias_desc = tl.make_tensor_descriptor(base=bias_ptr, shape=(N,), strides=(1,), block_shape=(BLOCK_SIZE_N,))
bias = bias_desc.load([off_n])[None, :].to(tl.float32)
accumulator = tl.fma(accumulator, scale_b, bias)
elif bias_ndim == 2:
accumulator = tl.mul(accumulator.to(tl.float32), scale_a)
bias_desc = tl.make_tensor_descriptor(base=bias_ptr, shape=(M, N), strides=(N, 1), block_shape=(BLOCK_SIZE_M, BLOCK_SIZE_N))
bias = bias_desc.load([off_m, off_n]).to(tl.float32)
accumulator = tl.fma(accumulator, scale_b, bias)
else:
accumulator = tl.mul(tl.mul(accumulator.to(tl.float32), scale_a), scale_b)
accumulator = accumulator.to(c_ptr.type.element_ty)
c_desc = tl.make_tensor_descriptor(base=c_ptr, shape=(M, N), strides=(N, 1), block_shape=(BLOCK_SIZE_M, BLOCK_SIZE_N))
c_desc.store([off_m, off_n], accumulator)
def sdnq_scaled_mm(
a: torch.Tensor,
b: torch.Tensor,
scale_a: torch.Tensor,
scale_b: torch.Tensor,
bias: torch.FloatTensor | None = None,
out_dtype: torch.dtype = torch.float32,
) -> torch.Tensor:
assert a.shape[1] == b.shape[0], "Incompatible dimensions"
assert a.is_contiguous(), "Matrix A must be contiguous"
assert scale_a.is_contiguous(), "Matrix A scale must be contiguous"
assert scale_b.is_contiguous(), "Matrix B scale must be contiguous"
if bias is not None:
assert bias.is_contiguous(), "Bias must be contiguous"
assert bias.ndim in {1, 2}, "Bias must be 1D or 2D"
M, K = a.shape
K, N = b.shape
c = torch.empty((M, N), device=a.device, dtype=out_dtype)
def grid(META):
return (triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]), )
sdnq_scaled_mm_kernel[grid](
a, b, c, bias,
scale_a, scale_b,
M, N, K,
(1 if b.is_contiguous() else 0),
(0 if bias is None else bias.ndim),
math.ceil(M / min_block_size),
math.ceil(N / min_block_size),
math.ceil(K / min_block_size),
str(a.dtype), str(c.dtype),
)
return c
modules/sdnq/layers/__init__.py
-93
View File
@@ -1,93 +0,0 @@
from collections.abc import Callable
import torch
class SDNQLayer(torch.nn.Module):
def __init__(self, original_layer: torch.nn.Module, forward_func: Callable):
torch.nn.Module.__init__(self)
for key, value in original_layer.__dict__.items():
if key not in {"forward", "forward_func", "original_class", "state_dict", "load_state_dict"}:
setattr(self, key, value)
self.original_class = original_layer.__class__
self.forward_func = forward_func
@property
def dtype(self: torch.nn.Module) -> torch.dtype:
return self.sdnq_dequantizer.result_dtype if hasattr(self, "sdnq_dequantizer") else self.weight.dtype
def dequantize(self: torch.nn.Module):
if self.weight.__class__.__name__ == "SDNQTensor": # pylint: disable=access-member-before-definition
self.weight = torch.nn.Parameter(self.weight.dequantize(), requires_grad=True) # pylint: disable=attribute-defined-outside-init
elif hasattr(self, "sdnq_dequantizer"):
self.weight = torch.nn.Parameter(self.sdnq_dequantizer(self.weight, self.scale, zero_point=self.zero_point, svd_up=self.svd_up, svd_down=self.svd_down, skip_quantized_matmul=self.sdnq_dequantizer.use_quantized_matmul), requires_grad=True) # pylint: disable=attribute-defined-outside-init
del self.sdnq_dequantizer, self.scale, self.zero_point, self.svd_up, self.svd_down
self.__class__ = self.original_class # pylint: disable=attribute-defined-outside-init
del self.original_class, self.forward_func
return self
def forward(self, *args, **kwargs) -> torch.Tensor:
return self.forward_func(self, *args, **kwargs)
def __repr__(self) -> str:
return f"{self.__class__.__name__}(original_class={self.original_class} forward_func={self.forward_func} sdnq_dequantizer={getattr(self, 'sdnq_dequantizer', None)})"
class SDNQLinear(SDNQLayer, torch.nn.Linear):
original_class: torch.nn.Linear
class SDNQEmbedding(SDNQLayer, torch.nn.Embedding):
original_class: torch.nn.Embedding
class SDNQConv1d(SDNQLayer, torch.nn.Conv1d):
original_class: torch.nn.Conv1d
class SDNQConv2d(SDNQLayer, torch.nn.Conv2d):
original_class: torch.nn.Conv2d
class SDNQConv3d(SDNQLayer, torch.nn.Conv3d):
original_class: torch.nn.Conv3d
class SDNQConvTranspose1d(SDNQLayer, torch.nn.ConvTranspose1d):
original_class: torch.nn.ConvTranspose1d
class SDNQConvTranspose2d(SDNQLayer, torch.nn.ConvTranspose2d):
original_class: torch.nn.ConvTranspose2d
class SDNQConvTranspose3d(SDNQLayer, torch.nn.ConvTranspose3d):
original_class: torch.nn.ConvTranspose3d
torch.serialization.add_safe_globals([SDNQLayer])
torch.serialization.add_safe_globals([SDNQLinear])
torch.serialization.add_safe_globals([SDNQEmbedding])
torch.serialization.add_safe_globals([SDNQConv1d])
torch.serialization.add_safe_globals([SDNQConv2d])
torch.serialization.add_safe_globals([SDNQConv3d])
torch.serialization.add_safe_globals([SDNQConvTranspose1d])
torch.serialization.add_safe_globals([SDNQConvTranspose2d])
torch.serialization.add_safe_globals([SDNQConvTranspose3d])
def get_sdnq_wrapper_class(original_layer: torch.nn.Module, forward_func: Callable) -> SDNQLayer:
match original_layer.__class__.__name__:
case "Linear":
return SDNQLinear(original_layer, forward_func)
case "Embedding":
return SDNQEmbedding(original_layer, forward_func)
case "Gemma4TextScaledWordEmbedding":
return SDNQEmbedding(original_layer, forward_func)
case "Conv1d":
return SDNQConv1d(original_layer, forward_func)
case "Conv2d":
return SDNQConv2d(original_layer, forward_func)
case "Conv3d":
return SDNQConv3d(original_layer, forward_func)
case "ConvTranspose1d":
return SDNQConvTranspose1d(original_layer, forward_func)
case "ConvTranspose2d":
return SDNQConvTranspose2d(original_layer, forward_func)
case "ConvTranspose3d":
return SDNQConvTranspose3d(original_layer, forward_func)
case _:
return SDNQLayer(original_layer, forward_func)
modules/sdnq/layers/conv/conv_fp16.py
-107
View File
@@ -1,107 +0,0 @@
# pylint: disable=relative-beyond-top-level,redefined-builtin,protected-access
import torch
from ...common import compile_func
from ...kernel_wrappers import fp_mm_func, fp_scaled_mm_func
from ...dequantizer import dequantize_symmetric, dequantize_asymmetric
from ...quant_utils import rotate_hadamard, get_hadamard
from ...packed_float import unpack_float
from .forward import get_conv_args, process_conv_input
from ..linear.linear_fp8 import quantize_fp_mm_input
from ..linear.forward import check_mats
def conv_fp16_matmul(
input: torch.FloatTensor,
weight: torch.Tensor,
scale: torch.FloatTensor,
result_shape: torch.Size,
reversed_padding_repeated_twice: list[int],
padding_mode: str, conv_type: int,
groups: int, stride: list[int],
padding: list[int], dilation: list[int],
bias: torch.FloatTensor | None = None,
svd_up: torch.FloatTensor | None = None,
svd_down: torch.FloatTensor | None = None,
hadamard: torch.FloatTensor | None = None,
quantized_weight_shape: torch.Size | None = None,
weights_dtype: str | None = None,
) -> torch.FloatTensor:
return_dtype = input.dtype
input, mm_output_shape = process_conv_input(conv_type, input, reversed_padding_repeated_twice, padding_mode, result_shape, stride, padding, dilation)
if quantized_weight_shape is not None:
weight = unpack_float(weight, weights_dtype, quantized_weight_shape).to(dtype=torch.float16).t_()
scale = scale.t()
elif weight.dtype != torch.float16:
weight = weight.to(dtype=torch.float16) # fp8 weights
if hadamard is not None:
input = rotate_hadamard(input, hadamard=hadamard)
if svd_up is not None:
input = input.flatten(0,-2)
if bias is not None:
bias = torch.addmm(bias.to(dtype=svd_down.dtype), torch.mm(input.to(dtype=svd_down.dtype), svd_down), svd_up)
else:
bias = torch.mm(torch.mm(input.to(dtype=svd_down.dtype), svd_down), svd_up)
input, input_scale = quantize_fp_mm_input(input, dtype=scale.dtype, matmul_dtype="float16")
input, weight = check_mats(input, weight, matmul_dtype="float16")
if groups == 1:
result = fp_scaled_mm_func(input, weight, input_scale, scale, bias=bias, out_dtype=return_dtype).view(mm_output_shape)
else:
weight = weight.view(weight.shape[0], groups, weight.shape[1] // groups)
input = input.view(input.shape[0], groups, input.shape[1] // groups)
result = []
for i in range(groups):
result.append(fp_mm_func(input[:, i], weight[:, i]))
result = torch.cat(result, dim=-1).to(dtype=input_scale.dtype).mul_(input_scale)
if bias is not None:
result = dequantize_asymmetric(result, scale, bias, dtype=return_dtype, result_shape=mm_output_shape)
else:
result = dequantize_symmetric(result, scale, dtype=return_dtype, result_shape=mm_output_shape)
if conv_type == 1:
result = result.transpose_(1,2)
elif conv_type == 2:
result = result.permute(0,3,1,2)
elif conv_type == 3:
result = result.permute(0,4,1,2,3)
result = result.contiguous()
return result
def quantized_conv_forward_fp16_matmul(self, input) -> torch.FloatTensor:
if torch.numel(input) / input.shape[2] < 32:
return self._conv_forward(input, self.sdnq_dequantizer(self.weight, self.scale, zero_point=self.zero_point, svd_up=self.svd_up, svd_down=self.svd_down, skip_quantized_matmul=True), self.bias)
if self.sdnq_dequantizer.re_quantize_for_matmul:
weight, scale = self.sdnq_dequantizer.re_quantize_matmul(self.weight, self.scale, zero_point=self.zero_point)
quantized_weight_shape = None
else:
weight, scale = self.weight, self.scale
quantized_weight_shape = self.sdnq_dequantizer.quantized_weight_shape if self.sdnq_dequantizer.is_packed else None
if self.sdnq_dequantizer.use_hadamard:
hadamard = get_hadamard(self.sdnq_dequantizer.hadamard_group_size, dtype=input.dtype, device=input.device)
else:
hadamard = None
conv_type, stride, padding, dilation = get_conv_args(input.ndim, self.stride, self.padding, self.dilation)
return conv_fp16_matmul(
input, weight, scale,
self.sdnq_dequantizer.result_shape,
self._reversed_padding_repeated_twice,
self.padding_mode, conv_type,
self.groups, stride, padding, dilation,
bias=self.bias,
svd_up=self.svd_up,
svd_down=self.svd_down,
hadamard=hadamard,
quantized_weight_shape=quantized_weight_shape,
weights_dtype=self.sdnq_dequantizer.weights_dtype,
)
conv_fp16_matmul = compile_func(conv_fp16_matmul)
modules/sdnq/layers/conv/conv_fp8.py
-105
View File
@@ -1,105 +0,0 @@
# pylint: disable=relative-beyond-top-level,redefined-builtin,protected-access
import torch
from ...common import compile_func
from ...kernel_wrappers import fp8_mm_func, fp8_scaled_mm_func
from ...dequantizer import dequantize_symmetric, dequantize_asymmetric
from ...quant_utils import rotate_hadamard, get_hadamard
from ...packed_float import unpack_float
from .forward import get_conv_args, process_conv_input
from ..linear.linear_fp8 import quantize_fp_mm_input
from ..linear.forward import check_mats
def conv_fp8_matmul(
input: torch.FloatTensor,
weight: torch.Tensor,
scale: torch.FloatTensor,
result_shape: torch.Size,
reversed_padding_repeated_twice: list[int],
padding_mode: str, conv_type: int,
groups: int, stride: list[int],
padding: list[int], dilation: list[int],
bias: torch.FloatTensor | None = None,
svd_up: torch.FloatTensor | None = None,
svd_down: torch.FloatTensor | None = None,
hadamard: torch.FloatTensor | None = None,
quantized_weight_shape: torch.Size | None = None,
weights_dtype: str | None = None,
) -> torch.FloatTensor:
return_dtype = input.dtype
input, mm_output_shape = process_conv_input(conv_type, input, reversed_padding_repeated_twice, padding_mode, result_shape, stride, padding, dilation)
if quantized_weight_shape is not None:
weight = unpack_float(weight, weights_dtype, quantized_weight_shape).to(dtype=torch.float8_e4m3fn).t_()
scale = scale.t()
if hadamard is not None:
input = rotate_hadamard(input, hadamard=hadamard)
if svd_up is not None:
input = input.flatten(0,-2)
if bias is not None:
bias = torch.addmm(bias.to(dtype=svd_down.dtype), torch.mm(input.to(dtype=svd_down.dtype), svd_down), svd_up)
else:
bias = torch.mm(torch.mm(input.to(dtype=svd_down.dtype), svd_down), svd_up)
input, input_scale = quantize_fp_mm_input(input, dtype=scale.dtype)
input, weight = check_mats(input, weight, matmul_dtype="float8_e4m3fn")
if groups == 1:
result = fp8_scaled_mm_func(input, weight, input_scale, scale, bias=bias, out_dtype=return_dtype).view(mm_output_shape)
else:
weight = weight.view(weight.shape[0], groups, weight.shape[1] // groups)
input = input.view(input.shape[0], groups, input.shape[1] // groups)
result = []
for i in range(groups):
result.append(fp8_mm_func(input[:, i], weight[:, i]))
result = torch.cat(result, dim=-1).to(dtype=input_scale.dtype).mul_(input_scale)
if bias is not None:
result = dequantize_asymmetric(result, scale, bias, dtype=return_dtype, result_shape=mm_output_shape)
else:
result = dequantize_symmetric(result, scale, dtype=return_dtype, result_shape=mm_output_shape)
if conv_type == 1:
result = result.transpose_(1,2)
elif conv_type == 2:
result = result.permute(0,3,1,2)
elif conv_type == 3:
result = result.permute(0,4,1,2,3)
result = result.contiguous()
return result
def quantized_conv_forward_fp8_matmul(self, input) -> torch.FloatTensor:
if torch.numel(input) / input.shape[2] < 32:
return self._conv_forward(input, self.sdnq_dequantizer(self.weight, self.scale, zero_point=self.zero_point, svd_up=self.svd_up, svd_down=self.svd_down, skip_quantized_matmul=True), self.bias)
if self.sdnq_dequantizer.re_quantize_for_matmul:
weight, scale = self.sdnq_dequantizer.re_quantize_matmul(self.weight, self.scale, zero_point=self.zero_point)
quantized_weight_shape = None
else:
weight, scale = self.weight, self.scale
quantized_weight_shape = self.sdnq_dequantizer.quantized_weight_shape if self.sdnq_dequantizer.is_packed else None
if self.sdnq_dequantizer.use_hadamard:
hadamard = get_hadamard(self.sdnq_dequantizer.hadamard_group_size, dtype=input.dtype, device=input.device)
else:
hadamard = None
conv_type, stride, padding, dilation = get_conv_args(input.ndim, self.stride, self.padding, self.dilation)
return conv_fp8_matmul(
input, weight, scale,
self.sdnq_dequantizer.result_shape,
self._reversed_padding_repeated_twice,
self.padding_mode, conv_type,
self.groups, stride, padding, dilation,
bias=self.bias,
svd_up=self.svd_up,
svd_down=self.svd_down,
hadamard=hadamard,
quantized_weight_shape=quantized_weight_shape,
weights_dtype=self.sdnq_dequantizer.weights_dtype,
)
conv_fp8_matmul = compile_func(conv_fp8_matmul)
modules/sdnq/layers/conv/conv_int8.py
-123
View File
@@ -1,123 +0,0 @@
# pylint: disable=relative-beyond-top-level,redefined-builtin,protected-access
import torch
from ...common import compile_func
from ...kernel_wrappers import int_mm_func, int_scaled_mm_func
from ...dequantizer import dequantize_symmetric, dequantize_asymmetric
from ...quant_utils import rotate_hadamard, get_hadamard
from ...packed_int import unpack_int
from .forward import get_conv_args, process_conv_input
from ..linear.linear_int8 import quantize_int_mm_input
from ..linear.forward import check_mats
def conv_int8_matmul(
input: torch.FloatTensor,
weight: torch.Tensor,
scale: torch.FloatTensor,
result_shape: torch.Size,
reversed_padding_repeated_twice: list[int],
padding_mode: str, conv_type: int,
groups: int, stride: list[int],
padding: list[int], dilation: list[int],
bias: torch.FloatTensor | None = None,
svd_up: torch.FloatTensor | None = None,
svd_down: torch.FloatTensor | None = None,
zero_point: torch.FloatTensor | None = None,
hadamard: torch.FloatTensor | None = None,
quantized_weight_shape: torch.Size | None = None,
weights_dtype: str | None = None,
) -> torch.FloatTensor:
return_dtype = input.dtype
input, mm_output_shape = process_conv_input(conv_type, input, reversed_padding_repeated_twice, padding_mode, result_shape, stride, padding, dilation)
if quantized_weight_shape is not None:
weight = unpack_int(weight, weights_dtype, quantized_weight_shape, dtype=torch.int8).t_()
scale = scale.t()
if zero_point is not None:
zero_point = zero_point.t()
if weight.dtype == torch.uint8:
weight = weight.view(dtype=torch.int8)
elif weight.dtype == torch.uint8:
weight = weight.bitwise_xor(128).view(torch.int8)
if zero_point is not None:
zero_point = torch.add(zero_point, scale, alpha=128)
else:
zero_point = torch.mul(scale, 128)
if hadamard is not None:
input = rotate_hadamard(input, hadamard=hadamard)
if svd_up is not None:
input = input.flatten(0,-2)
if bias is not None:
bias = torch.addmm(bias.to(dtype=svd_down.dtype), torch.mm(input.to(dtype=svd_down.dtype), svd_down), svd_up)
else:
bias = torch.mm(torch.mm(input.to(dtype=svd_down.dtype), svd_down), svd_up)
input, input_scale = quantize_int_mm_input(input, dtype=scale.dtype)
if zero_point is not None:
zero_bias = torch.sum(input, dim=-1, keepdim=True, dtype=torch.int32).to(input_scale.dtype).mul_(input_scale).mul(zero_point)
if bias is not None:
zero_bias.add_(bias)
bias = zero_bias
input, weight = check_mats(input, weight, matmul_dtype="int8")
if groups == 1:
result = int_scaled_mm_func(input, weight, input_scale, scale, bias=bias, out_dtype=return_dtype).view(mm_output_shape)
else:
weight = weight.view(weight.shape[0], groups, weight.shape[1] // groups)
input = input.view(input.shape[0], groups, input.shape[1] // groups)
result = []
for i in range(groups):
result.append(int_mm_func(input[:, i], weight[:, i]))
result = torch.cat(result, dim=-1).to(dtype=input_scale.dtype).mul_(input_scale)
if bias is not None:
result = dequantize_asymmetric(result, scale, bias, dtype=return_dtype, result_shape=mm_output_shape)
else:
result = dequantize_symmetric(result, scale, dtype=return_dtype, result_shape=mm_output_shape)
if conv_type == 1:
result = result.transpose_(1,2)
elif conv_type == 2:
result = result.permute(0,3,1,2)
elif conv_type == 3:
result = result.permute(0,4,1,2,3)
result = result.contiguous()
return result
def quantized_conv_forward_int8_matmul(self, input) -> torch.FloatTensor:
if torch.numel(input) / input.shape[2] < 32:
return self._conv_forward(input, self.sdnq_dequantizer(self.weight, self.scale, zero_point=self.zero_point, svd_up=self.svd_up, svd_down=self.svd_down, skip_quantized_matmul=True), self.bias)
if self.sdnq_dequantizer.re_quantize_for_matmul:
weight, scale = self.sdnq_dequantizer.re_quantize_matmul(self.weight, self.scale, zero_point=self.zero_point)
quantized_weight_shape = None
zero_point = None
else:
weight, scale, zero_point = self.weight, self.scale, self.zero_point
quantized_weight_shape = self.sdnq_dequantizer.quantized_weight_shape if self.sdnq_dequantizer.is_packed else None
if self.sdnq_dequantizer.use_hadamard:
hadamard = get_hadamard(self.sdnq_dequantizer.hadamard_group_size, dtype=input.dtype, device=input.device)
else:
hadamard = None
conv_type, stride, padding, dilation = get_conv_args(input.ndim, self.stride, self.padding, self.dilation)
return conv_int8_matmul(
input, weight, scale,
self.sdnq_dequantizer.result_shape,
self._reversed_padding_repeated_twice,
self.padding_mode, conv_type,
self.groups, stride, padding, dilation,
bias=self.bias,
svd_up=self.svd_up,
svd_down=self.svd_down,
zero_point=zero_point,
hadamard=hadamard,
quantized_weight_shape=quantized_weight_shape,
weights_dtype=self.sdnq_dequantizer.weights_dtype,
)
conv_int8_matmul = compile_func(conv_int8_matmul)
modules/sdnq/layers/conv/conv_uint8.py
-122
View File
@@ -1,122 +0,0 @@
# pylint: disable=relative-beyond-top-level,redefined-builtin,protected-access
import torch
from ...common import compile_func
from ...kernel_wrappers import int_mm_func, int_scaled_mm_func
from ...dequantizer import dequantize_asymmetric
from ...quant_utils import rotate_hadamard, get_hadamard
from ...packed_int import unpack_int
from .forward import get_conv_args, process_conv_input
from ..linear.linear_uint8 import quantize_uint_mm_input
from ..linear.forward import check_mats
def conv_uint8_matmul(
input: torch.FloatTensor,
weight: torch.Tensor,
scale: torch.FloatTensor,
zero_point: torch.FloatTensor,
result_shape: torch.Size,
reversed_padding_repeated_twice: list[int],
padding_mode: str, conv_type: int,
groups: int, stride: list[int],
padding: list[int], dilation: list[int],
bias: torch.FloatTensor | None = None,
svd_up: torch.FloatTensor | None = None,
svd_down: torch.FloatTensor | None = None,
hadamard: torch.FloatTensor | None = None,
quantized_weight_shape: torch.Size | None = None,
weights_dtype: str | None = None,
) -> torch.FloatTensor:
return_dtype = input.dtype
input, mm_output_shape = process_conv_input(conv_type, input, reversed_padding_repeated_twice, padding_mode, result_shape, stride, padding, dilation)
if quantized_weight_shape is not None:
weight = unpack_int(weight, weights_dtype, quantized_weight_shape, dtype=torch.int8).t_()
scale = scale.t()
if zero_point is not None:
zero_point = zero_point.t()
if weight.dtype == torch.uint8:
weight = weight.view(dtype=torch.int8)
elif weight.dtype == torch.uint8:
weight = weight.bitwise_xor(128).view(torch.int8)
if zero_point is not None:
zero_point = torch.add(zero_point, scale, alpha=128)
else:
zero_point = torch.mul(scale, 128)
if hadamard is not None:
input = rotate_hadamard(input, hadamard=hadamard)
if svd_up is not None:
input = input.flatten(0,-2)
if bias is not None:
bias = torch.addmm(bias.to(dtype=svd_down.dtype), torch.mm(input.to(dtype=svd_down.dtype), svd_down), svd_up)
else:
bias = torch.mm(torch.mm(input.to(dtype=svd_down.dtype), svd_down), svd_up)
input, input_scale, input_zero_point = quantize_uint_mm_input(input, dtype=scale.dtype)
if zero_point is not None:
zero_bias = torch.sum(input, dim=-1, keepdim=True, dtype=torch.int32).to(input_scale.dtype).mul_(input_scale).mul(zero_point)
zero_bias.add_(torch.sum(weight, dim=0, keepdim=True, dtype=torch.int32).to(scale.dtype).mul_(scale).mul(input_zero_point))
zero_bias.add_(torch.mul(input_zero_point.mul_(input.shape[-1]), zero_point))
else:
zero_bias = torch.sum(weight, dim=0, keepdim=True, dtype=torch.int32).to(scale.dtype).mul_(scale).mul(input_zero_point)
if bias is not None:
zero_bias.add_(bias)
input, weight = check_mats(input, weight, matmul_dtype="uint8")
if groups == 1:
result = int_scaled_mm_func(input, weight, input_scale, scale, bias=zero_bias, out_dtype=return_dtype).view(mm_output_shape)
else:
weight = weight.view(weight.shape[0], groups, weight.shape[1] // groups)
input = input.view(input.shape[0], groups, input.shape[1] // groups)
result = []
for i in range(groups):
result.append(int_mm_func(input[:, i], weight[:, i]))
result = torch.cat(result, dim=-1).to(dtype=input_scale.dtype).mul_(input_scale)
result = dequantize_asymmetric(result, scale, zero_bias, dtype=return_dtype, result_shape=mm_output_shape)
if conv_type == 1:
result = result.transpose_(1,2)
elif conv_type == 2:
result = result.permute(0,3,1,2)
elif conv_type == 3:
result = result.permute(0,4,1,2,3)
result = result.contiguous()
return result
def quantized_conv_forward_uint8_matmul(self, input) -> torch.FloatTensor:
if torch.numel(input) / input.shape[2] < 32:
return self._conv_forward(input, self.sdnq_dequantizer(self.weight, self.scale, zero_point=self.zero_point, svd_up=self.svd_up, svd_down=self.svd_down, skip_quantized_matmul=True), self.bias)
if self.sdnq_dequantizer.re_quantize_for_matmul:
weight, scale, zero_point = self.sdnq_dequantizer.re_quantize_matmul(self.weight, self.scale, zero_point=self.zero_point)
quantized_weight_shape = None
else:
weight, scale, zero_point = self.weight, self.scale, self.zero_point
quantized_weight_shape = self.sdnq_dequantizer.quantized_weight_shape if self.sdnq_dequantizer.is_packed else None
if self.sdnq_dequantizer.use_hadamard:
hadamard = get_hadamard(self.sdnq_dequantizer.hadamard_group_size, dtype=input.dtype, device=input.device)
else:
hadamard = None
conv_type, stride, padding, dilation = get_conv_args(input.ndim, self.stride, self.padding, self.dilation)
return conv_uint8_matmul(
input, weight,
scale, zero_point,
self.sdnq_dequantizer.result_shape,
self._reversed_padding_repeated_twice,
self.padding_mode, conv_type,
self.groups, stride, padding, dilation,
bias=self.bias,
svd_up=self.svd_up,
svd_down=self.svd_down,
hadamard=hadamard,
quantized_weight_shape=quantized_weight_shape,
weights_dtype=self.sdnq_dequantizer.weights_dtype,
)
conv_uint8_matmul = compile_func(conv_uint8_matmul)
modules/sdnq/layers/conv/forward.py
-91
View File
@@ -1,91 +0,0 @@
# pylint: disable=relative-beyond-top-level,redefined-builtin,protected-access
import torch
def get_conv_args(input_ndim: int, stride, padding, dilation):
if input_ndim == 3:
conv_type = 1
elif input_ndim == 4:
conv_type = 2
else:
conv_type = 3
if isinstance(stride, int):
stride = (stride,) * conv_type
if isinstance(padding, int):
padding = (padding,) * conv_type
if isinstance(dilation, int):
dilation = (dilation,) * conv_type
if conv_type == 1:
stride = (1, stride[0])
padding = (0, padding[0])
dilation = (1, dilation[0])
return conv_type, stride, padding, dilation
def process_conv_input(conv_type, input, reversed_padding_repeated_twice, padding_mode, result_shape, stride, padding, dilation):
if conv_type == 1:
batch_size, _, L_in = input.shape
C_out, _, K_l = result_shape
L_out = (L_in + 2 * padding[1] - dilation[1] * (K_l - 1) - 1) // stride[1] + 1
mm_output_shape = (batch_size, L_out, C_out)
kernel_size = (1, K_l)
elif conv_type == 2:
batch_size, _, H_in, W_in = input.shape
C_out, _, K_h, K_w = result_shape
H_out = (H_in + 2 * padding[0] - dilation[0] * (K_h - 1) - 1) // stride[0] + 1
W_out = (W_in + 2 * padding[1] - dilation[1] * (K_w - 1) - 1) // stride[1] + 1
mm_output_shape = (batch_size, H_out, W_out, C_out)
kernel_size = (K_h, K_w)
else:
batch_size, _, D_in, H_in, W_in = input.shape
C_out, _, K_d, K_h, K_w = result_shape
D_out = (D_in + 2 * padding[0] - dilation[0] * (K_d - 1) - 1) // stride[0] + 1
H_out = (H_in + 2 * padding[1] - dilation[1] * (K_h - 1) - 1) // stride[1] + 1
W_out = (W_in + 2 * padding[2] - dilation[2] * (K_w - 1) - 1) // stride[2] + 1
mm_output_shape = (batch_size, D_out, H_out, W_out, C_out)
kernel_size = (K_d, K_h, K_w)
if padding_mode != "zeros":
input = torch.nn.functional.pad(input, reversed_padding_repeated_twice, mode=padding_mode)
padding = (0,) * (conv_type if conv_type != 1 else 2)
elif conv_type == 3:
input = torch.nn.functional.pad(input, reversed_padding_repeated_twice)
if conv_type == 1:
input = input.unsqueeze(2)
if conv_type == 3:
K_D_eff = K_d + (K_d - 1) * (dilation[0] - 1)
K_H_eff = K_h + (K_h - 1) * (dilation[0] - 1)
K_W_eff = K_w + (K_w - 1) * (dilation[0] - 1)
input = input.unfold(2, K_D_eff, stride[0]).unfold(3, K_H_eff, stride[1]).unfold(4, K_W_eff, stride[2])
if dilation[0] > 1:
input = input[..., ::dilation[0], :, :]
if dilation[1] > 1:
input = input[..., ::dilation[1], :]
if dilation[2] > 1:
input = input[..., ::dilation[2]]
input = input.permute(0, 2, 3, 4, 1, 5, 6, 7).reshape(batch_size, D_out * H_out * W_out, -1)
else:
input = torch.nn.functional.unfold(input, kernel_size=kernel_size, padding=padding, stride=stride, dilation=dilation).transpose(1,2)
return input, mm_output_shape
def quantized_conv_forward(self, input) -> torch.FloatTensor:
return self._conv_forward(input, self.sdnq_dequantizer(self.weight, self.scale, self.zero_point, self.svd_up, self.svd_down), self.bias)
def quantized_conv_transpose_1d_forward(self, input: torch.FloatTensor, output_size: list[int] | None = None) -> torch.FloatTensor:
output_padding = self._output_padding(input, output_size, self.stride, self.padding, self.kernel_size, 1, self.dilation)
return torch.nn.functional.conv_transpose1d(input, self.sdnq_dequantizer(self.weight, self.scale, self.zero_point, self.svd_up, self.svd_down), self.bias, self.stride, self.padding, output_padding, self.groups, self.dilation)
def quantized_conv_transpose_2d_forward(self, input: torch.FloatTensor, output_size: list[int] | None = None) -> torch.FloatTensor:
output_padding = self._output_padding(input, output_size, self.stride, self.padding, self.kernel_size, 2, self.dilation)
return torch.nn.functional.conv_transpose2d(input, self.sdnq_dequantizer(self.weight, self.scale, self.zero_point, self.svd_up, self.svd_down), self.bias, self.stride, self.padding, output_padding, self.groups, self.dilation)
def quantized_conv_transpose_3d_forward(self, input: torch.FloatTensor, output_size: list[int] | None = None) -> torch.FloatTensor:
output_padding = self._output_padding(input, output_size, self.stride, self.padding, self.kernel_size, 3, self.dilation)
return torch.nn.functional.conv_transpose3d(input, self.sdnq_dequantizer(self.weight, self.scale, self.zero_point, self.svd_up, self.svd_down), self.bias, self.stride, self.padding, output_padding, self.groups, self.dilation)
modules/sdnq/layers/embedding/forward.py
-83
View File
@@ -1,83 +0,0 @@
# pylint: disable=relative-beyond-top-level,redefined-builtin,protected-access
import torch
from ...common import dtype_dict, compile_func
from ...dequantizer import dequantize_symmetric, dequantize_asymmetric
from ...quant_utils import get_hadamard
from ...packed_int import unpack_int
from ...packed_float import unpack_float
def quantized_embedding(
input: torch.Tensor,
weight: torch.Tensor,
scale: torch.FloatTensor,
zero_point: torch.FloatTensor | None = None,
svd_up: torch.FloatTensor | None = None,
svd_down: torch.FloatTensor | None = None,
hadamard: torch.FloatTensor | None = None,
embed_scale: torch.FloatTensor | float | None = None,
result_dtype: torch.dtype | None = None,
weight_shape: torch.Size | None = None,
quantized_weight_shape: torch.Size | None = None,
weights_dtype: str | None = None,
) -> torch.FloatTensor:
return_shape = list(input.shape) + [weight_shape[-1] if weight_shape is not None else quantized_weight_shape[-1] if quantized_weight_shape is not None else weight.shape[-1]]
input = input.flatten()
if weights_dtype is not None and dtype_dict[weights_dtype]["is_packed"]:
if dtype_dict[weights_dtype]["is_integer"]:
weight = unpack_int(weight, weights_dtype, quantized_weight_shape)
else:
weight = unpack_float(weight, weights_dtype, quantized_weight_shape)
if zero_point is not None:
result = dequantize_asymmetric(
weight[input], scale[input], zero_point[input],
svd_up=svd_up[input] if svd_up is not None else svd_up,
svd_down=svd_down,
hadamard=hadamard,
dtype=result_dtype,
)
else:
result = dequantize_symmetric(
weight[input], scale[input],
svd_up=svd_up[input] if svd_up is not None else svd_up,
svd_down=svd_down,
hadamard=hadamard,
dtype=result_dtype,
)
del input
result = result.view(return_shape).contiguous()
if embed_scale is not None:
result = result.mul_(embed_scale)
return result
def quantized_embedding_forward(self: torch.nn.Module, input: torch.Tensor) -> torch.FloatTensor:
if self.sdnq_dequantizer.use_hadamard:
hadamard = get_hadamard(self.sdnq_dequantizer.hadamard_group_size, dtype=self.sdnq_dequantizer.result_dtype, device=input.device)
else:
hadamard = None
return quantized_embedding(
input,
self.weight,
self.scale,
zero_point=self.zero_point,
svd_up=self.svd_up,
svd_down=self.svd_down,
hadamard=hadamard,
embed_scale=getattr(self, "scalar_embed_scale", None),
result_dtype=self.sdnq_dequantizer.result_dtype,
weight_shape=self.sdnq_dequantizer.result_shape,
quantized_weight_shape=self.sdnq_dequantizer.quantized_weight_shape,
weights_dtype=self.sdnq_dequantizer.weights_dtype,
)
quantized_embedding = compile_func(quantized_embedding)
modules/sdnq/layers/linear/forward.py
-23
View File
@@ -1,23 +0,0 @@
# pylint: disable=relative-beyond-top-level,redefined-builtin,protected-access
import torch
from ...kernel_wrappers import use_contiguous_int8_mm, use_contiguous_fp16_mm, use_contiguous_fp8_mm
def check_mats(input: torch.Tensor, weight: torch.Tensor, matmul_dtype: str = "int8") -> tuple[torch.Tensor, torch.Tensor]:
if input is not None:
input = input.contiguous()
if (
(use_contiguous_int8_mm and matmul_dtype in {"int8", "uint8"})
or (use_contiguous_fp16_mm and matmul_dtype in {"fp16", "float16"})
or (use_contiguous_fp8_mm and matmul_dtype in {"fp8", "float8_e4m3fn"})
):
weight = weight.contiguous()
elif weight.is_contiguous():
weight = weight.t().contiguous().t()
return input, weight
def quantized_linear_forward(self, input: torch.FloatTensor) -> torch.FloatTensor:
return torch.nn.functional.linear(input, self.sdnq_dequantizer(self.weight, self.scale, self.zero_point, self.svd_up, self.svd_down), self.bias)
modules/sdnq/layers/linear/linear_fp16.py
-98
View File
@@ -1,98 +0,0 @@
# pylint: disable=relative-beyond-top-level,redefined-builtin,protected-access
import torch
from ...common import compile_func
from ...kernel_wrappers import fp_scaled_mm_func, include_mm_kernel_in_compile
from ...quant_utils import rotate_hadamard, get_hadamard
from ...packed_float import unpack_float
from .forward import check_mats
from .linear_fp8 import quantize_fp_mm_input
def get_fp16_matmul_inputs(
input: torch.FloatTensor,
weight: torch.Tensor,
scale: torch.FloatTensor,
bias: torch.FloatTensor | None = None,
svd_up: torch.FloatTensor | None = None,
svd_down: torch.FloatTensor | None = None,
hadamard: torch.FloatTensor | None = None,
quantized_weight_shape: torch.Size | None = None,
weights_dtype: str | None = None,
) -> torch.FloatTensor:
if quantized_weight_shape is not None:
weight = unpack_float(weight, weights_dtype, quantized_weight_shape).to(dtype=torch.float16).t_()
scale = scale.t()
elif weight.dtype != torch.float16:
weight = weight.to(dtype=torch.float16) # fp8 weights
return_dtype = input.dtype
output_shape = (*input.shape[:-1], weight.shape[-1])
if hadamard is not None:
input = rotate_hadamard(input, hadamard=hadamard)
if svd_up is not None:
input = input.flatten(0,-2)
if bias is not None:
bias = torch.addmm(bias.to(dtype=svd_down.dtype), torch.mm(input.to(dtype=svd_down.dtype), svd_down), svd_up)
else:
bias = torch.mm(torch.mm(input.to(dtype=svd_down.dtype), svd_down), svd_up)
input, input_scale = quantize_fp_mm_input(input, dtype=scale.dtype, matmul_dtype="float16")
input, weight = check_mats(input, weight, matmul_dtype="float16")
return input, weight, input_scale, scale, bias, return_dtype, output_shape
def fp16_matmul(
input: torch.FloatTensor,
weight: torch.Tensor,
scale: torch.FloatTensor,
bias: torch.FloatTensor | None = None,
svd_up: torch.FloatTensor | None = None,
svd_down: torch.FloatTensor | None = None,
hadamard: torch.FloatTensor | None = None,
quantized_weight_shape: torch.Size | None = None,
weights_dtype: str | None = None,
) -> torch.FloatTensor:
input, weight, input_scale, scale, bias, return_dtype, output_shape = get_fp16_matmul_inputs(
input, weight, scale,
bias=bias,
svd_up=svd_up,
svd_down=svd_down,
hadamard=hadamard,
quantized_weight_shape=quantized_weight_shape,
weights_dtype=weights_dtype,
)
return fp_scaled_mm_func(input, weight, input_scale, scale, bias=bias, out_dtype=return_dtype).view(output_shape)
def quantized_linear_forward_fp16_matmul(self, input: torch.FloatTensor) -> torch.FloatTensor:
if torch.numel(input) / input.shape[-1] < 32:
return torch.nn.functional.linear(input, self.sdnq_dequantizer(self.weight, self.scale, zero_point=self.zero_point, svd_up=self.svd_up, svd_down=self.svd_down, skip_quantized_matmul=True), self.bias)
if self.sdnq_dequantizer.re_quantize_for_matmul:
weight, scale = self.sdnq_dequantizer.re_quantize_matmul(self.weight, self.scale, zero_point=self.zero_point)
quantized_weight_shape = None
else:
weight, scale = self.weight, self.scale
quantized_weight_shape = self.sdnq_dequantizer.quantized_weight_shape if self.sdnq_dequantizer.is_packed else None
if self.sdnq_dequantizer.use_hadamard:
hadamard = get_hadamard(self.sdnq_dequantizer.hadamard_group_size, dtype=input.dtype, device=input.device)
else:
hadamard = None
return fp16_matmul(
input, weight, scale,
bias=self.bias,
svd_up=self.svd_up,
svd_down=self.svd_down,
hadamard=hadamard,
quantized_weight_shape=quantized_weight_shape,
weights_dtype=self.sdnq_dequantizer.weights_dtype,
)
if not include_mm_kernel_in_compile:
get_fp16_matmul_inputs = compile_func(get_fp16_matmul_inputs)
else:
fp16_matmul = compile_func(fp16_matmul)
modules/sdnq/layers/linear/linear_fp8.py
-105
View File
@@ -1,105 +0,0 @@
# pylint: disable=relative-beyond-top-level,redefined-builtin,protected-access
import torch
from ...common import compile_func
from ...kernel_wrappers import fp8_scaled_mm_func, is_fp8_mm_supported, include_mm_kernel_in_compile
from ...quant_utils import quantize_fp_mm, rotate_hadamard, get_hadamard
from ...packed_float import unpack_float
from .forward import check_mats
def quantize_fp_mm_input(input: torch.FloatTensor, dtype: torch.dtype | None = None, matmul_dtype: str = "float8_e4m3fn") -> tuple[torch.Tensor, torch.FloatTensor]:
input = input.flatten(0,-2)
if dtype is not None:
input = input.to(dtype=dtype)
input, input_scale = quantize_fp_mm(input, dim=-1, matmul_dtype=matmul_dtype)
if input_scale.dtype == torch.float16: # fp16 will overflow
input_scale = input_scale.to(dtype=torch.float32)
return input, input_scale
def get_fp8_matmul_inputs(
input: torch.FloatTensor,
weight: torch.Tensor,
scale: torch.FloatTensor,
bias: torch.FloatTensor | None = None,
svd_up: torch.FloatTensor | None = None,
svd_down: torch.FloatTensor | None = None,
hadamard: torch.FloatTensor | None = None,
quantized_weight_shape: torch.Size | None = None,
weights_dtype: str | None = None,
) -> torch.FloatTensor:
if quantized_weight_shape is not None:
weight = unpack_float(weight, weights_dtype, quantized_weight_shape).to(dtype=torch.float8_e4m3fn).t_()
scale = scale.t()
return_dtype = input.dtype
output_shape = (*input.shape[:-1], weight.shape[-1])
if hadamard is not None:
input = rotate_hadamard(input, hadamard=hadamard)
if svd_up is not None:
input = input.flatten(0,-2)
if bias is not None:
bias = torch.addmm(bias.to(dtype=svd_down.dtype), torch.mm(input.to(dtype=svd_down.dtype), svd_down), svd_up)
else:
bias = torch.mm(torch.mm(input.to(dtype=svd_down.dtype), svd_down), svd_up)
input, input_scale = quantize_fp_mm_input(input, dtype=scale.dtype)
input, weight = check_mats(input, weight, matmul_dtype="float8_e4m3fn")
return input, weight, input_scale, scale, bias, return_dtype, output_shape
def fp8_matmul(
input: torch.FloatTensor,
weight: torch.Tensor,
scale: torch.FloatTensor,
bias: torch.FloatTensor | None = None,
svd_up: torch.FloatTensor | None = None,
svd_down: torch.FloatTensor | None = None,
hadamard: torch.FloatTensor | None = None,
quantized_weight_shape: torch.Size | None = None,
weights_dtype: str | None = None,
) -> torch.FloatTensor:
input, weight, input_scale, scale, bias, return_dtype, output_shape = get_fp8_matmul_inputs(
input, weight, scale,
bias=bias,
svd_up=svd_up,
svd_down=svd_down,
hadamard=hadamard,
quantized_weight_shape=quantized_weight_shape,
weights_dtype=weights_dtype,
)
return fp8_scaled_mm_func(input, weight, input_scale, scale, bias=bias, out_dtype=return_dtype).view(output_shape)
def quantized_linear_forward_fp8_matmul(self, input: torch.FloatTensor) -> torch.FloatTensor:
if torch.numel(input) / input.shape[-1] < 32:
return torch.nn.functional.linear(input, self.sdnq_dequantizer(self.weight, self.scale, zero_point=self.zero_point, svd_up=self.svd_up, svd_down=self.svd_down, skip_quantized_matmul=True), self.bias)
if self.sdnq_dequantizer.re_quantize_for_matmul:
weight, scale = self.sdnq_dequantizer.re_quantize_matmul(self.weight, self.scale, zero_point=self.zero_point)
quantized_weight_shape = None
else:
weight, scale = self.weight, self.scale
quantized_weight_shape = self.sdnq_dequantizer.quantized_weight_shape if self.sdnq_dequantizer.is_packed else None
if self.sdnq_dequantizer.use_hadamard:
hadamard = get_hadamard(self.sdnq_dequantizer.hadamard_group_size, dtype=input.dtype, device=input.device)
else:
hadamard = None
return fp8_matmul(
input, weight, scale,
bias=self.bias,
svd_up=self.svd_up,
svd_down=self.svd_down,
hadamard=hadamard,
quantized_weight_shape=quantized_weight_shape,
weights_dtype=self.sdnq_dequantizer.weights_dtype,
)
if is_fp8_mm_supported and not include_mm_kernel_in_compile:
get_fp8_matmul_inputs = compile_func(get_fp8_matmul_inputs)
else:
fp8_matmul = compile_func(fp8_matmul)
modules/sdnq/layers/linear/linear_int8.py
-126
View File
@@ -1,126 +0,0 @@
# pylint: disable=relative-beyond-top-level,redefined-builtin,protected-access
import torch
from ...common import compile_func
from ...kernel_wrappers import int_scaled_mm_func, include_mm_kernel_in_compile
from ...quant_utils import quantize_int_mm, rotate_hadamard, get_hadamard
from ...packed_int import unpack_int
from .forward import check_mats
def quantize_int_mm_input(input: torch.FloatTensor, dtype: torch.dtype | None = None, matmul_dtype: str = "int8") -> tuple[torch.Tensor, torch.FloatTensor]:
input = input.flatten(0,-2)
if dtype is not None:
input = input.to(dtype=dtype)
input, input_scale = quantize_int_mm(input, dim=-1, matmul_dtype=matmul_dtype)
if input_scale.dtype == torch.float16: # fp16 will overflow
input_scale = input_scale.to(dtype=torch.float32)
return input, input_scale
def get_int8_matmul_inputs(
input: torch.FloatTensor,
weight: torch.Tensor,
scale: torch.FloatTensor,
bias: torch.FloatTensor | None = None,
svd_up: torch.FloatTensor | None = None,
svd_down: torch.FloatTensor | None = None,
zero_point: torch.FloatTensor | None = None,
hadamard: torch.FloatTensor | None = None,
quantized_weight_shape: torch.Size | None = None,
weights_dtype: str | None = None,
) -> torch.FloatTensor:
if quantized_weight_shape is not None:
weight = unpack_int(weight, weights_dtype, quantized_weight_shape, dtype=torch.int8).t_()
scale = scale.t()
if zero_point is not None:
zero_point = zero_point.t()
if weight.dtype == torch.uint8:
weight = weight.view(dtype=torch.int8)
elif weight.dtype == torch.uint8:
weight = weight.bitwise_xor(128).view(torch.int8)
if zero_point is not None:
zero_point = torch.add(zero_point, scale, alpha=128)
else:
zero_point = torch.mul(scale, 128)
return_dtype = input.dtype
output_shape = (*input.shape[:-1], weight.shape[-1])
if hadamard is not None:
input = rotate_hadamard(input, hadamard=hadamard)
if svd_up is not None:
input = input.flatten(0,-2)
if bias is not None:
bias = torch.addmm(bias.to(dtype=svd_down.dtype), torch.mm(input.to(dtype=svd_down.dtype), svd_down), svd_up)
else:
bias = torch.mm(torch.mm(input.to(dtype=svd_down.dtype), svd_down), svd_up)
input, input_scale = quantize_int_mm_input(input, dtype=scale.dtype)
if zero_point is not None:
zero_bias = torch.sum(input, dim=-1, keepdim=True, dtype=torch.int32).to(dtype=input_scale.dtype).mul_(input_scale).mul(zero_point)
if bias is not None:
zero_bias.add_(bias)
bias = zero_bias
input, weight = check_mats(input, weight, matmul_dtype="int8")
return input, weight, input_scale, scale, bias, return_dtype, output_shape
def int8_matmul(
input: torch.FloatTensor,
weight: torch.Tensor,
scale: torch.FloatTensor,
bias: torch.FloatTensor | None = None,
svd_up: torch.FloatTensor | None = None,
svd_down: torch.FloatTensor | None = None,
zero_point: torch.FloatTensor | None = None,
hadamard: torch.FloatTensor | None = None,
quantized_weight_shape: torch.Size | None = None,
weights_dtype: str | None = None,
) -> torch.FloatTensor:
input, weight, input_scale, scale, bias, return_dtype, output_shape = get_int8_matmul_inputs(
input, weight, scale,
bias=bias,
svd_up=svd_up,
svd_down=svd_down,
zero_point=zero_point,
hadamard=hadamard,
quantized_weight_shape=quantized_weight_shape,
weights_dtype=weights_dtype,
)
return int_scaled_mm_func(input, weight, input_scale, scale, bias=bias, out_dtype=return_dtype).view(output_shape)
def quantized_linear_forward_int8_matmul(self, input: torch.FloatTensor) -> torch.FloatTensor:
if torch.numel(input) / input.shape[-1] < 32:
return torch.nn.functional.linear(input, self.sdnq_dequantizer(self.weight, self.scale, zero_point=self.zero_point, svd_up=self.svd_up, svd_down=self.svd_down, skip_quantized_matmul=True), self.bias)
if self.sdnq_dequantizer.re_quantize_for_matmul:
weight, scale = self.sdnq_dequantizer.re_quantize_matmul(self.weight, self.scale, zero_point=self.zero_point)
quantized_weight_shape = None
zero_point = None
else:
weight, scale, zero_point = self.weight, self.scale, self.zero_point
quantized_weight_shape = self.sdnq_dequantizer.quantized_weight_shape if self.sdnq_dequantizer.is_packed else None
if self.sdnq_dequantizer.use_hadamard:
hadamard = get_hadamard(self.sdnq_dequantizer.hadamard_group_size, dtype=input.dtype, device=input.device)
else:
hadamard = None
return int8_matmul(
input, weight, scale,
bias=self.bias,
svd_up=self.svd_up,
svd_down=self.svd_down,
zero_point=zero_point,
hadamard=hadamard,
quantized_weight_shape=quantized_weight_shape,
weights_dtype=self.sdnq_dequantizer.weights_dtype,
)
if not include_mm_kernel_in_compile:
get_int8_matmul_inputs = compile_func(get_int8_matmul_inputs)
else:
int8_matmul = compile_func(int8_matmul)
modules/sdnq/layers/linear/linear_uint8.py
-130
View File
@@ -1,130 +0,0 @@
# pylint: disable=relative-beyond-top-level,redefined-builtin,protected-access
import torch
from ...common import compile_func
from ...kernel_wrappers import int_scaled_mm_func, include_mm_kernel_in_compile
from ...quant_utils import quantize_uint_mm, rotate_hadamard, get_hadamard
from ...packed_int import unpack_int
from .forward import check_mats
def quantize_uint_mm_input(input: torch.FloatTensor, dtype: torch.dtype | None = None, matmul_dtype: str = "uint8") -> tuple[torch.Tensor, torch.FloatTensor, torch.FloatTensor]:
input = input.flatten(0,-2)
if dtype is not None:
input = input.to(dtype=dtype)
input, input_scale, input_zero_point = quantize_uint_mm(input, dim=-1, matmul_dtype=matmul_dtype)
if input_scale.dtype == torch.float16: # fp16 will overflow
input_scale = input_scale.to(dtype=torch.float32)
input_zero_point = input_zero_point.to(dtype=torch.float32)
return input, input_scale, input_zero_point
def get_uint8_matmul_inputs(
input: torch.FloatTensor,
weight: torch.Tensor,
scale: torch.FloatTensor,
zero_point: torch.FloatTensor,
bias: torch.FloatTensor | None = None,
svd_up: torch.FloatTensor | None = None,
svd_down: torch.FloatTensor | None = None,
hadamard: torch.FloatTensor | None = None,
quantized_weight_shape: torch.Size | None = None,
weights_dtype: str | None = None,
) -> torch.FloatTensor:
if quantized_weight_shape is not None:
weight = unpack_int(weight, weights_dtype, quantized_weight_shape, dtype=torch.int8).t_()
scale = scale.t()
if zero_point is not None:
zero_point = zero_point.t()
if weight.dtype == torch.uint8:
weight = weight.view(dtype=torch.int8)
elif weight.dtype == torch.uint8:
weight = weight.bitwise_xor(128).view(torch.int8)
if zero_point is not None:
zero_point = torch.add(zero_point, scale, alpha=128)
else:
zero_point = torch.mul(scale, 128)
return_dtype = input.dtype
output_shape = (*input.shape[:-1], weight.shape[-1])
if hadamard is not None:
input = rotate_hadamard(input, hadamard=hadamard)
if svd_up is not None:
input = input.flatten(0,-2)
if bias is not None:
bias = torch.addmm(bias.to(dtype=svd_down.dtype), torch.mm(input.to(dtype=svd_down.dtype), svd_down), svd_up)
else:
bias = torch.mm(torch.mm(input.to(dtype=svd_down.dtype), svd_down), svd_up)
input, input_scale, input_zero_point = quantize_uint_mm_input(input, dtype=scale.dtype)
if zero_point is not None:
zero_bias = torch.sum(input, dim=-1, keepdim=True, dtype=torch.int32).to(dtype=input_scale.dtype).mul_(input_scale).mul(zero_point)
zero_bias.add_(torch.sum(weight, dim=0, keepdim=True, dtype=torch.int32).to(dtype=scale.dtype).mul_(scale).mul(input_zero_point))
zero_bias.add_(torch.mul(input_zero_point, zero_point), alpha=input.shape[-1])
else:
zero_bias = torch.sum(weight, dim=0, keepdim=True, dtype=torch.int32).to(dtype=scale.dtype).mul_(scale).mul(input_zero_point)
if bias is not None:
zero_bias.add_(bias)
input, weight = check_mats(input, weight, matmul_dtype="uint8")
return input, weight, input_scale, scale, zero_bias, return_dtype, output_shape
def uint8_matmul(
input: torch.FloatTensor,
weight: torch.Tensor,
scale: torch.FloatTensor,
zero_point: torch.FloatTensor,
bias: torch.FloatTensor | None = None,
svd_up: torch.FloatTensor | None = None,
svd_down: torch.FloatTensor | None = None,
hadamard: torch.FloatTensor | None = None,
quantized_weight_shape: torch.Size | None = None,
weights_dtype: str | None = None,
) -> torch.FloatTensor:
input, weight, input_scale, scale, zero_bias, return_dtype, output_shape = get_uint8_matmul_inputs(
input, weight,
scale, zero_point,
bias=bias,
svd_up=svd_up,
svd_down=svd_down,
hadamard=hadamard,
quantized_weight_shape=quantized_weight_shape,
weights_dtype=weights_dtype,
)
return int_scaled_mm_func(input, weight, input_scale, scale, bias=zero_bias, out_dtype=return_dtype).view(output_shape)
def quantized_linear_forward_uint8_matmul(self, input: torch.FloatTensor) -> torch.FloatTensor:
if torch.numel(input) / input.shape[-1] < 32:
return torch.nn.functional.linear(input, self.sdnq_dequantizer(self.weight, self.scale, zero_point=self.zero_point, svd_up=self.svd_up, svd_down=self.svd_down, skip_quantized_matmul=True), self.bias)
if self.sdnq_dequantizer.re_quantize_for_matmul:
weight, scale, zero_point = self.sdnq_dequantizer.re_quantize_matmul(self.weight, self.scale, zero_point=self.zero_point)
quantized_weight_shape = None
else:
weight, scale, zero_point = self.weight, self.scale, self.zero_point
quantized_weight_shape = self.sdnq_dequantizer.quantized_weight_shape if self.sdnq_dequantizer.is_packed else None
if self.sdnq_dequantizer.use_hadamard:
hadamard = get_hadamard(self.sdnq_dequantizer.hadamard_group_size, dtype=input.dtype, device=input.device)
else:
hadamard = None
return uint8_matmul(
input, weight,
scale, zero_point,
bias=self.bias,
svd_up=self.svd_up,
svd_down=self.svd_down,
hadamard=hadamard,
quantized_weight_shape=quantized_weight_shape,
weights_dtype=self.sdnq_dequantizer.weights_dtype,
)
if not include_mm_kernel_in_compile:
get_uint8_matmul_inputs = compile_func(get_uint8_matmul_inputs)
else:
uint8_matmul = compile_func(uint8_matmul)
modules/sdnq/loader.py
-346
View File
@@ -1,346 +0,0 @@
import os
import json
import torch
from modules import shared
from .common import dtype_dict, check_torch_compile, linear_types
from .kernel_wrappers import is_fp8_mm_supported, use_tensorwise_fp8_matmul
from .quantizer import QuantizationMethod, SDNQConfig, SDNQQuantizer, sdnq_post_load_quant
from .quant_utils import prepare_weight_for_matmul, prepare_svd_for_matmul
from .utils import get_quant_args_from_config, check_param_name_in
from .forward import get_forward_func
from .file_loader import load_files
def get_module_names(model: torch.nn.Module) -> list:
modules_names = model._internal_dict.keys() # pylint: disable=protected-access
modules_names = [m for m in modules_names if not m.startswith("_")]
modules_names = [m for m in modules_names if isinstance(getattr(model, m, None), torch.nn.Module)]
modules_names = sorted(set(modules_names))
return modules_names
def normalize_tied_weights_keys_for_save(model: torch.nn.Module, is_pipeline: bool = False) -> list[tuple[torch.nn.Module, object]]:
normalized_modules = []
modules_to_walk = []
if is_pipeline:
for module_name in get_module_names(model):
module = getattr(model, module_name, None)
if isinstance(module, torch.nn.Module):
modules_to_walk.append(module)
elif isinstance(model, torch.nn.Module):
modules_to_walk.append(model)
for root_module in modules_to_walk:
for submodule in root_module.modules():
tied_weights_keys = getattr(submodule, "_tied_weights_keys", None)
if isinstance(tied_weights_keys, list):
normalized_modules.append((submodule, tied_weights_keys))
submodule._tied_weights_keys = {key: key for key in tied_weights_keys} # pylint: disable=protected-access
return normalized_modules
def restore_tied_weights_keys_after_save(normalized_modules: list[tuple[torch.nn.Module, object]]) -> None:
for submodule, tied_weights_keys in normalized_modules:
submodule._tied_weights_keys = tied_weights_keys # pylint: disable=protected-access
def save_sdnq_model(
model: torch.nn.Module,
model_path: str,
max_shard_size: str = "5GB",
is_pipeline: bool = False,
sdnq_config: SDNQConfig | None = None,
) -> None:
normalized_modules = normalize_tied_weights_keys_for_save(model, is_pipeline=is_pipeline)
try:
model.save_pretrained(model_path, max_shard_size=max_shard_size) # actual save
finally:
restore_tied_weights_keys_after_save(normalized_modules)
quantization_config_path = os.path.join(model_path, "quantization_config.json")
if sdnq_config is not None: # if provided, save global config
sdnq_config.to_json_file(quantization_config_path)
if is_pipeline:
for module_name in get_module_names(model): # save per-module config if available
module = getattr(model, module_name, None)
if module is None:
continue
module_quantization_config_path = os.path.join(model_path, module_name, "quantization_config.json")
if hasattr(module, "quantization_config") and isinstance(module.quantization_config, SDNQConfig):
module.quantization_config.to_json_file(module_quantization_config_path)
elif hasattr(module, "config") and hasattr(module.config, "quantization_config") and isinstance(module.config.quantization_config, SDNQConfig):
module.config.quantization_config.to_json_file(module_quantization_config_path)
elif sdnq_config is None:
if hasattr(model, "quantization_config") and isinstance(model.quantization_config, SDNQConfig):
model.quantization_config.to_json_file(quantization_config_path)
elif hasattr(model, "config") and hasattr(model.config, "quantization_config") and isinstance(model.config.quantization_config, SDNQConfig):
model.config.quantization_config.to_json_file(quantization_config_path)
def load_sdnq_model(
model_path: str,
model_cls: torch.nn.Module | None = None,
file_name: str | None = None,
dtype: torch.dtype | None = None,
device: torch.device = "cpu",
dequantize_fp32: bool | None = None,
use_quantized_matmul: bool | None = None,
model_config: dict | None = None,
quantization_config: dict | None = None,
load_method: str = "safetensors",
) -> torch.nn.Module:
from accelerate import init_empty_weights
with init_empty_weights():
model_config_path = os.path.join(model_path, "config.json")
quantization_config_path = os.path.join(model_path, "quantization_config.json")
if model_config is None:
if os.path.exists(model_config_path):
with open(model_config_path, encoding="utf-8") as f:
model_config = json.load(f)
else:
model_config = {}
if quantization_config is None:
if os.path.exists(quantization_config_path):
with open(quantization_config_path, encoding="utf-8") as f:
quantization_config = json.load(f)
else:
quantization_config = model_config.get("quantization_config", None)
if quantization_config is None:
raise ValueError(f"Cannot determine quantization_config for {model_path}, please provide quantization_config argument")
if not isinstance(quantization_config, SDNQConfig):
quantization_config = SDNQConfig.from_dict(quantization_config)
if model_cls is None:
import transformers
import diffusers
class_name = model_config.get("_class_name", None) or model_config.get("architectures", None)
if isinstance(class_name, list):
class_name = class_name[0]
if class_name is not None:
model_cls = getattr(diffusers, class_name, None) or getattr(transformers, class_name, None)
if model_cls is None:
raise ValueError(f"Cannot determine model class for {model_path}, please provide model_cls argument")
if hasattr(model_cls, "load_config") and hasattr(model_cls, "from_config"):
config = model_cls.load_config(model_path)
if hasattr(config, "quantization_config"):
del config.quantization_config
if hasattr(config, "pop"):
config.pop("quantization_config", None)
model = model_cls.from_config(config)
elif hasattr(model_cls, "_from_config"):
config = transformers.AutoConfig.from_pretrained(model_path)
if hasattr(config, "quantization_config"):
del config.quantization_config
if hasattr(config, "pop"):
config.pop("quantization_config", None)
model = model_cls(config)
else:
if hasattr(model_config, "quantization_config"):
del model_config.quantization_config
if hasattr(model_config, "pop"):
model_config.pop("quantization_config", None)
model = model_cls(**model_config)
model = sdnq_post_load_quant(model, torch_dtype=dtype, pre_quantized=True, **get_quant_args_from_config(quantization_config))
key_mapping = getattr(model, "_checkpoint_conversion_mapping", None)
files = []
if file_name:
files.append(os.path.join(model_path, file_name))
else:
all_files = os.listdir(model_path)
files = sorted([os.path.join(model_path, f) for f in all_files if f.endswith(".safetensors")]) # pylint: disable=not-an-iterable
state_dict = load_files(files, key_mapping=key_mapping, device=device, method=load_method)
if isinstance(getattr(model, "_tied_weights_keys", None), dict):
for key, value in model._tied_weights_keys.items(): # pylint: disable=protected-access
if value in state_dict and key not in state_dict:
state_dict[key] = state_dict[value]
else:
# older transformers case, handle known models manually
if model.__class__.__name__ in {"T5EncoderModel", "UMT5EncoderModel"} and "encoder.embed_tokens.weight" not in state_dict:
state_dict["encoder.embed_tokens.weight"] = state_dict["shared.weight"]
elif model.__class__.__name__ in {"Qwen3ForCausalLM"} and "lm_head.weight" not in state_dict and "model.embed_tokens.weight" in state_dict:
state_dict["lm_head.weight"] = state_dict["model.embed_tokens.weight"]
model.load_state_dict(state_dict, assign=True)
del state_dict
model.quantization_config = quantization_config
model.quantization_method = QuantizationMethod.SDNQ
if hasattr(model, "config"):
try:
model.config.quantization_config = quantization_config
except Exception:
pass
try:
model.config["quantization_config"] = quantization_config.to_dict()
except Exception:
pass
if hasattr(model, "hf_quantizer"):
model.hf_quantizer.quantization_config = quantization_config
else:
model.hf_quantizer = SDNQQuantizer(quantization_config)
model = post_process_model(model)
if (dtype is not None) or (dequantize_fp32 is not None) or (use_quantized_matmul is not None):
model = apply_sdnq_options_to_model(model, dtype=dtype, dequantize_fp32=dequantize_fp32, use_quantized_matmul=use_quantized_matmul)
return model
def post_process_model(model: torch.nn.Module) -> torch.nn.Module:
has_children = list(model.children())
if not has_children:
return model
for module_name, module in model.named_children():
if hasattr(module, "sdnq_dequantizer"):
module.weight.requires_grad_(False)
module.scale.requires_grad_(False)
if module.zero_point is not None:
module.zero_point.requires_grad_(False)
if module.sdnq_dequantizer.use_quantized_matmul and not module.sdnq_dequantizer.re_quantize_for_matmul:
module.weight.data = prepare_weight_for_matmul(module.weight, matmul_dtype=module.sdnq_dequantizer.quantized_matmul_dtype)
if module.svd_up is not None:
module.svd_up.requires_grad_(False)
module.svd_down.requires_grad_(False)
module.svd_up.data, module.svd_down.data = prepare_svd_for_matmul(module.svd_up, module.svd_down, module.sdnq_dequantizer.use_quantized_matmul)
setattr(model, module_name, module)
else:
setattr(model, module_name, post_process_model(module))
return model
def apply_sdnq_options_to_module(
model: torch.nn.Module,
quantization_config: SDNQConfig,
dtype: torch.dtype | None = None,
dequantize_fp32: bool | None = None,
use_quantized_matmul: bool | None = None,
full_param_name: str = "",
) -> torch.nn.Module:
has_children = list(model.children())
if not has_children:
if dtype is not None and getattr(model, "dtype", torch.float32) not in {torch.float32, torch.float64}:
model = model.to(dtype=dtype)
return model
for module_name, module in model.named_children():
if full_param_name:
param_name = full_param_name + "." + module_name
else:
param_name = module_name
if hasattr(module, "sdnq_dequantizer"):
param_name = param_name + ".weight"
layer_class_name = module.original_class.__name__
current_use_quantized_matmul = use_quantized_matmul
if layer_class_name in linear_types:
if not is_fp8_mm_supported and module.sdnq_dequantizer.quantized_matmul_dtype in {"fp8", "float8_e4m3fn"}:
current_use_quantized_matmul = False
elif check_param_name_in(param_name, quantization_config.modules_to_not_use_matmul) is not None:
current_use_quantized_matmul = None
if current_use_quantized_matmul:
output_channel_size, channel_size = module.sdnq_dequantizer.original_shape
current_use_quantized_matmul = current_use_quantized_matmul and channel_size >= 32 and output_channel_size >= 32
current_use_quantized_matmul = current_use_quantized_matmul and output_channel_size % 16 == 0 and channel_size % 16 == 0
else:
current_use_quantized_matmul = None
if dtype is not None and module.sdnq_dequantizer.result_dtype not in {torch.float32, torch.float64}:
module.sdnq_dequantizer.result_dtype = dtype
if module.svd_up is not None:
module.svd_up.data = module.svd_up.to(dtype=dtype)
module.svd_down.data = module.svd_down.to(dtype=dtype)
upcast_scale = bool(
dequantize_fp32
or dtype_dict[module.sdnq_dequantizer.weights_dtype]["num_bits"] > 8
or (
(current_use_quantized_matmul or (current_use_quantized_matmul is None and module.sdnq_dequantizer.use_quantized_matmul))
and not dtype_dict[module.sdnq_dequantizer.quantized_matmul_dtype]["is_integer"]
and (not use_tensorwise_fp8_matmul or dtype_dict[module.sdnq_dequantizer.quantized_matmul_dtype]["num_bits"] == 16)
)
)
if upcast_scale or dequantize_fp32:
if module.scale.dtype in {torch.float32, torch.float64}:
scale_dtype = module.scale.dtype
else:
scale_dtype = torch.float32 if module.sdnq_dequantizer.result_dtype != torch.float64 else torch.float64
elif dequantize_fp32 is None and module.scale.dtype in {torch.float32, torch.float64}:
scale_dtype = module.scale.dtype
else:
scale_dtype = module.sdnq_dequantizer.result_dtype
module.scale.data = module.scale.to(dtype=scale_dtype)
if module.zero_point is not None:
module.zero_point.data = module.zero_point.to(dtype=scale_dtype)
if current_use_quantized_matmul is not None:
if current_use_quantized_matmul != module.sdnq_dequantizer.use_quantized_matmul:
if not module.sdnq_dequantizer.re_quantize_for_matmul and not dtype_dict[module.sdnq_dequantizer.weights_dtype]["is_packed"]:
module.scale.data = module.scale.t_().contiguous()
module.weight.data = module.weight.t_()
if module.zero_point is not None:
module.zero_point.data = module.zero_point.t_().contiguous()
if current_use_quantized_matmul:
module.weight.data = prepare_weight_for_matmul(module.weight, matmul_dtype=module.sdnq_dequantizer.quantized_matmul_dtype)
else:
module.scale.data = module.scale.contiguous()
module.weight.data = module.weight.contiguous()
if module.svd_up is not None:
module.svd_up.data, module.svd_down.data = prepare_svd_for_matmul(module.svd_up.t_(), module.svd_down.t_(), current_use_quantized_matmul)
module.sdnq_dequantizer.use_quantized_matmul = current_use_quantized_matmul
module.forward_func = get_forward_func(module.original_class.__name__, module.sdnq_dequantizer.quantized_matmul_dtype, current_use_quantized_matmul)
if (
not module.sdnq_dequantizer.use_quantized_matmul
and (use_quantized_matmul or (use_quantized_matmul is None and quantization_config.use_quantized_matmul))
and check_param_name_in(param_name, quantization_config.modules_to_not_use_matmul) is None
):
quantization_config.modules_to_not_use_matmul.append(param_name)
setattr(model, module_name, module)
else:
setattr(model, module_name, apply_sdnq_options_to_module(module, quantization_config, dtype=dtype, dequantize_fp32=dequantize_fp32, use_quantized_matmul=use_quantized_matmul, full_param_name=param_name))
return model
def apply_sdnq_options_to_model(model: torch.nn.Module, dtype: torch.dtype | None = None, dequantize_fp32: bool | None = None, use_quantized_matmul: bool | None = None) -> torch.nn.Module:
if use_quantized_matmul and not check_torch_compile():
shared.log.warning("SDNQ: Quantized MatMul requires a working Triton install for best performance.")
model = apply_sdnq_options_to_module(model, model.quantization_config, dtype=dtype, dequantize_fp32=dequantize_fp32, use_quantized_matmul=use_quantized_matmul)
if hasattr(model, "quantization_config"):
if use_quantized_matmul is not None:
model.quantization_config.use_quantized_matmul = use_quantized_matmul
if dequantize_fp32 is not None:
model.quantization_config.dequantize_fp32 = dequantize_fp32
if hasattr(model, "config"):
try:
if hasattr(model.config, "quantization_config"):
if use_quantized_matmul is not None:
model.config.quantization_config.use_quantized_matmul = use_quantized_matmul
if dequantize_fp32 is not None:
model.config.quantization_config.dequantize_fp32 = dequantize_fp32
except Exception:
pass
try:
if hasattr(model.config, "get") and model.config.get("quantization_config", None) is not None:
if use_quantized_matmul is not None:
model.config["quantization_config"].use_quantized_matmul = use_quantized_matmul
if dequantize_fp32 is not None:
model.config["quantization_config"].dequantize_fp32 = dequantize_fp32
except Exception:
pass
if hasattr(model, "hf_quantizer"):
if use_quantized_matmul is not None:
model.hf_quantizer.quantization_config.use_quantized_matmul = use_quantized_matmul
if dequantize_fp32 is not None:
model.hf_quantizer.quantization_config.dequantize_fp32 = dequantize_fp32
return model
modules/sdnq/packed_float.py
-129
View File
@@ -1,129 +0,0 @@
import torch
from .common import dtype_dict
from .packed_int import pack_int, unpack_int
float_bits_to_uint_dict = {
1: "uint1",
2: "uint2",
3: "uint3",
4: "uint4",
5: "uint5",
6: "uint6",
7: "uint7",
9: "uint9",
10: "uint10",
11: "uint11",
12: "uint12",
13: "uint13",
14: "uint14",
15: "uint15",
}
def pack_float(x: torch.FloatTensor, weights_dtype: str) -> torch.Tensor:
exponent_bits = dtype_dict[weights_dtype]["exponent"]
mantissa_bits = dtype_dict[weights_dtype]["mantissa"]
total_bits = dtype_dict[weights_dtype]["num_bits"]
if dtype_dict[weights_dtype]["is_unsigned"]:
sign_mask = (1 << (total_bits-1)) # pylint: disable=superfluous-parens
else:
sign_mask = (1 << (total_bits-1)) + (1 << (total_bits-2))
mantissa_difference = 23 - mantissa_bits
exponent_difference = 8 - exponent_bits
mantissa_mask = (1 << mantissa_difference) # pylint: disable=superfluous-parens
x = x.to(dtype=torch.float32).view(torch.int32)
x = torch.where(
torch.gt(
torch.bitwise_and(x, -(1 << (mantissa_difference-4)) & ~(-mantissa_mask)),
(1 << (mantissa_difference-1)),
),
torch.add(x, mantissa_mask),
x,
)
if exponent_bits < 8:
min_normal = 2.0 ** (2 - (1 << (exponent_bits - 1)))
x_f32_abs = x.view(dtype=torch.float32).abs()
is_subnormal = torch.lt(x_f32_abs, min_normal)
x = torch.where(
is_subnormal,
torch.bitwise_or(
torch.bitwise_and(x, -2147483648),
torch.bitwise_left_shift(
x_f32_abs.mul_((1 << mantissa_bits) / min_normal).round_().to(torch.int32),
mantissa_difference,
),
),
x,
)
del x_f32_abs, is_subnormal
x = torch.bitwise_right_shift(x, mantissa_difference)
x = torch.bitwise_and(
torch.bitwise_or(
torch.bitwise_and(torch.bitwise_right_shift(x, exponent_difference), sign_mask),
torch.bitwise_and(x, ~sign_mask),
),
~(-(1 << total_bits)),
).view(torch.uint32)
if total_bits not in {8, 16}:
x = pack_int(x, float_bits_to_uint_dict[total_bits])
else:
x = x.to(dtype=dtype_dict[weights_dtype]["storage_dtype"])
return x
def unpack_float(x: torch.Tensor, weights_dtype: str, shape: torch.Size) -> torch.FloatTensor:
exponent_bits = dtype_dict[weights_dtype]["exponent"]
mantissa_bits = dtype_dict[weights_dtype]["mantissa"]
total_bits = dtype_dict[weights_dtype]["num_bits"]
if dtype_dict[weights_dtype]["is_unsigned"]:
sign_mask = (1 << (total_bits-1)) # pylint: disable=superfluous-parens
else:
sign_mask = (1 << (total_bits-1)) + (1 << (total_bits-2))
mantissa_difference = 23 - mantissa_bits
exponent_difference = 8 - exponent_bits
if total_bits not in {8, 16}:
x = unpack_int(x, float_bits_to_uint_dict[total_bits], shape)
x = x.to(dtype=torch.uint32).view(torch.int32)
x = torch.bitwise_left_shift(
torch.bitwise_or(
torch.bitwise_left_shift(torch.bitwise_and(x, sign_mask), exponent_difference),
torch.bitwise_and(x, ~sign_mask),
),
mantissa_difference,
)
x = torch.bitwise_or(
x,
torch.bitwise_and(
torch.bitwise_right_shift(
-torch.bitwise_and(torch.bitwise_not(x), 1073741824),
exponent_difference,
),
1065353216,
),
)
overflow_mask = (~(-(1 << (22 + exponent_bits))) | 1090519039)
x = torch.where(torch.bitwise_and(x, overflow_mask).to(dtype=torch.bool), x, 0)
x = x.view(torch.float32)
if exponent_bits < 8:
min_normal = 2.0 ** (2 - (1 << (exponent_bits - 1)))
x = torch.where(
torch.lt(x.abs(), min_normal),
torch.sign(x).mul_(-min_normal).add_(x, alpha=2.0),
x,
)
return x
modules/sdnq/packed_int/__init__.py
-85
View File
@@ -1,85 +0,0 @@
import torch
from ..common import dtype_dict
from .pack import (
pack_uint15,
pack_uint14,
pack_uint13,
pack_uint12,
pack_uint11,
pack_uint10,
pack_uint9,
pack_uint7,
pack_uint6,
pack_uint5,
pack_uint4,
pack_uint3,
pack_uint2,
pack_uint1,
)
from .unpack import (
unpack_uint15,
unpack_uint14,
unpack_uint13,
unpack_uint12,
unpack_uint11,
unpack_uint10,
unpack_uint9,
unpack_uint7,
unpack_uint6,
unpack_uint5,
unpack_uint4,
unpack_uint3,
unpack_uint2,
unpack_uint1,
)
packed_int_function_dict = {
"uint15": {"pack": pack_uint15, "unpack": unpack_uint15},
"uint14": {"pack": pack_uint14, "unpack": unpack_uint14},
"uint13": {"pack": pack_uint13, "unpack": unpack_uint13},
"uint12": {"pack": pack_uint12, "unpack": unpack_uint12},
"uint11": {"pack": pack_uint11, "unpack": unpack_uint11},
"uint10": {"pack": pack_uint10, "unpack": unpack_uint10},
"uint9": {"pack": pack_uint9, "unpack": unpack_uint9},
"uint7": {"pack": pack_uint7, "unpack": unpack_uint7},
"uint6": {"pack": pack_uint6, "unpack": unpack_uint6},
"uint5": {"pack": pack_uint5, "unpack": unpack_uint5},
"uint4": {"pack": pack_uint4, "unpack": unpack_uint4},
"uint3": {"pack": pack_uint3, "unpack": unpack_uint3},
"uint2": {"pack": pack_uint2, "unpack": unpack_uint2},
"uint1": {"pack": pack_uint1, "unpack": unpack_uint1},
}
packed_int_function_dict["int15"] = packed_int_function_dict["uint15"]
packed_int_function_dict["int14"] = packed_int_function_dict["uint14"]
packed_int_function_dict["int13"] = packed_int_function_dict["uint13"]
packed_int_function_dict["int12"] = packed_int_function_dict["uint12"]
packed_int_function_dict["int11"] = packed_int_function_dict["uint11"]
packed_int_function_dict["int10"] = packed_int_function_dict["uint10"]
packed_int_function_dict["int9"] = packed_int_function_dict["uint9"]
packed_int_function_dict["int7"] = packed_int_function_dict["uint7"]
packed_int_function_dict["int6"] = packed_int_function_dict["uint6"]
packed_int_function_dict["int5"] = packed_int_function_dict["uint5"]
packed_int_function_dict["int4"] = packed_int_function_dict["uint4"]
packed_int_function_dict["int3"] = packed_int_function_dict["uint3"]
packed_int_function_dict["int2"] = packed_int_function_dict["uint2"]
packed_int_function_dict["int1"] = packed_int_function_dict["uint1"]
packed_int_function_dict["bool"] = packed_int_function_dict["uint1"]
def pack_int(tensor: torch.Tensor, weights_dtype: str) -> torch.Tensor:
if not dtype_dict[weights_dtype]["is_unsigned"]:
tensor = tensor.sub(dtype_dict[weights_dtype]["min"])
return packed_int_function_dict[weights_dtype]["pack"](tensor.to(dtype=dtype_dict[weights_dtype]["storage_dtype"]))
def unpack_int(packed_tensor: torch.Tensor, weights_dtype: str, shape: torch.Size, dtype: torch.dtype = None) -> torch.Tensor:
packed_tensor = packed_int_function_dict[weights_dtype]["unpack"](packed_tensor, shape)
if not dtype_dict[weights_dtype]["is_unsigned"]:
packed_tensor = packed_tensor.to(dtype=dtype_dict[weights_dtype]["torch_dtype"] if dtype is None else dtype).add_(dtype_dict[weights_dtype]["min"])
return packed_tensor
modules/sdnq/packed_int/pack.py
-305
View File
@@ -1,305 +0,0 @@
import torch
def pack_uint15(tensor: torch.Tensor) -> torch.Tensor:
packed_tensor = tensor.contiguous().view(-1, 16)
packed_tensor = torch.bitwise_or(
packed_tensor[:, :15],
torch.bitwise_and(
torch.stack(
(
torch.bitwise_left_shift(packed_tensor[:, 15], 1),
torch.bitwise_left_shift(packed_tensor[:, 15], 2),
torch.bitwise_left_shift(packed_tensor[:, 15], 3),
torch.bitwise_left_shift(packed_tensor[:, 15], 4),
torch.bitwise_left_shift(packed_tensor[:, 15], 5),
torch.bitwise_left_shift(packed_tensor[:, 15], 6),
torch.bitwise_left_shift(packed_tensor[:, 15], 7),
torch.bitwise_left_shift(packed_tensor[:, 15], 8),
torch.bitwise_left_shift(packed_tensor[:, 15], 9),
torch.bitwise_left_shift(packed_tensor[:, 15], 10),
torch.bitwise_left_shift(packed_tensor[:, 15], 11),
torch.bitwise_left_shift(packed_tensor[:, 15], 12),
torch.bitwise_left_shift(packed_tensor[:, 15], 13),
torch.bitwise_left_shift(packed_tensor[:, 15], 14),
torch.bitwise_left_shift(packed_tensor[:, 15], 15),
),
dim=-1
),
32768
),
)
return packed_tensor
def pack_uint14(tensor: torch.Tensor) -> torch.Tensor:
packed_tensor = tensor.contiguous().view(-1, 8)
packed_tensor = torch.bitwise_or(
packed_tensor[:, :7],
torch.bitwise_and(
torch.stack(
(
torch.bitwise_left_shift(packed_tensor[:, 7], 2),
torch.bitwise_left_shift(packed_tensor[:, 7], 4),
torch.bitwise_left_shift(packed_tensor[:, 7], 6),
torch.bitwise_left_shift(packed_tensor[:, 7], 8),
torch.bitwise_left_shift(packed_tensor[:, 7], 10),
torch.bitwise_left_shift(packed_tensor[:, 7], 12),
torch.bitwise_left_shift(packed_tensor[:, 7], 14),
),
dim=-1
),
49152
),
)
return packed_tensor
def pack_uint13(tensor: torch.Tensor) -> torch.Tensor:
packed_tensor = tensor.contiguous().view(-1, 16)
packed_tensor = torch.bitwise_or(
packed_tensor[:, :13],
torch.bitwise_and(
torch.cat(
(
torch.bitwise_left_shift(packed_tensor[:, 13:], 13),
torch.bitwise_left_shift(packed_tensor[:, 13:], 10),
torch.bitwise_left_shift(packed_tensor[:, 13:], 7),
torch.bitwise_left_shift(packed_tensor[:, 13:], 4),
torch.bitwise_or(
torch.bitwise_or(
torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 13], 1), 8192),
torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 14], 2), 16384),
),
torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 15], 3), 32768),
).unsqueeze(-1),
),
dim=-1,
),
57344,
),
)
return packed_tensor
def pack_uint12(tensor: torch.Tensor) -> torch.Tensor:
packed_tensor = tensor.contiguous().view(-1, 4)
packed_tensor = torch.bitwise_or(
packed_tensor[:, :3],
torch.bitwise_and(
torch.stack(
(
torch.bitwise_left_shift(packed_tensor[:, 3], 4),
torch.bitwise_left_shift(packed_tensor[:, 3], 8),
torch.bitwise_left_shift(packed_tensor[:, 3], 12),
),
dim=-1
),
61440
)
)
return packed_tensor
def pack_uint11(tensor: torch.Tensor) -> torch.Tensor:
packed_tensor = tensor.contiguous().view(-1, 16)
packed_tensor = torch.cat(
(
torch.bitwise_or(packed_tensor[:, :8], torch.bitwise_left_shift(packed_tensor[:, 8:], 11)),
torch.bitwise_or(
torch.bitwise_or(
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 8:11], 5), 63),
torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 11:14], 1), 4032),
),
torch.cat(
(
torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 14:], 7), -4096),
torch.bitwise_or(
torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 14], 3), 12288),
torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 15], 5), -16384),
).unsqueeze(-1),
),
dim=-1,
),
),
),
dim=-1,
)
return packed_tensor
def pack_uint10(tensor: torch.Tensor) -> torch.Tensor:
packed_tensor = tensor.contiguous().view(-1, 8)
packed_tensor = torch.cat(
(
torch.bitwise_or(packed_tensor[:, :3], torch.bitwise_left_shift(packed_tensor[:, 5:8], 10)),
torch.bitwise_or(
packed_tensor[:, 3:5],
torch.bitwise_or(
torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 5:7], 4), 15360),
torch.bitwise_and(
torch.stack(
(
torch.bitwise_left_shift(packed_tensor[:, 7], 6),
torch.bitwise_left_shift(packed_tensor[:, 7], 8),
),
dim=-1,
),
49152
),
),
),
),
dim=-1
)
return packed_tensor
def pack_uint9(tensor: torch.Tensor) -> torch.Tensor:
packed_tensor = tensor.contiguous().view(-1, 16)
packed_tensor = torch.cat(
(
torch.bitwise_or(packed_tensor[:, :8], torch.bitwise_left_shift(packed_tensor[:, 8:], 9)),
torch.bitwise_or(
torch.bitwise_or(
torch.bitwise_or(
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 8], 7), 3),
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 9], 5), 12),
),
torch.bitwise_or(
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 10], 3), 48),
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 11], 1), 192),
),
),
torch.bitwise_or(
torch.bitwise_or(
torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 12], 1), 768),
torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 13], 3), 3072),
),
torch.bitwise_or(
torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 14], 5), 12288),
torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 15], 7), 49152),
),
),
).unsqueeze(-1),
),
dim=-1,
)
return packed_tensor
def pack_uint7(tensor: torch.ByteTensor) -> torch.ByteTensor:
packed_tensor = tensor.contiguous().view(-1, 8)
packed_tensor = torch.bitwise_or(
packed_tensor[:, :7],
torch.bitwise_and(
torch.stack(
(
torch.bitwise_left_shift(packed_tensor[:, 7], 1),
torch.bitwise_left_shift(packed_tensor[:, 7], 2),
torch.bitwise_left_shift(packed_tensor[:, 7], 3),
torch.bitwise_left_shift(packed_tensor[:, 7], 4),
torch.bitwise_left_shift(packed_tensor[:, 7], 5),
torch.bitwise_left_shift(packed_tensor[:, 7], 6),
torch.bitwise_left_shift(packed_tensor[:, 7], 7),
),
dim=-1
),
128
),
)
return packed_tensor
def pack_uint6(tensor: torch.ByteTensor) -> torch.ByteTensor:
packed_tensor = tensor.contiguous().view(-1, 4)
packed_tensor = torch.bitwise_or(
packed_tensor[:, :3],
torch.bitwise_and(
torch.stack(
(
torch.bitwise_left_shift(packed_tensor[:, 3], 2),
torch.bitwise_left_shift(packed_tensor[:, 3], 4),
torch.bitwise_left_shift(packed_tensor[:, 3], 6),
),
dim=-1
),
192
)
)
return packed_tensor
def pack_uint5(tensor: torch.ByteTensor) -> torch.ByteTensor:
packed_tensor = tensor.contiguous().view(-1, 8)
packed_tensor = torch.cat(
(
torch.bitwise_or(packed_tensor[:, :3], torch.bitwise_left_shift(packed_tensor[:, 5:8], 5)),
torch.bitwise_or(
packed_tensor[:, 3:5],
torch.bitwise_or(
torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 5:7], 2), 96),
torch.bitwise_and(
torch.stack(
(
torch.bitwise_left_shift(packed_tensor[:, 7], 3),
torch.bitwise_left_shift(packed_tensor[:, 7], 4),
),
dim=-1,
),
128,
),
),
),
),
dim=-1
)
return packed_tensor
def pack_uint4(tensor: torch.ByteTensor) -> torch.ByteTensor:
packed_tensor = tensor.contiguous().view(-1, 2)
packed_tensor = torch.bitwise_or(packed_tensor[:, 0], torch.bitwise_left_shift(packed_tensor[:, 1], 4))
return packed_tensor
def pack_uint3(tensor: torch.ByteTensor) -> torch.ByteTensor:
packed_tensor = tensor.contiguous().view(-1, 8)
packed_tensor = torch.bitwise_or(
torch.bitwise_or(packed_tensor[:, :3], torch.bitwise_left_shift(packed_tensor[:, 3:6], 3)),
torch.cat(
(
torch.bitwise_left_shift(packed_tensor[:, 6:8], 6),
torch.bitwise_or(
torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 6], 4), 64),
torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 7], 5), 128),
).unsqueeze(-1),
),
dim=-1
)
)
return packed_tensor
def pack_uint2(tensor: torch.ByteTensor) -> torch.ByteTensor:
packed_tensor = tensor.contiguous().view(-1, 4)
packed_tensor = torch.bitwise_or(
torch.bitwise_or(packed_tensor[:, 0], torch.bitwise_left_shift(packed_tensor[:, 1], 2)),
torch.bitwise_or(torch.bitwise_left_shift(packed_tensor[:, 2], 4), torch.bitwise_left_shift(packed_tensor[:, 3], 6)),
)
return packed_tensor
def pack_uint1(tensor: torch.Tensor) -> torch.Tensor:
packed_tensor = tensor.contiguous().view(-1, 8)
packed_tensor = torch.bitwise_or(
torch.bitwise_or(
torch.bitwise_or(packed_tensor[:, 0], torch.bitwise_left_shift(packed_tensor[:, 1], 1)),
torch.bitwise_or(torch.bitwise_left_shift(packed_tensor[:, 2], 2), torch.bitwise_left_shift(packed_tensor[:, 3], 3))
),
torch.bitwise_or(
torch.bitwise_or(torch.bitwise_left_shift(packed_tensor[:, 4], 4), torch.bitwise_left_shift(packed_tensor[:, 5], 5)),
torch.bitwise_or(torch.bitwise_left_shift(packed_tensor[:, 6], 6), torch.bitwise_left_shift(packed_tensor[:, 7], 7))
),
)
return packed_tensor
modules/sdnq/packed_int/unpack.py
-356
View File
@@ -1,356 +0,0 @@
import torch
def unpack_uint15(packed_tensor: torch.Tensor, shape: torch.Size) -> torch.Tensor:
result = torch.cat(
(
torch.bitwise_and(packed_tensor[:, :15], 32767),
torch.bitwise_or(
torch.bitwise_or(
torch.bitwise_or(
torch.bitwise_or(
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 0], 1), 16384),
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 1], 2), 8192),
),
torch.bitwise_or(
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 2], 3), 4096),
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 3], 4), 2048),
),
),
torch.bitwise_or(
torch.bitwise_or(
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 4], 5), 1024),
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 5], 6), 512),
),
torch.bitwise_or(
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 6], 7), 256),
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 7], 8), 128),
),
),
),
torch.bitwise_or(
torch.bitwise_or(
torch.bitwise_or(
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 8], 9), 64),
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 9], 10), 32),
),
torch.bitwise_or(
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 10], 11), 16),
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 11], 12), 8),
),
),
torch.bitwise_or(
torch.bitwise_or(
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 12], 13), 4),
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 13], 14), 2),
),
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 14], 15), 1),
),
)
).unsqueeze(-1)
),
dim=-1
).view(shape)
return result
def unpack_uint14(packed_tensor: torch.Tensor, shape: torch.Size) -> torch.Tensor:
result = torch.cat(
(
torch.bitwise_and(packed_tensor[:, :7], 16383),
torch.bitwise_or(
torch.bitwise_or(
torch.bitwise_or(
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 0], 2), 12288),
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 1], 4), 3072),
),
torch.bitwise_or(
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 2], 6), 768),
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 3], 8), 192),
),
),
torch.bitwise_or(
torch.bitwise_or(
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 4], 10), 48),
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 5], 12), 12),
),
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 6], 14), 3),
),
).unsqueeze(-1)
),
dim=-1
).view(shape)
return result
def unpack_uint13(packed_tensor: torch.Tensor, shape: torch.Size) -> torch.Tensor:
result = torch.cat(
(
torch.bitwise_and(packed_tensor[:, :13], 8191),
torch.bitwise_or(
torch.bitwise_or(
torch.bitwise_or(
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, :3], 13), 7),
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 3:6], 10), 56),
),
torch.bitwise_or(
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 6:9], 7), 448),
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 9:12], 4), 3584),
),
),
torch.bitwise_and(
torch.stack(
(
torch.bitwise_right_shift(packed_tensor[:, 12], 1),
torch.bitwise_right_shift(packed_tensor[:, 12], 2),
torch.bitwise_right_shift(packed_tensor[:, 12], 3),
),
dim=-1,
),
4096,
)
),
),
dim=-1
).view(shape)
return result
def unpack_uint12(packed_tensor: torch.Tensor, shape: torch.Size) -> torch.Tensor:
result = torch.cat(
(
torch.bitwise_and(packed_tensor[:, :3], 4095),
torch.bitwise_or(
torch.bitwise_or(
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 0], 4), 3840),
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 1], 8), 240),
),
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 2], 12), 15)
).unsqueeze(-1)
),
dim=-1
).view(shape)
return result
def unpack_uint11(packed_tensor: torch.Tensor, shape: torch.Size) -> torch.Tensor:
result = torch.cat(
(
torch.bitwise_and(packed_tensor[:, :8], 2047),
torch.bitwise_or(
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, :8], 11), 31),
torch.bitwise_and(
torch.cat(
(
torch.bitwise_left_shift(packed_tensor[:, 8:], 5),
torch.bitwise_right_shift(packed_tensor[:, 8:], 1),
torch.bitwise_or(
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 8:10], 7), 480),
torch.bitwise_and(
torch.stack(
(
torch.bitwise_right_shift(packed_tensor[:, 10], 3),
torch.bitwise_right_shift(packed_tensor[:, 10], 5),
),
dim=-1,
),
1536,
),
),
),
dim=-1,
),
2016,
),
),
),
dim=-1
).view(shape)
return result
def unpack_uint10(packed_tensor: torch.Tensor, shape: torch.Size) -> torch.Tensor:
result_bitwise_right_shift = torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, :3], 10), 63)
result = torch.cat(
(
torch.bitwise_and(packed_tensor[:, :5], 1023),
torch.bitwise_or(
result_bitwise_right_shift[:, :2],
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 3:5], 4), 960),
),
torch.bitwise_or(
result_bitwise_right_shift[:, 2],
torch.bitwise_or(
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 3], 6), 768),
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 4], 8), 192),
),
).unsqueeze(-1),
),
dim=-1
).view(shape)
return result
def unpack_uint9(packed_tensor: torch.Tensor, shape: torch.Size) -> torch.Tensor:
result = torch.cat(
(
torch.bitwise_and(packed_tensor[:, :8], 511),
torch.bitwise_or(
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, :8], 9), 127),
torch.bitwise_and(
torch.stack(
(
torch.bitwise_left_shift(packed_tensor[:, 8], 7),
torch.bitwise_left_shift(packed_tensor[:, 8], 5),
torch.bitwise_left_shift(packed_tensor[:, 8], 3),
torch.bitwise_left_shift(packed_tensor[:, 8], 1),
torch.bitwise_right_shift(packed_tensor[:, 8], 1),
torch.bitwise_right_shift(packed_tensor[:, 8], 3),
torch.bitwise_right_shift(packed_tensor[:, 8], 5),
torch.bitwise_right_shift(packed_tensor[:, 8], 7),
),
dim=-1,
),
384
)
)
),
dim=-1
).view(shape)
return result
def unpack_uint7(packed_tensor: torch.ByteTensor, shape: torch.Size) -> torch.ByteTensor:
result = torch.cat(
(
torch.bitwise_and(packed_tensor[:, :7], 127),
torch.bitwise_or(
torch.bitwise_or(
torch.bitwise_or(
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 0], 1), 64),
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 1], 2), 32),
),
torch.bitwise_or(
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 2], 3), 16),
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 3], 4), 8),
),
),
torch.bitwise_or(
torch.bitwise_or(
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 4], 5), 4),
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 5], 6), 2),
),
torch.bitwise_right_shift(packed_tensor[:, 6], 7),
),
).unsqueeze(-1)
),
dim=-1
).view(shape)
return result
def unpack_uint6(packed_tensor: torch.ByteTensor, shape: torch.Size) -> torch.ByteTensor:
result = torch.cat(
(
torch.bitwise_and(packed_tensor[:, :3], 63),
torch.bitwise_or(
torch.bitwise_or(
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 0], 2), 48),
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 1], 4), 12),
),
torch.bitwise_right_shift(packed_tensor[:, 2], 6)
).unsqueeze(-1)
),
dim=-1
).view(shape)
return result
def unpack_uint5(packed_tensor: torch.ByteTensor, shape: torch.Size) -> torch.ByteTensor:
result_bitwise_right_shift = torch.bitwise_right_shift(packed_tensor[:, :3], 5)
result = torch.cat(
(
torch.bitwise_and(packed_tensor[:, :5], 31),
torch.bitwise_or(
result_bitwise_right_shift[:, :2],
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 3:5], 2), 24),
),
torch.bitwise_or(
result_bitwise_right_shift[:, 2],
torch.bitwise_or(
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 3], 3), 16),
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 4], 4), 8),
),
).unsqueeze(-1),
),
dim=-1
).view(shape)
return result
def unpack_uint4(packed_tensor: torch.ByteTensor, shape: torch.Size) -> torch.ByteTensor:
result = torch.stack((torch.bitwise_and(packed_tensor, 15), torch.bitwise_right_shift(packed_tensor, 4)), dim=-1).view(shape)
return result
def unpack_uint3(packed_tensor: torch.ByteTensor, shape: torch.Size) -> torch.ByteTensor:
result = torch.bitwise_and(
torch.cat(
(
packed_tensor[:, :3],
torch.bitwise_right_shift(packed_tensor[:, :3], 3),
torch.bitwise_or(
torch.bitwise_right_shift(packed_tensor[:, :2], 6),
torch.bitwise_and(
torch.stack(
(
torch.bitwise_right_shift(packed_tensor[:, 2], 4),
torch.bitwise_right_shift(packed_tensor[:, 2], 5),
),
dim=-1
),
4
),
),
),
dim=-1
),
7
).view(shape)
return result
def unpack_uint2(packed_tensor: torch.ByteTensor, shape: torch.Size) -> torch.ByteTensor:
result = torch.bitwise_and(
torch.stack(
(
packed_tensor,
torch.bitwise_right_shift(packed_tensor, 2),
torch.bitwise_right_shift(packed_tensor, 4),
torch.bitwise_right_shift(packed_tensor, 6),
),
dim=-1
),
3
).view(shape)
return result
def unpack_uint1(packed_tensor: torch.Tensor, shape: torch.Size) -> torch.Tensor:
result = torch.bitwise_and(
torch.stack(
(
packed_tensor,
torch.bitwise_right_shift(packed_tensor, 1),
torch.bitwise_right_shift(packed_tensor, 2),
torch.bitwise_right_shift(packed_tensor, 3),
torch.bitwise_right_shift(packed_tensor, 4),
torch.bitwise_right_shift(packed_tensor, 5),
torch.bitwise_right_shift(packed_tensor, 6),
torch.bitwise_right_shift(packed_tensor, 7),
),
dim=-1
),
1
).view(shape)
return result
modules/sdnq/quant_utils.py
-236
View File
@@ -1,236 +0,0 @@
# pylint: disable=redefined-builtin
import torch
from modules import devices
from .common import dtype_dict, compile_func, conv_types, conv_transpose_types
from .kernel_wrappers import use_contiguous_int8_mm, use_contiguous_fp16_mm, use_contiguous_fp8_mm
from .utils import is_pow2, is_pow4, next_power_of_2
@devices.inference_context()
def get_scale_asymmetric(weight: torch.FloatTensor, dim: int | list[int], weights_dtype: str) -> tuple[torch.FloatTensor, torch.FloatTensor]:
zero_point, scale = torch.aminmax(weight, dim=dim, keepdims=True)
scale = scale.sub_(zero_point).div_(dtype_dict[weights_dtype]["max"] - dtype_dict[weights_dtype]["min"])
if dtype_dict[weights_dtype]["min"] != 0:
zero_point.sub_(scale, alpha=dtype_dict[weights_dtype]["min"])
return scale, zero_point
@devices.inference_context()
def get_scale_symmetric(weight: torch.FloatTensor, dim: int | list[int], weights_dtype: str) -> torch.FloatTensor:
return torch.amax(weight.abs(), dim=dim, keepdims=True).div_(dtype_dict[weights_dtype]["max"])
@devices.inference_context()
def quantize_weight(weight: torch.FloatTensor, dim: int | list[int], weights_dtype: str, dtype: torch.dtype = None, use_stochastic_rounding: bool = False) -> tuple[torch.Tensor, torch.FloatTensor, torch.FloatTensor]:
if weight.dtype != torch.float64:
weight = weight.to(dtype=torch.float32, copy=False)
if dtype_dict[weights_dtype]["is_unsigned"]:
scale, zero_point = get_scale_asymmetric(weight, dim, weights_dtype)
if dtype is not None:
scale = scale.to(dtype=dtype)
zero_point = zero_point.to(dtype=dtype)
quantized_weight = torch.sub(weight, zero_point).div_(scale)
else:
scale = get_scale_symmetric(weight, dim, weights_dtype)
zero_point = None
if dtype is not None:
scale = scale.to(dtype=dtype)
quantized_weight = torch.div(weight, scale)
if dtype_dict[weights_dtype]["is_integer"]:
if use_stochastic_rounding:
quantized_weight.add_(torch.randn_like(quantized_weight), alpha=0.1)
quantized_weight.round_()
else:
if use_stochastic_rounding:
mantissa_difference = 1 << (23 - dtype_dict[weights_dtype]["mantissa"])
quantized_weight = quantized_weight.to(dtype=torch.float32).view(dtype=torch.int32)
quantized_weight = quantized_weight.add_(torch.randint_like(quantized_weight, low=0, high=mantissa_difference, dtype=torch.int32)).bitwise_and_(-mantissa_difference).view(dtype=torch.float32)
quantized_weight.nan_to_num_()
quantized_weight = quantized_weight.clamp_(dtype_dict[weights_dtype]["min"], dtype_dict[weights_dtype]["max"]).to(dtype_dict[weights_dtype]["torch_dtype"])
return quantized_weight, scale, zero_point
@devices.inference_context()
def apply_svdquant(weight: torch.FloatTensor, rank: int = 32, niter: int = 8, dtype: torch.dtype = None) -> tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]:
reshape_weight = False
if weight.ndim > 2: # convs
reshape_weight = True
weight_shape = weight.shape
weight = weight.flatten(1,-1)
if weight.dtype != torch.float64:
weight = weight.to(dtype=torch.float32)
U, S, svd_down = torch.svd_lowrank(weight, q=rank, niter=niter)
svd_up = torch.mul(U, S.unsqueeze(0))
svd_down = svd_down.t_()
if dtype is not None:
svd_up = svd_up.to(dtype=dtype)
svd_down = svd_down.to(dtype=dtype)
weight = weight.sub(torch.mm(svd_up, svd_down))
if reshape_weight:
weight = weight.unflatten(-1, (*weight_shape[1:],)) # pylint: disable=possibly-used-before-assignment
return weight, svd_up, svd_down
@devices.inference_context()
def build_hadamard_n2(n: int, dtype: torch.dtype | None = None, device: torch.device | None = None) -> torch.FloatTensor:
current_size = 2
H = H_N2 = torch.tensor([[1, 1], [1, -1]], dtype=dtype, device=device)
while current_size < n:
H = torch.kron(H, H_N2)
current_size *= 2
H = H.div_(n**0.5)
H = prepare_weight_for_matmul(H, matmul_dtype="float16")
return H
@devices.inference_context()
def build_hadamard_n4(n: int, dtype: torch.dtype | None = None, device: torch.device | None = None) -> torch.FloatTensor:
current_size = 4
H = H_N4 = torch.tensor([[ 1, 1, 1, -1], [ 1, 1, -1, 1], [ 1, -1, 1, 1], [-1, 1, 1, 1]], dtype=dtype, device=device)
while current_size < n:
H = torch.kron(H, H_N4)
current_size *= 4
H = H.div_(n**0.5)
H = prepare_weight_for_matmul(H, matmul_dtype="float16")
return H
@devices.inference_context()
def build_hadamard(n: int, dtype: torch.dtype | None = None, device: torch.device | None = None) -> torch.FloatTensor:
if is_pow4(n):
return build_hadamard_n4(n, device=device, dtype=dtype)
elif is_pow2(n):
return build_hadamard_n2(n, device=device, dtype=dtype)
else:
raise RuntimeError(f"Hadamard Group Size must be a power of 2 but got {n}.")
# 256x256 Hadamard matrix is just 256 KB at FP32
# And is the exact same matrix on all model layers
# So we can safely cache a single one
HADAMARD_MATRIX_CACHE: dict[tuple[int, torch.device, torch.dtype], torch.FloatTensor] = {}
@devices.inference_context()
def get_hadamard(n: int, dtype: torch.dtype | None = None, device: torch.device | None = None) -> torch.FloatTensor:
device = devices.normalize_device(device)
H_key = (n, device, dtype)
H = HADAMARD_MATRIX_CACHE.get(H_key, None)
if H is None:
H = build_hadamard(n, dtype=dtype, device=device)
HADAMARD_MATRIX_CACHE[H_key] = H
return H
@devices.inference_context()
def rotate_hadamard(weight: torch.Tensor, group_size: int = 256, hadamard: torch.FloatTensor | None = None, is_conv: bool = False) -> torch.Tensor:
if hadamard is None:
hadamard = get_hadamard(group_size, dtype=weight.dtype, device=weight.device)
else:
group_size = hadamard.shape[-1]
if hadamard.dtype != weight.dtype:
hadamard = hadamard.to(dtype=weight.dtype)
if is_conv:
weight_shape = list(weight.shape)[1:]
weight = weight.flatten(1,-1)
weight = weight.unflatten(-1, (-1,group_size))
result = torch.matmul(weight, hadamard).flatten(-2,-1)
del hadamard
if is_conv:
result = result.unflatten(-1, weight_shape)
return result
def get_hadamard_group_size(channel_size: int, group_size: int) -> tuple[bool, int]:
group_size = next_power_of_2(min(channel_size, group_size))
if channel_size % group_size != 0:
while channel_size % group_size != 0:
group_size = group_size // 2
use_hadamard = group_size >= 4
return use_hadamard, group_size
@devices.inference_context()
def apply_hadamard(weight: torch.Tensor, group_size: int = 256, hadamard: torch.FloatTensor | None = None, layer_class_name: str | None = None) -> tuple[torch.Tensor, bool, int]:
is_conv = False
if hadamard is not None:
group_size = hadamard.shape[-1]
if layer_class_name in conv_types or layer_class_name in conv_transpose_types:
is_conv = True
channel_size = weight.shape[1]
else:
channel_size = weight.shape[-1]
use_hadamard, group_size = get_hadamard_group_size(channel_size, group_size)
if use_hadamard:
if hadamard is not None and group_size != hadamard.shape[-1]:
hadamard = None
weight = rotate_hadamard(weight, group_size=group_size, hadamard=hadamard, is_conv=is_conv)
return weight, use_hadamard, group_size
@devices.inference_context()
def prepare_weight_for_matmul(weight: torch.Tensor, matmul_dtype: str | None = "int8") -> torch.Tensor:
if (
(use_contiguous_int8_mm and matmul_dtype in {"int8", "uint8"})
or (use_contiguous_fp16_mm and matmul_dtype in {"fp16", "float16"})
or (use_contiguous_fp8_mm and matmul_dtype in {"fp8", "float8_e4m3fn"})
):
weight = weight.contiguous()
elif weight.is_contiguous():
weight = weight.t_().contiguous().t_()
return weight
@devices.inference_context()
def prepare_svd_for_matmul(svd_up: torch.FloatTensor, svd_down: torch.FloatTensor, use_quantized_matmul: bool) -> tuple[torch.FloatTensor, torch.FloatTensor]:
if svd_up is not None:
if use_quantized_matmul:
svd_up = prepare_weight_for_matmul(svd_up, matmul_dtype="float16")
else:
svd_up = svd_up.contiguous()
if svd_down is not None:
svd_down = prepare_weight_for_matmul(svd_down, matmul_dtype="float16")
return svd_up, svd_down
@devices.inference_context()
def quantize_int_mm(weight: torch.FloatTensor, dim: int = -1, hadamard: torch.FloatTensor | None = None, matmul_dtype: str = "int8", use_sr: bool = False) -> tuple[torch.Tensor, torch.FloatTensor]:
if hadamard is not None:
weight = rotate_hadamard(weight, hadamard=hadamard)
scale = get_scale_symmetric(weight, dim, matmul_dtype)
weight = torch.div(weight, scale)
if use_sr:
weight = weight.add_(torch.randn_like(weight), alpha=0.1)
weight = weight.round_().clamp_(dtype_dict[matmul_dtype]["min"], dtype_dict[matmul_dtype]["max"]).to(dtype=dtype_dict[matmul_dtype]["torch_dtype"])
return weight, scale
@devices.inference_context()
def quantize_uint_mm(weight: torch.FloatTensor, dim: int = -1, hadamard: torch.FloatTensor | None = None, matmul_dtype: str = "uint8", use_sr: bool = False) -> tuple[torch.FloatTensor, torch.FloatTensor]:
if hadamard is not None:
weight = rotate_hadamard(weight, hadamard=hadamard)
matmul_dtype = matmul_dtype.removeprefix("u")
scale, zero_point = get_scale_asymmetric(weight, dim, matmul_dtype)
weight = torch.sub(weight, zero_point).div_(scale)
if use_sr:
weight = weight.add_(torch.randn_like(weight), alpha=0.1)
weight = weight.round_().clamp_(dtype_dict[matmul_dtype]["min"], dtype_dict[matmul_dtype]["max"]).to(dtype=dtype_dict[matmul_dtype]["torch_dtype"])
return weight, scale, zero_point
@devices.inference_context()
def quantize_fp_mm(weight: torch.FloatTensor, dim: int = -1, hadamard: torch.FloatTensor | None = None, matmul_dtype: str = "float8_e4m3fn", use_sr: bool = False) -> tuple[torch.Tensor, torch.FloatTensor]:
if hadamard is not None:
weight = rotate_hadamard(weight, hadamard=hadamard)
scale = get_scale_symmetric(weight, dim, matmul_dtype)
if use_sr:
mantissa_difference = 1 << (23 - dtype_dict[matmul_dtype]["mantissa"])
weight = weight.to(dtype=torch.float32).view(dtype=torch.int32)
weight = weight.add_(torch.randint_like(weight, low=0, high=mantissa_difference, dtype=torch.int32)).bitwise_and_(-mantissa_difference).view(dtype=torch.float32)
weight = torch.div(weight, scale).nan_to_num_().clamp_(dtype_dict[matmul_dtype]["min"], dtype_dict[matmul_dtype]["max"]).to(dtype=dtype_dict[matmul_dtype]["torch_dtype"])
return weight, scale
rotate_hadamard_compiled = compile_func(rotate_hadamard)
modules/sdnq/quantizer.py
-1065
View File
File diff suppressed because it is too large Load Diff
modules/sdnq/utils.py
-220
View File
@@ -1,220 +0,0 @@
import re
import torch
from .common import (
dtype_dict,
common_skip_keys,
module_skip_keys_dict,
allowed_types,
embedding_types,
conv_types,
conv_transpose_types,
)
def is_pow2(n: int) -> bool:
return (n & (n - 1)) == 0
def is_pow4(n: int) -> bool:
return is_pow2(n) and (n.bit_length() & 1 == 1)
def next_power_of_2(n: int) -> int:
if is_pow2(n):
return n
return 2 ** n.bit_length()
def check_param_name_in(param_name: str, param_list: list[str]) -> str:
split_param_name = param_name.split(".")
for param in param_list:
if param.startswith("."):
if param_name.startswith(param[1:]):
return param
else:
continue
if (
param_name == param
or param in split_param_name
or ("*" in param and re.match(param.replace(".*", "\\.*").replace("*", ".*"), param_name))
):
return param
return None
def check_quant_is_allowed(layer_class_name: str, weight: torch.Tensor, quantization_config, pre_quantized: bool = False) -> bool:
if (
layer_class_name in allowed_types
and weight.dtype in {torch.float64, torch.float32, torch.float16, torch.bfloat16}
and not (layer_class_name in embedding_types and not quantization_config.quant_embedding)
and not ((layer_class_name in conv_types or layer_class_name in conv_transpose_types) and not quantization_config.quant_conv)
):
if pre_quantized:
return True
if layer_class_name in conv_types:
channel_size = weight.shape[1]
elif layer_class_name in conv_transpose_types:
channel_size = weight.shape[0]
else:
channel_size = weight.shape[-1]
if channel_size >= quantization_config.minimum_allowed_channel_size and weight.numel() >= quantization_config.minimum_allowed_numel:
return True
return False
def check_quantized_matmul_is_allowed(use_quantized_matmul: bool, output_channel_size: int, channel_size: int) -> bool:
return bool(
use_quantized_matmul
and output_channel_size >= 32 and channel_size >= 32
and output_channel_size % 16 == 0 and channel_size % 16 == 0
)
def get_quant_args_from_config(quantization_config: dict) -> dict:
from .quantizer import SDNQConfig
if isinstance(quantization_config, SDNQConfig):
quantization_config_dict = quantization_config.to_dict()
else:
quantization_config_dict = quantization_config.copy()
quantization_config_dict.pop("is_integer", None)
quantization_config_dict.pop("quant_method", None)
quantization_config_dict.pop("quantization_device", None)
quantization_config_dict.pop("return_device", None)
quantization_config_dict.pop("non_blocking", None)
quantization_config_dict.pop("add_skip_keys", None)
quantization_config_dict.pop("use_dynamic_quantization", None)
quantization_config_dict.pop("use_static_quantization", None)
quantization_config_dict.pop("use_stochastic_rounding", None)
quantization_config_dict.pop("use_grad_ckpt", None)
quantization_config_dict.pop("is_training", None)
quantization_config_dict.pop("sdnq_version", None)
if quantization_config_dict.get("modules_quant_config", None) is not None:
for key in quantization_config_dict["modules_quant_config"]:
quantization_config_dict["modules_quant_config"][key] = get_quant_args_from_config(quantization_config_dict["modules_quant_config"][key])
return quantization_config_dict
def get_minimum_dtype(weights_dtype: str, param_name: str, modules_dtype_dict: dict[str, list[str]]):
if len(modules_dtype_dict.keys()) > 0:
for key, value in modules_dtype_dict.items():
if check_param_name_in(param_name, value) is not None:
key = key.lower()
if key.startswith("minimum") or key.endswith(("bit", "bits")):
minimum_bits_str = key.removeprefix("minimum").removeprefix("-").removeprefix("_").removesuffix("bits").removesuffix("bit").removesuffix("-").removesuffix("_")
if minimum_bits_str.startswith("uint"):
is_unsigned = True
minimum_bits_str = minimum_bits_str.removeprefix("uint")
else:
is_unsigned = False
minimum_bits_str = minimum_bits_str.removeprefix("int")
minimum_bits = int(minimum_bits_str)
if dtype_dict[weights_dtype]["num_bits"] < minimum_bits:
if is_unsigned or minimum_bits <= 4:
return "uint" + minimum_bits_str
else:
return "int" + minimum_bits_str
else:
return key
return weights_dtype
def get_quant_kwargs(layer: torch.nn.Module, quantization_config, torch_dtype: torch.dtype | None = None, param_name: str = "", **kwargs) -> dict:
from .quantizer import SDNQConfig
if not isinstance(quantization_config, SDNQConfig):
quantization_config = SDNQConfig(**quantization_config)
layer_class_name = layer.__class__.__name__
quant_kwargs = {
"weights_dtype": quantization_config.weights_dtype,
"quantized_matmul_dtype": quantization_config.quantized_matmul_dtype,
"hadamard_group_size": quantization_config.hadamard_group_size,
"group_size": quantization_config.group_size,
"svd_rank": quantization_config.svd_rank,
"svd_steps": quantization_config.svd_steps,
"dynamic_loss_threshold": quantization_config.dynamic_loss_threshold,
"use_svd": quantization_config.use_svd,
"use_hadamard": quantization_config.use_hadamard,
"use_quantized_matmul": quantization_config.use_quantized_matmul,
"use_quantized_matmul_conv": quantization_config.use_quantized_matmul_conv,
"use_dynamic_quantization": quantization_config.use_dynamic_quantization,
"use_stochastic_rounding": quantization_config.use_stochastic_rounding,
"dequantize_fp32": quantization_config.dequantize_fp32,
"non_blocking": quantization_config.non_blocking,
"quantization_device": quantization_config.quantization_device,
"return_device": quantization_config.return_device,
"layer_class_name": layer_class_name,
"torch_dtype": torch_dtype,
"param_name": param_name,
}
for key, value in kwargs.items():
quant_kwargs[key] = value
param_key = check_param_name_in(quant_kwargs["param_name"], quantization_config.modules_quant_config.keys())
if param_key is not None:
for key, value in quantization_config.modules_quant_config[param_key].items():
quant_kwargs[key] = value
if layer_class_name in conv_transpose_types or layer_class_name in conv_types:
quant_kwargs["use_quantized_matmul"] = quant_kwargs.pop("use_quantized_matmul_conv")
else:
quant_kwargs.pop("use_quantized_matmul_conv")
if not quant_kwargs["use_dynamic_quantization"]:
quant_kwargs.pop("dynamic_loss_threshold")
quant_kwargs["weights_dtype"] = get_minimum_dtype(quant_kwargs["weights_dtype"], quant_kwargs["param_name"], quantization_config.modules_dtype_dict)
if check_param_name_in(quant_kwargs["param_name"], quantization_config.modules_to_not_use_matmul) is not None:
quant_kwargs["use_quantized_matmul"] = False
return quant_kwargs
def get_quantized_matmul_dtype(weights_dtype: str, quantized_matmul_dtype: str | None = None) -> str:
if quantized_matmul_dtype is None:
if dtype_dict[weights_dtype]["is_integer"]:
if weights_dtype == "uint8":
quantized_matmul_dtype = "uint8"
else:
quantized_matmul_dtype = "int8"
elif dtype_dict[weights_dtype]["num_bits"] < 16:
quantized_matmul_dtype = "float8_e4m3fn"
else:
quantized_matmul_dtype = "float16"
return quantized_matmul_dtype
def add_module_skip_keys(model: torch.nn.Module, quantization_config):
if getattr(model, "_keep_in_fp32_modules", None) is not None:
quantization_config.modules_to_not_convert.extend(model._keep_in_fp32_modules) # pylint: disable=protected-access
if getattr(model, "_tied_weights_keys", None) is not None:
if isinstance(model._tied_weights_keys, dict): # pylint: disable=protected-access
quantization_config.modules_to_not_convert.extend(model._tied_weights_keys.keys()) # pylint: disable=protected-access
quantization_config.modules_to_not_convert.extend(model._tied_weights_keys.values()) # pylint: disable=protected-access
else:
quantization_config.modules_to_not_convert.extend(model._tied_weights_keys) # pylint: disable=protected-access
skip_key_list = module_skip_keys_dict.get(model.__class__.__name__, None)
if skip_key_list is not None:
quantization_config.modules_to_not_convert.extend(skip_key_list[0])
for key, value in skip_key_list[1].items():
if key in quantization_config.modules_dtype_dict:
quantization_config.modules_dtype_dict[key].extend(value)
else:
quantization_config.modules_dtype_dict[key] = value
quantized_matmul_dtype = get_quantized_matmul_dtype(quantization_config.weights_dtype, quantization_config.quantized_matmul_dtype)
quantization_config.modules_to_not_use_matmul.extend(skip_key_list[2].get(quantized_matmul_dtype, []))
else:
quantization_config.modules_to_not_convert.extend(common_skip_keys)
if getattr(model, "_skip_layerwise_casting_patterns", None) is not None:
quantization_config.modules_to_not_convert.extend(model._skip_layerwise_casting_patterns) # pylint: disable=protected-access
# dedupe
quantization_config.modules_to_not_convert = list(set(quantization_config.modules_to_not_convert))
quantization_config.modules_to_not_use_matmul = list(set(quantization_config.modules_to_not_use_matmul))
for key, value in quantization_config.modules_dtype_dict.items():
quantization_config.modules_dtype_dict[key] = list(set(value))
return model, quantization_config
modules/timer_sdnq.py
+1 -1
View File
@@ -1,7 +1,7 @@
def update_sdnq_attention_timers():
from modules.timer import autotune
autotune.reset()
from modules.sdnq.kernels import triton_atten, triton_mm, triton_scaled_mm
from sdnq.kernels import triton_atten, triton_mm, triton_scaled_mm
if getattr(triton_atten.sdnq_attn_kernel, 'bench_time', None) is not None:
autotune.add('sdnq_attn_kernel', getattr(triton_atten.sdnq_attn_kernel, 'bench_time', 0))
triton_atten.sdnq_attn_kernel.bench_time = 0
modules/video_models/video_load.py
+1 -1
View File
@@ -31,7 +31,7 @@ def load_custom(model_name: str):
def load_model(selected: models_def.Model):
from modules import sdnq # pylint: disable=unused-import
import sdnq # pylint: disable=unused-import
if selected is None or selected.repo is None:
return ''
if isinstance(selected.repo_cls, str):
pipelines/native_transformer.py
+9 -9
View File
@@ -727,7 +727,7 @@ def build_component_quantized(
import rich.progress as rp
from accelerate import init_empty_weights
from accelerate.utils import set_module_tensor_to_device
from modules.sdnq.quantizer import SDNQQuantizer
from sdnq.quantizer import SDNQQuantizer
quantization_config = quant_args.get("quantization_config")
if quantization_config is None:
@@ -819,19 +819,19 @@ def build_component_prequantized(
target dtype (the fp32 scales must survive). Layers are assembled in
canonical dequant layout; ``apply_sdnq_options_to_model`` then applies
the user's quantized-matmul settings, matching
``modules.sdnq.loader.load_sdnq_model``.
``sdnq.loader.load_sdnq_model``.
"""
import rich.progress as rp
from accelerate import init_empty_weights
from accelerate.utils import set_module_tensor_to_device
from diffusers.utils import get_module_from_name
from modules.sdnq.common import dtype_dict, check_torch_compile
from modules.sdnq.kernel_wrappers import is_fp8_compile_supported
from modules.sdnq.quantizer import SDNQConfig, SDNQQuantizer
from modules.sdnq.dequantizer import SDNQDequantizer
from modules.sdnq.layers import get_sdnq_wrapper_class
from modules.sdnq.forward import get_forward_func
from modules.sdnq.loader import apply_sdnq_options_to_model
from sdnq.common import dtype_dict, check_torch_compile
from sdnq.kernel_wrappers import is_fp8_compile_supported
from sdnq.quantizer import SDNQConfig, SDNQQuantizer
from sdnq.dequantizer import SDNQDequantizer
from sdnq.layers import get_sdnq_wrapper_class
from sdnq.forward import get_forward_func
from sdnq.loader import apply_sdnq_options_to_model
weights_dtype = COMFY_QUANT_FORMATS[comfy_format]
matmul_dtype = "int8" if dtype_dict[weights_dtype]["is_integer"] else "float8_e4m3fn"
test/test-ideogram4-smoke.py
+1 -1
View File
@@ -52,7 +52,7 @@ from transformers import AutoTokenizer
from transformers.models.qwen3_vl import Qwen3VLModel
from modules import devices
from modules.sdnq import SDNQConfig
from sdnq import SDNQConfig
TE_REPO = "Qwen/Qwen3-VL-8B-Instruct"
test/test-native-transformer.py
+4 -4
View File
@@ -580,7 +580,7 @@ def test_unswizzle_block_scales_roundtrip():
def test_nvfp4_codec_ocp_table():
"""SDNQ's float4_e2m1fn decodes OCP FP4 E2M1 exactly, including the
subnormal codes 1/9 as +/-0.5; nvfp4 containers adopt it directly."""
from modules.sdnq.packed_float import unpack_float
from sdnq.packed_float import unpack_float
packed = torch.tensor([(2 * j) | (((2 * j) + 1) << 4) for j in range(8)], dtype=torch.uint8)
dec = unpack_float(packed, 'float4_e2m1fn', torch.Size([16]))
for code in range(16):
@@ -592,7 +592,7 @@ def test_nvfp4_pack_ocp_grid_roundtrip():
(subnormals included), and off-grid values land inside the value set.
Exact nearest-rounding near the grid midpoints is not asserted: the
packer's staged rounding may resolve boundary values to either side."""
from modules.sdnq.packed_float import pack_float, unpack_float
from sdnq.packed_float import pack_float, unpack_float
grid = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0, 0.0]
vals = torch.tensor(grid, dtype=torch.float32)
out = unpack_float(pack_float(vals, 'float4_e2m1fn'), 'float4_e2m1fn', vals.shape)
@@ -1370,8 +1370,8 @@ class ComfyTestEnv:
def __enter__(self):
from modules import model_quant, shared
from modules.sdnq import common as sdnq_common
from modules.sdnq import kernel_wrappers
from sdnq import common as sdnq_common
from sdnq import kernel_wrappers
self.shared = shared
self.sdnq_common = sdnq_common
self.kernel_wrappers = kernel_wrappers