From 712a13f1a09059e3d2dd6cd32b583db705e0d89c Mon Sep 17 00:00:00 2001 From: Dity0 Date: Mon, 10 Aug 2026 12:13:10 +0300 Subject: [PATCH] Move SDNQ to upstream SDNQ repo --- cli/sdnq-attention-benchmark.py | 56 +- installer.py | 26 +- launch.py | 1 + modules/attention.py | 2 +- modules/lora/lora_apply.py | 2 +- modules/model_quant.py | 4 +- modules/sd_models.py | 10 +- modules/sdnq/__init__.py | 16 - modules/sdnq/common.py | 524 ---------- modules/sdnq/dequantizer.py | 364 ------- modules/sdnq/file_loader.py | 64 -- modules/sdnq/forward.py | 59 -- modules/sdnq/kernel_wrappers.py | 215 ---- modules/sdnq/kernels/openvino_mm.py | 155 --- modules/sdnq/kernels/triton_atten.py | 355 ------- modules/sdnq/kernels/triton_mm.py | 131 --- modules/sdnq/kernels/triton_scaled_mm.py | 142 --- modules/sdnq/layers/__init__.py | 93 -- modules/sdnq/layers/conv/conv_fp16.py | 107 -- modules/sdnq/layers/conv/conv_fp8.py | 105 -- modules/sdnq/layers/conv/conv_int8.py | 123 --- modules/sdnq/layers/conv/conv_uint8.py | 122 --- modules/sdnq/layers/conv/forward.py | 91 -- modules/sdnq/layers/embedding/forward.py | 83 -- modules/sdnq/layers/linear/forward.py | 23 - modules/sdnq/layers/linear/linear_fp16.py | 98 -- modules/sdnq/layers/linear/linear_fp8.py | 105 -- modules/sdnq/layers/linear/linear_int8.py | 126 --- modules/sdnq/layers/linear/linear_uint8.py | 130 --- modules/sdnq/loader.py | 346 ------- modules/sdnq/packed_float.py | 129 --- modules/sdnq/packed_int/__init__.py | 85 -- modules/sdnq/packed_int/pack.py | 305 ------ modules/sdnq/packed_int/unpack.py | 356 ------- modules/sdnq/quant_utils.py | 236 ----- modules/sdnq/quantizer.py | 1065 -------------------- modules/sdnq/utils.py | 220 ---- modules/timer_sdnq.py | 2 +- modules/video_models/video_load.py | 2 +- pipelines/native_transformer.py | 18 +- test/test-ideogram4-smoke.py | 2 +- test/test-native-transformer.py | 8 +- 42 files changed, 79 insertions(+), 6027 deletions(-) delete mode 100644 modules/sdnq/__init__.py delete mode 100644 modules/sdnq/common.py delete mode 100644 modules/sdnq/dequantizer.py delete mode 100644 modules/sdnq/file_loader.py delete mode 100644 modules/sdnq/forward.py delete mode 100644 modules/sdnq/kernel_wrappers.py delete mode 100644 modules/sdnq/kernels/openvino_mm.py delete mode 100644 modules/sdnq/kernels/triton_atten.py delete mode 100644 modules/sdnq/kernels/triton_mm.py delete mode 100644 modules/sdnq/kernels/triton_scaled_mm.py delete mode 100644 modules/sdnq/layers/__init__.py delete mode 100644 modules/sdnq/layers/conv/conv_fp16.py delete mode 100644 modules/sdnq/layers/conv/conv_fp8.py delete mode 100644 modules/sdnq/layers/conv/conv_int8.py delete mode 100644 modules/sdnq/layers/conv/conv_uint8.py delete mode 100644 modules/sdnq/layers/conv/forward.py delete mode 100644 modules/sdnq/layers/embedding/forward.py delete mode 100644 modules/sdnq/layers/linear/forward.py delete mode 100644 modules/sdnq/layers/linear/linear_fp16.py delete mode 100644 modules/sdnq/layers/linear/linear_fp8.py delete mode 100644 modules/sdnq/layers/linear/linear_int8.py delete mode 100644 modules/sdnq/layers/linear/linear_uint8.py delete mode 100644 modules/sdnq/loader.py delete mode 100644 modules/sdnq/packed_float.py delete mode 100644 modules/sdnq/packed_int/__init__.py delete mode 100644 modules/sdnq/packed_int/pack.py delete mode 100644 modules/sdnq/packed_int/unpack.py delete mode 100644 modules/sdnq/quant_utils.py delete mode 100644 modules/sdnq/quantizer.py delete mode 100644 modules/sdnq/utils.py diff --git a/cli/sdnq-attention-benchmark.py b/cli/sdnq-attention-benchmark.py index e774d0344..614eb68fe 100755 --- a/cli/sdnq-attention-benchmark.py +++ b/cli/sdnq-attention-benchmark.py @@ -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 diff --git a/installer.py b/installer.py index 595510e2b..d6755a336 100644 --- a/installer.py +++ b/installer.py @@ -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') diff --git a/launch.py b/launch.py index e3b0b995e..418bfdc3b 100755 --- a/launch.py +++ b/launch.py @@ -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() diff --git a/modules/attention.py b/modules/attention.py index 7380b5210..5213da012 100644 --- a/modules/attention.py +++ b/modules/attention.py @@ -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: diff --git a/modules/lora/lora_apply.py b/modules/lora/lora_apply.py index da250b860..467330e3e 100644 --- a/modules/lora/lora_apply.py +++ b/modules/lora/lora_apply.py @@ -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) diff --git a/modules/model_quant.py b/modules/model_quant.py index bd990bcc1..031e986cb 100644 --- a/modules/model_quant.py +++ b/modules/model_quant.py @@ -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") diff --git a/modules/sd_models.py b/modules/sd_models.py index e55fdf88c..c4a829509 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -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()) diff --git a/modules/sdnq/__init__.py b/modules/sdnq/__init__.py deleted file mode 100644 index c69f4887f..000000000 --- a/modules/sdnq/__init__.py +++ /dev/null @@ -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", -] diff --git a/modules/sdnq/common.py b/modules/sdnq/common.py deleted file mode 100644 index 915e266d1..000000000 --- a/modules/sdnq/common.py +++ /dev/null @@ -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"] diff --git a/modules/sdnq/dequantizer.py b/modules/sdnq/dequantizer.py deleted file mode 100644 index 26f6add91..000000000 --- a/modules/sdnq/dequantizer.py +++ /dev/null @@ -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]) diff --git a/modules/sdnq/file_loader.py b/modules/sdnq/file_loader.py deleted file mode 100644 index 5e921da0d..000000000 --- a/modules/sdnq/file_loader.py +++ /dev/null @@ -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 diff --git a/modules/sdnq/forward.py b/modules/sdnq/forward.py deleted file mode 100644 index 513be32bb..000000000 --- a/modules/sdnq/forward.py +++ /dev/null @@ -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 diff --git a/modules/sdnq/kernel_wrappers.py b/modules/sdnq/kernel_wrappers.py deleted file mode 100644 index 9a9751fbc..000000000 --- a/modules/sdnq/kernel_wrappers.py +++ /dev/null @@ -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) diff --git a/modules/sdnq/kernels/openvino_mm.py b/modules/sdnq/kernels/openvino_mm.py deleted file mode 100644 index b0832b009..000000000 --- a/modules/sdnq/kernels/openvino_mm.py +++ /dev/null @@ -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) diff --git a/modules/sdnq/kernels/triton_atten.py b/modules/sdnq/kernels/triton_atten.py deleted file mode 100644 index e23b994c8..000000000 --- a/modules/sdnq/kernels/triton_atten.py +++ /dev/null @@ -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) diff --git a/modules/sdnq/kernels/triton_mm.py b/modules/sdnq/kernels/triton_mm.py deleted file mode 100644 index be49a94fe..000000000 --- a/modules/sdnq/kernels/triton_mm.py +++ /dev/null @@ -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 diff --git a/modules/sdnq/kernels/triton_scaled_mm.py b/modules/sdnq/kernels/triton_scaled_mm.py deleted file mode 100644 index c2c1d8689..000000000 --- a/modules/sdnq/kernels/triton_scaled_mm.py +++ /dev/null @@ -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 diff --git a/modules/sdnq/layers/__init__.py b/modules/sdnq/layers/__init__.py deleted file mode 100644 index 2e66f7718..000000000 --- a/modules/sdnq/layers/__init__.py +++ /dev/null @@ -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) diff --git a/modules/sdnq/layers/conv/conv_fp16.py b/modules/sdnq/layers/conv/conv_fp16.py deleted file mode 100644 index 259a02dff..000000000 --- a/modules/sdnq/layers/conv/conv_fp16.py +++ /dev/null @@ -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) diff --git a/modules/sdnq/layers/conv/conv_fp8.py b/modules/sdnq/layers/conv/conv_fp8.py deleted file mode 100644 index 3f211bb55..000000000 --- a/modules/sdnq/layers/conv/conv_fp8.py +++ /dev/null @@ -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) diff --git a/modules/sdnq/layers/conv/conv_int8.py b/modules/sdnq/layers/conv/conv_int8.py deleted file mode 100644 index a265fbc98..000000000 --- a/modules/sdnq/layers/conv/conv_int8.py +++ /dev/null @@ -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) diff --git a/modules/sdnq/layers/conv/conv_uint8.py b/modules/sdnq/layers/conv/conv_uint8.py deleted file mode 100644 index c291fb5d3..000000000 --- a/modules/sdnq/layers/conv/conv_uint8.py +++ /dev/null @@ -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) diff --git a/modules/sdnq/layers/conv/forward.py b/modules/sdnq/layers/conv/forward.py deleted file mode 100644 index 38295d3cf..000000000 --- a/modules/sdnq/layers/conv/forward.py +++ /dev/null @@ -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) diff --git a/modules/sdnq/layers/embedding/forward.py b/modules/sdnq/layers/embedding/forward.py deleted file mode 100644 index 50370ec7e..000000000 --- a/modules/sdnq/layers/embedding/forward.py +++ /dev/null @@ -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) diff --git a/modules/sdnq/layers/linear/forward.py b/modules/sdnq/layers/linear/forward.py deleted file mode 100644 index 38ef65f26..000000000 --- a/modules/sdnq/layers/linear/forward.py +++ /dev/null @@ -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) diff --git a/modules/sdnq/layers/linear/linear_fp16.py b/modules/sdnq/layers/linear/linear_fp16.py deleted file mode 100644 index 07afeb27f..000000000 --- a/modules/sdnq/layers/linear/linear_fp16.py +++ /dev/null @@ -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) diff --git a/modules/sdnq/layers/linear/linear_fp8.py b/modules/sdnq/layers/linear/linear_fp8.py deleted file mode 100644 index 71ce7e404..000000000 --- a/modules/sdnq/layers/linear/linear_fp8.py +++ /dev/null @@ -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) diff --git a/modules/sdnq/layers/linear/linear_int8.py b/modules/sdnq/layers/linear/linear_int8.py deleted file mode 100644 index 4a2c0d3ff..000000000 --- a/modules/sdnq/layers/linear/linear_int8.py +++ /dev/null @@ -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) diff --git a/modules/sdnq/layers/linear/linear_uint8.py b/modules/sdnq/layers/linear/linear_uint8.py deleted file mode 100644 index 12fb23756..000000000 --- a/modules/sdnq/layers/linear/linear_uint8.py +++ /dev/null @@ -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) diff --git a/modules/sdnq/loader.py b/modules/sdnq/loader.py deleted file mode 100644 index 60235af69..000000000 --- a/modules/sdnq/loader.py +++ /dev/null @@ -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 diff --git a/modules/sdnq/packed_float.py b/modules/sdnq/packed_float.py deleted file mode 100644 index 1d75afb31..000000000 --- a/modules/sdnq/packed_float.py +++ /dev/null @@ -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 diff --git a/modules/sdnq/packed_int/__init__.py b/modules/sdnq/packed_int/__init__.py deleted file mode 100644 index 465b98e86..000000000 --- a/modules/sdnq/packed_int/__init__.py +++ /dev/null @@ -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 diff --git a/modules/sdnq/packed_int/pack.py b/modules/sdnq/packed_int/pack.py deleted file mode 100644 index 839398eda..000000000 --- a/modules/sdnq/packed_int/pack.py +++ /dev/null @@ -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 diff --git a/modules/sdnq/packed_int/unpack.py b/modules/sdnq/packed_int/unpack.py deleted file mode 100644 index ff76482f7..000000000 --- a/modules/sdnq/packed_int/unpack.py +++ /dev/null @@ -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 diff --git a/modules/sdnq/quant_utils.py b/modules/sdnq/quant_utils.py deleted file mode 100644 index a711cfe64..000000000 --- a/modules/sdnq/quant_utils.py +++ /dev/null @@ -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) diff --git a/modules/sdnq/quantizer.py b/modules/sdnq/quantizer.py deleted file mode 100644 index a0f216538..000000000 --- a/modules/sdnq/quantizer.py +++ /dev/null @@ -1,1065 +0,0 @@ -# pylint: disable=redefined-builtin,no-member,protected-access - -from dataclasses import dataclass -from enum import Enum - -import os -import torch - -from transformers.quantizers import HfQuantizer -from diffusers.quantizers.base import DiffusersQuantizer -from diffusers.quantizers.quantization_config import QuantizationConfigMixin - -from diffusers.utils import get_module_from_name -from accelerate import init_empty_weights - -from modules import devices, shared -from .common import ( - sdnq_keys, - dtype_dict, - accepted_weight_dtypes, - accepted_matmul_dtypes, - linear_types, - embedding_types, - conv_types, - conv_transpose_types, - weights_dtype_order, - check_torch_compile, - compile_func, -) - -from .quant_utils import ( - quantize_weight, - apply_svdquant, - apply_hadamard, - prepare_weight_for_matmul, - prepare_svd_for_matmul, -) -from .utils import ( - check_param_name_in, - check_quant_is_allowed, - check_quantized_matmul_is_allowed, - get_quant_args_from_config, - get_quant_kwargs, - get_quantized_matmul_dtype, - add_module_skip_keys, -) - -from .kernel_wrappers import is_fp8_mm_supported, use_tensorwise_fp8_matmul -from .dequantizer import SDNQDequantizer, dequantize_sdnq_model -from .packed_int import pack_int -from .packed_float import pack_float -from .forward import get_forward_func -from .layers import get_sdnq_wrapper_class -from .common import sdnq_version as current_sdnq_version - - -from diffusers import __version__ as diffusers_version_str # pylint: disable=ungrouped-imports,wrong-import-order -diffusers_version = [int(i) for i in diffusers_version_str.split(".")[:3]] - - -class QuantizationMethod(str, Enum): - SDNQ = "sdnq" - SDNQ_TRAINING = "sdnq_training" - - -@devices.inference_context() -def sdnq_quantize_layer_weight( - weight: torch.FloatTensor, - layer_class_name: str | None = None, - weights_dtype: str = "int8", - quantized_matmul_dtype: str | None = None, - group_size: int = 0, - hadamard_group_size: int = 256, - svd_rank: int = 32, - svd_steps: int = 8, - use_svd: bool = False, - use_hadamard: bool = False, - use_quantized_matmul: bool = False, - use_stochastic_rounding: bool = False, - dequantize_fp32: bool = True, - hadamard: torch.FloatTensor | None = None, - using_pre_calculated_svd: bool = False, - using_pre_rotated_hadamard: bool = False, - skip_sr: bool = False, - param_name: str | None = None, # pylint: disable=unused-argument - torch_dtype: torch.dtype | None = None, -) -> tuple[SDNQDequantizer, dict[str, torch.Tensor]]: - num_of_groups = 1 - is_conv_type = False - is_conv_transpose_type = False - is_linear_type = False - result_shape = None - scale_dtype = None - - weight = weight.detach() - original_shape = weight.shape - original_stride = weight.stride() - - if torch_dtype is None: - torch_dtype = weight.dtype - quantized_matmul_dtype = get_quantized_matmul_dtype(weights_dtype, quantized_matmul_dtype) - - re_quantize_for_matmul = bool( - dtype_dict[weights_dtype]["num_bits"] > dtype_dict[quantized_matmul_dtype]["num_bits"] - or dtype_dict[weights_dtype]["is_integer"] != dtype_dict[quantized_matmul_dtype]["is_integer"] - or (dtype_dict[weights_dtype]["is_unsigned"] and not dtype_dict[quantized_matmul_dtype]["is_integer"]) - or ( - dtype_dict[weights_dtype]["is_packed"] - and not dtype_dict[weights_dtype]["is_integer"] - and not dtype_dict[quantized_matmul_dtype]["is_integer"] - and ( - dtype_dict[weights_dtype]["num_bits"] >= dtype_dict[quantized_matmul_dtype]["num_bits"] - or dtype_dict[weights_dtype]["max"] > dtype_dict[quantized_matmul_dtype]["max"] - ) - ) - ) - - if layer_class_name in conv_types: - is_conv_type = True - reduction_axes = 1 - output_channel_size, channel_size = weight.shape[:2] - use_quantized_matmul = check_quantized_matmul_is_allowed(use_quantized_matmul, output_channel_size, channel_size) - if use_quantized_matmul and not re_quantize_for_matmul and not dtype_dict[weights_dtype]["is_packed"]: - result_shape = weight.shape - weight = weight.flatten(1,-1) - reduction_axes = -1 - elif layer_class_name in conv_transpose_types: - is_conv_transpose_type = True - reduction_axes = 0 - channel_size, output_channel_size = weight.shape[:2] - use_quantized_matmul = False - elif layer_class_name in linear_types: - is_linear_type = True - reduction_axes = -1 - output_channel_size, channel_size = weight.shape - use_quantized_matmul = check_quantized_matmul_is_allowed(use_quantized_matmul, output_channel_size, channel_size) - else: - if weight.ndim > 1: - output_channel_size, channel_size = weight.shape[-2:] - else: - output_channel_size, channel_size = 1, weight.shape[-1] - reduction_axes = -1 - use_quantized_matmul = False - - if ( - not dequantize_fp32 - and dtype_dict[weights_dtype]["max"] <= 16384 # 1/fp16_min_normal - and not ( - use_quantized_matmul - and not dtype_dict[quantized_matmul_dtype]["is_integer"] - and (not use_tensorwise_fp8_matmul or dtype_dict[quantized_matmul_dtype]["num_bits"] == 16) - ) - ): - scale_dtype = torch_dtype - - if use_hadamard: - weight, use_hadamard, hadamard_group_size = apply_hadamard(weight, group_size=hadamard_group_size, hadamard=hadamard, layer_class_name=layer_class_name) - - if use_svd: - try: - weight, svd_up, svd_down = apply_svdquant(weight, rank=svd_rank, niter=svd_steps, dtype=torch_dtype) - if use_quantized_matmul: - svd_up = svd_up.t_() - svd_down = svd_down.t_() - svd_up, svd_down = prepare_svd_for_matmul(svd_up, svd_down, use_quantized_matmul) - except Exception: - svd_up, svd_down = None, None - else: - svd_up, svd_down = None, None - - if group_size == 0: - if use_quantized_matmul and not re_quantize_for_matmul and dtype_dict[weights_dtype]["num_bits"] >= 6: - group_size = -1 - elif is_linear_type: - group_size = 2 ** ((3 if (svd_up is not None or using_pre_calculated_svd) else 2) + dtype_dict[weights_dtype]["num_bits"]) - else: - group_size = 2 ** ((2 if (svd_up is not None or using_pre_calculated_svd) else 1) + dtype_dict[weights_dtype]["num_bits"]) - - if group_size > 0: - if group_size >= channel_size: - group_size = channel_size - num_of_groups = 1 - else: - num_of_groups = channel_size // group_size - while num_of_groups * group_size != channel_size: # find something divisible - num_of_groups -= 1 - if num_of_groups <= 1: - group_size = channel_size - num_of_groups = 1 - break - group_size = channel_size // num_of_groups - group_size = int(group_size) - num_of_groups = int(num_of_groups) - - if num_of_groups > 1: - if result_shape is None: - result_shape = weight.shape - if is_conv_type: - # output_channel_size, channel_size, X, X - # output_channel_size, num_of_groups, group_size, X, X - reduction_axes = 2 - weight = weight.unflatten(1, (num_of_groups, group_size)) - elif is_conv_transpose_type: - #channel_size, output_channel_size, X, X - #num_of_groups, group_size, output_channel_size, X, X - reduction_axes = 1 - weight = weight.unflatten(1, (group_size, num_of_groups)) - else: - # output_channel_size, channel_size - # output_channel_size, num_of_groups, group_size - reduction_axes = -1 - weight = weight.unflatten(-1, (num_of_groups, group_size)) - else: - group_size = -1 - - cast_scale = True - transpose_weights = False - re_quantize_for_matmul = re_quantize_for_matmul or num_of_groups > 1 - if use_quantized_matmul and not re_quantize_for_matmul and not dtype_dict[weights_dtype]["is_packed"]: - transpose_weights = True - if not use_tensorwise_fp8_matmul and not dtype_dict[quantized_matmul_dtype]["is_integer"]: - cast_scale = False - - weight, scale, zero_point = quantize_weight(weight, reduction_axes, weights_dtype, dtype=(scale_dtype if cast_scale else None), use_stochastic_rounding=(use_stochastic_rounding and not skip_sr)) - - if transpose_weights: - scale = scale.t_().contiguous() - weight = weight.t_() - if zero_point is not None: - zero_point = zero_point.t_().contiguous() - weight = prepare_weight_for_matmul(weight, matmul_dtype=quantized_matmul_dtype) - - quantized_weight_shape = weight.shape - if dtype_dict[weights_dtype]["is_packed"]: - if dtype_dict[weights_dtype]["is_integer"]: - weight = pack_int(weight, weights_dtype) - else: - weight = pack_float(weight, weights_dtype) - else: - weight = weight.to(dtype=dtype_dict[weights_dtype]["torch_dtype"]) - - sdnq_dequantizer = SDNQDequantizer( - result_dtype=torch_dtype, - result_shape=result_shape, - original_shape=original_shape, - original_stride=original_stride, - quantized_weight_shape=quantized_weight_shape, - weights_dtype=weights_dtype, - quantized_matmul_dtype=quantized_matmul_dtype, - hadamard_group_size=hadamard_group_size, - group_size=group_size, - svd_rank=svd_rank, - svd_steps=svd_steps, - use_quantized_matmul=use_quantized_matmul, - re_quantize_for_matmul=re_quantize_for_matmul, - use_stochastic_rounding=use_stochastic_rounding, - use_hadamard=bool(use_hadamard or using_pre_rotated_hadamard), - layer_class_name=layer_class_name, - ) - - return (sdnq_dequantizer, {"weight": weight, "scale": scale, "zero_point": zero_point, "svd_up": svd_up, "svd_down": svd_down}) - - -@devices.inference_context() -def sdnq_quantize_layer_weight_dynamic( - weight: torch.FloatTensor, - layer_class_name: str | None = None, - weights_dtype: str = "uint4", - quantized_matmul_dtype: str | None = None, - group_size: int = 0, - hadamard_group_size: int = 256, - svd_rank: int = 32, - svd_steps: int = 8, - dynamic_loss_threshold: float | None =None, - use_svd: bool = False, - use_hadamard: bool = False, - use_quantized_matmul: bool = False, - use_stochastic_rounding: bool = False, - dequantize_fp32: bool = True, - hadamard: torch.FloatTensor | None = None, - param_name: str | None = None, - torch_dtype: torch.dtype | None = None, - quantization_config: "SDNQConfig" = None, -) -> tuple[SDNQDequantizer, dict[str, torch.Tensor]] | None: - if torch_dtype is None: - torch_dtype = weight.dtype - if dynamic_loss_threshold is None or dynamic_loss_threshold < 0: - dynamic_loss_threshold = 10 ** -(dtype_dict[weights_dtype]["num_bits"] / 2) - - weight = weight.detach() - if weight.dtype != torch.float64: - weight = weight.to(dtype=torch.float32, copy=False) - weight_std = weight.std().square_().clamp_(min=1e-8) - original_weight_fp32 = weight - - if use_hadamard: - weight, use_hadamard, hadamard_group_size = apply_hadamard(weight, group_size=hadamard_group_size, hadamard=hadamard, layer_class_name=layer_class_name) - - if use_svd: - try: - weight, svd_up, svd_down = apply_svdquant(weight, rank=svd_rank, niter=svd_steps, dtype=torch_dtype) - svd_up, svd_down = prepare_svd_for_matmul(svd_up, svd_down, False) - if use_quantized_matmul: - svd_up_t, svd_down_t = svd_up.clone().t_(), svd_down.clone().t_() - svd_up_t, svd_down_t = prepare_svd_for_matmul(svd_up_t, svd_down_t, True) - else: - svd_up_t, svd_down_t = None, None - except Exception: - svd_up, svd_down = None, None - svd_up_t, svd_down_t = None, None - else: - svd_up, svd_down = None, None - svd_up_t, svd_down_t = None, None - - quantization_loss = None - for i in range(weights_dtype_order.index(weights_dtype), len(weights_dtype_order)): - add_param_to_not_use_matmul = False - current_weights_dtype = weights_dtype_order[i] - current_quantized_matmul_dtype = get_quantized_matmul_dtype(current_weights_dtype, quantized_matmul_dtype) - if quantized_matmul_dtype is None and not is_fp8_mm_supported and current_quantized_matmul_dtype in {"fp8", "float8_e4m3fn", "float8_e5m2"}: - current_use_quantized_matmul = False - add_param_to_not_use_matmul = use_quantized_matmul - else: - current_use_quantized_matmul = use_quantized_matmul - - if ( - (dtype_dict[current_quantized_matmul_dtype]["is_integer"] and not dtype_dict[current_weights_dtype]["is_integer"]) - or ( - dtype_dict[current_weights_dtype]["num_bits"] == dtype_dict[current_quantized_matmul_dtype]["num_bits"] - and dtype_dict[current_weights_dtype]["is_unsigned"] and not dtype_dict[current_quantized_matmul_dtype]["is_integer"] - ) - or ( - dtype_dict[weights_dtype]["num_bits"] <= dtype_dict[current_quantized_matmul_dtype]["num_bits"] - and dtype_dict[current_weights_dtype]["num_bits"] > dtype_dict[current_quantized_matmul_dtype]["num_bits"] - ) - ): - current_use_quantized_matmul = False - add_param_to_not_use_matmul = True - - sdnq_dequantizer, weight_data = sdnq_quantize_layer_weight( - weight, - layer_class_name=layer_class_name, - weights_dtype=current_weights_dtype, - quantized_matmul_dtype=current_quantized_matmul_dtype, - torch_dtype=torch_dtype, - hadamard_group_size=hadamard_group_size, - group_size=group_size, - svd_rank=svd_rank, - svd_steps=svd_steps, - use_svd=False, - use_hadamard=False, - use_quantized_matmul=current_use_quantized_matmul, - use_stochastic_rounding=use_stochastic_rounding, - dequantize_fp32=dequantize_fp32, - using_pre_calculated_svd=use_svd, - using_pre_rotated_hadamard=use_hadamard, - param_name=param_name, - ) - - if sdnq_dequantizer.use_quantized_matmul: - weight_data["svd_up"] = svd_up_t - weight_data["svd_down"] = svd_down_t - else: - weight_data["svd_up"] = svd_up - weight_data["svd_down"] = svd_down - - quantization_loss = torch.nn.functional.mse_loss( - original_weight_fp32, - sdnq_dequantizer( - weight_data["weight"], - weight_data["scale"], - zero_point=weight_data["zero_point"], - svd_up=weight_data["svd_up"], - svd_down=weight_data["svd_down"], - skip_quantized_matmul=sdnq_dequantizer.use_quantized_matmul, - dtype=weight.dtype, - skip_compile=True, - ), - ).div_(weight_std) - if quantization_loss <= dynamic_loss_threshold: - del original_weight_fp32 - if quantization_config is not None: - if sdnq_dequantizer.weights_dtype not in quantization_config.modules_dtype_dict: - quantization_config.modules_dtype_dict[sdnq_dequantizer.weights_dtype] = [param_name] - else: - quantization_config.modules_dtype_dict[sdnq_dequantizer.weights_dtype].append(param_name) - if add_param_to_not_use_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) - return (sdnq_dequantizer, weight_data), quantization_config - else: - return (sdnq_dequantizer, weight_data) - - del original_weight_fp32 - if quantization_config is not None: - quantization_config.modules_to_not_convert.append(param_name) - return None, quantization_config - else: - return None - - -@devices.inference_context() -def sdnq_quantize_layer(layer: torch.nn.Module, quantization_config: "SDNQConfig", torch_dtype: torch.dtype | None = None, param_name: str = "", quant_kwargs: dict | None = None) -> tuple[torch.nn.Module, "SDNQConfig"]: # pylint: disable=unused-argument - if torch_dtype is None: - torch_dtype = layer.weight.dtype - if quant_kwargs is None: - quant_kwargs = get_quant_kwargs(layer, quantization_config, torch_dtype=torch_dtype, param_name=param_name) - - layer_class_name = layer.__class__.__name__ - if ( - (layer_class_name in embedding_types and not quantization_config.quant_embedding) - or ((layer_class_name in conv_transpose_types or layer_class_name in conv_types) and not quantization_config.quant_conv) - ): - quantization_config.modules_to_not_convert.append(param_name) - return layer, quantization_config - - return_device = quant_kwargs.pop("return_device") - quantization_device = quant_kwargs.pop("quantization_device") - non_blocking = quant_kwargs.pop("non_blocking") - use_dynamic_quantization = quant_kwargs.pop("use_dynamic_quantization") - - layer.weight.requires_grad_(False) - if return_device is None: - return_device = layer.weight.device - if quantization_device is not None: - layer.weight.data = layer.weight.to(quantization_device, non_blocking=non_blocking, copy=False) - - if use_dynamic_quantization: - weight_data, quantization_config = sdnq_quantize_layer_weight_dynamic(layer.weight, quantization_config=quantization_config, **quant_kwargs) - else: - weight_data = sdnq_quantize_layer_weight(layer.weight, **quant_kwargs) - - if weight_data is not None: - layer.sdnq_dequantizer, weight_data = weight_data - layer = get_sdnq_wrapper_class(layer, get_forward_func(layer_class_name, layer.sdnq_dequantizer.quantized_matmul_dtype, layer.sdnq_dequantizer.use_quantized_matmul)) - - for key, value in weight_data.items(): - if isinstance(value, (torch.Tensor, torch.nn.Parameter)): - setattr(layer, key, torch.nn.Parameter(value.to(return_device, non_blocking=non_blocking, copy=False), requires_grad=False)) - setattr(getattr(layer, key), "_is_hf_initialized", True) # noqa: B010 - else: - setattr(layer, key, value) - del weight_data - - if ( - quant_kwargs["use_quantized_matmul"] and not layer.sdnq_dequantizer.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) - else: - layer.weight = torch.nn.Parameter(layer.weight.to(return_device, dtype=torch_dtype, non_blocking=non_blocking, copy=False), requires_grad=False) - - return layer, quantization_config - - -@devices.inference_context() -def apply_sdnq_to_module(model: torch.nn.Module, quantization_config: "SDNQConfig", torch_dtype: torch.dtype | None = None, pre_quantized: bool = False, full_param_name: str = "") -> tuple[torch.nn.Module, "SDNQConfig"]: # pylint: disable=unused-argument - if not list(model.children()): - return model, quantization_config - 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, "weight") and module.weight is not None: - param_name = param_name + ".weight" - if check_param_name_in(param_name, quantization_config.modules_to_not_convert) is None: - if check_quant_is_allowed(module.__class__.__name__, module.weight, quantization_config, pre_quantized=pre_quantized): - module, quantization_config = sdnq_quantize_layer(module, quantization_config, torch_dtype=torch_dtype, param_name=param_name) - setattr(model, module_name, module) - else: - quantization_config.modules_to_not_convert.append(param_name) - module, quantization_config = apply_sdnq_to_module(module, quantization_config, torch_dtype=torch_dtype, pre_quantized=pre_quantized, full_param_name=param_name) - setattr(model, module_name, module) - return model, quantization_config - - -@devices.inference_context() -def sdnq_post_load_quant( - model: torch.nn.Module, - weights_dtype: str = "int8", - quantized_matmul_dtype: str | None = None, - hadamard_group_size: int = 256, - group_size: int = 0, - svd_rank: int = 32, - svd_steps: int = 8, - dynamic_loss_threshold: float | None = None, - use_svd: bool = False, - use_hadamard: bool = False, - quant_conv: bool = False, - quant_embedding: bool = False, - use_quantized_matmul: bool = False, - use_quantized_matmul_conv: bool = False, - use_dynamic_quantization: bool = False, - use_stochastic_rounding: bool = False, - dequantize_fp32: bool = True, - non_blocking: bool = False, - add_skip_keys:bool = True, - minimum_allowed_numel: int = 16384, - minimum_allowed_channel_size: int = 32, - modules_to_not_convert: list[str] | None = None, - modules_to_not_use_matmul: list[str] | None = None, - modules_dtype_dict: dict[str, list[str]] | None = None, - modules_quant_config: dict[str, dict] | None = None, - quantization_device: torch.device | None = None, - return_device: torch.device | None = None, - torch_dtype: torch.dtype | None = None, - quantization_config: "SDNQConfig" = None, - pre_quantized: bool = False, -) -> torch.nn.Module: - if pre_quantized: - add_skip_keys = False - use_dynamic_quantization = False - if quantization_config is not None: - quantization_config.add_skip_keys = False - quantization_config.use_dynamic_quantization = False - elif ( - hasattr(model, "quantization_config") - or (hasattr(model, "config") and hasattr(model.config, "quantization_config")) - or (hasattr(model, "config") and isinstance(model.config, dict) and "quantization_config" in model.config) - ): - raise RuntimeError("Quantizing a pre-quantized model is not supported!") - - if quantization_config is None: - quantization_config = SDNQConfig( - weights_dtype=weights_dtype, - quantized_matmul_dtype=quantized_matmul_dtype, - hadamard_group_size=hadamard_group_size, - group_size=group_size, - svd_rank=svd_rank, - svd_steps=svd_steps, - dynamic_loss_threshold=dynamic_loss_threshold, - use_svd=use_svd, - use_hadamard=use_hadamard, - quant_conv=quant_conv, - quant_embedding=quant_embedding, - use_quantized_matmul=use_quantized_matmul, - use_quantized_matmul_conv=use_quantized_matmul_conv, - use_dynamic_quantization=use_dynamic_quantization, - use_stochastic_rounding=use_stochastic_rounding, - dequantize_fp32=dequantize_fp32, - non_blocking=non_blocking, - add_skip_keys=add_skip_keys, - minimum_allowed_numel=minimum_allowed_numel, - minimum_allowed_channel_size=minimum_allowed_channel_size, - modules_to_not_convert=modules_to_not_convert, - modules_to_not_use_matmul=modules_to_not_use_matmul, - modules_dtype_dict=modules_dtype_dict, - modules_quant_config=modules_quant_config, - quantization_device=quantization_device, - return_device=return_device, - ) - if quantization_config.add_skip_keys: - model, quantization_config = add_module_skip_keys(model, quantization_config) - - model.eval() - model, quantization_config = apply_sdnq_to_module(model, quantization_config, torch_dtype=torch_dtype, pre_quantized=pre_quantized) - - model.quantization_config = quantization_config - if hasattr(model, "config"): - try: - model.config.quantization_config = model.quantization_config - except Exception: - pass - try: - model.config["quantization_config"] = model.quantization_config.to_dict() - except Exception: - pass - model.quantization_method = QuantizationMethod.SDNQ - - return model - - -class SDNQQuantize: - def __init__(self, hf_quantizer: "SDNQQuantizer"): - self.hf_quantizer = hf_quantizer - - def convert( # pylint: disable=unused-argument - self, - input_dict: dict[str, list[torch.Tensor]], - model: torch.nn.Module | None = None, - full_layer_name: str | None = None, - missing_keys: list[str] | None = None, - **kwargs, - ) -> dict[str, torch.Tensor]: - _module_name, value = next(iter(input_dict.items())) - value = value[0] - self.hf_quantizer.create_quantized_param(model, value, full_layer_name, value.device) - param, name = get_module_from_name(model, full_layer_name) - param = getattr(param, name) - return {full_layer_name: param} - - @property - def reverse_op(self): - raise NotImplementedError - - -class SDNQQuantizer(DiffusersQuantizer, HfQuantizer): - r""" - Diffusers and Transformers Quantizer for SDNQ - """ - - requires_parameters_quantization = True - use_keep_in_fp32_modules = True - requires_calibration = False - required_packages = None - torch_dtype = None - - def __str__(self) -> str: - return f"SDNQQuantizer(torch_dtype={self.torch_dtype}, requires_parameters_quantization={self.requires_parameters_quantization}, use_keep_in_fp32_modules={self.use_keep_in_fp32_modules}, requires_calibration={self.requires_calibration}, required_packages={self.required_packages})" - - def check_if_quantized_param( # pylint: disable=unused-argument - self, - model: torch.nn.Module, - param_value: torch.Tensor, - param_name: str, - *args, **kwargs, - ) -> bool: - if self.pre_quantized: - layer, _tensor_name = get_module_from_name(model, param_name) - if hasattr(layer, "sdnq_dequantizer") and param_name.rsplit(".", maxsplit=1)[-1] in sdnq_keys: - return True - elif param_name.endswith(".weight"): - if check_param_name_in(param_name, self.quantization_config.modules_to_not_convert) is None: - layer = get_module_from_name(model, param_name)[0] - if check_quant_is_allowed(layer.__class__.__name__, layer.weight, self.quantization_config, pre_quantized=self.pre_quantized): - return True - else: - self.quantization_config.modules_to_not_convert.append(param_name) - return False - - @devices.inference_context() - def create_quantized_param( # pylint: disable=unused-argument,arguments-differ - self, - model: torch.nn.Module, - param_value: torch.FloatTensor, - param_name: str, - target_device: torch.device, - *args, - **kwargs, - ) -> None: - layer, tensor_name = get_module_from_name(model, param_name) - - if self.pre_quantized: - if param_value is not None: - if tensor_name == "weight": - return_dtype = param_value.dtype - elif self.quantization_config.dequantize_fp32 and tensor_name in sdnq_keys: - if torch.float64 not in {param_value.dtype, self.torch_dtype}: - return_dtype = torch.float32 - else: - return_dtype = torch.float64 - else: - return_dtype = kwargs.get("dtype", param_value.dtype if self.torch_dtype is None else self.torch_dtype) - - param_value = param_value.detach() - param_value = param_value.to(target_device, dtype=return_dtype, copy=False) - - if tensor_name == "weight" and layer.sdnq_dequantizer.use_quantized_matmul and not layer.sdnq_dequantizer.re_quantize_for_matmul: - param_value = prepare_weight_for_matmul(param_value, matmul_dtype=layer.sdnq_dequantizer.quantized_matmul_dtype) - elif tensor_name == "svd_up": - param_value, _ = prepare_svd_for_matmul(param_value, None, layer.sdnq_dequantizer.use_quantized_matmul) - elif tensor_name == "svd_down": - _, param_value = prepare_svd_for_matmul(None, param_value, layer.sdnq_dequantizer.use_quantized_matmul) - - param_value = torch.nn.Parameter(param_value, requires_grad=False) - param_value._is_hf_initialized = True # pylint: disable=protected-access - setattr(layer, tensor_name, param_value) - return - - torch_dtype = kwargs.get("dtype", param_value.dtype if self.torch_dtype is None else self.torch_dtype) - quant_kwargs = get_quant_kwargs(layer, self.quantization_config, torch_dtype=torch_dtype, param_name=param_name) - if quant_kwargs["return_device"] is None: - quant_kwargs["return_device"] = target_device - if quant_kwargs["quantization_device"] is not None: - target_device = quant_kwargs["quantization_device"] - quant_kwargs["quantization_device"] = None - - param_value = param_value.detach() - param_value = param_value.to(target_device, non_blocking=self.quantization_config.non_blocking, copy=False).to(dtype=torch.float32 if param_value.dtype != torch.float64 else torch.float64) - - layer.weight = torch.nn.Parameter(param_value, requires_grad=False) - layer, self.quantization_config = sdnq_quantize_layer(layer, self.quantization_config, torch_dtype=torch_dtype, param_name=param_name, quant_kwargs=quant_kwargs) # pylint: disable=attribute-defined-outside-init - - parent_module, tensor_name = get_module_from_name(model, param_name.removesuffix(tensor_name).removesuffix(".")) - setattr(parent_module, tensor_name, layer) - - def _process_model_before_weight_loading( # pylint: disable=unused-argument,arguments-differ - self, - model: torch.nn.Module, - device_map, - keep_in_fp32_modules: list[str] | None = None, - **kwargs, - ) -> None: - if self.pre_quantized: - self.quantization_config.quantization_device = None - self.quantization_config.return_device = None - self.quantization_config.non_blocking = False - self.quantization_config.add_skip_keys = False - - with init_empty_weights(): - model = sdnq_post_load_quant(model, torch_dtype=self.torch_dtype, pre_quantized=True, **get_quant_args_from_config(self.quantization_config)) - - if self.quantization_config.add_skip_keys: - if keep_in_fp32_modules is not None: - self.quantization_config.modules_to_not_convert.extend(keep_in_fp32_modules) - if hasattr(self, "get_modules_to_not_convert") and hasattr(model, "tie_weights"): - self.quantization_config.modules_to_not_convert.extend(self.get_modules_to_not_convert(model, add_default_skips=True)) - model, self.quantization_config = add_module_skip_keys(model, self.quantization_config) # pylint: disable=attribute-defined-outside-init - - - def _process_model_after_weight_loading(self, model: torch.nn.Module, **kwargs) -> torch.nn.Module: # pylint: disable=unused-argument - model.quantization_config = self.quantization_config - model.quantization_method = QuantizationMethod.SDNQ - if hasattr(model, "config"): - try: - model.config.quantization_config = self.quantization_config - except Exception: - pass - try: - model.config["quantization_config"] = self.quantization_config.to_dict() - except Exception: - pass - - if self.pre_quantized: - from .loader import post_process_model - model = post_process_model(model) - if self.quantization_config.is_training: - from .training import convert_sdnq_model_to_training - model = convert_sdnq_model_to_training( - model, - dtype=self.torch_dtype, - quantized_matmul_dtype=self.quantization_config.quantized_matmul_dtype, - use_grad_ckpt=self.quantization_config.use_grad_ckpt, - use_quantized_matmul=self.quantization_config.use_quantized_matmul, - use_stochastic_rounding=self.quantization_config.use_stochastic_rounding, - dequantize_fp32=self.quantization_config.dequantize_fp32, - ) - - if shared.opts.diffusers_offload_mode != "none": - try: - model = model.to(device=devices.cpu) - except Exception: - model = model.to_empty(device=devices.cpu) - devices.torch_gc(force=True, reason="sdnq") - return model - - def get_quantize_ops(self) -> SDNQQuantize: - return SDNQQuantize(self) - - def adjust_max_memory(self, max_memory: dict[str, int | str]) -> dict[str, int | str]: - max_memory = {key: val * 0.80 for key, val in max_memory.items()} - return max_memory - - def adjust_target_dtype(self, target_dtype: torch.dtype) -> torch.dtype: # pylint: disable=unused-argument,arguments-renamed - return dtype_dict[self.quantization_config.weights_dtype]["target_dtype"] - - def update_torch_dtype(self, torch_dtype: torch.dtype) -> torch.dtype: - self.torch_dtype = torch_dtype - return torch_dtype - - def get_state_dict_and_metadata(self, state_dict: dict | torch.nn.Module, **kwargs) -> tuple[dict | None, dict]: # pylint: disable=unused-argument, arguments-differ - # transformers - if isinstance(state_dict, torch.nn.Module): - return None, {} - # diffusers - return state_dict, {} - - def get_accelerator_warm_up_factor(self) -> int: - return 32 // dtype_dict[self.quantization_config.weights_dtype]["num_bits"] - - def _dequantize(self, model: torch.nn.Module) -> torch.nn.Module: - return dequantize_sdnq_model(model) - - def is_serializable(self, *args, **kwargs) -> bool: # pylint: disable=unused-argument, invalid-overridden-method - return not self.quantization_config.is_training - - @property - def supports_safetensors_serialization(self) -> bool: - return self.is_serializable() - - @property - def is_trainable(self) -> bool: - return self.quantization_config.is_training - - @property - def is_qat_trainable(self) -> bool: - return self.is_trainable() - - @property - def is_compileable(self) -> bool: - return True - - def check_quantized_param(self, *args, **kwargs) -> bool: - """ - needed for transformers compatibility, returns self.check_if_quantized_param - """ - return self.check_if_quantized_param(*args, **kwargs) - - def param_needs_quantization(self, model: torch.nn.Module, param_name: str, *args, **kwargs) -> bool: - """ - needed for transformers compatibility, returns self.check_if_quantized_param - """ - return self.check_if_quantized_param(model, None, param_name, *args, **kwargs) - - def get_cuda_warm_up_factor(self) -> int: - """ - needed for transformers compatibility, returns self.get_accelerator_warm_up_factor - """ - return self.get_accelerator_warm_up_factor() - - def update_dtype(self, dtype: torch.dtype) -> torch.dtype: - """ - needed for transformers compatibility, returns self.update_torch_dtype - """ - return self.update_torch_dtype(dtype) - - -@dataclass -class SDNQConfig(QuantizationConfigMixin): - """ - This is a wrapper class about all possible attributes and features that you can play with a model that has been - loaded using `sdnq`. - - Args: - weights_dtype (`str`, *optional*, defaults to `"int8"`): - The target dtype for the weights after quantization. - See `sdnq.common.accepted_weight_dtypes` for all the supported values. - These are some of the recommended values to use: ("int8", "int7", "int6", "uint5", "uint4", "uint3", "uint2", "float8_e4m3fn", "float7_e3m3fn", "float6_e3m2fn", "float5_e2m2fn", "float4_e2m1fn", "float3_e1m1fn", "float2_e1m0fn") - quantized_matmul_dtype (`str`, *optional*, defaults to `None`): - The target dtype for quantized matmul. - `None` will use "int8" with integer weight dtypes and "float8_e4m3fn" or "float16" with float weight dtypes. - Supported values are: ("int8", "float8_e4m3fn", "float16") - hadamard_group_size (`int`, *optional*, defaults to `0`): - Used to decide how many elements of a tensor will share the same Hadamard rotation group. - group_size (`int`, *optional*, defaults to `0`): - Used to decide how many elements of a tensor will share the same quantization group. - group_size = `0` will automatically select a group size based on weights_dtype. - group_size = `-1` will disable group sizes and use simple row-wise quantization instead. - svd_rank (`int`, *optional*, defaults to `32`): - The rank size used for the SVDQuant algorithm. - svd_steps (`int`, *optional*, defaults to `8`): - The number of iterations to use in svd lowrank estimation. - dynamic_loss_threshold (`float`, *optional*, defaults to `None`): - The target quantization mse loss threshold to use for dynamic quantization. - The value `None` or negative values means auto select a threshold based on the weights_dtype. - use_svd (`bool`, *optional*, defaults to `False`): - Enabling this option will use SVDQuant algorithm on top of SDNQ quantization. - use_hadamard (`bool`, *optional*, defaults to `False`): - Enabling this option will use Hadamard rotation on top of SDNQ quantization. - use_grad_ckpt (`bool`, *optional*, defaults to `True`): - This option is only used for training models when `is_training` is enabled or with `sdnq.training.sdnq_training_post_load_quant`. - Disabling this option will quantize the tensors needed for the backward pass before saving in the forward pass. - Disable this option only if you are not using gradient checkpointing and enable it when you are using gradient checkpointing. - quant_conv (`bool`, *optional*, defaults to `False`): - Enabling this option will quantize the nn.Conv layers in UNet models too. - Note: nn.Conv layers doesn't support training. - quant_embedding (`bool`, *optional*, defaults to `False`): - Enabling this option will quantize the nn.Embedding layers in text models too. - Note: nn.Embedding layers doesn't support training. - use_quantized_matmul (`bool`, *optional*, defaults to `False`): - Enabling this option will use quantized INT8, FP8 or FP16 MatMul on the forward pass and the backward pass instead of BF16 / FP16. - `False` means no quantized matmul and quantized_matmul_dtype will be ignored. - use_quantized_matmul_conv (`bool`, *optional*, defaults to `False`): - Same as use_quantized_matmul but for the nn.Conv layers with UNets like SDXL. - use_stochastic_rounding (`bool`, *optional*, defaults to `False`): - Enabling this option will use stochastic rounding on the quantization step. - Recommended for full finetuning / full rank training. - use_static_quantization (`bool`, *optional*, defaults to `True`): - This option is only used for training models with `sdnq.training.sdnq_training_post_load_quant`. - Disabling this option means model weights will be kept unquantized in their original precision and only the use_quantized_matmul (if enabled) will be used. - use_dynamic_quantization (`bool`, *optional*, defaults to `False`): - Enabling this option will dynamically select a per layer quantization type based on the dynamic_loss_threshold. - weights_dtype will be used as the minimum allowed quantization type when this option is enabled. - dequantize_fp32 (`bool`, *optional*, defaults to `True`): - Enabling this option will keep the quant scales in FP32 and compute the de-quant steps in FP32. - This option is highly recommended to be enabled. - non_blocking (`bool`, *optional*, defaults to `False`): - Enabling this option will use non blocking ops when moving layers between the quantization device and the return device. - add_skip_keys (`bool`, *optional*, defaults to `True`): - Disabling this option won't add model specific keys to modules_to_not_convert, modules_to_not_use_matmul and modules_dtype_dict. - minimum_allowed_numel (`int`, *optional*, defaults to `16384`): - Layers that have less than `minimum_allowed_numel` elements in them will be skipped and added to `modules_to_not_convert`. - minimum_allowed_channel_size (`int`, *optional*, defaults to `32`): - Layers that have less than `minimum_allowed_channel_size` channels in them will be skipped and added to `modules_to_not_convert`. - modules_to_not_convert (`list`, *optional*, default to `None`): - The list of modules to not quantize. Useful for quantizing models that explicitly require to have some - modules left in their original precision (e.g. Whisper encoder, Llava encoder, Mixtral gate layers). - modules_to_not_use_matmul (`list`, *optional*, default to `None`): - The list of modules to not use quantized matmul. Useful for quantizing models that explicitly require to have some - modules left in their original activation precision (e.g. AdaLN layers). - modules_dtype_dict (`dict`, *optional*, default to `None`): - The dict of dtypes and list of modules. Useful for quantizing some modules with a different dtype. - modules_quant_config (`dict`, *optional*, default to `None`): - The dict of modules and a dict of quantization kwargs to use for that module. - Useful for quantizing some modules with a different quantization config. - quantization_device (`torch.device`, *optional*, defaults to `None`): - Used to set which device will be used for the quantization calculation on model load. - return_device (`torch.device`, *optional*, defaults to `None`): - Used to set which device will the quantized weights be sent back to. - is_training (`bool`, *optional*, defaults to `False`): - This option is auto enabled when using `sdnq.training.sdnq_training_post_load_quant`. - Enabling this option with SDNQQuantizer will convert the model to SDNQ Training format after quantization. - Note: Safetensors serialization is not supported with SDNQ training. - """ - - def __init__( # pylint: disable=super-init-not-called,unused-argument - self, - weights_dtype: str = "int8", - quantized_matmul_dtype: str | None = None, - hadamard_group_size: int = 256, - group_size: int = 0, - svd_rank: int = 32, - svd_steps: int = 8, - dynamic_loss_threshold: float | None = None, - use_svd: bool = False, - use_hadamard: bool = False, - use_grad_ckpt: bool = True, - quant_conv: bool = False, - quant_embedding: bool = False, - use_quantized_matmul: bool = False, - use_quantized_matmul_conv: bool = False, - use_static_quantization: bool = True, - use_dynamic_quantization: bool = False, - use_stochastic_rounding: bool = False, - dequantize_fp32: bool = True, - non_blocking: bool = False, - add_skip_keys: bool = True, - minimum_allowed_numel: int = 16384, - minimum_allowed_channel_size: int = 32, - modules_to_not_convert: list[str] | None = None, - modules_to_not_use_matmul: list[str] | None = None, - modules_dtype_dict: dict[str, list[str]] | None = None, - modules_quant_config: dict[str, dict] | None = None, - quantization_device: torch.device | None = None, - return_device: torch.device | None = None, - sdnq_version: str | None = None, - is_training: bool = False, - **kwargs, - ): - self.weights_dtype = weights_dtype - self.quantized_matmul_dtype = quantized_matmul_dtype - self.is_training = is_training - self.hadamard_group_size = hadamard_group_size - self.group_size = group_size - self.svd_rank = svd_rank - self.dynamic_loss_threshold = dynamic_loss_threshold - self.svd_steps = svd_steps - self.use_svd = use_svd - self.use_hadamard = use_hadamard - self.use_grad_ckpt = use_grad_ckpt - self.quant_conv = quant_conv - self.quant_embedding = quant_embedding - self.use_quantized_matmul = use_quantized_matmul - self.use_quantized_matmul_conv = use_quantized_matmul_conv - self.use_static_quantization = use_static_quantization - self.use_dynamic_quantization = use_dynamic_quantization - self.use_stochastic_rounding = use_stochastic_rounding - self.dequantize_fp32 = dequantize_fp32 - self.non_blocking = non_blocking - self.add_skip_keys = add_skip_keys - self.minimum_allowed_numel = minimum_allowed_numel - self.minimum_allowed_channel_size = minimum_allowed_channel_size - self.modules_to_not_convert = modules_to_not_convert - self.modules_to_not_use_matmul = modules_to_not_use_matmul - self.modules_dtype_dict = modules_dtype_dict - self.modules_quant_config = modules_quant_config - self.quantization_device = quantization_device - self.return_device = return_device - self.sdnq_version = current_sdnq_version if sdnq_version is None else sdnq_version - self.is_integer = dtype_dict[self.weights_dtype]["is_integer"] - if self.is_training: - self.quant_method = QuantizationMethod.SDNQ_TRAINING - else: - self.quant_method = QuantizationMethod.SDNQ - self.post_init() - - def post_init(self) -> None: - r""" - Safety checker that arguments are correct - """ - if self.use_quantized_matmul and not check_torch_compile(): - shared.log.warning("SDNQ: Quantized MatMul requires a working Triton install for best performance.") - if self.weights_dtype not in accepted_weight_dtypes: - raise ValueError(f"SDNQ only support weight dtypes in {accepted_weight_dtypes} but found {self.weights_dtype}") - if self.quantized_matmul_dtype is not None and self.quantized_matmul_dtype not in accepted_matmul_dtypes: - raise ValueError(f"SDNQ only support quantized matmul dtypes in {accepted_matmul_dtypes} but found {self.quantized_matmul_dtype}") - - if self.modules_to_not_convert is None: - self.modules_to_not_convert = [] - elif isinstance(self.modules_to_not_convert, str): - self.modules_to_not_convert = [self.modules_to_not_convert] - elif isinstance(self.modules_to_not_convert, tuple): - self.modules_to_not_convert = list(self.modules_to_not_convert) - elif not isinstance(self.modules_to_not_convert, list): - raise ValueError(f"modules_to_not_convert must be a list but got {type(self.modules_to_not_convert)}") - - if self.modules_to_not_use_matmul is None: - self.modules_to_not_use_matmul = [] - elif isinstance(self.modules_to_not_use_matmul, str): - self.modules_to_not_use_matmul = [self.modules_to_not_use_matmul] - elif isinstance(self.modules_to_not_use_matmul, tuple): - self.modules_to_not_use_matmul = list(self.modules_to_not_use_matmul) - elif not isinstance(self.modules_to_not_use_matmul, list): - raise ValueError(f"modules_to_not_use_matmul must be a list but got {type(self.modules_to_not_use_matmul)}") - - if self.modules_dtype_dict is None: - self.modules_dtype_dict = {} - elif not isinstance(self.modules_dtype_dict, dict): - raise ValueError(f"modules_dtype_dict must be a dict but got {type(self.modules_dtype_dict)}") - elif len(self.modules_dtype_dict.keys()) > 0: - self.modules_dtype_dict = self.modules_dtype_dict.copy() - for key, value in self.modules_dtype_dict.items(): - if isinstance(value, str): - value = [value] - self.modules_dtype_dict[key] = value - elif isinstance(value, tuple): - value = list(value) - self.modules_dtype_dict[key] = value - if not isinstance(key, str) or not isinstance(value, list): - raise TypeError(f"modules_dtype_dict must be a dictionary of strings and lists but got {type(key)} and {type(value)}") - - if self.modules_quant_config is None: - self.modules_quant_config = {} - - self.modules_to_not_convert = self.modules_to_not_convert.copy() - self.modules_to_not_use_matmul = self.modules_to_not_use_matmul.copy() - self.modules_dtype_dict = self.modules_dtype_dict.copy() - self.modules_quant_config = self.modules_quant_config.copy() - - # dedupe - self.modules_to_not_convert = list(set(self.modules_to_not_convert)) - self.modules_to_not_use_matmul = list(set(self.modules_to_not_use_matmul)) - for key, value in self.modules_dtype_dict.items(): - self.modules_dtype_dict[key] = list(set(value)) - - def to_dict(self) -> dict: - quantization_config_dict = self.__dict__.copy() # make serializable - quantization_config_dict["quantization_device"] = str(quantization_config_dict["quantization_device"]) if quantization_config_dict["quantization_device"] is not None else None - quantization_config_dict["return_device"] = str(quantization_config_dict["return_device"]) if quantization_config_dict["return_device"] is not None else None - return quantization_config_dict - - def __str__(self) -> str: - return f"SDNQConfig(weights_dtype={self.weights_dtype} quantization_device={self.quantization_device} return_device={self.return_device} group_size={self.group_size} use_quantized_matmul={self.use_quantized_matmul} quantized_matmul_dtype={self.quantized_matmul_dtype} quant_conv={self.quant_conv} quant_embedding={self.quant_embedding} use_quantized_matmul_conv={self.use_quantized_matmul_conv} use_static_quantization={self.use_static_quantization} use_dynamic_quantization={self.use_dynamic_quantization} dynamic_loss_threshold={self.dynamic_loss_threshold} use_stochastic_rounding={self.use_stochastic_rounding} use_hadamard={self.use_hadamard} hadamard_group_size={self.hadamard_group_size} use_svd={self.use_svd} svd_rank={self.svd_rank} svd_steps={self.svd_steps} dequantize_fp32={self.dequantize_fp32} non_blocking={self.non_blocking} add_skip_keys={self.add_skip_keys} modules_to_not_convert={self.modules_to_not_convert} modules_to_not_use_matmul={self.modules_to_not_use_matmul} modules_dtype_dict={self.modules_dtype_dict} modules_quant_config={self.modules_quant_config} )" - - -if ( - os.environ.get("SDNQ_REGISTER_DIFFUSERS", "0").lower() not in {"0", "false", "no"} - or (diffusers_version[0] == 0 and diffusers_version[1] < 40) -): - import diffusers.quantizers.auto # noqa: E402,RUF100 # pylint: disable=wrong-import-order,wrong-import-position - diffusers.quantizers.auto.AUTO_QUANTIZER_MAPPING["sdnq"] = SDNQQuantizer - diffusers.quantizers.auto.AUTO_QUANTIZATION_CONFIG_MAPPING["sdnq"] = SDNQConfig - diffusers.quantizers.auto.AUTO_QUANTIZER_MAPPING["sdnq_training"] = SDNQQuantizer - diffusers.quantizers.auto.AUTO_QUANTIZATION_CONFIG_MAPPING["sdnq_training"] = SDNQConfig - - -if os.environ.get("SDNQ_REGISTER_TRANSFORMERS", "1").lower() not in {"0", "false", "no"}: - import transformers.quantizers.auto # noqa: E402,RUF100 # pylint: disable=wrong-import-order,wrong-import-position - transformers.quantizers.auto.AUTO_QUANTIZER_MAPPING["sdnq"] = SDNQQuantizer - transformers.quantizers.auto.AUTO_QUANTIZATION_CONFIG_MAPPING["sdnq"] = SDNQConfig - transformers.quantizers.auto.AUTO_QUANTIZER_MAPPING["sdnq_training"] = SDNQQuantizer - transformers.quantizers.auto.AUTO_QUANTIZATION_CONFIG_MAPPING["sdnq_training"] = SDNQConfig - - -sdnq_quantize_layer_weight_compiled = compile_func(sdnq_quantize_layer_weight) diff --git a/modules/sdnq/utils.py b/modules/sdnq/utils.py deleted file mode 100644 index bdc865087..000000000 --- a/modules/sdnq/utils.py +++ /dev/null @@ -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 diff --git a/modules/timer_sdnq.py b/modules/timer_sdnq.py index 97345a430..79f694f6a 100644 --- a/modules/timer_sdnq.py +++ b/modules/timer_sdnq.py @@ -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 diff --git a/modules/video_models/video_load.py b/modules/video_models/video_load.py index 5744e68e6..66419d6c5 100644 --- a/modules/video_models/video_load.py +++ b/modules/video_models/video_load.py @@ -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): diff --git a/pipelines/native_transformer.py b/pipelines/native_transformer.py index 325c72e3f..59f655403 100644 --- a/pipelines/native_transformer.py +++ b/pipelines/native_transformer.py @@ -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" diff --git a/test/test-ideogram4-smoke.py b/test/test-ideogram4-smoke.py index 971671dff..5b441844d 100644 --- a/test/test-ideogram4-smoke.py +++ b/test/test-ideogram4-smoke.py @@ -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" diff --git a/test/test-native-transformer.py b/test/test-native-transformer.py index 0bf7d8b77..84775b51e 100644 --- a/test/test-native-transformer.py +++ b/test/test-native-transformer.py @@ -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