mirror of
https://github.com/Comfy-Org/ComfyUI.git
synced 2026-08-05 18:05:08 +08:00
Support convrot int4 models. (#14859)
linear_dtype in comfy_quant metadata can be used to set if the int4 op does the matrix multiplication in int8 or int4, the default is int4 on GPUs that support it with fallback to int8 for GPUs that don't.
This commit is contained in:
@@ -10,6 +10,7 @@ try:
|
||||
QuantizedLayout,
|
||||
TensorCoreFP8Layout as _CKFp8Layout,
|
||||
TensorCoreNVFP4Layout as _CKNvfp4Layout,
|
||||
TensorCoreConvRotW4A4Layout as _CKTensorCoreConvRotW4A4Layout,
|
||||
TensorWiseINT8Layout as _CKTensorWiseINT8Layout,
|
||||
register_layout_op,
|
||||
register_layout_class,
|
||||
@@ -51,6 +52,9 @@ except ImportError as e:
|
||||
class _CKTensorWiseINT8Layout:
|
||||
pass
|
||||
|
||||
class _CKTensorCoreConvRotW4A4Layout:
|
||||
pass
|
||||
|
||||
def register_layout_class(name, cls):
|
||||
pass
|
||||
|
||||
@@ -179,6 +183,7 @@ class TensorCoreFP8E5M2Layout(_TensorCoreFP8LayoutBase):
|
||||
# Backward compatibility alias - default to E4M3
|
||||
TensorCoreFP8Layout = TensorCoreFP8E4M3Layout
|
||||
TensorWiseINT8Layout = _CKTensorWiseINT8Layout
|
||||
TensorCoreConvRotW4A4Layout = _CKTensorCoreConvRotW4A4Layout
|
||||
|
||||
|
||||
# ==============================================================================
|
||||
@@ -190,6 +195,7 @@ register_layout_class("TensorCoreFP8E4M3Layout", TensorCoreFP8E4M3Layout)
|
||||
register_layout_class("TensorCoreFP8E5M2Layout", TensorCoreFP8E5M2Layout)
|
||||
register_layout_class("TensorCoreNVFP4Layout", TensorCoreNVFP4Layout)
|
||||
register_layout_class("TensorWiseINT8Layout", _CKTensorWiseINT8Layout)
|
||||
register_layout_class("TensorCoreConvRotW4A4Layout", _CKTensorCoreConvRotW4A4Layout)
|
||||
if _CK_MXFP8_AVAILABLE:
|
||||
register_layout_class("TensorCoreMXFP8Layout", TensorCoreMXFP8Layout)
|
||||
|
||||
@@ -227,6 +233,13 @@ QUANT_ALGOS["int8_tensorwise"] = {
|
||||
"quantize_input": False,
|
||||
}
|
||||
|
||||
QUANT_ALGOS["convrot_w4a4"] = {
|
||||
"storage_t": torch.int8,
|
||||
"parameters": {"weight_scale"},
|
||||
"comfy_tensor_layout": "TensorCoreConvRotW4A4Layout",
|
||||
"quantize_input": False,
|
||||
}
|
||||
|
||||
|
||||
# ==============================================================================
|
||||
# Re-exports for backward compatibility
|
||||
@@ -239,6 +252,7 @@ __all__ = [
|
||||
"TensorCoreFP8E4M3Layout",
|
||||
"TensorCoreFP8E5M2Layout",
|
||||
"TensorCoreNVFP4Layout",
|
||||
"TensorCoreConvRotW4A4Layout",
|
||||
"TensorWiseINT8Layout",
|
||||
"QUANT_ALGOS",
|
||||
"register_layout_op",
|
||||
|
||||
Reference in New Issue
Block a user