lint

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))

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

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)

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)

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)

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)

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:

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