tico 0.1.0.dev250904__py3-none-any.whl → 0.1.0.dev250908__py3-none-any.whl

tico/__init__.py

@@ -29,7 +29,7 @@ __all__ = [
 ]
 
 # THIS LINE IS AUTOMATICALLY GENERATED BY setup.py
-__version__ = "0.1.0.dev250904"
+__version__ = "0.1.0.dev250908"
 
 MINIMUM_SUPPORTED_VERSION = "2.5.0"
 SECURE_TORCH_VERSION = "2.6.0"

tico/experimental/quantization/algorithm/smoothquant/observer.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 import functools
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Literal
 
 import torch
 
@@ -21,18 +21,24 @@ import torch
 class ChannelwiseMaxActsObserver:
     """
     Observer to calcuate channelwise maximum activation
+    It supports collecting activations from either module inputs or outputs.
     """
 
-    def __init__(self, model):
+    def __init__(
+        self, model: torch.nn.Module, acts_from: Literal["input", "output"] = "input"
+    ):
         """
         model
             A torch module whose activations are to be analyzed.
+        acts_from
+            Where to hook: "input" for forward-pre-hook, "output" for forward-hook.
         hooks
-            A list to store the hooks which are registered to collect activation statistics.
+            A list to store the hooks registered to collect activation statistics.
         max_acts
-            A dictionary to store the maximum activation values
+            A dictionary to store the per-channel maxima.
         """
         self.model = model
+        self.acts_from: Literal["input", "output"] = acts_from
         self.hooks: List[Any] = []
         self.max_acts: Dict[str, torch.Tensor] = {}
 
@@ -62,13 +68,25 @@ class ChannelwiseMaxActsObserver:
                 input = input[0]
             stat_tensor(name, input)
 
+        def stat_output_hook(m, input, output, name):
+            if isinstance(output, tuple):
+                output = output[0]
+            stat_tensor(name, output)
+
         for name, m in self.model.named_modules():
             if isinstance(m, torch.nn.Linear):
-                self.hooks.append(
-                    m.register_forward_pre_hook(
-                        functools.partial(stat_input_hook, name=name)
+                if self.acts_from == "input":
+                    self.hooks.append(
+                        m.register_forward_pre_hook(
+                            functools.partial(stat_input_hook, name=name)
+                        )
+                    )
+                else:  # "output"
+                    self.hooks.append(
+                        m.register_forward_hook(
+                            functools.partial(stat_output_hook, name=name)
+                        )
                     )
-                )
 
     def remove(self):
         for hook in self.hooks:

tico/experimental/quantization/algorithm/smoothquant/quantizer.py

@@ -30,6 +30,25 @@ from tico.experimental.quantization.quantizer import BaseQuantizer
 class SmoothQuantQuantizer(BaseQuantizer):
     """
     Quantizer for applying the SmoothQuant algorithm
+
+    Q) Why allow choosing between input and output activations?
+
+    SmoothQuant relies on channel-wise activation statistics to balance
+    weights and activations. In practice, there are two natural sources:
+
+    - "input": captures the tensor right before a Linear layer
+      (forward-pre-hook). This matches the original SmoothQuant paper
+      and focuses on scaling the raw hidden state.
+
+    - "output": captures the tensor right after a Linear layer
+      (forward-hook). This can better reflect post-weight dynamics,
+      especially when subsequent operations (bias, activation functions)
+      dominate the dynamic range.
+
+    Allowing both options provides flexibility: depending on model
+    architecture and calibration data, one may yield lower error than
+    the other. The default remains "input" for compatibility, but "output"
+    can be selected to empirically reduce error or runtime overhead.
     """
 
     def __init__(self, config: SmoothQuantConfig):
@@ -37,6 +56,7 @@ class SmoothQuantQuantizer(BaseQuantizer):
 
         self.alpha = config.alpha
         self.custom_alpha_map = config.custom_alpha_map
+        self.acts_from = config.acts_from  # "input" (default) or "output"
         self.observer: Optional[ChannelwiseMaxActsObserver] = None
 
     @torch.no_grad()
@@ -55,7 +75,8 @@ class SmoothQuantQuantizer(BaseQuantizer):
         Returns:
             The model prepared for SmoothQuant quantization.
         """
-        self.observer = ChannelwiseMaxActsObserver(model)
+        # Attach hooks according to `config.acts_from`
+        self.observer = ChannelwiseMaxActsObserver(model, acts_from=self.acts_from)
         self.observer.attach()
 
         return model

tico/experimental/quantization/algorithm/smoothquant/smooth_quant.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Dict, List, Optional
+from typing import Callable, Dict, List, Optional
 
 import torch
 
@@ -110,6 +110,185 @@ def smooth_weights(
         back_m.weight.mul_(scales.view(1, -1))  # type: ignore[operator]
 
 
+# TODO Split the files per model
+# ────────────────────────────────────────────────────────────
+# fairseq ReLU bridge (input-hook stats) helpers
+# ────────────────────────────────────────────────────────────
+
+
+@torch.no_grad()
+def _compute_s_for_linear(
+    linear_like: torch.nn.Module,  # 2D weight [out, in]
+    activation_max: torch.Tensor,  # shape [in]
+    alpha: float,
+) -> torch.Tensor:
+    """
+    s = (amax^alpha / w_col_max^(1-alpha))
+      - amax: channel-wise max of the input to this module
+      - w_col_max: max(|W|) per input column
+    """
+    if not hasattr(linear_like, "weight"):
+        raise RuntimeError(f"{type(linear_like).__name__} has no 'weight' attribute.")
+    W = linear_like.weight  # [out, in]
+    assert isinstance(W, torch.Tensor)
+    if W.ndim != 2:
+        raise RuntimeError(
+            f"Expected 2D weight, got {W.ndim}D for {type(linear_like).__name__}"
+        )
+
+    device, dtype = W.device, W.dtype
+    amax = activation_max.to(device=device, dtype=dtype)
+
+    if amax.numel() != W.shape[1]:
+        raise ValueError(
+            f"activation_max numel({amax.numel()}) != in_features({W.shape[1]})"
+        )
+
+    w_col_max = W.abs().max(dim=0)[0].clamp(min=1e-5)  # [in]
+    s = (amax.pow(alpha) / w_col_max.pow(1.0 - alpha)).clamp(min=1e-5)  # [in]
+    return s
+
+
+@torch.no_grad()
+def _fuse_relu_bridge_no_runtime_mul(
+    fc1: torch.nn.Module,
+    fc2: torch.nn.Module,
+    s_hidden: torch.Tensor,
+):
+    """
+    Fuse scaling across fc1 → ReLU → fc2 without runtime multiplies:
+      - fc1 rows *= 1/s, (fc1.bias *= 1/s)
+      - fc2 cols *= s
+    Assumes middle activation is ReLU (positive homogeneous).
+    """
+    if not hasattr(fc1, "weight") or not hasattr(fc2, "weight"):
+        raise RuntimeError("fc1/fc2 must have 'weight' attributes.")
+
+    W1, W2 = fc1.weight, fc2.weight
+    assert isinstance(W1, torch.Tensor) and isinstance(W2, torch.Tensor)
+    if W1.ndim != 2 or W2.ndim != 2:
+        raise RuntimeError("fc1/fc2 weights must be 2D.")
+
+    hidden = W1.shape[0]
+    if W2.shape[1] != hidden or s_hidden.numel() != hidden:
+        raise ValueError(
+            f"Dimension mismatch: hidden={hidden}, W2.in={W2.shape[1]}, s={s_hidden.numel()}"
+        )
+
+    s = s_hidden.to(device=W1.device, dtype=W1.dtype).clamp(min=1e-5)  # [hidden]
+    inv_s = (1.0 / s).clamp(min=1e-5)
+
+    # fc1: row-wise scale
+    W1.mul_(inv_s.view(-1, 1))
+    if hasattr(fc1, "bias") and getattr(fc1, "bias") is not None:
+        assert isinstance(fc1.bias, torch.Tensor)
+        fc1.bias.mul_(inv_s)
+
+    # fc2: column-wise scale
+    W2.mul_(s.view(1, -1))
+
+
+# ────────────────────────────────────────────────────────────
+# Per-layer appliers (uniform protocol): return True if applied, else False
+# ────────────────────────────────────────────────────────────
+
+
+@torch.no_grad()
+def _apply_if_llama_decoder(
+    name: str,
+    module: torch.nn.Module,
+    activation_max: Dict[str, torch.Tensor],
+    alpha_to_apply: float,
+) -> bool:
+    """
+    Apply LLaMA decoder-layer smoothing (input-hook stats).
+    Returns True if this handler applied smoothing to `module`.
+    """
+    try:
+        from transformers.models.llama.modeling_llama import (  # type: ignore
+            LlamaDecoderLayer,
+        )
+    except Exception:
+        return False
+
+    if not isinstance(module, LlamaDecoderLayer):
+        return False
+
+    attn_ln = module.input_layernorm
+    qkv = [
+        module.self_attn.q_proj,
+        module.self_attn.k_proj,
+        module.self_attn.v_proj,
+    ]
+    # Input-hook stats for q_proj input
+    qkv_input_scales = activation_max[name + ".self_attn.q_proj"]
+    smooth_weights(attn_ln, qkv, qkv_input_scales, alpha_to_apply)
+
+    ffn_ln = module.post_attention_layernorm
+    fcs = [module.mlp.gate_proj, module.mlp.up_proj]
+    # Input-hook stats for gate_proj input
+    fcs_input_scales = activation_max[name + ".mlp.gate_proj"]
+    smooth_weights(ffn_ln, fcs, fcs_input_scales, alpha_to_apply)
+
+    return True
+
+
+@torch.no_grad()
+def _apply_if_fairseq_relu_bridge(
+    name: str,
+    module: torch.nn.Module,
+    activation_max: Dict[str, torch.Tensor],
+    alpha_to_apply: float,
+) -> bool:
+    """
+    Apply fairseq Transformer (Encoder/Decoder) ReLU-FFN bridge fusion
+    using input-hook stats at '{name}.fc1'. Returns True if applied.
+    """
+    try:
+        from fairseq.modules.transformer_layer import (
+            TransformerDecoderLayerBase,
+            TransformerEncoderLayerBase,
+        )  # type: ignore
+    except Exception:
+        return False
+
+    if not isinstance(
+        module, (TransformerEncoderLayerBase, TransformerDecoderLayerBase)
+    ):
+        return False
+
+    # Only when FFN activation is ReLU (positive homogeneity)
+    act_fn = getattr(module, "activation_fn", None)
+    is_relu = (act_fn is torch.nn.functional.relu) or getattr(
+        act_fn, "__name__", ""
+    ) == "relu"
+    if not is_relu:
+        return False
+
+    fc1_key = f"{name}.fc1"
+    amax2 = activation_max.get(fc1_key)
+    if amax2 is None:
+        return False
+
+    fc1 = getattr(module, "fc1", None)
+    fc2 = getattr(module, "fc2", None)
+    if fc1 is None or fc2 is None or not hasattr(fc2, "weight") or fc2.weight.ndim != 2:
+        return False
+
+    s_hidden = _compute_s_for_linear(fc2, amax2, alpha_to_apply)  # [hidden]
+    _fuse_relu_bridge_no_runtime_mul(fc1, fc2, s_hidden)
+    return True
+
+
+# Registry of appliers (order matters: try LLaMA first, then fairseq)
+_APPLIERS: List[
+    Callable[[str, torch.nn.Module, Dict[str, torch.Tensor], float], bool]
+] = [
+    _apply_if_llama_decoder,
+    _apply_if_fairseq_relu_bridge,
+]
+
+
 @torch.no_grad()
 def apply_smoothing(
     model: torch.nn.Module,
@@ -133,32 +312,16 @@ def apply_smoothing(
             Layers specified in this dictionary will use the corresponding alpha
              value instead of the default.
     """
-    from transformers.models.llama.modeling_llama import LlamaDecoderLayer
-
     for name, module in model.named_modules():
-        alpha_to_apply = alpha
-        if custom_alpha_map and name in custom_alpha_map:
-            alpha_to_apply = custom_alpha_map[name]
+        alpha_to_apply = (
+            custom_alpha_map.get(name, alpha) if custom_alpha_map else alpha
+        )
         if alpha_to_apply > 1.0:
             raise RuntimeError(
                 f"Alpha value cannot exceed 1.0. Given alpha: {alpha_to_apply}"
             )
-        # SmoothQuant is applied before capturing the graph. Therefore, it needs to know
-        #  specific module information.
-        # TODO Suport more modules.
-        if isinstance(module, LlamaDecoderLayer):
-            attn_ln = module.input_layernorm
-            qkv = [
-                module.self_attn.q_proj,
-                module.self_attn.k_proj,
-                module.self_attn.v_proj,
-            ]
-
-            qkv_input_scales = activation_max[name + ".self_attn.q_proj"]
-            smooth_weights(attn_ln, qkv, qkv_input_scales, alpha_to_apply)
-
-            ffn_ln = module.post_attention_layernorm
-            fcs = [module.mlp.gate_proj, module.mlp.up_proj]
-            fcs_input_scales = activation_max[name + ".mlp.gate_proj"]
-
-            smooth_weights(ffn_ln, fcs, fcs_input_scales, alpha_to_apply)
+
+        # Try each applier until one succeeds.
+        for applier in _APPLIERS:
+            if applier(name, module, activation_max, alpha_to_apply):
+                break  # applied → stop trying others

tico/experimental/quantization/config.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 from abc import ABC, abstractmethod
-from typing import Dict, Optional
+from typing import Dict, Literal, Optional
 
 
 class BaseConfig(ABC):
@@ -60,9 +60,15 @@ class SmoothQuantConfig(BaseConfig):
         self,
         alpha: float = 0.5,
         custom_alpha_map: Optional[Dict[str, float]] = None,
+        acts_from: Literal["input", "output"] = "input",
     ):
         self.alpha = alpha
         self.custom_alpha_map = custom_alpha_map
+        # Where to collect activation statistics from:
+        # - "input": use forward-pre-hook (Tensor before the Linear op)
+        # - "output": use forward-hook    (Tensor after the Linear op)
+        # Default is "input".
+        self.acts_from = acts_from
 
     @property
     def name(self) -> str:

tico/experimental/quantization/ptq/wrappers/fairseq/__init__.py

@@ -0,0 +1,5 @@
+from tico.experimental.quantization.ptq.wrappers.fairseq.quant_mha import (
+    QuantFairseqMultiheadAttention,
+)
+
+__all__ = ["QuantFairseqMultiheadAttention"]

tico/experimental/quantization/ptq/wrappers/fairseq/quant_mha.py

@@ -0,0 +1,383 @@
+# Copyright (c) 2025 Samsung Electronics Co., Ltd. All Rights Reserved
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# -----------------------------------------------------------------------------
+# This file includes modifications based on fairseq
+#  (https://github.com/facebookresearch/fairseq), originally licensed under
+# the MIT License. See the LICENSE file in the fairseq repository for details.
+# -----------------------------------------------------------------------------
+
+from typing import Dict, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from tico.experimental.quantization.ptq.quant_config import QuantConfig
+from tico.experimental.quantization.ptq.wrappers.ptq_wrapper import PTQWrapper
+from tico.experimental.quantization.ptq.wrappers.quant_module_base import (
+    QuantModuleBase,
+)
+from tico.experimental.quantization.ptq.wrappers.registry import try_register
+
+
+@try_register("fairseq.modules.multihead_attention.MultiheadAttention")
+class QuantFairseqMultiheadAttention(QuantModuleBase):
+    """
+    Quant-aware drop-in for Fairseq MultiheadAttention.
+
+    - No xFormers / no torch F.multi_head_attention_forward fast-path.
+    - Self/cross attention + minimal incremental KV cache.
+    - Causal mask is pre-built statically; `key_padding_mask` is additive float.
+    - I/O shape: [T, B, C]
+
+    Runtime optimization flags
+    --------------------------
+    use_static_causal : bool
+        If True, reuse a precomputed upper-triangular causal mask template
+        instead of rebuilding it each forward step. Reduces per-step mask
+        construction overhead during incremental decoding.
+
+    assume_additive_key_padding : bool
+        If True, assume the `key_padding_mask` is already an additive float
+        tensor (large negative values at padded positions). Skips conversion
+        from boolean masks, reducing runtime overhead.
+    """
+
+    def __init__(
+        self,
+        fp_attn: nn.Module,
+        *,
+        qcfg: Optional[QuantConfig] = None,
+        fp_name: Optional[str] = None,
+        max_seq: int = 4096,
+        use_static_causal: bool = False,
+        mask_fill_value: float = -120.0,
+        assume_additive_key_padding: bool = False,
+    ):
+        super().__init__(qcfg, fp_name=fp_name)
+
+        self.use_static_causal = use_static_causal
+        self.mask_fill_value = mask_fill_value
+        self.assume_additive_key_padding = assume_additive_key_padding
+        self.embed_dim: int = int(fp_attn.embed_dim)  # type: ignore[arg-type]
+        self.num_heads: int = int(fp_attn.num_heads)  # type: ignore[arg-type]
+        self.head_dim: int = self.embed_dim // self.num_heads
+        assert self.head_dim * self.num_heads == self.embed_dim
+
+        self.self_attention: bool = bool(getattr(fp_attn, "self_attention", False))
+        self.encoder_decoder_attention: bool = bool(
+            getattr(fp_attn, "encoder_decoder_attention", False)
+        )
+        assert self.self_attention != self.encoder_decoder_attention
+
+        # PTQ-wrapped projections
+        qc = qcfg.child("q_proj") if qcfg else None
+        kc = qcfg.child("k_proj") if qcfg else None
+        vc = qcfg.child("v_proj") if qcfg else None
+        oc = qcfg.child("out_proj") if qcfg else None
+        assert hasattr(fp_attn, "q_proj") and hasattr(fp_attn, "k_proj")
+        assert hasattr(fp_attn, "v_proj") and hasattr(fp_attn, "out_proj")
+        assert isinstance(fp_attn.q_proj, nn.Module) and isinstance(
+            fp_attn.k_proj, nn.Module
+        )
+        assert isinstance(fp_attn.v_proj, nn.Module) and isinstance(
+            fp_attn.out_proj, nn.Module
+        )
+        self.q_proj = PTQWrapper(fp_attn.q_proj, qcfg=qc, fp_name=f"{fp_name}.q_proj")
+        self.k_proj = PTQWrapper(fp_attn.k_proj, qcfg=kc, fp_name=f"{fp_name}.k_proj")
+        self.v_proj = PTQWrapper(fp_attn.v_proj, qcfg=vc, fp_name=f"{fp_name}.v_proj")
+        self.out_proj = PTQWrapper(
+            fp_attn.out_proj, qcfg=oc, fp_name=f"{fp_name}.out_proj"
+        )
+
+        # scale & static causal mask
+        self.register_buffer(
+            "scale_const", torch.tensor(self.head_dim**-0.5), persistent=False
+        )
+        mask = torch.full((1, 1, max_seq, max_seq), float(self.mask_fill_value))
+        mask.triu_(1)
+        self.register_buffer("causal_mask_template", mask, persistent=False)
+
+        # observers (no *_proj_out here; PTQWrapper handles module outputs)
+        mk = self._make_obs
+        self.obs_query_in = mk("query_in")
+        self.obs_key_in = mk("key_in")
+        self.obs_value_in = mk("value_in")
+        self.obs_kpm_in = mk("kpm_in")
+        self.obs_causal_mask = mk("causal_mask")
+        self.obs_q_fold = mk("q_fold")
+        self.obs_k_fold = mk("k_fold")
+        self.obs_v_fold = mk("v_fold")
+        self.obs_scale = mk("scale")
+        self.obs_logits_raw = mk("logits_raw")
+        self.obs_logits = mk("logits_scaled")
+        self.obs_attn_mask_add = mk("obs_attn_mask_add")
+        self.obs_kp_mask_add = mk("obs_kp_mask_add")
+        self.obs_softmax = mk("softmax")
+        self.obs_attn_out = mk("attn_out")
+
+        safe_name = (
+            fp_name if (fp_name not in (None, "", "None")) else f"QuantFsMHA_{id(self)}"
+        )
+        assert safe_name is not None
+        self._state_key = safe_name + ".attn_state"
+
+    def _get_input_buffer(
+        self,
+        incremental_state: Optional[Dict[str, Dict[str, Optional[torch.Tensor]]]],
+    ) -> Optional[Dict[str, Optional[torch.Tensor]]]:
+        """Return saved KV/mask dict or None."""
+        if incremental_state is None:
+            return None
+        return incremental_state.get(self._state_key, None)
+
+    def _set_input_buffer(
+        self,
+        incremental_state: Optional[Dict[str, Dict[str, Optional[torch.Tensor]]]],
+        buffer: Dict[str, Optional[torch.Tensor]],
+    ):
+        """Store KV/mask dict in incremental_state."""
+        if incremental_state is not None:
+            incremental_state[self._state_key] = buffer
+        return incremental_state
+
+    # ---- utils ----
+    def _fold_heads(self, x: torch.Tensor, B: int) -> torch.Tensor:
+        # [T,B,E] -> [B*H, T, Dh]
+        T = x.size(0)
+        x = x.view(T, B, self.num_heads, self.head_dim).permute(1, 2, 0, 3).contiguous()
+        return x.view(B * self.num_heads, T, self.head_dim)
+
+    def _unfold_heads(self, x: torch.Tensor, B: int, T: int) -> torch.Tensor:
+        # [B*H, T, Dh] -> [T,B,E]
+        x = x.view(B, self.num_heads, T, self.head_dim).permute(2, 0, 1, 3).contiguous()
+        return x.view(T, B, self.embed_dim)
+
+    def forward(
+        self,
+        query: torch.Tensor,  # [Tq,B,C]
+        key: Optional[torch.Tensor],
+        value: Optional[torch.Tensor],
+        key_padding_mask: Optional[
+            torch.Tensor
+        ] = None,  # additive float (e.g. -120 at pads)
+        incremental_state: Optional[
+            Dict[str, Dict[str, Optional[torch.Tensor]]]
+        ] = None,
+        need_weights: bool = False,
+        static_kv: bool = False,
+        attn_mask: Optional[torch.Tensor] = None,  # if None -> internal causal
+        before_softmax: bool = False,
+        need_head_weights: bool = False,
+        return_new_kv: bool = False,
+    ) -> Union[
+        Tuple[torch.Tensor, Optional[torch.Tensor]],
+        Tuple[
+            torch.Tensor,
+            Optional[torch.Tensor],
+            Optional[torch.Tensor],
+            Optional[torch.Tensor],
+        ],
+    ]:
+
+        if need_head_weights:
+            need_weights = True
+
+        Tq, B, _ = query.shape
+        if self.self_attention:
+            key = query if key is None else key
+            value = query if value is None else value
+        else:
+            assert key is not None and value is not None
+
+        Tk, Bk, _ = key.shape
+        Tv, Bv, _ = value.shape
+        assert B == Bk == Bv
+
+        q = self.q_proj(self._fq(query, self.obs_query_in))
+        k = self.k_proj(self._fq(key, self.obs_key_in))
+        v = self.v_proj(self._fq(value, self.obs_value_in))
+
+        state = self._get_input_buffer(incremental_state)
+        if incremental_state is not None and state is None:
+            state = {}
+
+        # Capture "new" K/V for this call BEFORE concatenating with cache
+        new_k_bh: Optional[torch.Tensor] = None
+        new_v_bh: Optional[torch.Tensor] = None
+
+        # Fold heads
+        q = self._fq(self._fold_heads(q, B), self.obs_q_fold)
+        if state is not None and "prev_key" in state and static_kv:
+            # Cross-attention static_kv path: reuse cached KV; there is no new KV this call.
+            k = None
+            v = None
+        if k is not None:
+            k = self._fq(self._fold_heads(k, B), self.obs_k_fold)  # [B*H, Tnew, Dh]
+            if return_new_kv:
+                new_k_bh = k.contiguous()
+        if v is not None:
+            v = self._fq(self._fold_heads(v, B), self.obs_v_fold)  # [B*H, Tnew, Dh]
+            if return_new_kv:
+                new_v_bh = v.contiguous()
+
+        # Append/reuse cache
+        if state is not None:
+            pk = state.get("prev_key")
+            pv = state.get("prev_value")
+            if pk is not None:
+                pk = pk.view(B * self.num_heads, -1, self.head_dim)
+                k = pk if static_kv else torch.cat([pk, k], dim=1)
+            if pv is not None:
+                pv = pv.view(B * self.num_heads, -1, self.head_dim)
+                v = pv if static_kv else torch.cat([pv, v], dim=1)
+
+        assert k is not None and v is not None
+        Ts = k.size(1)
+
+        # Scaled dot-product
+        scale = self._fq(self.scale_const, self.obs_scale).to(q.dtype)
+        logits_raw = self._fq(
+            torch.bmm(q, k.transpose(1, 2)), self.obs_logits_raw
+        )  # [B*H,Tq,Ts]
+        logits = self._fq(logits_raw * scale, self.obs_logits)
+
+        assert isinstance(self.causal_mask_template, torch.Tensor)
+        # Masks
+        device = logits.device
+        if attn_mask is None and self.use_static_causal:
+            # Incremental decoding aware slicing:
+            # align the causal row(s) to the current time indices
+            start_q = max(Ts - Tq, 0)
+            cm = self.causal_mask_template[..., start_q : start_q + Tq, :Ts].to(
+                device=device, dtype=logits.dtype
+            )
+            attn_mask = cm.squeeze(0).squeeze(0)  # [Tq,Ts]
+
+        if attn_mask is not None:
+            # Bool/byte mask -> additive float with large negatives
+            if not torch.is_floating_point(attn_mask):
+                fill = self.causal_mask_template.new_tensor(self.mask_fill_value)
+                attn_mask = torch.where(
+                    attn_mask.to(torch.bool), fill, fill.new_zeros(())
+                )
+            attn_mask = self._fq(attn_mask, self.obs_causal_mask)
+            assert isinstance(attn_mask, torch.Tensor)
+
+            if not self.assume_additive_key_padding:
+                # attn_mask -> [B*H,Tq,Ts]
+                if attn_mask.dim() == 2:
+                    add_mask = attn_mask.unsqueeze(0).expand(logits.size(0), -1, -1)
+                elif attn_mask.dim() == 3:
+                    add_mask = (
+                        attn_mask.unsqueeze(1)
+                        .expand(B, self.num_heads, Tq, Ts)
+                        .contiguous()
+                    )
+                    add_mask = add_mask.view(B * self.num_heads, Tq, Ts)
+                else:
+                    raise RuntimeError("attn_mask must be [T,S] or [B,T,S]")
+            else:
+                add_mask = attn_mask
+            logits = self._fq(logits + add_mask, self.obs_attn_mask_add)
+
+        if key_padding_mask is not None:
+            if not torch.is_floating_point(key_padding_mask):
+                fill = self.causal_mask_template.new_tensor(self.mask_fill_value)
+                kpm = torch.where(
+                    key_padding_mask.to(torch.bool), fill, fill.new_zeros(())
+                )
+            else:
+                kpm = key_padding_mask
+            kpm = self._fq(kpm, self.obs_kpm_in)
+
+            if not self.assume_additive_key_padding:
+                # key_padding_mask: additive float already
+                kpm = kpm.to(dtype=logits.dtype, device=device)
+                if kpm.dim() == 2:  # [B,S]
+                    kpm = (
+                        kpm.view(B, 1, 1, Ts)
+                        .expand(B, self.num_heads, Tq, Ts)
+                        .contiguous()
+                    )
+                    kpm = kpm.view(B * self.num_heads, Tq, Ts)
+                elif kpm.dim() == 3:  # [B,T,S]
+                    kpm = (
+                        kpm.unsqueeze(1).expand(B, self.num_heads, Tq, Ts).contiguous()
+                    )
+                    kpm = kpm.view(B * self.num_heads, Tq, Ts)
+                else:
+                    raise RuntimeError(
+                        "key_padding_mask must be [B,S] or [B,T,S] (additive)"
+                    )
+            logits = self._fq(logits + kpm, self.obs_kp_mask_add)
+
+        if before_softmax:
+            if return_new_kv:
+                return logits, v, new_k_bh, new_v_bh
+            return logits, v
+
+        # Softmax (float32) -> back to q.dtype
+        attn_probs = torch.softmax(logits, dim=-1, dtype=torch.float32).to(q.dtype)
+        attn_probs = self._fq(attn_probs, self.obs_softmax)
+
+        # Context + output proj
+        ctx = self._fq(torch.bmm(attn_probs, v), self.obs_attn_out)  # [B*H,Tq,Dh]
+        ctx = self._unfold_heads(ctx, B, Tq)  # [Tq,B,E]
+        out = self.out_proj(ctx)
+
+        # Weights (optional)
+        attn_weights_out: Optional[torch.Tensor] = None
+        if need_weights:
+            aw = (
+                torch.softmax(logits, dim=-1, dtype=torch.float32)
+                .view(B, self.num_heads, Tq, Ts)
+                .transpose(1, 0)
+            )
+            if not need_head_weights:
+                aw = aw.mean(dim=1)  # [B,Tq,Ts]
+            attn_weights_out = aw
+
+        # Cache write
+        if state is not None:
+            state["prev_key"] = k.view(B, self.num_heads, -1, self.head_dim).detach()
+            state["prev_value"] = v.view(B, self.num_heads, -1, self.head_dim).detach()
+            self._set_input_buffer(incremental_state, state)
+
+        if return_new_kv:
+            return out, attn_weights_out, new_k_bh, new_v_bh
+        return out, attn_weights_out
+
+    def _all_observers(self):
+        yield from (
+            self.obs_query_in,
+            self.obs_key_in,
+            self.obs_value_in,
+            self.obs_kpm_in,
+            self.obs_causal_mask,
+            self.obs_q_fold,
+            self.obs_k_fold,
+            self.obs_v_fold,
+            self.obs_scale,
+            self.obs_logits_raw,
+            self.obs_logits,
+            self.obs_attn_mask_add,
+            self.obs_kp_mask_add,
+            self.obs_softmax,
+            self.obs_attn_out,
+        )
+        for m in (self.q_proj, self.k_proj, self.v_proj, self.out_proj):
+            if isinstance(m, QuantModuleBase):
+                yield from m._all_observers()

tico/experimental/quantization/ptq/wrappers/registry.py

@@ -28,9 +28,12 @@ _CORE_MODULES = (
     "tico.experimental.quantization.ptq.wrappers.nn.quant_layernorm",
     "tico.experimental.quantization.ptq.wrappers.nn.quant_linear",
     "tico.experimental.quantization.ptq.wrappers.nn.quant_silu",
+    # llama
     "tico.experimental.quantization.ptq.wrappers.llama.quant_attn",
     "tico.experimental.quantization.ptq.wrappers.llama.quant_decoder_layer",
     "tico.experimental.quantization.ptq.wrappers.llama.quant_mlp",
+    # fairseq
+    "tico.experimental.quantization.ptq.wrappers.fairseq.quant_mha",
     # add future core wrappers here
 )
 

{tico-0.1.0.dev250904.dist-info → tico-0.1.0.dev250908.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: tico
-Version: 0.1.0.dev250904
+Version: 0.1.0.dev250908
 Summary: Convert exported Torch module to circle
 Home-page: UNKNOWN
 License: UNKNOWN

{tico-0.1.0.dev250904.dist-info → tico-0.1.0.dev250908.dist-info}/RECORD

@@ -1,4 +1,4 @@
-tico/__init__.py,sha256=Q7A_sZgX9kbJSpo2ndJB19BlbFSGnajl9HAFGt8D2Q0,1883
+tico/__init__.py,sha256=KNsLIqvVA_qyf9VOIt_PmHk0L_QUjJqnldhE8_X4n6I,1883
 tico/pt2_to_circle.py,sha256=gu3MD4Iqc0zMZcCZ2IT8oGbyj21CTSbT3Rgd9s2B_9A,2767
 tico/config/__init__.py,sha256=xZzCXjZ84qE-CsBi-dfaL05bqpQ3stKKfTXhnrJRyVs,142
 tico/config/base.py,sha256=q5xMqGxTUZs4mFqt5c7i_y9U00fYgdMGl9nUqIVMlCo,1248
@@ -6,7 +6,7 @@ tico/config/factory.py,sha256=il0zqB6Lm5NX2LnG-TUhmiP9vVeZ_3TucJMorVZIodY,1324
 tico/config/v1.py,sha256=O1jzpUBDwoWpLohEpI08pJNwVB-yz3ufPrQm2_XWq4Y,1108
 tico/experimental/__init__.py,sha256=IO6FP_xYbGy0dW0HL26GXD3ouxARaxCK7bz9dn4blPQ,26
 tico/experimental/quantization/__init__.py,sha256=IaJPZegVJp0P3luutBo907Kp5sOJensE1Mm-XBG_jBs,122
-tico/experimental/quantization/config.py,sha256=1bCSAUI043Kbq08j59mb-K1cP2lmBMbekh8p3hNK6b8,1675
+tico/experimental/quantization/config.py,sha256=nMepa_H471t7f3bKMvR8cZUZgruy_8kdb147rBkTWCQ,2004
 tico/experimental/quantization/public_interface.py,sha256=4-v9VXsokRG2-UUYYHd_MlbHxChqdGI5iuySyYDY_Pw,4420
 tico/experimental/quantization/quantizer.py,sha256=_2pDtWFKDCuKfYF2bptOwIYsa0VFNFM1ZNgi8_OGvHM,2365
 tico/experimental/quantization/algorithm/__init__.py,sha256=IO6FP_xYbGy0dW0HL26GXD3ouxARaxCK7bz9dn4blPQ,26
@@ -37,9 +37,9 @@ tico/experimental/quantization/algorithm/pt2e/annotation/op/sub.py,sha256=4z8HoY
 tico/experimental/quantization/algorithm/pt2e/transformation/__init__.py,sha256=IO6FP_xYbGy0dW0HL26GXD3ouxARaxCK7bz9dn4blPQ,26
 tico/experimental/quantization/algorithm/pt2e/transformation/convert_scalars_to_attrs.py,sha256=Idtoya2RcGKlgUJgC9WqNz0jH3gf6ViuPmsD9ySHbls,2253
 tico/experimental/quantization/algorithm/smoothquant/__init__.py,sha256=IO6FP_xYbGy0dW0HL26GXD3ouxARaxCK7bz9dn4blPQ,26
-tico/experimental/quantization/algorithm/smoothquant/observer.py,sha256=yi9nR_BEuxKVjgFcYPeldhXlEbE9V-0r4kRRPcI3C70,2639
-tico/experimental/quantization/algorithm/smoothquant/quantizer.py,sha256=rQMtnqM1dBzjhY-KJRp3TWZSW-dzXrs5N5FagXzh0IQ,2564
-tico/experimental/quantization/algorithm/smoothquant/smooth_quant.py,sha256=O1h7IojcsJaprFvRM9tsAuR3q7vTjKKRi88jD-nH79Y,6175
+tico/experimental/quantization/algorithm/smoothquant/observer.py,sha256=OWBKQ3ox6PqeqgevxOjpXvb7uApoqE4YbUBelGhVSN8,3435
+tico/experimental/quantization/algorithm/smoothquant/quantizer.py,sha256=QuZBi24L-LYI26nwZd6JmTdokxr6-l_vIgZvWVdqx_o,3637
+tico/experimental/quantization/algorithm/smoothquant/smooth_quant.py,sha256=fxCy4m-BsSjraciSVPFlPhgsOT46RjrOgczQGb7B9TA,11561
 tico/experimental/quantization/evaluation/__init__.py,sha256=IO6FP_xYbGy0dW0HL26GXD3ouxARaxCK7bz9dn4blPQ,26
 tico/experimental/quantization/evaluation/backend.py,sha256=CZL9rZOA0t8cH7PHp6u9l7dGqWNvTj9bKOvwo0PVul0,692
 tico/experimental/quantization/evaluation/evaluate.py,sha256=kfa_GvFaX6DoSTAmuCImMJqF2jgqtnor5UpC7wVmGPI,7877
@@ -84,7 +84,9 @@ tico/experimental/quantization/ptq/wrappers/__init__.py,sha256=IO6FP_xYbGy0dW0HL
 tico/experimental/quantization/ptq/wrappers/ptq_wrapper.py,sha256=F9sK_DiRaXiGNHULcwIbs5EUtHz6ZJ7N4r5CWTTfhsM,2442
 tico/experimental/quantization/ptq/wrappers/quant_elementwise.py,sha256=LhEoobfvto6zKrBOKL4gmxfFFc31jHzyQV_zfps-iQM,3604
 tico/experimental/quantization/ptq/wrappers/quant_module_base.py,sha256=vkcDos_knGSS29rIZuEIWkAJLHrENbGz8nCH2-iara8,5969
-tico/experimental/quantization/ptq/wrappers/registry.py,sha256=wauoZdZBR15bGj1Upt9owEfFDT-Tj6HzciG9HDM1BHo,4845
+tico/experimental/quantization/ptq/wrappers/registry.py,sha256=cblE6dpLlK1lh61Xpqxcr2Vwi9XuqdrV5Y7TPSjpETQ,4940
+tico/experimental/quantization/ptq/wrappers/fairseq/__init__.py,sha256=Mc8FLd9DusyB_IT1vk1OYrRkngOYnYd05IvtA9ORVQc,160
+tico/experimental/quantization/ptq/wrappers/fairseq/quant_mha.py,sha256=HsigmOLeacLXc46QNeFqwQ0DwKQhNrtWTKEtLJoqXoc,15562
 tico/experimental/quantization/ptq/wrappers/llama/__init__.py,sha256=IO6FP_xYbGy0dW0HL26GXD3ouxARaxCK7bz9dn4blPQ,26
 tico/experimental/quantization/ptq/wrappers/llama/quant_attn.py,sha256=-K1COLHIHfJZhQu-RE6KfJIkaL7S6yR4iUj48QkjMTw,8652
 tico/experimental/quantization/ptq/wrappers/llama/quant_decoder_layer.py,sha256=2XsIf5rcabDXXkahqriSxfo2curFq0Y5bnRPcYkJPg8,7187
@@ -247,9 +249,9 @@ tico/utils/mx/__init__.py,sha256=IO6FP_xYbGy0dW0HL26GXD3ouxARaxCK7bz9dn4blPQ,26
 tico/utils/mx/elemwise_ops.py,sha256=V6glyAHsVR1joqpsgnNytatCD_ew92xNWZ19UFDoMTA,10281
 tico/utils/mx/formats.py,sha256=uzNWyu-1onUlwQfX5cZ6fZSUfHMRqorper7_T1k3jfk,3404
 tico/utils/mx/mx_ops.py,sha256=RcfUTYVi-wilGB2sC35OeARdwDqnixv7dG5iyZ-fQT8,8555
-tico-0.1.0.dev250904.dist-info/LICENSE,sha256=kp4JLII7bzRhPb0CPD5XTDZMh22BQ7h3k3B7t8TiSbw,12644
-tico-0.1.0.dev250904.dist-info/METADATA,sha256=jCLzVEpVnwflcRRdt3AZElB_Lw1_lWSzusXsiA5RAig,8450
-tico-0.1.0.dev250904.dist-info/WHEEL,sha256=G16H4A3IeoQmnOrYV4ueZGKSjhipXx8zc8nu9FGlvMA,92
-tico-0.1.0.dev250904.dist-info/entry_points.txt,sha256=kBKYSS_IYrSXmUYevmmepqIVPScq5vF8ulQRu3I_Zf0,59
-tico-0.1.0.dev250904.dist-info/top_level.txt,sha256=oqs7UPoNSKZEwqsX8B-KAWdQwfAa7i60pbxW_Jk7P3w,5
-tico-0.1.0.dev250904.dist-info/RECORD,,
+tico-0.1.0.dev250908.dist-info/LICENSE,sha256=kp4JLII7bzRhPb0CPD5XTDZMh22BQ7h3k3B7t8TiSbw,12644
+tico-0.1.0.dev250908.dist-info/METADATA,sha256=SBFNzxkvYrsWshqSGGkmtVoHtDk1Rp9wVNfhdIRsFdg,8450
+tico-0.1.0.dev250908.dist-info/WHEEL,sha256=G16H4A3IeoQmnOrYV4ueZGKSjhipXx8zc8nu9FGlvMA,92
+tico-0.1.0.dev250908.dist-info/entry_points.txt,sha256=kBKYSS_IYrSXmUYevmmepqIVPScq5vF8ulQRu3I_Zf0,59
+tico-0.1.0.dev250908.dist-info/top_level.txt,sha256=oqs7UPoNSKZEwqsX8B-KAWdQwfAa7i60pbxW_Jk7P3w,5
+tico-0.1.0.dev250908.dist-info/RECORD,,