From 16a284eecf23a148d9dcf40c79bd2d7dcc541437 Mon Sep 17 00:00:00 2001 From: Disty0 Date: Sat, 4 Jul 2026 21:53:40 +0300 Subject: [PATCH] lint --- modules/sdnq/kernels/openvino_mm.py | 2 +- modules/sdnq/kernels/triton_atten.py | 33 ++++++++++++++-------- modules/sdnq/kernels/triton_mm.py | 24 ++++++++-------- modules/sdnq/layers/linear/linear_int8.py | 2 +- modules/sdnq/layers/linear/linear_uint8.py | 2 +- modules/sdnq/loader.py | 8 +++--- modules/sdnq/quant_utils.py | 3 +- modules/sdnq/quantizer.py | 33 +++++++++++----------- 8 files changed, 58 insertions(+), 49 deletions(-) diff --git a/modules/sdnq/kernels/openvino_mm.py b/modules/sdnq/kernels/openvino_mm.py index 51b5f392a..a6cf9255e 100644 --- a/modules/sdnq/kernels/openvino_mm.py +++ b/modules/sdnq/kernels/openvino_mm.py @@ -15,7 +15,7 @@ for ov_device in core.get_available_devices(): core.set_property(ov_device, {ov_hints.execution_mode: ov_hints.ExecutionMode.ACCURACY}) -def ov_mm(A: torch.CharTensor, B: torch.CharTensor, infer_request: ov.InferRequest, out_name: str) -> torch.FloatTensor: +def ov_mm(A: torch.Tensor, B: torch.Tensor, infer_request: ov.InferRequest, out_name: str) -> torch.FloatTensor: C = torch.empty((A.shape[0], B.shape[-1]), device="cpu", dtype=torch.float32) infer_request.set_tensor("A", ov.Tensor(A.detach().contiguous().to("cpu").numpy(), shared_memory=True)) infer_request.set_tensor("B", ov.Tensor(B.detach().contiguous().to("cpu").numpy(), shared_memory=True)) diff --git a/modules/sdnq/kernels/triton_atten.py b/modules/sdnq/kernels/triton_atten.py index 3a15dc1af..93afee23f 100644 --- a/modules/sdnq/kernels/triton_atten.py +++ b/modules/sdnq/kernels/triton_atten.py @@ -34,19 +34,28 @@ matmul_configs = [ ) @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, + 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, KN_AT: tl.constexpr, VN_AT: tl.constexpr, - qk_is_quantized: tl.constexpr, pv_is_quantized: tl.constexpr, - q_dtype: tl.constexpr, v_dtype: tl.constexpr, - out_dtype: tl.constexpr, mask_dtype: tl.constexpr, - BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, -) -> None: # pylint: disable=unused-argument + 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) @@ -76,10 +85,10 @@ def sdnq_attn_kernel( 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) - tl.assume(is_causal == 0 or is_causal == 1) - tl.assume(qk_is_quantized == 0 or qk_is_quantized == 1) - tl.assume(pv_is_quantized == 0 or pv_is_quantized == 1) + 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: tl.constexpr = KN % BLOCK_SIZE_N != 0 start_m_block = start_m * BLOCK_SIZE_M diff --git a/modules/sdnq/kernels/triton_mm.py b/modules/sdnq/kernels/triton_mm.py index bbdd7a692..839d54a00 100644 --- a/modules/sdnq/kernels/triton_mm.py +++ b/modules/sdnq/kernels/triton_mm.py @@ -34,19 +34,19 @@ def triton_mm_kernel( M: tl.constexpr, N: tl.constexpr, K: tl.constexpr, - M_AT: tl.constexpr, - N_AT: tl.constexpr, - K_AT: 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 stride_am: tl.constexpr, stride_ak: tl.constexpr, stride_bk: tl.constexpr, stride_bn: tl.constexpr, stride_cm: tl.constexpr, stride_cn: tl.constexpr, - a_dtype: tl.constexpr, - out_dtype: tl.constexpr, + 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: # pylint: disable=unused-argument +) -> None: pid = tl.program_id(axis=0) num_pid_m: tl.constexpr = tl.cdiv(M, BLOCK_SIZE_M) num_pid_n: tl.constexpr = tl.cdiv(N, BLOCK_SIZE_N) @@ -104,16 +104,16 @@ def triton_mm_td_kernel( M: tl.constexpr, N: tl.constexpr, K: tl.constexpr, - M_AT: tl.constexpr, - N_AT: tl.constexpr, - K_AT: tl.constexpr, - a_dtype: tl.constexpr, - out_dtype: 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: # pylint: disable=unused-argument +) -> None: pid = tl.program_id(axis=0) num_pid_m: tl.constexpr = tl.cdiv(M, BLOCK_SIZE_M) num_pid_n: tl.constexpr = tl.cdiv(N, BLOCK_SIZE_N) diff --git a/modules/sdnq/layers/linear/linear_int8.py b/modules/sdnq/layers/linear/linear_int8.py index ef9d33430..9c72e22e6 100644 --- a/modules/sdnq/layers/linear/linear_int8.py +++ b/modules/sdnq/layers/linear/linear_int8.py @@ -10,7 +10,7 @@ from ...packed_int import unpack_int from .forward import check_mats -def quantize_int_mm_input(input: torch.FloatTensor, dtype: torch.dtype | None = None) -> tuple[torch.CharTensor, torch.FloatTensor]: +def quantize_int_mm_input(input: torch.FloatTensor, dtype: torch.dtype | None = None) -> tuple[torch.Tensor, torch.FloatTensor]: input = input.flatten(0,-2) if dtype is not None: input = input.to(dtype=dtype) diff --git a/modules/sdnq/layers/linear/linear_uint8.py b/modules/sdnq/layers/linear/linear_uint8.py index dd5f0a9e9..64693abe6 100644 --- a/modules/sdnq/layers/linear/linear_uint8.py +++ b/modules/sdnq/layers/linear/linear_uint8.py @@ -10,7 +10,7 @@ from ...packed_int import unpack_int from .forward import check_mats -def quantize_uint_mm_input(input: torch.FloatTensor, dtype: torch.dtype | None = None) -> tuple[torch.CharTensor, torch.FloatTensor]: +def quantize_uint_mm_input(input: torch.FloatTensor, dtype: torch.dtype | None = None) -> tuple[torch.Tensor, torch.FloatTensor, torch.FloatTensor]: input = input.flatten(0,-2) if dtype is not None: input = input.to(dtype=dtype) diff --git a/modules/sdnq/loader.py b/modules/sdnq/loader.py index 9d881c8d8..0fa376676 100644 --- a/modules/sdnq/loader.py +++ b/modules/sdnq/loader.py @@ -161,14 +161,14 @@ def load_sdnq_model( 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.keys() and key not in state_dict.keys(): + 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.keys(): + 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.keys(): - if "model.embed_tokens.weight" in state_dict.keys(): + elif model.__class__.__name__ in {"Qwen3ForCausalLM"} and "lm_head.weight" not in state_dict: + if "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) diff --git a/modules/sdnq/quant_utils.py b/modules/sdnq/quant_utils.py index 814d83497..5ea676fac 100644 --- a/modules/sdnq/quant_utils.py +++ b/modules/sdnq/quant_utils.py @@ -151,8 +151,7 @@ def apply_hadamard(weight: torch.Tensor, group_size: int = 256, hadamard: torch. channel_size = weight.shape[1] else: channel_size = weight.shape[-1] - if channel_size < group_size: - group_size = channel_size + group_size = min(group_size, channel_size) if channel_size % group_size != 0: hadamard_pow2 = int(math.log2(group_size)) while channel_size % group_size != 0: diff --git a/modules/sdnq/quantizer.py b/modules/sdnq/quantizer.py index eba4b56c9..952fb23c1 100644 --- a/modules/sdnq/quantizer.py +++ b/modules/sdnq/quantizer.py @@ -555,13 +555,13 @@ class SDNQQuantize: def __init__(self, hf_quantizer: "SDNQQuantizer"): self.hf_quantizer = hf_quantizer - def convert( + 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, # pylint: disable=unused-argument - **kwargs, # pylint: disable=unused-argument + missing_keys: list[str] | None = None, + **kwargs, ) -> dict[str, torch.Tensor]: _module_name, value = tuple(input_dict.items())[0] value = value[0] @@ -589,12 +589,12 @@ class SDNQQuantizer(DiffusersQuantizer, HfQuantizer): 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( + def check_if_quantized_param( # pylint: disable=unused-argument self, model: torch.nn.Module, - param_value: torch.Tensor, # pylint: disable=unused-argument + param_value: torch.Tensor, param_name: str, - *args, **kwargs, # pylint: disable=unused-argument + *args, **kwargs, ) -> bool: if self.pre_quantized: layer, _tensor_name = get_module_from_name(model, param_name) @@ -616,13 +616,14 @@ class SDNQQuantizer(DiffusersQuantizer, HfQuantizer): return False @devices.inference_context() - def create_quantized_param( # pylint: disable=arguments-differ + def create_quantized_param( # pylint: disable=unused-argument self, model: torch.nn.Module, param_value: torch.FloatTensor, param_name: str, target_device: torch.device, - *args, **kwargs, # pylint: disable=unused-argument + *args, + **kwargs, ) -> None: layer, tensor_name = get_module_from_name(model, param_name) @@ -631,7 +632,7 @@ class SDNQQuantizer(DiffusersQuantizer, HfQuantizer): if tensor_name == "weight": return_dtype = param_value.dtype elif self.quantization_config.dequantize_fp32 and tensor_name in sdnq_keys: - if param_value.dtype != torch.float64 and self.torch_dtype != torch.float64: + if torch.float64 not in {param_value.dtype, self.torch_dtype}: return_dtype = torch.float32 else: return_dtype = torch.float64 @@ -670,12 +671,12 @@ class SDNQQuantizer(DiffusersQuantizer, HfQuantizer): 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=arguments-differ + def _process_model_before_weight_loading( # pylint: disable=unused-argument self, model: torch.nn.Module, - device_map, # pylint: disable=unused-argument + device_map, keep_in_fp32_modules: list[str] | None = None, - **kwargs, # pylint: disable=unused-argument + **kwargs, ) -> None: if self.pre_quantized: self.quantization_config.quantization_device = None @@ -885,7 +886,7 @@ class SDNQConfig(QuantizationConfigMixin): Note: Safetensors serialization is not supported with SDNQ training. """ - def __init__( # pylint: disable=super-init-not-called + def __init__( # pylint: disable=super-init-not-called,unused-argument self, weights_dtype: str = "int8", quantized_matmul_dtype: str | None = None, @@ -914,7 +915,7 @@ class SDNQConfig(QuantizationConfigMixin): modules_dtype_dict: dict[str, list[str]] | None = None, modules_quant_config: dict[str, dict] | None = None, is_training: bool = False, - **kwargs, # pylint: disable=unused-argument + **kwargs, ): self.weights_dtype = weights_dtype self.quantized_matmul_dtype = quantized_matmul_dtype @@ -1020,14 +1021,14 @@ class SDNQConfig(QuantizationConfigMixin): 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} )" -import diffusers.quantizers.auto # noqa: E402,RUF100 # pylint: disable=wrong-import-order +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 -import transformers.quantizers.auto # noqa: E402,RUF100 # pylint: disable=wrong-import-order +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