feat: add per-head fp8 kv attention calibration

auto_round/autoround.py

@@ -336,6 +336,8 @@ def _select_rtn_compressor_base_cls(quant_config: "RTNConfig", scheme, format, b
     "enable_deterministic_algorithms": "base",
     "static_kv_dtype": "base",
     "static_attention_dtype": "base",
+    "static_kv_granularity": "base",
+    "static_attention_granularity": "base",
     "processor": "mllm",
     "image_processor": "mllm",
     "template": "mllm",

auto_round/cli/main.py

@@ -66,6 +66,10 @@ def _build_entry_base_kwargs(args, *, low_cpu_mem_usage, enable_torch_compile, l
         "layer_config": layer_config,
         "model_dtype": args.model_dtype,
         "trust_remote_code": not args.disable_trust_remote_code,
+        "static_kv_dtype": args.static_kv_dtype,
+        "static_kv_granularity": args.static_kv_granularity,
+        "static_attention_dtype": args.static_attention_dtype,
+        "static_attention_granularity": args.static_attention_granularity,
     }
 
 

auto_round/cli/parser.py

@@ -184,13 +184,27 @@ def build_quantize_parser(*, prog: str = "auto_round quantize") -> argparse.Argu
         choices=["fp8", "float8_e4m3fn"],
         help="Static KV-cache quantization data type.",
     )
+    rt.add_argument(
+        "--static_kv_granularity",
+        default="tensor",
+        type=str,
+        choices=["tensor", "head"],
+        help="Static KV-cache FP8 calibration granularity.",
+    )
     rt.add_argument(
         "--static_attention_dtype",
         default=None,
         type=str,
         choices=["fp8", "float8_e4m3fn"],
         help="Static attention quantization data type.",
     )
+    rt.add_argument(
+        "--static_attention_granularity",
+        default="tensor",
+        type=str,
+        choices=["tensor", "head"],
+        help="Static attention FP8 calibration granularity.",
+    )
 
     # ---- Evaluation ----
     ev = parser.add_argument_group("Evaluation Arguments")

auto_round/compressors/base.py

@@ -110,6 +110,8 @@ class SerializedCompressorConfig:
     supported_types: Optional[list[str]] = SUPPORTED_LAYER_TYPES
     static_attention_dtype: Optional[str] = None
     static_kv_dtype: Optional[str] = None
+    static_attention_granularity: Optional[str] = "tensor"
+    static_kv_granularity: Optional[str] = "tensor"
     super_bits: Optional[int] = None
     super_group_size: Optional[int] = None
     to_quant_block_names: Optional[list[str]] = None
@@ -314,12 +316,18 @@ def __init__(
         if device is not None:
             logger.warning("`device` is deprecated, please use `device_map` instead")
 
+        from auto_round.experimental.utils import normalize_fp8_granularity
+
         self.static_attention_dtype = kwargs.pop("static_attention_dtype", None)
+        self.static_attention_granularity = normalize_fp8_granularity(
+            kwargs.pop("static_attention_granularity", "tensor")
+        )
         # Attention static dtype
         if self.static_attention_dtype is not None:
             logger.warning("The static attention dtype is experimental and currently has limited support.")
         # KV cache, this one does not affect tuning but will collect some infos during tuning
         self.static_kv_dtype = kwargs.pop("static_kv_dtype", None)
+        self.static_kv_granularity = normalize_fp8_granularity(kwargs.pop("static_kv_granularity", "tensor"))
         if self.static_kv_dtype is not None:
             logger.warning("The static kv is experimental and currently has limited support.")
 
@@ -432,6 +440,8 @@ def __init__(
             formats=self.formats,
             static_kv_dtype=self.static_kv_dtype,
             static_attention_dtype=self.static_attention_dtype,
+            static_kv_granularity=self.static_kv_granularity,
+            static_attention_granularity=self.static_attention_granularity,
         )
         self.shard_writer = None
         # Resumability state deferred from Orchestrator._quantize_data_driven() until
@@ -1832,13 +1842,21 @@ def quantize_and_save(
         if self.static_attention_dtype is not None:
             from auto_round.experimental.attention import attention_quant_ctx
 
-            with attention_quant_ctx(self.model_context.model, static_attention_dtype=self.static_attention_dtype):
+            with attention_quant_ctx(
+                self.model_context.model,
+                static_attention_dtype=self.static_attention_dtype,
+                static_attention_granularity=self.static_attention_granularity,
+            ):
                 self.quantize()
                 self.model_context.quantized = True
         elif self.static_kv_dtype is not None:
             from auto_round.experimental.kv_cache import kvcache_quant_context
 
-            with kvcache_quant_context(self.model_context.model, static_kv_dtype=self.static_kv_dtype):
+            with kvcache_quant_context(
+                self.model_context.model,
+                static_kv_dtype=self.static_kv_dtype,
+                static_kv_granularity=self.static_kv_granularity,
+            ):
                 self.quantize()
                 self.model_context.quantized = True
         else:

auto_round/context/compress.py

@@ -40,6 +40,8 @@ def __init__(
         output_dir: str = "./compressed_models",
         static_kv_dtype: Optional[torch.dtype] = None,  # TODO later this should be scheme wenhuach
         static_attention_dtype: Optional[torch.dtype] = None,
+        static_kv_granularity: str = "tensor",
+        static_attention_granularity: str = "tensor",
         **kwargs,
     ) -> None:
         super().__init__()
@@ -58,6 +60,8 @@ def __init__(
         self.is_immediate_saving = is_immediate_saving
         self.static_kv_dtype = static_kv_dtype
         self.static_attention_dtype = static_attention_dtype
+        self.static_kv_granularity = static_kv_granularity
+        self.static_attention_granularity = static_attention_granularity
 
     def clear_memory(self, tensor=None):
         """Clear GPU/CPU memory only when ``low_gpu_mem_usage`` is enabled."""

auto_round/experimental/attention.py

@@ -32,8 +32,9 @@
 from auto_round.experimental.kv_cache import kvcache_quant_context
 from auto_round.experimental.utils import (
     clean_model_parameters_and_buffers_,
+    fp8_qdq,
     is_attention_module,
-    per_tensor_fp8_qdq,
+    normalize_fp8_granularity,
     update_parameter_data,
 )
 from auto_round.utils import logger
@@ -75,9 +76,10 @@ class QuantizedAttentionImpl(torch.nn.Module):
 
     _original_impl = "sdpa"
 
-    def __init__(self, config: PretrainedConfig, attn_module: Module):
+    def __init__(self, config: PretrainedConfig, attn_module: Module, granularity: str = "tensor"):
         super().__init__()
         self.config = config
+        self.granularity = normalize_fp8_granularity(granularity)
         self.attn_module = ref(attn_module)  # avoid circular references
         # register query max
         device = next(attn_module.parameters()).device
@@ -95,14 +97,17 @@ def forward(
         *args,
         **kwargs,
     ):
-        cur_query_max = query.abs().max()
+        if self.granularity == "head":
+            cur_query_max = query.abs().amax(dim=(0, 2, 3))
+        else:
+            cur_query_max = query.abs().max()
         query_max = torch.max(
             getattr(module, QUERY_MAX_NAME).data,
             cur_query_max.detach().to(getattr(module, QUERY_MAX_NAME).data.device),
         )
         update_parameter_data(module, query_max, QUERY_MAX_NAME)
-        _, query_scale = per_tensor_fp8_qdq(query, tensor_max=query_max)
-        update_parameter_data(module, query_scale.squeeze(0).detach(), QUERY_SCALE_NAME)
+        _, query_scale = fp8_qdq(query, tensor_max=query_max, granularity=self.granularity)
+        update_parameter_data(module, query_scale.reshape(-1).detach(), QUERY_SCALE_NAME)
         # original attention
         return ALL_ATTENTION_FUNCTIONS[self._original_impl](
             module,
@@ -124,7 +129,7 @@ def _ct_hooked_attention(module: Module, *args, **kwargs):
         return ALL_ATTENTION_FUNCTIONS[_original_impl](module, *args, **kwargs)  # pylint: disable=E0601
 
 
-def init_hooked_attention(module: Module, config):
+def init_hooked_attention(module: Module, config, granularity: str = "tensor"):
     """
     Initialize `QuantizedAttentionImpl` and `QuantizedKVCache` instances
     attached to attention
@@ -133,7 +138,7 @@ def init_hooked_attention(module: Module, config):
     :param module: attention module to initialize with
     """
     if not hasattr(module, ATTN_IMPL_ATTR_NAME):
-        module.register_module(ATTN_IMPL_ATTR_NAME, QuantizedAttentionImpl(config, module))
+        module.register_module(ATTN_IMPL_ATTR_NAME, QuantizedAttentionImpl(config, module, granularity=granularity))
         if config._attn_implementation != HOOKED_ATTENTION_NAME:
             # assumes only one model at a time
             global _original_impl
@@ -145,10 +150,10 @@ def init_hooked_attention(module: Module, config):
     # initialize_hooked_kv_cache(model, module)
 
 
-def prep_attention_module_for_calibration(module: torch.nn.Module, config):
+def prep_attention_module_for_calibration(module: torch.nn.Module, config, granularity: str = "tensor"):
     if is_attention_module(module):
         logger.trace(f"Preparing attention module {module.__class__.__name__} for calibration")
-        init_hooked_attention(module, config)
+        init_hooked_attention(module, config, granularity=granularity)
 
 
 def clean_up_hooked_attention(module, model):
@@ -161,12 +166,25 @@ def clean_up_hooked_attention(module, model):
 
 
 @contextlib.contextmanager
-def attention_quant_ctx(model: PreTrainedModel, static_attention_dtype=torch.float8_e4m3fn):
+def attention_quant_ctx(
+    model: PreTrainedModel,
+    static_attention_dtype=torch.float8_e4m3fn,
+    static_attention_granularity: str = "tensor",
+):
     try:
         # Setup phase: Initialize hooked attention
-        prepare_fn = partial(prep_attention_module_for_calibration, config=model.config)
+        static_attention_granularity = normalize_fp8_granularity(static_attention_granularity)
+        prepare_fn = partial(
+            prep_attention_module_for_calibration,
+            config=model.config,
+            granularity=static_attention_granularity,
+        )
         model.apply(prepare_fn)
-        with kvcache_quant_context(model, static_kv_dtype=static_attention_dtype):
+        with kvcache_quant_context(
+            model,
+            static_kv_dtype=static_attention_dtype,
+            static_kv_granularity=static_attention_granularity,
+        ):
             yield model
     finally:
         clean_fn = partial(clean_up_hooked_attention, model=model)

auto_round/experimental/kv_cache.py

@@ -25,9 +25,10 @@
 from transformers.cache_utils import DynamicCache
 
 from auto_round.experimental.utils import (
+    fp8_qdq,
     is_attention_module,
+    normalize_fp8_granularity,
     normalize_static_kv_dtype,
-    per_tensor_fp8_qdq,
     update_parameter_data,
 )
 from auto_round.utils import logger
@@ -106,9 +107,10 @@ def __new__(cls, *args, **kwargs):
             cls._instance = super(QuantizedKVParameterCache, cls).__new__(cls)
         return cls._instance
 
-    def __init__(self, dtype: torch.dtype = torch.float8_e4m3fn):
+    def __init__(self, dtype: torch.dtype = torch.float8_e4m3fn, granularity: str = "tensor"):
 
         assert dtype == torch.float8_e4m3fn, "Only fp8_e4m3fn is supported for now."
+        self.granularity = normalize_fp8_granularity(granularity)
         if not self._initialized:
             super().__init__()
 
@@ -172,13 +174,13 @@ def _quant_dequant(self, tensor: torch.Tensor, kv_type: KVCacheScaleType, layer_
             assert kv_type == KVCacheScaleType.VALUE
             scales = self.v_scales
 
-        qdq_tensor, scale = per_tensor_fp8_qdq(tensor)
+        qdq_tensor, scale = fp8_qdq(tensor, granularity=self.granularity)
         # Detach scale to prevent holding computation graph references
-        _pad_and_append_at_idx_(scales, layer_idx, scale.squeeze(0).detach())
+        _pad_and_append_at_idx_(scales, layer_idx, scale.reshape(-1).detach())
         return qdq_tensor
 
 
-def initialize_quantized_kv_cache(module: torch.nn.Module, dtype=torch.float8_e4m3fn):
+def initialize_quantized_kv_cache(module: torch.nn.Module, dtype=torch.float8_e4m3fn, granularity: str = "tensor"):
     """
     Initialize a quantized kv_cache on a module (analogous to initializing an observer)
     """
@@ -189,7 +191,7 @@ def initialize_quantized_kv_cache(module: torch.nn.Module, dtype=torch.float8_e4
     if isinstance(existing_kv_cache, QuantizedKVParameterCache):
         return
 
-    quantized_kv_cache = QuantizedKVParameterCache(dtype=dtype)
+    quantized_kv_cache = QuantizedKVParameterCache(dtype=dtype, granularity=granularity)
     setattr(module, "kv_cache", quantized_kv_cache)
     logger.debug(f"Initialized quantized kv_cache for {module.__class__.__name__} {getattr(module, 'layer_idx', None)}")
     init_scale = torch.tensor([0.0], device=next(module.parameters()).device)
@@ -234,15 +236,20 @@ def prep_attention_module_for_calibration(module: torch.nn.Module):
 
 
 @contextlib.contextmanager
-def kvcache_quant_context(model: torch.nn.Module, static_kv_dtype=torch.float8_e4m3fn):
+def kvcache_quant_context(
+    model: torch.nn.Module, static_kv_dtype=torch.float8_e4m3fn, static_kv_granularity: str = "tensor"
+):
     """Context manager for FP8 KV cache quantization operations."""
     try:
         # Setup phase: Initialize KV cache for quantization
         static_kv_dtype = normalize_static_kv_dtype(static_kv_dtype)
+        static_kv_granularity = normalize_fp8_granularity(static_kv_granularity)
         if static_kv_dtype != torch.float8_e4m3fn:
             logger.warning(f"Ignoring static kv dtype {static_kv_dtype}, only fp8_e4m3fn is supported.")
         else:
-            initialize_fn = partial(initialize_quantized_kv_cache, dtype=static_kv_dtype)
+            initialize_fn = partial(
+                initialize_quantized_kv_cache, dtype=static_kv_dtype, granularity=static_kv_granularity
+            )
             model.apply(initialize_fn)
             model.apply(prep_attention_module_for_calibration)
 

auto_round/experimental/utils.py

@@ -23,6 +23,22 @@
 SUPPORTED_QUANTIZATION_SCHEMES = ["MXFP8", "MXFP4", "NVFP4"]
 
 
+FP8_GRANULARITY_TENSOR = "tensor"
+FP8_GRANULARITY_HEAD = "head"
+
+
+def normalize_fp8_granularity(granularity: str | None) -> str:
+    if granularity is None:
+        return FP8_GRANULARITY_TENSOR
+    granularity = granularity.lower()
+    if granularity in (FP8_GRANULARITY_TENSOR, FP8_GRANULARITY_HEAD):
+        return granularity
+    raise ValueError(
+        f"Invalid FP8 calibration granularity: {granularity}. "
+        f"Supported granularities are: {FP8_GRANULARITY_TENSOR}, {FP8_GRANULARITY_HEAD}."
+    )
+
+
 def per_tensor_fp8_qdq(
     tensor: torch.Tensor, tensor_max: None | torch.Tensor = None
 ) -> tuple[torch.Tensor, torch.Tensor]:
@@ -32,6 +48,41 @@ def per_tensor_fp8_qdq(
     return qdq_tensor, scale
 
 
+def per_head_fp8_qdq(tensor: torch.Tensor, tensor_max: None | torch.Tensor = None) -> tuple[torch.Tensor, torch.Tensor]:
+    """Quantize/dequantize an attention tensor with one FP8 scale per attention head."""
+    if tensor.dim() < 4:
+        raise ValueError(f"Per-head FP8 calibration expects [batch, heads, seq, head_dim], got {tuple(tensor.shape)}")
+    if tensor_max is not None and tensor_max.shape != (tensor.shape[1],):
+        raise ValueError(
+            f"Per-head FP8 calibration expects tensor_max shape ({tensor.shape[1]},), " f"got {tuple(tensor_max.shape)}"
+        )
+
+    info = torch.finfo(torch.float8_e4m3fn)
+    orig_dtype = tensor.dtype
+    if tensor.dtype == torch.float16:
+        tensor = tensor.to(torch.bfloat16)
+
+    if tensor_max is None:
+        max_tensor = tensor.abs().amax(dim=(0, 2, 3)).to(torch.float32)
+    else:
+        max_tensor = tensor_max.to(tensor.device, dtype=torch.float32)
+
+    scale = torch.clip(max_tensor / info.max, min=float(1.0 / (info.max * 512.0)))
+    qdq_scale = scale.view(1, -1, 1, 1)
+    fp8_res = torch.clip(tensor / qdq_scale, info.min, info.max)
+    fp8_res = fp8_res.to(torch.float8_e4m3fn)
+    return (fp8_res.to(tensor.dtype) * qdq_scale).to(orig_dtype), scale
+
+
+def fp8_qdq(
+    tensor: torch.Tensor, tensor_max: None | torch.Tensor = None, granularity: str | None = None
+) -> tuple[torch.Tensor, torch.Tensor]:
+    granularity = normalize_fp8_granularity(granularity)
+    if granularity == FP8_GRANULARITY_HEAD:
+        return per_head_fp8_qdq(tensor, tensor_max=tensor_max)
+    return per_tensor_fp8_qdq(tensor, tensor_max=tensor_max)
+
+
 @torch.compiler.disable
 def update_parameter_data(module: torch.nn.Module, new_val: torch.Tensor, name: str):
     """
@@ -41,7 +92,17 @@ def update_parameter_data(module: torch.nn.Module, new_val: torch.Tensor, name:
     if hasattr(module, name):
         param = getattr(module, name)
         if isinstance(param, torch.nn.Parameter):
-            param.data.copy_(new_val)
+            if param.shape == new_val.shape:
+                param.data.copy_(new_val)
+            else:
+                logger.warning(
+                    "Replacing parameter %s in module %s with a new shape: %s -> %s",
+                    name,
+                    module.__class__.__name__,
+                    tuple(param.shape),
+                    tuple(new_val.shape),
+                )
+                setattr(module, name, torch.nn.Parameter(new_val))
         else:
             module.register_parameter(name, torch.nn.Parameter(new_val))
     else: