From 72a4d2c2bdab668701b5a2c5475f16ddda7aea5e Mon Sep 17 00:00:00 2001 From: Disty0 Date: Tue, 23 Jun 2026 14:17:34 +0300 Subject: [PATCH] Add Hadamard support to SDNQ Atten --- modules/attention.py | 6 ++- modules/sdnq/kernels/triton_atten.py | 65 ++++++++++++++++++++++++---- modules/sdnq/quant_utils.py | 4 ++ modules/ui_definitions.py | 2 + 4 files changed, 66 insertions(+), 11 deletions(-) diff --git a/modules/attention.py b/modules/attention.py index d91159f79..90cd5fd39 100644 --- a/modules/attention.py +++ b/modules/attention.py @@ -30,8 +30,10 @@ def set_sdnq_attention(): attn_mask=attn_mask, scale=scale, enable_gqa=enable_gqa, quant_group_size=shared.opts.sdnq_attention_quant_group_size, quant_group_size_kv=shared.opts.sdnq_attention_quant_group_size_kv, - matmul_dtype = "int8" if shared.opts.sdnq_attention_matmul_type == "auto" else shared.opts.sdnq_attention_matmul_type, - pv_matmul_dtype = None if shared.opts.sdnq_attention_pv_matmul_type == "auto" else shared.opts.sdnq_attention_pv_matmul_type, + matmul_dtype="int8" if shared.opts.sdnq_attention_matmul_type == "auto" else shared.opts.sdnq_attention_matmul_type, + pv_matmul_dtype=None if shared.opts.sdnq_attention_pv_matmul_type == "auto" else shared.opts.sdnq_attention_pv_matmul_type, + use_hadamard=shared.opts.sdnq_attention_use_hadamard, + smooth_k=shared.opts.sdnq_attention_smooth_k, ) else: if enable_gqa: diff --git a/modules/sdnq/kernels/triton_atten.py b/modules/sdnq/kernels/triton_atten.py index 3c51abd29..a15147a59 100644 --- a/modules/sdnq/kernels/triton_atten.py +++ b/modules/sdnq/kernels/triton_atten.py @@ -1,7 +1,3 @@ -""" -Heavily modified from SageAttention Triton kernel to run a lot faster. -This one also supports Intel and AMD on top of Nvidia. -""" import os import math @@ -10,7 +6,7 @@ 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# pylint: disable=relative-beyond-top-level +from ..quant_utils import quantize_int_mm, quantize_fp_mm, get_hadamard, apply_hadamard # pylint: disable=relative-beyond-top-level matmul_configs = [ @@ -41,7 +37,16 @@ def quantize_tensor(tensor: torch.FloatTensor, group_size: int = 128, matmul_dty return tensor, scale -def quantize_attn(q, k, v, group_size: int = 128, group_size_kv: int = 32, scale: float | None = None, smooth_k: bool = False, matmul_dtype: str = "int8", pv_matmul_dtype: str | None = "float16"): +def quantize_attn( + q, k, v, + hadamard: torch.FloatTensor | None = None, + group_size: int = 128, + group_size_kv: int = 32, + scale: float | None = None, + smooth_k: bool = False, + matmul_dtype: str = "int8", + pv_matmul_dtype: str | None = "float16", +): if pv_matmul_dtype is None: pv_matmul_dtype = matmul_dtype if scale is None: @@ -52,6 +57,10 @@ def quantize_attn(q, k, v, group_size: int = 128, group_size_kv: int = 32, scale k = k.sub_(k.mean(dim=2, keepdim=True)) else: k = k.sub(k.mean(dim=2, keepdim=True)) + if hadamard is not None: + q, use_hadamard, hadamard_group_size = apply_hadamard(q, group_size=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] q_q, q_scale = quantize_tensor(q, group_size=group_size, matmul_dtype=matmul_dtype) k_q, k_scale = quantize_tensor(k, group_size=group_size_kv, matmul_dtype=matmul_dtype) v_q, v_scale = quantize_tensor(v, group_size=group_size_kv, matmul_dtype=pv_matmul_dtype) @@ -157,11 +166,12 @@ def triton_attn_kernel( O_desc.store([start_m * BLOCK_M, 0], acc) -def sdnq_triton_atten( +def sdnq_triton_atten_forward( query: torch.FloatTensor, key: torch.FloatTensor, value: torch.FloatTensor, - attn_mask: torch.Tensor = None, + hadamard: torch.FloatTensor | None = None, + attn_mask: torch.Tensor | None = None, dropout_p: float = 0.0, # pylint: disable=unused-argument is_causal: bool = False, scale: float | None = None, @@ -195,6 +205,7 @@ def sdnq_triton_atten( query, query_scale, key, key_scale, value, value_scale = quantize_attn( query, key, value, + hadamard=hadamard, group_size=quant_group_size, group_size_kv=quant_group_size_kv, scale=scale, smooth_k=smooth_k, @@ -216,4 +227,40 @@ def sdnq_triton_atten( return out[..., :qhd] -sdnq_triton_atten = compile_func(sdnq_triton_atten) +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, + smooth_k: bool = False, + use_hadamard: bool = False, + quant_group_size: int = 128, + quant_group_size_kv: int = 32, + matmul_dtype: str = "int8", + pv_matmul_dtype: str | None = "float16", +) -> torch.FloatTensor: + if use_hadamard: + hadamard = get_hadamard(min(quant_group_size, query.shape[-1], key.shape[-1]), dtype=query.dtype, device=query.device) + else: + hadamard = None + return sdnq_triton_atten_forward( + query, key, 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, + quant_group_size=quant_group_size, + quant_group_size_kv=quant_group_size_kv, + matmul_dtype=matmul_dtype, + pv_matmul_dtype=pv_matmul_dtype, + ) + + +sdnq_triton_atten_forward = compile_func(sdnq_triton_atten_forward) diff --git a/modules/sdnq/quant_utils.py b/modules/sdnq/quant_utils.py index 94ba60c45..25ef88a96 100644 --- a/modules/sdnq/quant_utils.py +++ b/modules/sdnq/quant_utils.py @@ -149,6 +149,8 @@ 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 if channel_size % group_size != 0: hadamard_pow2 = int(math.log2(group_size)) while channel_size % group_size != 0: @@ -157,6 +159,8 @@ def apply_hadamard(weight: torch.Tensor, group_size: int = 256, hadamard: torch. if group_size < 4: use_hadamard = False 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 diff --git a/modules/ui_definitions.py b/modules/ui_definitions.py index f425e9123..e02c495cb 100644 --- a/modules/ui_definitions.py +++ b/modules/ui_definitions.py @@ -254,6 +254,8 @@ def create_settings(cmd_opts): "xformers_options": OptionInfo(['Flash attention'], "xFormers options", gr.CheckboxGroup, {"choices": ['Flash attention'] }), "dynamic_attention_slice_rate": OptionInfo(0.5, "Dynamic Attention slicing rate", gr.Slider, {"minimum": 0.01, "maximum": max(gpu_memory,4), "step": 0.01}), "dynamic_attention_trigger_rate": OptionInfo(1, "Dynamic Attention trigger rate", gr.Slider, {"minimum": 0.01, "maximum": max(gpu_memory,4)*2, "step": 0.01}), + "sdnq_attention_use_hadamard": OptionInfo(False, "SDNQ Attention use Hadamard", gr.Checkbox), + "sdnq_attention_smooth_k": OptionInfo(False, "SDNQ Attention use Smooth K", gr.Checkbox), "sdnq_attention_matmul_type": OptionInfo("auto", "SDNQ Attention MatMul type", gr.Radio, {"choices": sdnq_matmul_modes}), "sdnq_attention_pv_matmul_type": OptionInfo("auto", "SDNQ Attention PV MatMul type", gr.Radio, {"choices": sdnq_matmul_modes}), "sdnq_attention_quant_group_size": OptionInfo(128, "SDNQ Attention Quantization Group Size", gr.Number, {"minimum": 32, "maximum": 1024, "step": 1}),