|
| 1 | +# Copyright (c) 2025 Red Hat AI, vLLM Project and Intel Corporation |
| 2 | +# |
| 3 | +# Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | +# you may not use this file except in compliance with the License. |
| 5 | +# You may obtain a copy of the License at |
| 6 | +# |
| 7 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 8 | +# |
| 9 | +# Unless required by applicable law or agreed to in writing, software |
| 10 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | +# See the License for the specific language governing permissions and |
| 13 | +# limitations under the License. |
| 14 | + |
| 15 | +# NOTICE: The design adapted from: |
| 16 | +# https://github.com/vllm-project/compressed-tensors/pull/491 |
| 17 | + |
| 18 | + |
| 19 | +import contextlib |
| 20 | +import inspect |
| 21 | +from functools import partial |
| 22 | +from typing import Callable, Optional |
| 23 | +from weakref import ref |
| 24 | + |
| 25 | +import torch |
| 26 | +from torch import Tensor |
| 27 | +from torch.nn import Module |
| 28 | +from torch.utils.hooks import RemovableHandle |
| 29 | +from transformers import AttentionInterface, PretrainedConfig, PreTrainedModel |
| 30 | +from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS |
| 31 | + |
| 32 | +from auto_round.experimental.kv_cache import kvcache_quant_context |
| 33 | +from auto_round.experimental.utils import ( |
| 34 | + is_attention_module, |
| 35 | + per_tensor_fp8_qdq, |
| 36 | + update_parameter_data, |
| 37 | +) |
| 38 | +from auto_round.utils import logger |
| 39 | + |
| 40 | +__all__ = [ |
| 41 | + "QuantizedAttentionImpl", |
| 42 | + "initialize_hooked_attention", |
| 43 | + "IMPL_ATTR", |
| 44 | + "attention_quant_ctx", |
| 45 | +] |
| 46 | + |
| 47 | + |
| 48 | +IMPL_ATTR = "impl" |
| 49 | +HOOKED_ATTENTION_NAME = "ct_hooked_attention" |
| 50 | +QUERY_SCALE_NAME = "q_scale" |
| 51 | +QUERY_MAX_NAME = "q_max" |
| 52 | + |
| 53 | + |
| 54 | +class QuantizedAttentionImpl(torch.nn.Module): |
| 55 | + """ |
| 56 | + QuantizedAttentionImpl module which wraps the functionality of the original |
| 57 | + attention implementation. Unlike the original attention function, this |
| 58 | + implementation is a `torch.nn.Module` which can be hooked to trigger |
| 59 | + transforms and calibration hooks. |
| 60 | +
|
| 61 | + This module works by being registered as a submodule to attention modules via |
| 62 | + `initialize_hooked_attention`, registering a new attention implementation function |
| 63 | + which calls this module, then setting the model attention implementation to the new |
| 64 | + function. After triggering hooks and quantization, this module calls the original |
| 65 | + attention implementation function. |
| 66 | +
|
| 67 | + :param attn_module: parent attention module |
| 68 | + """ |
| 69 | + |
| 70 | + _original_impl = "sdpa" |
| 71 | + |
| 72 | + def __init__(self, config: PretrainedConfig, attn_module: Module): |
| 73 | + super().__init__() |
| 74 | + self.config = config |
| 75 | + self.attn_module = ref(attn_module) # avoid circular references |
| 76 | + # register query max |
| 77 | + device = next(attn_module.parameters()).device |
| 78 | + initial_max = torch.tensor([float("-inf")], device=device) |
| 79 | + update_parameter_data(attn_module, initial_max, QUERY_MAX_NAME) |
| 80 | + initial_scale = torch.tensor([0.0], device=device) |
| 81 | + update_parameter_data(attn_module, initial_scale, QUERY_SCALE_NAME) |
| 82 | + |
| 83 | + def forward( |
| 84 | + self, |
| 85 | + module: Module, |
| 86 | + query: Tensor, |
| 87 | + key: Tensor, |
| 88 | + value: Tensor, |
| 89 | + *args, |
| 90 | + **kwargs, |
| 91 | + ): |
| 92 | + cur_query_max = query.abs().max() |
| 93 | + query_max = torch.max( |
| 94 | + getattr(module, QUERY_MAX_NAME).data, |
| 95 | + cur_query_max.detach().to(getattr(module, QUERY_MAX_NAME).data.device), |
| 96 | + ) |
| 97 | + update_parameter_data(module, query_max, QUERY_MAX_NAME) |
| 98 | + query, query_scale = per_tensor_fp8_qdq(query, tensor_max=query_max) |
| 99 | + update_parameter_data(module, query_scale.squeeze(0), QUERY_SCALE_NAME) |
| 100 | + # original attention |
| 101 | + return ALL_ATTENTION_FUNCTIONS[self._original_impl]( |
| 102 | + module, |
| 103 | + query, |
| 104 | + key, |
| 105 | + value, |
| 106 | + *args, |
| 107 | + **kwargs, |
| 108 | + ) |
| 109 | + |
| 110 | + |
| 111 | +# ----- initialize ----- # |
| 112 | + |
| 113 | + |
| 114 | +def _ct_hooked_attention(module: Module, *args, **kwargs): |
| 115 | + if hasattr(module, IMPL_ATTR): |
| 116 | + return module.impl(module, *args, **kwargs) |
| 117 | + else: |
| 118 | + return ALL_ATTENTION_FUNCTIONS[_original_impl](module, *args, **kwargs) # pylint: disable=E0601 |
| 119 | + |
| 120 | + |
| 121 | +def initialize_hooked_attention(module: Module, config): |
| 122 | + """ |
| 123 | + Initialize `QuantizedAttentionImpl` and `QuantizedKVCache` instances |
| 124 | + attached to attention |
| 125 | +
|
| 126 | + :param model: parent model of attention module |
| 127 | + :param module: attention module to initialize with |
| 128 | + """ |
| 129 | + if not hasattr(module, IMPL_ATTR): |
| 130 | + module.register_module(IMPL_ATTR, QuantizedAttentionImpl(config, module)) |
| 131 | + if config._attn_implementation != HOOKED_ATTENTION_NAME: |
| 132 | + # assumes only one model at a time |
| 133 | + global _original_impl |
| 134 | + _original_impl = config._attn_implementation |
| 135 | + # Add new implementation to AttentionInterface(mapping) |
| 136 | + AttentionInterface.register(HOOKED_ATTENTION_NAME, _ct_hooked_attention) |
| 137 | + config._attn_implementation = HOOKED_ATTENTION_NAME |
| 138 | + |
| 139 | + # initialize_hooked_kv_cache(model, module) |
| 140 | + |
| 141 | + |
| 142 | +def prep_attention_module_for_calibration(module: torch.nn.Module, config): |
| 143 | + if is_attention_module(module): |
| 144 | + logger.trace(f"Preparing attention module {module.__class__.__name__} for calibration") |
| 145 | + initialize_hooked_attention(module, config) |
| 146 | + |
| 147 | + |
| 148 | +# # ----- hooks ----- # |
| 149 | + |
| 150 | + |
| 151 | +# def register_query_hook(module: Module, hook: Callable[[Module, Tensor], Optional[Tensor]]) -> RemovableHandle: |
| 152 | +# """ |
| 153 | +# Register a hook which takes post-rope query states as an argument and |
| 154 | +# returns the modified query states or `None` |
| 155 | + |
| 156 | +# :param module: attention module to add hook to |
| 157 | +# :param hook: query hook function |
| 158 | +# """ |
| 159 | +# impl = getattr(module, IMPL_ATTR) |
| 160 | + |
| 161 | +# def _hook(impl: QuantizedAttentionImpl, args, kwargs): |
| 162 | +# bound = inspect.signature(impl.forward).bind(*args, **kwargs) |
| 163 | +# value = hook(module, bound.arguments["query"]) |
| 164 | +# if value is not None: |
| 165 | +# bound.arguments["query"] = value |
| 166 | + |
| 167 | +# return bound.args, bound.kwargs |
| 168 | + |
| 169 | +# return impl.register_forward_pre_hook(_hook, with_kwargs=True) |
| 170 | + |
| 171 | + |
| 172 | +def clean_up_hooked_attention(module, model): |
| 173 | + if is_attention_module(module): |
| 174 | + # Cleanup phase: Restore the original attention implementation |
| 175 | + if hasattr(model.config, "_attn_implementation") and hasattr(model, "_original_impl"): |
| 176 | + model.config._attn_implementation = model._original_impl |
| 177 | + del model._original_impl |
| 178 | + |
| 179 | + |
| 180 | +@contextlib.contextmanager |
| 181 | +def attention_quant_ctx(model: PreTrainedModel, static_attention_dtype=torch.float8_e4m3fn): |
| 182 | + try: |
| 183 | + # Setup phase: Initialize hooked attention |
| 184 | + prepare_fn = partial(prep_attention_module_for_calibration, config=model.config) |
| 185 | + model.apply(prepare_fn) |
| 186 | + with kvcache_quant_context(model, static_kv_dtype=static_attention_dtype): |
| 187 | + yield model |
| 188 | + finally: |
| 189 | + clean_fn = partial(clean_up_hooked_attention, model=model) |
| 190 | + model.apply(clean_fn) |
0 commit comments