| """
|
| SageAttention patching functions for QwenVL models.
|
| This module contains the safe SageAttention implementation with proper error handling.
|
| """
|
|
|
| import torch
|
|
|
|
|
| def set_sage_attention(model):
|
| """
|
| Apply SageAttention patching to the model.
|
| Patches Qwen2Attention and Qwen3VLTextAttention modules to use SageAttention kernels.
|
| """
|
| from AILab_QwenVL import sage_attn_available, get_sage_attention_config
|
|
|
| if not sage_attn_available():
|
| raise ImportError("SageAttention library is not installed or GPU doesn't support it.")
|
|
|
| SAGE_ATTN_FUNC, QK_QUANT_GRAN, PV_ACCUM_DTYPE = get_sage_attention_config()
|
| if SAGE_ATTN_FUNC is None:
|
| raise RuntimeError("No compatible SageAttention kernel found for this GPU.")
|
|
|
|
|
| attention_classes = []
|
|
|
|
|
| try:
|
| from transformers.models.qwen2.modeling_qwen2 import Qwen2Attention, apply_rotary_pos_emb as qwen2_apply_rotary
|
| attention_classes.append((Qwen2Attention, qwen2_apply_rotary))
|
| except ImportError:
|
| pass
|
|
|
|
|
| try:
|
| from transformers.models.qwen3.modeling_qwen3 import Qwen3Attention, apply_rotary_pos_emb as qwen3_apply_rotary
|
| attention_classes.append((Qwen3Attention, qwen3_apply_rotary))
|
| except ImportError:
|
| pass
|
|
|
|
|
| try:
|
| from transformers.models.qwen3_vl.modeling_qwen3_vl import Qwen3VLTextAttention, apply_rotary_pos_emb as qwen3vl_apply_rotary
|
| attention_classes.append((Qwen3VLTextAttention, qwen3vl_apply_rotary))
|
| except ImportError:
|
| pass
|
|
|
| if not attention_classes:
|
| print("[QwenVL] Could not import any attention classes for SageAttention patching")
|
| return
|
|
|
| def make_sage_forward(AttentionClass, apply_rotary_pos_emb_func):
|
| def sage_attention_forward(
|
| self,
|
| hidden_states: torch.Tensor,
|
| position_embeddings: tuple = None,
|
| attention_mask: torch.Tensor = None,
|
| past_key_values=None,
|
| cache_position: torch.LongTensor = None,
|
| position_ids: torch.LongTensor = None,
|
| **kwargs,
|
| ):
|
| original_dtype = hidden_states.dtype
|
|
|
|
|
| is_4bit = hasattr(self.q_proj, 'quant_state')
|
| if is_4bit:
|
| target_dtype = torch.bfloat16
|
| else:
|
| target_dtype = self.q_proj.weight.dtype
|
|
|
| if hidden_states.dtype != target_dtype:
|
| hidden_states = hidden_states.to(target_dtype)
|
|
|
| input_shape = hidden_states.shape[:-1]
|
| hidden_shape = (*input_shape, -1, self.head_dim)
|
| bsz, q_len = input_shape[0], input_shape[1] if len(input_shape) > 1 else hidden_states.size(1)
|
|
|
|
|
| query_states = self.q_proj(hidden_states)
|
| key_states = self.k_proj(hidden_states)
|
| value_states = self.v_proj(hidden_states)
|
|
|
|
|
| if hasattr(self, 'q_norm'):
|
| query_states = self.q_norm(query_states.view(hidden_shape)).transpose(1, 2)
|
| else:
|
| query_states = query_states.view(hidden_shape).transpose(1, 2)
|
|
|
| if hasattr(self, 'k_norm'):
|
| key_states = self.k_norm(key_states.view(hidden_shape)).transpose(1, 2)
|
| else:
|
| key_states = key_states.view(hidden_shape).transpose(1, 2)
|
|
|
| value_states = value_states.view(hidden_shape).transpose(1, 2)
|
|
|
|
|
| if position_embeddings is not None:
|
| cos, sin = position_embeddings
|
| query_states, key_states = apply_rotary_pos_emb_func(query_states, key_states, cos, sin)
|
|
|
| if past_key_values is not None:
|
| cache_kwargs = {"sin": sin if position_embeddings else None, "cos": cos if position_embeddings else None, "cache_position": cache_position}
|
| key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)
|
|
|
| is_causal = attention_mask is None and q_len > 1
|
|
|
| attn_output = SAGE_ATTN_FUNC(
|
| query_states.to(target_dtype),
|
| key_states.to(target_dtype),
|
| value_states.to(target_dtype),
|
| tensor_layout="HND",
|
| is_causal=is_causal,
|
| qk_quant_gran=QK_QUANT_GRAN,
|
| pv_accum_dtype=PV_ACCUM_DTYPE,
|
| )
|
|
|
| if isinstance(attn_output, tuple):
|
| attn_output = attn_output[0]
|
|
|
| attn_output = attn_output.transpose(1, 2).contiguous()
|
| attn_output = attn_output.reshape(*input_shape, -1)
|
|
|
| attn_output = self.o_proj(attn_output)
|
|
|
| if attn_output.dtype != original_dtype:
|
| attn_output = attn_output.to(original_dtype)
|
|
|
| return attn_output, None
|
|
|
| return sage_attention_forward
|
|
|
|
|
| patched_count = 0
|
| for AttentionClass, apply_rotary_func in attention_classes:
|
| sage_forward = make_sage_forward(AttentionClass, apply_rotary_func)
|
| for module in model.modules():
|
| if isinstance(module, AttentionClass):
|
| module.forward = sage_forward.__get__(module, AttentionClass)
|
| patched_count += 1
|
|
|
| if patched_count > 0:
|
| print(f"[QwenVL] SageAttention: Patched {patched_count} attention layers")
|
| else:
|
| print("[QwenVL] SageAttention: No compatible attention layers found to patch")
|
|
|
