From e1eeff7063d8083170d1396db164738b9357f148 Mon Sep 17 00:00:00 2001
From: WeiweiZhang1 <weiwei1.zhang@intel.com>
Date: Fri, 31 Jul 2026 16:02:38 +0800
Subject: [PATCH 1/4] refine awq alg, add smooth seqlen, bugfix, refine docs

Signed-off-by: WeiweiZhang1 <weiwei1.zhang@intel.com>
---
 auto_round/algorithms/transforms/awq/base.py  | 181 ++++++++++++++----
 .../algorithms/transforms/awq/config.py       |  16 ++
 .../algorithms/transforms/awq/mappings.py     |  56 +++++-
 docs/step_by_step.md                          |  66 +++++--
 docs/step_by_step_CN.md                       |  68 +++++--
 test/test_cpu/algorithms/test_awq.py          | 111 +++++++++++
 6 files changed, 423 insertions(+), 75 deletions(-)

diff --git a/auto_round/algorithms/transforms/awq/base.py b/auto_round/algorithms/transforms/awq/base.py
index 7b1bf54376..cf651209b0 100644
--- a/auto_round/algorithms/transforms/awq/base.py
+++ b/auto_round/algorithms/transforms/awq/base.py
@@ -38,9 +38,10 @@
 from auto_round.algorithms.registry import register_pipeline_member
 from auto_round.algorithms.transforms.awq.config import AWQConfig
 from auto_round.algorithms.transforms.awq.mappings import (
+    AWQ_DYNAMIC_MAPPING_REGISTRY,
+    AWQ_MAPPING_REGISTRY,
     ResolvedMapping,
     _extract_block_prefix,
-    check_model_compatibility,
     resolve_mappings,
 )
 from auto_round.algorithms.transforms.awq.qdq import QDQTool
@@ -102,6 +103,45 @@ def _rmsnorm_has_unit_offset(module: torch.nn.Module) -> bool:
     return result
 
 
+def _slice_seq_tensor(v: Any, actual_seq: int, seqlen: int) -> Any:
+    """Slice a value's sequence dimension to *seqlen*, recursing into containers."""
+    if isinstance(v, torch.Tensor):
+        if v.ndim == 3 and v.shape[1] == actual_seq:
+            return v[:, :seqlen]
+        if v.ndim == 2 and v.shape[1] == actual_seq:
+            return v[:, :seqlen]
+        if v.ndim == 4 and v.shape[2] == actual_seq and v.shape[3] == actual_seq:
+            return v[:, :, :seqlen, :seqlen]
+        if v.ndim == 4 and v.shape[2] == actual_seq:
+            return v[:, :, :seqlen]
+        return v
+    if isinstance(v, (tuple, list)):
+        return type(v)(_slice_seq_tensor(t, actual_seq, seqlen) for t in v)
+    return v
+
+
+def _detect_actual_seq(values) -> int | None:
+    """Infer sequence length from the first 3-D tensor, including nested containers."""
+    for v in values:
+        if isinstance(v, torch.Tensor) and v.ndim == 3:
+            return v.shape[1]
+        if isinstance(v, (tuple, list)):
+            nested = _detect_actual_seq(v)
+            if nested is not None:
+                return nested
+    return None
+
+
+def _truncate_args_kwargs(args: tuple, kwargs: dict, seqlen: int) -> tuple[tuple, dict]:
+    """Truncate parent-forward positional args and kwargs to at most *seqlen* tokens."""
+    actual_seq = _detect_actual_seq(list(args) + list(kwargs.values()))
+    if actual_seq is None or actual_seq <= seqlen:
+        return args, kwargs
+    new_args = tuple(_slice_seq_tensor(v, actual_seq, seqlen) for v in args)
+    new_kwargs = {k: _slice_seq_tensor(v, actual_seq, seqlen) for k, v in kwargs.items()}
+    return new_args, new_kwargs
+
+
 @register_pipeline_member(AWQConfig)
 class AWQTransform(BasePreprocessor):
     """AWQ transform: activation-aware weight smoothing pre-processor.
@@ -124,6 +164,11 @@ def __init__(self, config: AWQConfig) -> None:
         self.clip_max_shrink: float = getattr(config, "clip_max_shrink", 0.5)
         self.clip_n_sample_token: int = getattr(config, "clip_n_sample_token", 512)
 
+        # Cap parent-forward samples used by the expensive scale search.
+        # A value <= 0 disables truncation and uses the full calibration sequence.
+        smooth_seqlen = getattr(config, "smooth_seqlen", 512)
+        self._smooth_seqlen: int | None = smooth_seqlen if smooth_seqlen > 0 else None
+
         # Single source of truth for "QDQ a candidate weight under the target
         # block-quantizer scheme", used as AWQ's grid-search / clip loss. AWQ
         # composes this instead of re-implementing block-quantizer dispatch;
@@ -135,6 +180,7 @@ def __init__(self, config: AWQConfig) -> None:
         )
 
         self._user_mappings: list[dict] | None = config.mappings
+        self._skip_moe: bool = getattr(config, "skip_moe", False)
 
         # Set at runtime by the compressor's post_init() via ``pre.layer_config = self.layer_config``.
         self.layer_config: dict | None = None
@@ -143,8 +189,7 @@ def __init__(self, config: AWQConfig) -> None:
         self._block_mappings: dict[str, list[ResolvedMapping]] = {}
 
         self._activation_stats: dict[str, list] = {}
-        self._parent_args_cache: dict[torch.nn.Module, list[dict]] = {}
-        self._parent_signatures: dict[int, inspect.Signature] = {}
+        self._parent_args_cache: dict[torch.nn.Module, list[tuple[tuple, dict]]] = {}
         # Per-mapping balance-layer input features captured for the clip search
         # (keyed by smooth_name). Only populated when ``apply_clip`` is set.
         self._clip_input_feat: dict[str, torch.Tensor] = {}
@@ -165,14 +210,11 @@ def bind(self, compressor) -> None:
             exit(-1)
 
     def prepare_run(self, composer: "AlgorithmComposer" = None) -> None:
-        """Validate compatibility, resolve model-wide mappings, and group by block prefix."""
+        """Resolve model-wide mappings and group them by transformer block."""
         model = self.model
-        report = check_model_compatibility(model, self._user_mappings)
-        for warning in report["warnings"]:
-            logger.warning(warning)
 
         # ── Resolve all model-level mappings (name-only, no module caching) ──
-        self._resolved_mappings = resolve_mappings(model, self._user_mappings)
+        self._resolved_mappings = resolve_mappings(model, self._user_mappings, skip_moe=self._skip_moe)
         if not self._resolved_mappings:
             raise ValueError(
                 "AWQ: no layer mappings were resolved for this model. "
@@ -181,11 +223,38 @@ def prepare_run(self, composer: "AlgorithmComposer" = None) -> None:
                 "add an entry to auto_round/algorithms/transforms/awq/mappings.py."
             )
 
-        # Group mappings by block prefix for O(1) lookup during block iteration.
+        cls_name = type(model).__name__
+        if (
+            self._user_mappings is None
+            and cls_name not in AWQ_MAPPING_REGISTRY
+            and cls_name not in AWQ_DYNAMIC_MAPPING_REGISTRY
+        ):
+            logger.warning(
+                "AWQ: model class '%s' is not in any AWQ mapping registry; using "
+                "default Llama-like mappings. If quantization quality is poor, "
+                "provide explicit mappings via AWQConfig(mappings=[...]).",
+                cls_name,
+            )
+
+        from auto_round.utils.common import flatten_list
+        from auto_round.utils.model import get_block_names
+
+        try:
+            iter_block_names = [b for b in flatten_list(get_block_names(model)) if b]
+        except Exception:  # noqa: BLE001 - fall back to prefix heuristic if block discovery fails
+            iter_block_names = []
+        iter_block_names.sort(key=len, reverse=True)
+
         self._block_mappings = {}
         for m in self._resolved_mappings:
-            prefix = _extract_block_prefix(m.smooth_name)
-            self._block_mappings.setdefault(prefix, []).append(m)
+            key = None
+            for block_name in iter_block_names:
+                if m.smooth_name == block_name or m.smooth_name.startswith(block_name + "."):
+                    key = block_name
+                    break
+            if key is None:
+                key = _extract_block_prefix(m.smooth_name)
+            self._block_mappings.setdefault(key, []).append(m)
 
         if composer is not None:
             self._qdq_tool.configure(composer)
@@ -226,11 +295,22 @@ def pre_quantize_block(self, ctx: "BlockContext") -> None:
         # The compressor sets ``layer_config`` after ``prepare_run``; keep the
         # QDQ service in sync before it is used for the grid-search / clip loss.
         self._qdq_tool.layer_config = self.layer_config
-        self._smooth_block(block_name, block_mappings)
+        active_mappings = [m for m in block_mappings if not self._mapping_has_ignored_layer(m)]
+        skipped = len(block_mappings) - len(active_mappings)
+        if skipped:
+            logger.warning_once(
+                "AWQ: skipped %d smoothing mapping(s) in block '%s' that include "
+                "ignore_layers / full-precision layers (kept pure).",
+                skipped,
+                block_name,
+            )
+        if not active_mappings:
+            return
+        self._smooth_block(block_name, active_mappings)
         if self.apply_clip:
-            self._clip_block(block_name, block_mappings)
+            self._clip_block(block_name, active_mappings)
         modified = []
-        for mapping in block_mappings:
+        for mapping in active_mappings:
             modified.extend(mapping.balance_names)
             modified.append(mapping.smooth_name)
 
@@ -260,7 +340,6 @@ def finalize_run(self) -> None:
             return
         self._activation_stats.clear()
         self._parent_args_cache.clear()
-        self._parent_signatures.clear()
         self._clip_input_feat.clear()
         self._finalized = True
         logger.debug("AWQ: finalize_quantization complete.")
@@ -365,41 +444,34 @@ def hook_fn(mod, args):
             def _make_parent_hook(parent_module: torch.nn.Module):
 
                 def hook_fn(mod, args, kwargs):
-                    cls_id = id(type(mod))
-                    if cls_id not in self._parent_signatures:
-                        self._parent_signatures[cls_id] = inspect.signature(mod.forward)
-                    sig = self._parent_signatures[cls_id]
-                    try:
-                        bound = sig.bind(*args, **kwargs)
-                        bound.apply_defaults()
-                    except TypeError:
-                        return  # signature mismatch; skip this sample
-
                     param = next(mod.parameters(), None)
                     w_dtype = param.dtype if param is not None else None
 
-                    stored: dict[str, Any] = {}
-                    for k, v in bound.arguments.items():
+                    def _proc(v):
                         if isinstance(v, torch.Tensor):
                             v = v.detach()
                             if w_dtype and v.is_floating_point() and v.dtype != w_dtype:
                                 v = v.to(w_dtype)
-                            stored[k] = v
-                        elif isinstance(v, tuple) and any(isinstance(t, torch.Tensor) for t in v):
-                            stored[k] = tuple(
+                            return v.to("cpu", non_blocking=False)
+                        if isinstance(v, tuple) and any(isinstance(t, torch.Tensor) for t in v):
+                            return tuple(
                                 (
-                                    t.detach().to(w_dtype)
+                                    t.detach().to(w_dtype).to("cpu")
                                     if (w_dtype and isinstance(t, torch.Tensor) and t.is_floating_point())
-                                    else (t.detach() if isinstance(t, torch.Tensor) else t)
+                                    else (t.detach().to("cpu") if isinstance(t, torch.Tensor) else t)
                                 )
                                 for t in v
                             )
-                        elif hasattr(v, "key_cache"):
-                            stored[k] = None  # Null out KV cache objects
-                        else:
-                            stored[k] = v
+                        if hasattr(v, "key_cache"):
+                            return None
+                        return v
+
+                    proc_args = tuple(_proc(a) for a in args)
+                    proc_kwargs = {k: _proc(v) for k, v in kwargs.items()}
 
-                    self._parent_args_cache[parent_module].append(stored)
+                    if self._smooth_seqlen is not None:
+                        proc_args, proc_kwargs = _truncate_args_kwargs(proc_args, proc_kwargs, self._smooth_seqlen)
+                    self._parent_args_cache[parent_module].append((proc_args, proc_kwargs))
 
                 return hook_fn
 
@@ -410,6 +482,23 @@ def hook_fn(mod, args, kwargs):
 
     # ── Smoothing (grid search + scale apply) ─────────────────────────────────
 
+    def _mapping_has_ignored_layer(self, mapping: ResolvedMapping) -> bool:
+        """Return True if the smooth layer or any balance layer is kept full precision."""
+        layer_config = self._qdq_tool.layer_config or {}
+        if not layer_config:
+            return False
+
+        def _is_fp(layer: torch.nn.Module) -> bool:
+            name = getattr(layer, "global_name", None)
+            if not name or name not in layer_config:
+                return False
+            bits = layer_config[name].get("bits", None)
+            return bits is not None and bits >= 16
+
+        if _is_fp(mapping.smooth_layer):
+            return True
+        return any(_is_fp(bl) for bl in mapping.balance_layers)
+
     def _smooth_block(self, block_prefix: str, block_mappings: list) -> None:
         """Run grid search and apply AWQ scales for one block.
 
@@ -587,11 +676,23 @@ def _grid_search_scales(
     def _run_parent_samples(
         self,
         parent: torch.nn.Module,
-        kwargs_list: list[dict],
+        kwargs_list: list[tuple[tuple, dict]],
     ) -> list[torch.Tensor]:
+        param = next(parent.parameters(), None)
+        device = param.device if param is not None else torch.device("cpu")
+
+        def _to_device(v):
+            if isinstance(v, torch.Tensor):
+                return v.to(device)
+            if isinstance(v, (tuple, list)):
+                return type(v)(_to_device(t) for t in v)
+            return v
+
         outputs = []
-        for stored_kwargs in kwargs_list:
-            out = parent(**stored_kwargs)
+        for stored_args, stored_kwargs in kwargs_list:
+            call_args = tuple(_to_device(a) for a in stored_args)
+            call_kwargs = {k: _to_device(v) for k, v in stored_kwargs.items()}
+            out = parent(*call_args, **call_kwargs)
             if isinstance(out, tuple):
                 out = out[0]
             outputs.append(out)
diff --git a/auto_round/algorithms/transforms/awq/config.py b/auto_round/algorithms/transforms/awq/config.py
index b5b6207355..93ea115f01 100644
--- a/auto_round/algorithms/transforms/awq/config.py
+++ b/auto_round/algorithms/transforms/awq/config.py
@@ -47,6 +47,8 @@ def __init__(
         clip_n_grid: int = 20,
         clip_max_shrink: float = 0.5,
         clip_n_sample_token: int = 512,
+        smooth_seqlen: int = 512,
+        skip_moe: bool = True,
         mappings: list[dict] | None = None,
         **kwargs,
     ):
@@ -104,6 +106,17 @@ def __init__(
             clip_n_sample_token: Maximum number of calibration tokens used per
                 balance layer when searching the clip threshold (subsampled to
                 bound memory).
+            smooth_seqlen: Maximum sequence length (number of tokens) used per
+                calibration sample during the AWQ scale grid search. Defaults to
+                ``512``. Set a larger positive integer to use longer sequences,
+                or a value ``<= 0`` to disable truncation entirely.
+            skip_moe: Whether to exclude routed MoE experts from AWQ smoothing.
+                When True, balance layers belonging to routed experts (module
+                names matching ``.experts.<N>.``) are dropped from the resolved
+                mappings, so AWQ only smooths attention and dense/shared paths
+                and leaves routed experts to the downstream block quantizer. This
+                has no effect on dense models and is ignored when explicit
+                ``mappings`` are provided.
             mappings: Optional explicit AWQ smooth/balance mappings. Each
                 item should contain ``smooth_layer`` and
                 ``balance_layers`` entries. If None, mappings are inferred
@@ -138,6 +151,8 @@ def __init__(
         self.clip_n_grid = clip_n_grid
         self.clip_max_shrink = clip_max_shrink
         self.clip_n_sample_token = clip_n_sample_token
+        self.smooth_seqlen = smooth_seqlen
+        self.skip_moe = skip_moe
         self.mappings = mappings
         self.infer_bs_coeff = 1
         self.batch_dim = None
@@ -159,6 +174,7 @@ def __repr__(self) -> str:
             f"AWQConfig(duo_scaling={self.duo_scaling!r}, n_grid={self.n_grid}, "
             f"smooth_iters={self.smooth_iters}, "
             f"apply_clip={self.apply_clip}, clip_as_init={self.clip_as_init}, "
+            f"smooth_seqlen={self.smooth_seqlen}, skip_moe={self.skip_moe}, "
             f"bits={self.bits}, group_size={self.group_size}, sym={self.sym}, "
             f"mappings={'<explicit>' if self.mappings else 'auto'})"
         )
diff --git a/auto_round/algorithms/transforms/awq/mappings.py b/auto_round/algorithms/transforms/awq/mappings.py
index d7f6fd1474..124d3b1d79 100644
--- a/auto_round/algorithms/transforms/awq/mappings.py
+++ b/auto_round/algorithms/transforms/awq/mappings.py
@@ -70,6 +70,10 @@ class ResolvedMapping:
     activation_hook_target: str | None = None
 
 
+# Matches routed MoE expert modules, e.g. "...mlp.experts.3.gate_proj".
+_ROUTED_EXPERT_RE = re.compile(r"\.experts\.\d+\.")
+
+
 # ── Mapping definitions ─────────────────────────
 # Reference: vllm-project/llm-compressor src/llmcompressor/modifiers/awq/mappings.py
 
@@ -435,6 +439,7 @@ def _get_mappings_for_model(model: torch.nn.Module) -> list[AWQMapping]:
 def resolve_mappings(
     model: torch.nn.Module,
     user_mappings: list[dict] | None = None,
+    skip_moe: bool = False,
 ) -> list[ResolvedMapping]:
     """Resolve AWQ mappings for the given model.
 
@@ -444,6 +449,14 @@ def resolve_mappings(
       2. ``AWQ_MAPPING_REGISTRY`` — model-class-name lookup
       3. ``default_mappings`` — Llama-like fallback.
 
+    Args:
+        model: The model to resolve mappings against.
+        user_mappings: Optional explicit mappings; when provided they are used
+            verbatim and ``skip_moe`` is ignored.
+        skip_moe: When True, drop routed MoE experts (module names matching
+            ``.experts.<N>.``) from the resolved mappings so AWQ leaves each
+            routed expert to the downstream block quantizer.
+
     Returns:
         List of ``ResolvedMapping`` objects ready for AWQ grid search.
     """
@@ -452,7 +465,36 @@ def resolve_mappings(
     else:
         mapping_defs = _get_mappings_for_model(model)
 
-    return _resolve_mapping_defs(model, mapping_defs)
+    resolved = _resolve_mapping_defs(model, mapping_defs)
+
+    if skip_moe and user_mappings is None:
+        resolved = _drop_routed_experts(model, resolved)
+
+    return resolved
+
+
+def _drop_routed_experts(model: torch.nn.Module, resolved: list[ResolvedMapping]) -> list[ResolvedMapping]:
+    """Remove routed MoE experts from resolved mappings for ``skip_moe``."""
+    kept: list[ResolvedMapping] = []
+    dropped_experts = 0
+    for mapping in resolved:
+        if _ROUTED_EXPERT_RE.search(mapping.smooth_name):
+            dropped_experts += len(mapping.balance_names)
+            continue
+
+        keep_idx = [i for i, name in enumerate(mapping.balance_names) if not _ROUTED_EXPERT_RE.search(name)]
+        dropped_experts += len(mapping.balance_names) - len(keep_idx)
+        if not keep_idx:
+            continue
+        if len(keep_idx) < len(mapping.balance_names):
+            mapping.balance_names = [mapping.balance_names[i] for i in keep_idx]
+            mapping.balance_layers = [mapping.balance_layers[i] for i in keep_idx]
+            mapping.parent_name, mapping.parent = _find_parent(model, mapping.balance_names)
+        kept.append(mapping)
+
+    if dropped_experts:
+        logger.info(f"AWQ skip_moe: excluded {dropped_experts} routed-expert balance layer(s) from smoothing.")
+    return kept
 
 
 def _resolve_mapping_defs(
@@ -554,13 +596,7 @@ def _resolve_mapping_defs(
             "AWQConfig(mappings=[...])."
         )
     else:
-        first_prefix = next(iter(block_modules))
-        n_blocks = len(block_modules)
-        mappings_per_block = sum(1 for r in resolved if r.smooth_name.startswith(first_prefix))
-        logger.info(
-            f"AWQ resolved {matched_count} smooth-balance mappings "
-            f"({mappings_per_block} per block × {n_blocks} blocks)."
-        )
+        logger.info(f"AWQ resolved {matched_count} smooth-balance mappings.")
 
     return resolved
 
@@ -585,11 +621,11 @@ def check_model_compatibility(
     """
     warnings_list = []
     cls_name = _get_model_class_name(model)
-    in_registry = cls_name in AWQ_MAPPING_REGISTRY
+    in_registry = cls_name in AWQ_MAPPING_REGISTRY or cls_name in AWQ_DYNAMIC_MAPPING_REGISTRY
 
     if not in_registry and user_mappings is None:
         warnings_list.append(
-            f"Model class '{cls_name}' is not in AWQ_MAPPING_REGISTRY. "
+            f"Model class '{cls_name}' is not in any AWQ mapping registry. "
             f"Using default Llama-like mappings. If quantization quality is "
             f"poor, provide explicit mappings via AWQConfig(mappings=[...])."
         )
diff --git a/docs/step_by_step.md b/docs/step_by_step.md
index 74a178834c..1d5cc1046c 100644
--- a/docs/step_by_step.md
+++ b/docs/step_by_step.md
@@ -344,11 +344,11 @@ W2G64 Average Accuracy of 13 tasks and Time Cost Results(Testing was conducted o
 
 ### AWQ Algorithm
 
-**Experimental feature: our current implementation does not apply weight clipping yet, so accuracy may drop compared to the original AWQ algorithm.**
+**Experimental feature:** AWQ weight clipping is optional. Enable it with `--awq-apply-clip` when you want to match the original AWQ flow more closely.
 
 AWQ (Activation-Aware Weight Quantization) is available as an alternative quantization algorithm. AWQ protects salient weight channels by analyzing activation patterns and applying channel-wise scaling before standard RTN quantization.
 
-The canonical AWQ deployment path is **W4A16** served by vLLM's AWQ/Marlin CUDA kernels. **W8A8** with AWQ smoothing can also be served via vLLM's compressed_tensors backend (cutlass INT8 GEMM).
+The canonical AWQ deployment path is **W4A16** served by vLLM's AWQ/Marlin CUDA kernels. **INT8** is AutoRound's W8A8 scheme and can use AWQ smoothing before RTN quantization for vLLM's compressed_tensors backend (cutlass INT8 GEMM).
 
 #### CLI Usage
 
@@ -356,19 +356,41 @@ The canonical AWQ deployment path is **W4A16** served by vLLM's AWQ/Marlin CUDA
 auto-round --model Qwen/Qwen3-0.6B --scheme "W4A16" --algorithm awq --format "auto_round"
 ```
 
+INT8/W8A8 with AWQ smoothing and RTN:
+
+```bash
+auto-round \
+  --model meta-llama/Llama-3.1-8B-Instruct \
+  --scheme INT8 \
+  --algorithm awq,rtn \
+  --nsamples 256 \
+  --seqlen 512 \
+  --awq-apply-clip \
+  --format auto_round:llm_compressor
+```
+
+For `INT8`, `disable_opt_rtn` defaults to `True`, so the command above uses plain RTN without requiring `--disable_opt_rtn`.
+
 AWQ-specific options:
-- `--duo_scaling`: Use both activations and weights for scaling. Options: `true`, `false`, or `both` (searches both modes and picks the best). (default: True).
-- `--n_grid`: Number of grid points for scaling ratio search (default: 20).
+- `--awq-duo-scaling`: Use both activations and weights for scaling. Options: `true`, `false`, or `both` (searches both modes and picks the best). (default: True).
+- `--awq-n-grid`: Number of grid points for scaling ratio search (default: 20).
+- `--awq-apply-clip`: Search and apply AWQ weight clipping after smoothing.
+
+API-only AWQ options:
+- `AWQConfig(smooth_seqlen=512)`: Caps the parent-forward replay length used by AWQ scale search. Set a value `<= 0` to use the full calibration sequence.
+- `AWQConfig(skip_moe=True)`: Skips routed MoE experts during AWQ smoothing while keeping attention and dense/shared paths. Explicit `mappings` are used as provided.
 
 #### API Usage
 
 ```python
-from auto_round import AutoRound
+from auto_round import AWQConfig, AutoRound, RTNConfig
 
 ar = AutoRound(
-    "Qwen/Qwen3-0.6B",
+    "meta-llama/Llama-3.1-8B-Instruct",
     scheme="INT8",
-    algorithm="awq",
+    alg_configs=[AWQConfig(apply_clip=True), RTNConfig()],
+    nsamples=256,
+    seqlen=512,
 )
 
 output_dir = "./tmp_awq"
@@ -499,20 +521,40 @@ auto-round --model Qwen/Qwen3-0.6B --algorithm awq --scheme W4A16
 # AWQ + AutoRound optimization
 auto-round --model Qwen/Qwen3-0.6B --algorithm awq,auto_round --scheme W4A16
 
+# INT8/W8A8 + AWQ + RTN. disable_opt_rtn defaults to True for INT8.
+auto-round \
+  --model meta-llama/Llama-3.1-8B-Instruct \
+  --scheme INT8 \
+  --algorithm awq,rtn \
+  --nsamples 256 \
+  --seqlen 512 \
+  --awq-apply-clip \
+  --format auto_round:llm_compressor
+
 # AWQ flags
---duo-scaling true|false|both  (default: true)
---n-grid 20                    (default: 20)
+--awq-duo-scaling true|false|both  (default: true)
+--awq-n-grid 20                    (default: 20)
+--awq-apply-clip
 ```
 
+`AWQConfig` also supports `smooth_seqlen=512` to cap AWQ scale-search replay length and `skip_moe=True` to leave routed MoE experts to the downstream block quantizer.
+
 #### API Usage
 ```python
-from auto_round import AutoRound
-from auto_round.algorithms.quantization.awq.config import AWQConfig
-from auto_round.algorithms.quantization.sign_round.config import SignRoundConfig
+from auto_round import AWQConfig, AutoRound, RTNConfig, SignRoundConfig
 
 # AWQ + default RTN (simplest)
 ar = AutoRound(model, tokenizer, algorithm="awq", scheme="W4A16")
 
+# INT8/W8A8 + AWQ + RTN
+ar = AutoRound(
+    "meta-llama/Llama-3.1-8B-Instruct",
+    alg_configs=[AWQConfig(apply_clip=True), RTNConfig()],
+    scheme="INT8",
+    nsamples=256,
+    seqlen=512,
+)
+
 # AWQ + AutoRound via alg_configs (explicit pipeline)
 ar = AutoRound(model, tokenizer, alg_configs=[AWQConfig(), SignRoundConfig(iters=200)], scheme="W4A16")
 ar.quantize_and_save(output_dir="./qmodel")
diff --git a/docs/step_by_step_CN.md b/docs/step_by_step_CN.md
index 9163b93d8c..0d346bbb7e 100644
--- a/docs/step_by_step_CN.md
+++ b/docs/step_by_step_CN.md
@@ -339,11 +339,11 @@ W2G64 在 13 个任务上的平均精度与耗时
 
 ### AWQ 算法
 
-实验性功能：原始实现中未使用 weight clipping（权重裁剪）逻辑，因此相比原版 AWQ 算法，可能会存在一定精度下降
+实验性功能：AWQ weight clipping（权重裁剪）为可选功能。如需更接近原始 AWQ 流程，可使用 `--awq-apply-clip` 开启。
 
 AWQ（Activation-Aware Weight Quantization，激活感知权重量化）是一种可选的量化算法。AWQ 通过分析激活模式来保护关键权重通道，在标准量化前对权重施加通道级缩放，从而降低量化误差。
 
-AWQ 的标准部署路径是 **W4A16**，通过 vLLM 的 AWQ/Marlin CUDA 内核提供服务。**W8A8** 搭配 AWQ 平滑化也可通过 vLLM 的 compressed_tensors 后端（cutlass INT8 GEMM）提供服务。
+AWQ 的标准部署路径是 **W4A16**，通过 vLLM 的 AWQ/Marlin CUDA 内核提供服务。**INT8** 是 AutoRound 的 W8A8 scheme，可在 RTN 量化前使用 AWQ 平滑化，并通过 vLLM 的 compressed_tensors 后端（cutlass INT8 GEMM）提供服务。
 
 #### 命令行用法
 
@@ -351,21 +351,43 @@ AWQ 的标准部署路径是 **W4A16**，通过 vLLM 的 AWQ/Marlin CUDA 内核
 auto-round --model Qwen/Qwen3-0.6B --scheme "W4A16" --algorithm awq --format "auto_round"
 ```
 
+INT8/W8A8 搭配 AWQ 平滑化和 RTN：
+
+```bash
+auto-round \
+  --model meta-llama/Llama-3.1-8B-Instruct \
+  --scheme INT8 \
+  --algorithm awq,rtn \
+  --nsamples 256 \
+  --seqlen 512 \
+  --awq-apply-clip \
+  --format auto_round:llm_compressor
+```
+
+对于 `INT8`，`disable_opt_rtn` 默认就是 `True`，因此上面的命令无需额外传入 `--disable_opt_rtn` 即可使用原始 RTN。
+
 AWQ 专用选项：
-- `--duo_scaling`：同时使用激活和权重计算缩放因子。选项：`true`、`false` 或 `both`（搜索两种模式并选择最佳）。（默认：True）。
-- `--n_grid`：缩放比率搜索的网格点数（默认：20）。
+- `--awq-duo-scaling`：同时使用激活和权重计算缩放因子。选项：`true`、`false` 或 `both`（搜索两种模式并选择最佳）。（默认：True）。
+- `--awq-n-grid`：缩放比率搜索的网格点数（默认：20）。
+- `--awq-apply-clip`：在 AWQ 平滑后搜索并应用权重裁剪。
+
+仅 API 支持的 AWQ 选项：
+- `AWQConfig(smooth_seqlen=512)`：限制 AWQ scale search 中 parent-forward replay 使用的序列长度。设为 `<= 0` 时使用完整标定序列。
+- `AWQConfig(skip_moe=True)`：AWQ 平滑时跳过 routed MoE experts，仅保留 attention 和 dense/shared 路径。显式传入的 `mappings` 会按原样使用。
 
 #### API 用法
 
-W8A8 搭配 AWQ 平滑化：
+INT8/W8A8 搭配 AWQ 平滑化：
 
 ```python
-from auto_round import AutoRound
+from auto_round import AWQConfig, AutoRound, RTNConfig
 
 ar = AutoRound(
-    "Qwen/Qwen3-0.6B",
+    "meta-llama/Llama-3.1-8B-Instruct",
     scheme="INT8",
-    algorithm="awq",
+    alg_configs=[AWQConfig(apply_clip=True), RTNConfig()],
+    nsamples=256,
+    seqlen=512,
 )
 
 output_dir = "./tmp_awq"
@@ -497,20 +519,40 @@ auto-round --model Qwen/Qwen3-0.6B --algorithm awq --scheme W4A16
 # AWQ + AutoRound 优化
 auto-round --model Qwen/Qwen3-0.6B --algorithm awq,auto_round --scheme W4A16
 
+# INT8/W8A8 + AWQ + RTN。INT8 默认 disable_opt_rtn=True。
+auto-round \
+  --model meta-llama/Llama-3.1-8B-Instruct \
+  --scheme INT8 \
+  --algorithm awq,rtn \
+  --nsamples 256 \
+  --seqlen 512 \
+  --awq-apply-clip \
+  --format auto_round:llm_compressor
+
 # AWQ 相关参数
---duo-scaling true|false|both  (默认: true)
---n-grid 20                    (默认: 20)
+--awq-duo-scaling true|false|both  (默认: true)
+--awq-n-grid 20                    (默认: 20)
+--awq-apply-clip
 ```
 
+`AWQConfig` 还支持 `smooth_seqlen=512` 用于限制 AWQ scale search 的 replay 序列长度，并支持 `skip_moe=True` 将 routed MoE experts 交给后续 block quantizer 处理。
+
 #### API 用法
 ```python
-from auto_round import AutoRound
-from auto_round.algorithms.quantization.awq.config import AWQConfig
-from auto_round.algorithms.quantization.sign_round.config import SignRoundConfig
+from auto_round import AWQConfig, AutoRound, RTNConfig, SignRoundConfig
 
 # AWQ + 默认 RTN (最简用法)
 ar = AutoRound(model, tokenizer, algorithm="awq", scheme="W4A16")
 
+# INT8/W8A8 + AWQ + RTN
+ar = AutoRound(
+    "meta-llama/Llama-3.1-8B-Instruct",
+    alg_configs=[AWQConfig(apply_clip=True), RTNConfig()],
+    scheme="INT8",
+    nsamples=256,
+    seqlen=512,
+)
+
 # 通过 alg_configs 指定 AWQ + AutoRound (显式流水线)
 ar = AutoRound(model, tokenizer, alg_configs=[AWQConfig(), SignRoundConfig(iters=200)], scheme="W4A16")
 ar.quantize_and_save(output_dir="./qmodel")
diff --git a/test/test_cpu/algorithms/test_awq.py b/test/test_cpu/algorithms/test_awq.py
index 32966bd079..890e713c5c 100644
--- a/test/test_cpu/algorithms/test_awq.py
+++ b/test/test_cpu/algorithms/test_awq.py
@@ -241,6 +241,117 @@ def test_awq_moe_dynamic_smoothing(self, tiny_qwen_moe_model_path):
 
         del model
 
+    def test_awq_moe_skip_moe(self, tiny_qwen_moe_model_path):
+        """skip_moe should drop routed-expert balance layers/mappings but keep dense paths."""
+        import torch.nn as nn
+
+        from auto_round.algorithms.transforms.awq.mappings import ResolvedMapping, _drop_routed_experts
+
+        model = nn.Module()
+        ln = nn.LayerNorm(8)
+        q = nn.Linear(8, 8, bias=False)
+        shared_gate = nn.Linear(8, 16, bias=False)
+        e0_gate = nn.Linear(8, 16, bias=False)
+        e1_gate = nn.Linear(8, 16, bias=False)
+        e0_up = nn.Linear(8, 16, bias=False)
+        e0_down = nn.Linear(16, 8, bias=False)
+
+        resolved = [
+            ResolvedMapping(
+                smooth_name="model.layers.0.input_layernorm",
+                smooth_layer=ln,
+                balance_names=["model.layers.0.self_attn.q_proj"],
+                balance_layers=[q],
+                parent_name="model.layers.0.self_attn",
+                parent=nn.Module(),
+            ),
+            ResolvedMapping(
+                smooth_name="model.layers.0.post_attention_layernorm",
+                smooth_layer=ln,
+                balance_names=[
+                    "model.layers.0.mlp.shared_expert.gate_proj",
+                    "model.layers.0.mlp.experts.0.gate_proj",
+                    "model.layers.0.mlp.experts.1.gate_proj",
+                ],
+                balance_layers=[shared_gate, e0_gate, e1_gate],
+                parent_name="model.layers.0.mlp",
+                parent=nn.Module(),
+            ),
+            ResolvedMapping(
+                smooth_name="model.layers.0.mlp.experts.0.up_proj",
+                smooth_layer=e0_up,
+                balance_names=["model.layers.0.mlp.experts.0.down_proj"],
+                balance_layers=[e0_down],
+                parent_name="model.layers.0.mlp.experts.0",
+                parent=nn.Module(),
+            ),
+        ]
+
+        kept = _drop_routed_experts(model, resolved)
+        smooth_names = [m.smooth_name for m in kept]
+
+        assert "model.layers.0.mlp.experts.0.up_proj" not in smooth_names
+        assert "model.layers.0.input_layernorm" in smooth_names
+        mixed = next(m for m in kept if m.smooth_name.endswith("post_attention_layernorm"))
+        assert mixed.balance_names == ["model.layers.0.mlp.shared_expert.gate_proj"]
+
+        del model
+
+    def test_awq_ignored_layer_skips_mapping(self):
+        """A mapping containing an ignore_layers / bits>=16 layer is skipped so it stays pure."""
+        import torch.nn as nn
+
+        from auto_round.algorithms.transforms.awq.base import AWQTransform
+        from auto_round.algorithms.transforms.awq.config import AWQConfig
+        from auto_round.algorithms.transforms.awq.mappings import ResolvedMapping
+
+        transform = AWQTransform(AWQConfig(bits=4, group_size=128, sym=True, data_type="int"))
+
+        ln = nn.LayerNorm(8)
+        q = nn.Linear(8, 8, bias=False)
+        q.global_name = "model.layers.0.self_attn.q_proj"
+        k = nn.Linear(8, 8, bias=False)
+        k.global_name = "model.layers.0.self_attn.k_proj"
+
+        mapping = ResolvedMapping(
+            smooth_name="model.layers.0.input_layernorm",
+            smooth_layer=ln,
+            balance_names=[q.global_name, k.global_name],
+            balance_layers=[q, k],
+            parent_name="model.layers.0.self_attn",
+            parent=nn.Module(),
+        )
+
+        transform._qdq_tool.layer_config = {}
+        assert transform._mapping_has_ignored_layer(mapping) is False
+
+        transform._qdq_tool.layer_config = {q.global_name: {"bits": 4}, k.global_name: {"bits": 4}}
+        assert transform._mapping_has_ignored_layer(mapping) is False
+
+        transform._qdq_tool.layer_config = {q.global_name: {"bits": 4}, k.global_name: {"bits": 16}}
+        assert transform._mapping_has_ignored_layer(mapping) is True
+
+    def test_awq_smooth_seqlen_truncates_parent_forward_inputs(self):
+        """smooth_seqlen replay cache should truncate matching sequence dimensions consistently."""
+        from auto_round.algorithms.transforms.awq.base import _truncate_args_kwargs
+
+        hidden_states = torch.zeros(1, 16, 8)
+        position_ids = torch.arange(16).reshape(1, 16)
+        attention_mask = torch.zeros(1, 1, 16, 16)
+        rotary = (torch.zeros(1, 16, 8), torch.ones(1, 16, 8))
+
+        args, kwargs = _truncate_args_kwargs(
+            (hidden_states,),
+            {"position_ids": position_ids, "attention_mask": attention_mask, "position_embeddings": rotary},
+            seqlen=4,
+        )
+
+        assert args[0].shape == (1, 4, 8)
+        assert kwargs["position_ids"].shape == (1, 4)
+        assert kwargs["attention_mask"].shape == (1, 1, 4, 4)
+        assert kwargs["position_embeddings"][0].shape == (1, 4, 8)
+        assert kwargs["position_embeddings"][1].shape == (1, 4, 8)
+
     def test_awq_moe_quantized_layers_check(self, tiny_qwen_moe_model_path):
         """AWQ on MoE: expert layers should be quantized, gates/routers stay FP."""
         ar = AutoRound(

From b457c3aacf4ebdf84a8917c3e440e98fefd27a9e Mon Sep 17 00:00:00 2001
From: WeiweiZhang1 <weiwei1.zhang@intel.com>
Date: Fri, 7 Aug 2026 14:41:57 +0800
Subject: [PATCH 2/4] check mapping for mixed_config setting

Signed-off-by: WeiweiZhang1 <weiwei1.zhang@intel.com>
---
 auto_round/algorithms/transforms/awq/base.py  | 358 +++++++++++++-----
 .../algorithms/transforms/awq/config.py       |  12 +-
 .../algorithms/transforms/awq/mappings.py     |  26 +-
 auto_round/cli/algorithms.py                  |  16 +
 auto_round/compressors/entry.py               |   4 +
 docs/step_by_step.md                          |   4 +-
 docs/step_by_step_CN.md                       |   4 +-
 test/test_cpu/algorithms/test_awq.py          | 326 +++++++++++++++-
 8 files changed, 644 insertions(+), 106 deletions(-)

diff --git a/auto_round/algorithms/transforms/awq/base.py b/auto_round/algorithms/transforms/awq/base.py
index cf651209b0..dcb747d3e0 100644
--- a/auto_round/algorithms/transforms/awq/base.py
+++ b/auto_round/algorithms/transforms/awq/base.py
@@ -14,7 +14,7 @@
 """AWQ (Activation-Aware Weight Quantization) quantizer.
 
 Algorithm:
-1. Collect per-channel activation magnitudes during calibration.
+1. Collect per-channel balance-layer input magnitudes during calibration.
 2. For each smooth-balance mapping, perform a grid search over scaling ratios
    to find the one that minimises quantization error (output-based loss).
 3. Apply the best channel-wise scaling:
@@ -117,6 +117,8 @@ def _slice_seq_tensor(v: Any, actual_seq: int, seqlen: int) -> Any:
         return v
     if isinstance(v, (tuple, list)):
         return type(v)(_slice_seq_tensor(t, actual_seq, seqlen) for t in v)
+    if isinstance(v, dict):
+        return {k: _slice_seq_tensor(t, actual_seq, seqlen) for k, t in v.items()}
     return v
 
 
@@ -129,6 +131,10 @@ def _detect_actual_seq(values) -> int | None:
             nested = _detect_actual_seq(v)
             if nested is not None:
                 return nested
+        if isinstance(v, dict):
+            nested = _detect_actual_seq(v.values())
+            if nested is not None:
+                return nested
     return None
 
 
@@ -142,6 +148,50 @@ def _truncate_args_kwargs(args: tuple, kwargs: dict, seqlen: int) -> tuple[tuple
     return new_args, new_kwargs
 
 
+def _detect_actual_batch(values) -> int | None:
+    """Infer batch size from the first tensor with an explicit batch axis."""
+    for v in values:
+        if isinstance(v, torch.Tensor) and v.ndim >= 2:
+            return v.shape[0]
+        if isinstance(v, (tuple, list)):
+            nested = _detect_actual_batch(v)
+            if nested is not None:
+                return nested
+        if isinstance(v, dict):
+            nested = _detect_actual_batch(v.values())
+            if nested is not None:
+                return nested
+    return None
+
+
+def _detect_parent_batch(args: tuple, kwargs: dict) -> int | None:
+    """Infer parent replay batch size, preferring positional model inputs."""
+    actual_batch = _detect_actual_batch(args)
+    if actual_batch is not None:
+        return actual_batch
+    return _detect_actual_batch(kwargs.values())
+
+
+def _slice_batch_value(v: Any, actual_batch: int, start: int, end: int) -> Any:
+    """Slice a value's leading batch dimension, recursing into containers."""
+    if isinstance(v, torch.Tensor):
+        if v.ndim >= 2 and v.shape[0] == actual_batch:
+            return v[start:end]
+        return v
+    if isinstance(v, (tuple, list)):
+        return type(v)(_slice_batch_value(t, actual_batch, start, end) for t in v)
+    if isinstance(v, dict):
+        return {k: _slice_batch_value(t, actual_batch, start, end) for k, t in v.items()}
+    return v
+
+
+def _slice_batch_args_kwargs(args: tuple, kwargs: dict, actual_batch: int, start: int, end: int) -> tuple[tuple, dict]:
+    """Slice parent-forward positional args and kwargs to a batch microchunk."""
+    new_args = tuple(_slice_batch_value(v, actual_batch, start, end) for v in args)
+    new_kwargs = {k: _slice_batch_value(v, actual_batch, start, end) for k, v in kwargs.items()}
+    return new_args, new_kwargs
+
+
 @register_pipeline_member(AWQConfig)
 class AWQTransform(BasePreprocessor):
     """AWQ transform: activation-aware weight smoothing pre-processor.
@@ -168,6 +218,8 @@ def __init__(self, config: AWQConfig) -> None:
         # A value <= 0 disables truncation and uses the full calibration sequence.
         smooth_seqlen = getattr(config, "smooth_seqlen", 512)
         self._smooth_seqlen: int | None = smooth_seqlen if smooth_seqlen > 0 else None
+        smooth_batch_size = getattr(config, "smooth_batch_size", None)
+        self._smooth_batch_size: int | None = smooth_batch_size if smooth_batch_size and smooth_batch_size > 0 else None
 
         # Single source of truth for "QDQ a candidate weight under the target
         # block-quantizer scheme", used as AWQ's grid-search / clip loss. AWQ
@@ -295,12 +347,12 @@ def pre_quantize_block(self, ctx: "BlockContext") -> None:
         # The compressor sets ``layer_config`` after ``prepare_run``; keep the
         # QDQ service in sync before it is used for the grid-search / clip loss.
         self._qdq_tool.layer_config = self.layer_config
-        active_mappings = [m for m in block_mappings if not self._mapping_has_ignored_layer(m)]
+        active_mappings = [m for m in block_mappings if self._mapping_is_smoothable(m)]
         skipped = len(block_mappings) - len(active_mappings)
         if skipped:
             logger.warning_once(
                 "AWQ: skipped %d smoothing mapping(s) in block '%s' that include "
-                "ignore_layers / full-precision layers (kept pure).",
+                "ignore_layers / full-precision layers or incompatible per-layer quantization parameters.",
                 skipped,
                 block_name,
             )
@@ -355,85 +407,88 @@ def _register_awq_hooks(
         """Register activation-stats and parent-kwargs hooks for one block."""
         handles = []
         mappings = self._block_mappings.get(block_name, [])
-        smooth_names = {m.smooth_name for m in mappings}
-
-        # ── Smooth-layer activation-stats hooks ───────────────────────────────
-        # Priority: smooth source forward_hook (output stats).
-        # Each smooth source is hooked exactly once (set de-duplication via name).
-        for name, module in block.named_modules():
-            full_name = f"{block_name}.{name}" if name else block_name
-            if full_name not in smooth_names:
+        module_lookup = dict(model.named_modules())
+
+        def _resolve_activation_hook_layer(mapping: ResolvedMapping) -> torch.nn.Module | None:
+            if not mapping.balance_layers:
+                return None
+
+            hook_target = mapping.activation_hook_target
+            if not hook_target:
+                return mapping.balance_layers[0]
+
+            target_layer = module_lookup.get(hook_target)
+            if target_layer is None and mapping.parent_name:
+                target_layer = module_lookup.get(f"{mapping.parent_name}.{hook_target}")
+            if target_layer is None:
+                try:
+                    target_layer = mapping.parent.get_submodule(hook_target)
+                except AttributeError:
+                    target_layer = None
+            if target_layer is None:
+                logger.warning(
+                    "AWQ: activation_hook_target '%s' for '%s' was not found; using first balance layer '%s'.",
+                    hook_target,
+                    mapping.smooth_name,
+                    mapping.balance_names[0] if mapping.balance_names else "<unknown>",
+                )
+                return mapping.balance_layers[0]
+            return target_layer
+
+        # ── Balance-layer input activation hooks ─────────────────────────────
+        # AWQ scales are derived from the tensor entering the balance layer.  For
+        # gated MLPs, smooth-layer output (for example up_proj(x)) is not the same
+        # tensor consumed by down_proj (act(gate_proj(x)) * up_proj(x)).
+        for mapping in mappings:
+            target_layer = _resolve_activation_hook_layer(mapping)
+            if target_layer is None:
                 continue
 
-            def _make_stats_hook(layer_name: str):
+            def _make_activation_hook(smooth_name: str):
 
-                def hook_fn(mod, args, output):
-                    x = output[0] if isinstance(output, tuple) else output
-                    if x is None or x.numel() == 0:
+                def hook_fn(mod, args):
+                    x = args[0] if isinstance(args, tuple) else args
+                    if x is None or not isinstance(x, torch.Tensor) or x.numel() == 0:
                         return
-                    channel_sum = x.detach().float().flatten(0, -2).abs().sum(dim=0).cpu()
-                    count = x[..., 0].numel()
-                    if layer_name not in self._activation_stats:
-                        self._activation_stats[layer_name] = [
+
+                    feat = x.detach()
+                    if feat.ndim == 1:
+                        feat = feat.view(1, -1)
+                    else:
+                        feat = feat.flatten(0, -2)
+
+                    channel_sum = feat.float().abs().sum(dim=0).cpu()
+                    count = feat.shape[0]
+                    if smooth_name not in self._activation_stats:
+                        self._activation_stats[smooth_name] = [
                             torch.zeros_like(channel_sum),
                             0,
                         ]
-                    self._activation_stats[layer_name][0] += channel_sum
-                    self._activation_stats[layer_name][1] += count
-
-                return hook_fn
-
-            h = module.register_forward_hook(_make_stats_hook(full_name))
-            handles.append(h)
+                    self._activation_stats[smooth_name][0] += channel_sum
+                    self._activation_stats[smooth_name][1] += count
 
-        # ── Clip input-feature hooks (only when apply_clip is enabled) ────────
-        # The balance layers of a mapping share the same input; capture it once
-        # per mapping (keyed by smooth_name) for the post-smooth clip search.
-        if self.apply_clip:
-            for mapping in mappings:
-                if not mapping.balance_layers:
-                    continue
-                target_layer = mapping.balance_layers[0]
-
-                def _make_clip_hook(smooth_name: str):
-
-                    def hook_fn(mod, args):
-                        x = args[0] if isinstance(args, tuple) else args
-                        if x is None or not isinstance(x, torch.Tensor) or x.numel() == 0:
-                            return
-                        feat = x.detach().reshape(-1, x.shape[-1])
+                    if self.apply_clip:
+                        clip_feat = feat
                         # Subsample tokens to bound memory.
-                        if feat.shape[0] > self.clip_n_sample_token:
-                            step = max(1, feat.shape[0] // self.clip_n_sample_token)
-                            feat = feat[::step]
-                        feat = feat.float().cpu()
+                        if clip_feat.shape[0] > self.clip_n_sample_token:
+                            step = max(1, clip_feat.shape[0] // self.clip_n_sample_token)
+                            clip_feat = clip_feat[::step]
+                        clip_feat = clip_feat.float().cpu()
                         prev = self._clip_input_feat.get(smooth_name)
                         if prev is None:
-                            self._clip_input_feat[smooth_name] = feat
+                            self._clip_input_feat[smooth_name] = clip_feat
                         else:
-                            self._clip_input_feat[smooth_name] = torch.cat([prev, feat], dim=0)
+                            self._clip_input_feat[smooth_name] = torch.cat([prev, clip_feat], dim=0)
 
-                    return hook_fn
+                return hook_fn
 
-                h = target_layer.register_forward_pre_hook(_make_clip_hook(mapping.smooth_name))
-                handles.append(h)
+            h = target_layer.register_forward_pre_hook(_make_activation_hook(mapping.smooth_name))
+            handles.append(h)
 
         # One forward_pre_hook per unique parent module in the current block.
         parent_modules_hooked: set[int] = set()
         for mapping in mappings:
             parent = mapping.parent
-            hook_target = mapping.activation_hook_target
-            if hook_target:
-                target_parent = dict(model.named_modules()).get(hook_target)
-                if target_parent is None:
-                    logger.warning(
-                        "AWQ: activation_hook_target '%s' for '%s' was not found; using resolved parent '%s'.",
-                        hook_target,
-                        mapping.smooth_name,
-                        mapping.parent_name,
-                    )
-                else:
-                    parent = target_parent
             if id(parent) in parent_modules_hooked:
                 continue
             parent_modules_hooked.add(id(parent))
@@ -448,22 +503,19 @@ def hook_fn(mod, args, kwargs):
                     w_dtype = param.dtype if param is not None else None
 
                     def _proc(v):
+                        if hasattr(v, "key_cache"):
+                            return None
                         if isinstance(v, torch.Tensor):
                             v = v.detach()
                             if w_dtype and v.is_floating_point() and v.dtype != w_dtype:
                                 v = v.to(w_dtype)
                             return v.to("cpu", non_blocking=False)
-                        if isinstance(v, tuple) and any(isinstance(t, torch.Tensor) for t in v):
-                            return tuple(
-                                (
-                                    t.detach().to(w_dtype).to("cpu")
-                                    if (w_dtype and isinstance(t, torch.Tensor) and t.is_floating_point())
-                                    else (t.detach().to("cpu") if isinstance(t, torch.Tensor) else t)
-                                )
-                                for t in v
-                            )
-                        if hasattr(v, "key_cache"):
-                            return None
+                        if isinstance(v, tuple):
+                            return tuple(_proc(t) for t in v)
+                        if isinstance(v, list):
+                            return [_proc(t) for t in v]
+                        if isinstance(v, dict):
+                            return {k: _proc(t) for k, t in v.items()}
                         return v
 
                     proc_args = tuple(_proc(a) for a in args)
@@ -499,6 +551,50 @@ def _is_fp(layer: torch.nn.Module) -> bool:
             return True
         return any(_is_fp(bl) for bl in mapping.balance_layers)
 
+    @staticmethod
+    def _freeze_quant_param(value):
+        if isinstance(value, list):
+            return tuple(AWQTransform._freeze_quant_param(item) for item in value)
+        if isinstance(value, tuple):
+            return tuple(AWQTransform._freeze_quant_param(item) for item in value)
+        return value
+
+    def _balance_quant_signature(self, layer: torch.nn.Module) -> tuple:
+        """Return the resolved quantization signature that must match within one AWQ mapping."""
+        params = self._qdq_tool.resolve_params(layer)
+        keys = ("bits", "group_size", "sym", "data_type", "super_bits", "super_group_size")
+        return tuple((key, self._freeze_quant_param(params.get(key))) for key in keys)
+
+    def _mapping_has_mixed_quant_params(self, mapping: ResolvedMapping) -> bool:
+        """Return True when balance layers in one AWQ smoothing group do not share quant params."""
+        if len(mapping.balance_layers) <= 1:
+            return False
+
+        signatures = [self._balance_quant_signature(layer) for layer in mapping.balance_layers]
+        first = signatures[0]
+        if all(signature == first for signature in signatures[1:]):
+            return False
+
+        details = {
+            name: dict(signature)
+            for name, signature in zip(mapping.balance_names, signatures)
+        }
+        logger.warning(
+            "AWQ: skipping smoothing for '%s' because balance layers in the same mapping "
+            "have different quantization parameters: %s.",
+            mapping.smooth_name,
+            details,
+        )
+        return True
+
+    def _mapping_is_smoothable(self, mapping: ResolvedMapping) -> bool:
+        """AWQ smoothing is all-or-nothing for layers sharing one smooth scale."""
+        if self._mapping_has_ignored_layer(mapping):
+            return False
+        if self._mapping_has_mixed_quant_params(mapping):
+            return False
+        return True
+
     def _smooth_block(self, block_prefix: str, block_mappings: list) -> None:
         """Run grid search and apply AWQ scales for one block.
 
@@ -592,7 +688,8 @@ def _grid_search_scales(
         device = mapping.balance_layers[0].weight.device
         x_mean = x_mean.to(device)
 
-        group_size = self._normalize_group_size(self._qdq_tool.group_size, -1)
+        bl_params = {bl: self._qdq_tool.resolve_params(bl) for bl in mapping.balance_layers}
+        group_size = self._normalize_group_size(bl_params[mapping.balance_layers[0]]["group_size"], -1)
         if self.duo_scaling is not False:
             w_mean = self._compute_layer_means(mapping.balance_layers, group_size).to(device)
 
@@ -600,7 +697,11 @@ def _grid_search_scales(
         use_parent_forward = len(parent_kwargs_list) > 0
 
         if use_parent_forward:
-            fp16_outputs = self._run_parent_samples(mapping.parent, parent_kwargs_list)
+            fp16_outputs = self._run_parent_samples(
+                mapping.parent,
+                parent_kwargs_list,
+                offload_to_cpu=self._smooth_batch_size is not None,
+            )
             if not fp16_outputs or all(f.numel() == 0 for f in fp16_outputs):
                 use_parent_forward = False
 
@@ -608,11 +709,11 @@ def _grid_search_scales(
         if not use_parent_forward:
             orig_weights = orig_state  # same reference is fine
 
-        # Resolve each balance layer's scheme once, then pre-resolve the quant
-        # functions for the grid-search loop. ``opt_quant_func`` is non-None only
-        # when the SignRoundV2 optimized init-scale path applies for this mapping.
-        bl_params = {bl: self._qdq_tool.resolve_params(bl) for bl in mapping.balance_layers}
-        cached_quant_func, opt_quant_func = self._qdq_tool.resolve_quant_funcs(bl_params[mapping.balance_layers[0]])
+        # Resolve each balance layer's quant functions once, then reuse them in
+        # the grid-search loop. Normal AWQ flow requires one mapping to have
+        # compatible quant params, but keeping this per-layer avoids hidden
+        # coupling to the first layer and makes direct calls robust.
+        bl_quant_funcs = {bl: self._qdq_tool.resolve_quant_funcs(bl_params[bl]) for bl in mapping.balance_layers}
 
         best_error = float("inf")
         best_scales = None
@@ -633,28 +734,28 @@ def _grid_search_scales(
                 # de-smoothed result back, so the parent forward below sees the
                 # weights the layer would actually compute with.
                 for bl in mapping.balance_layers:
+                    quant_func, opt_quant_func = bl_quant_funcs[bl]
                     w_qdq = self._qdq_tool.qdq(
                         orig_state[bl] * scales_view,
                         bl_params[bl],
-                        quant_func=cached_quant_func,
+                        quant_func=quant_func,
                         opt_quant_func=opt_quant_func,
                         imatrix=getattr(bl, "imatrix", None),
                     )
                     bl.weight.data = (w_qdq / scales_view).to(bl.weight.dtype)
 
-                int_w_outputs = self._run_parent_samples(mapping.parent, parent_kwargs_list)
-                total_loss = self._compute_loss(fp16_outputs, int_w_outputs)
-                del int_w_outputs
+                total_loss = self._compute_parent_loss(mapping.parent, parent_kwargs_list, fp16_outputs)
                 for bl in mapping.balance_layers:
                     bl.weight.data.copy_(orig_state[bl])
             else:
                 total_loss = 0.0
                 for bl in mapping.balance_layers:
+                    quant_func, opt_quant_func = bl_quant_funcs[bl]
                     w_orig = orig_weights[bl].to(device)
                     w_qdq = self._qdq_tool.qdq(
                         w_orig * scales_view,
                         bl_params[bl],
-                        quant_func=cached_quant_func,
+                        quant_func=quant_func,
                         opt_quant_func=opt_quant_func,
                         imatrix=getattr(bl, "imatrix", None),
                     )
@@ -672,32 +773,91 @@ def _grid_search_scales(
         logger.debug("AWQ '%s': best_ratio=%.2f, best_error=%.3e", mapping.smooth_name, best_ratio, best_error)
         return best_scales
 
+    def _iter_parent_calls(self, stored_args: tuple, stored_kwargs: dict):
+        """Yield full or microbatched parent-call args from one cached calibration batch."""
+        actual_batch = _detect_parent_batch(stored_args, stored_kwargs)
+        if self._smooth_batch_size is None or actual_batch is None or actual_batch <= self._smooth_batch_size:
+            yield stored_args, stored_kwargs
+            return
+
+        for start in range(0, actual_batch, self._smooth_batch_size):
+            end = min(actual_batch, start + self._smooth_batch_size)
+            yield _slice_batch_args_kwargs(stored_args, stored_kwargs, actual_batch, start, end)
+
+    @staticmethod
+    def _move_parent_value_to_device(v: Any, device: torch.device | str) -> Any:
+        """Move a nested parent-call value to the parent execution device."""
+        if isinstance(v, torch.Tensor):
+            return v.to(device)
+        if isinstance(v, (tuple, list)):
+            return type(v)(AWQTransform._move_parent_value_to_device(t, device) for t in v)
+        if isinstance(v, dict):
+            return {k: AWQTransform._move_parent_value_to_device(t, device) for k, t in v.items()}
+        return v
+
+    @staticmethod
+    def _normalize_parent_output(out: Any) -> torch.Tensor:
+        """Extract the tensor output used by AWQ parent-output loss."""
+        if isinstance(out, tuple):
+            return out[0]
+        return out
+
     @torch.no_grad()
     def _run_parent_samples(
         self,
         parent: torch.nn.Module,
         kwargs_list: list[tuple[tuple, dict]],
+        offload_to_cpu: bool = False,
     ) -> list[torch.Tensor]:
         param = next(parent.parameters(), None)
         device = param.device if param is not None else torch.device("cpu")
 
-        def _to_device(v):
-            if isinstance(v, torch.Tensor):
-                return v.to(device)
-            if isinstance(v, (tuple, list)):
-                return type(v)(_to_device(t) for t in v)
-            return v
-
         outputs = []
         for stored_args, stored_kwargs in kwargs_list:
-            call_args = tuple(_to_device(a) for a in stored_args)
-            call_kwargs = {k: _to_device(v) for k, v in stored_kwargs.items()}
-            out = parent(*call_args, **call_kwargs)
-            if isinstance(out, tuple):
-                out = out[0]
-            outputs.append(out)
+            for micro_args, micro_kwargs in self._iter_parent_calls(stored_args, stored_kwargs):
+                call_args = tuple(self._move_parent_value_to_device(a, device) for a in micro_args)
+                call_kwargs = {k: self._move_parent_value_to_device(v, device) for k, v in micro_kwargs.items()}
+                out = self._normalize_parent_output(parent(*call_args, **call_kwargs)).detach()
+                if offload_to_cpu:
+                    out = out.to("cpu", non_blocking=False)
+                outputs.append(out)
         return outputs
 
+    @torch.no_grad()
+    def _compute_parent_loss(
+        self,
+        parent: torch.nn.Module,
+        kwargs_list: list[tuple[tuple, dict]],
+        fp16_outputs: list[torch.Tensor],
+    ) -> float:
+        """Replay parent samples and stream MSE loss without storing candidate outputs."""
+        param = next(parent.parameters(), None)
+        device = param.device if param is not None else torch.device("cpu")
+
+        loss = torch.tensor(0.0, device=device)
+        num_elements = torch.tensor(0, device=device, dtype=torch.long)
+        output_idx = 0
+        for stored_args, stored_kwargs in kwargs_list:
+            for micro_args, micro_kwargs in self._iter_parent_calls(stored_args, stored_kwargs):
+                if output_idx >= len(fp16_outputs):
+                    return float("inf")
+                call_args = tuple(self._move_parent_value_to_device(a, device) for a in micro_args)
+                call_kwargs = {k: self._move_parent_value_to_device(v, device) for k, v in micro_kwargs.items()}
+                out = self._normalize_parent_output(parent(*call_args, **call_kwargs))
+                fp16_out = fp16_outputs[output_idx].to(device, non_blocking=False)
+                loss += torch.nn.functional.mse_loss(
+                    fp16_out.float(),
+                    out.float(),
+                    reduction="sum",
+                )
+                num_elements += fp16_out.numel()
+                output_idx += 1
+                del out, fp16_out
+
+        if output_idx != len(fp16_outputs) or num_elements == 0:
+            return float("inf")
+        return (loss / num_elements).item()
+
     @staticmethod
     @torch.no_grad()
     def _compute_loss(
diff --git a/auto_round/algorithms/transforms/awq/config.py b/auto_round/algorithms/transforms/awq/config.py
index 93ea115f01..5aa10af290 100644
--- a/auto_round/algorithms/transforms/awq/config.py
+++ b/auto_round/algorithms/transforms/awq/config.py
@@ -48,6 +48,7 @@ def __init__(
         clip_max_shrink: float = 0.5,
         clip_n_sample_token: int = 512,
         smooth_seqlen: int = 512,
+        smooth_batch_size: int | None = None,
         skip_moe: bool = True,
         mappings: list[dict] | None = None,
         **kwargs,
@@ -110,6 +111,11 @@ def __init__(
                 calibration sample during the AWQ scale grid search. Defaults to
                 ``512``. Set a larger positive integer to use longer sequences,
                 or a value ``<= 0`` to disable truncation entirely.
+            smooth_batch_size: Optional microbatch size used when replaying AWQ
+                parent modules during scale grid search. Smaller values reduce
+                peak VRAM while preserving the AWQ parent-output loss, at the
+                cost of more parent forward calls. ``None`` or ``<= 0`` replays
+                the cached calibration batch as-is.
             skip_moe: Whether to exclude routed MoE experts from AWQ smoothing.
                 When True, balance layers belonging to routed experts (module
                 names matching ``.experts.<N>.``) are dropped from the resolved
@@ -148,10 +154,13 @@ def __init__(
             raise ValueError(f"`clip_max_shrink` must be in (0, 1), got {clip_max_shrink!r}")
         if clip_n_sample_token is None or clip_n_sample_token < 1:
             raise ValueError(f"`clip_n_sample_token` must be a positive integer, got {clip_n_sample_token!r}")
+        if smooth_batch_size is not None and smooth_batch_size < 0:
+            raise ValueError(f"`smooth_batch_size` must be a non-negative integer or None, got {smooth_batch_size!r}")
         self.clip_n_grid = clip_n_grid
         self.clip_max_shrink = clip_max_shrink
         self.clip_n_sample_token = clip_n_sample_token
         self.smooth_seqlen = smooth_seqlen
+        self.smooth_batch_size = smooth_batch_size
         self.skip_moe = skip_moe
         self.mappings = mappings
         self.infer_bs_coeff = 1
@@ -174,7 +183,8 @@ def __repr__(self) -> str:
             f"AWQConfig(duo_scaling={self.duo_scaling!r}, n_grid={self.n_grid}, "
             f"smooth_iters={self.smooth_iters}, "
             f"apply_clip={self.apply_clip}, clip_as_init={self.clip_as_init}, "
-            f"smooth_seqlen={self.smooth_seqlen}, skip_moe={self.skip_moe}, "
+            f"smooth_seqlen={self.smooth_seqlen}, smooth_batch_size={self.smooth_batch_size}, "
+            f"skip_moe={self.skip_moe}, "
             f"bits={self.bits}, group_size={self.group_size}, sym={self.sym}, "
             f"mappings={'<explicit>' if self.mappings else 'auto'})"
         )
diff --git a/auto_round/algorithms/transforms/awq/mappings.py b/auto_round/algorithms/transforms/awq/mappings.py
index 124d3b1d79..c04c79417f 100644
--- a/auto_round/algorithms/transforms/awq/mappings.py
+++ b/auto_round/algorithms/transforms/awq/mappings.py
@@ -177,6 +177,15 @@ class ResolvedMapping:
     AWQMapping(r"up_proj$", [r"down_proj$"]),
 ]
 
+_bagel_mappings = [
+    AWQMapping(
+        r"input_layernorm$",
+        [r"\.self_attn\.q_proj$", r"\.self_attn\.k_proj$", r"\.self_attn\.v_proj$"],
+    ),
+    AWQMapping(r"post_attention_layernorm$", [r"\.mlp\.gate_proj$", r"\.mlp\.up_proj$"]),
+    AWQMapping(r"\.mlp\.up_proj$", [r"\.mlp\.down_proj$"]),
+]
+
 # ── Model class name → mappings registry ──────────────────────────────────────
 # Aligned with llm-compressor AWQ_MAPPING_REGISTRY (llmcompressor v0.10.0).
 # Models not in this registry fall back to default_mappings.
@@ -224,6 +233,10 @@ class ResolvedMapping:
     # Other models using default mappings
     "SeedOssForCausalLM": default_mappings,
     "Ernie4_5_MoeForCausalLM": default_mappings,
+    # BAGEL wraps a Qwen2 language model and carries parallel *_moe_gen modules
+    # for image generation. Keep AWQ smoothing on the normal text path only.
+    "BagelForQuantization": _bagel_mappings,
+    "BagelForConditionalGeneration": _bagel_mappings,
 }
 
 
@@ -335,6 +348,14 @@ def _build_hybrid_attention_mappings(model: torch.nn.Module) -> list[AWQMapping]
 
     layer_types, num_layers = result
 
+    if len(layer_types) < num_layers:
+        logger.warning(
+            "Hybrid attention model config has num_hidden_layers=%d but only %d layer_types entries. Falling back.",
+            num_layers,
+            len(layer_types),
+        )
+        return None
+
     full_indices = [i for i in range(num_layers) if layer_types[i] == "full_attention"]
     linear_indices = [i for i in range(num_layers) if layer_types[i] == "linear_attention"]
 
@@ -461,7 +482,10 @@ def resolve_mappings(
         List of ``ResolvedMapping`` objects ready for AWQ grid search.
     """
     if user_mappings is not None:
-        mapping_defs = [AWQMapping(m["smooth_layer"], m["balance_layers"]) for m in user_mappings]
+        mapping_defs = [
+            AWQMapping(m["smooth_layer"], m["balance_layers"], m.get("activation_hook_target"))
+            for m in user_mappings
+        ]
     else:
         mapping_defs = _get_mappings_for_model(model)
 
diff --git a/auto_round/cli/algorithms.py b/auto_round/cli/algorithms.py
index 8291d1d821..9c4f7e2444 100644
--- a/auto_round/cli/algorithms.py
+++ b/auto_round/cli/algorithms.py
@@ -202,6 +202,20 @@ def register(self, group) -> None:
             type=int,
             help="Number of grid-search points for AWQ scaling ratio.",
         )
+        group.add_argument(
+            "--awq-smooth-seqlen",
+            dest="awq_smooth_seqlen",
+            default=None,
+            type=int,
+            help="Maximum sequence length used by AWQ scale-search parent replay.",
+        )
+        group.add_argument(
+            "--awq-smooth-batch-size",
+            dest="awq_smooth_batch_size",
+            default=None,
+            type=int,
+            help="Microbatch size for AWQ parent replay during scale search; <=0 disables microbatching.",
+        )
         group.add_argument(
             "--awq-apply-clip",
             dest="awq_apply_clip",
@@ -226,6 +240,8 @@ def build(self, args, common_kwargs: dict[str, Any]):
             n_grid=getattr(args, "n_grid", 20),
             apply_clip=getattr(args, "awq_apply_clip", False),
             clip_as_init=getattr(args, "awq_clip_as_init", False),
+            smooth_seqlen=getattr(args, "awq_smooth_seqlen", None) or 512,
+            smooth_batch_size=getattr(args, "awq_smooth_batch_size", None),
             **common_kwargs,
         )
 
diff --git a/auto_round/compressors/entry.py b/auto_round/compressors/entry.py
index 776a3049ad..a5e91cce0e 100644
--- a/auto_round/compressors/entry.py
+++ b/auto_round/compressors/entry.py
@@ -565,6 +565,10 @@ def _build_awq_config(
             seqlen=seqlen,
             nsamples=nsamples,
             batch_size=batch_size,
+            apply_clip=kwargs.pop("apply_clip", False),
+            clip_as_init=kwargs.pop("clip_as_init", False),
+            smooth_seqlen=kwargs.pop("smooth_seqlen", 512),
+            smooth_batch_size=kwargs.pop("smooth_batch_size", None),
             mappings=kwargs.pop("mappings", None),
             **common_config_kwargs,
         )
diff --git a/docs/step_by_step.md b/docs/step_by_step.md
index 1d5cc1046c..3f656aa208 100644
--- a/docs/step_by_step.md
+++ b/docs/step_by_step.md
@@ -388,7 +388,7 @@ from auto_round import AWQConfig, AutoRound, RTNConfig
 ar = AutoRound(
     "meta-llama/Llama-3.1-8B-Instruct",
     scheme="INT8",
-    alg_configs=[AWQConfig(apply_clip=True), RTNConfig()],
+    alg_configs=[AWQConfig(apply_clip=True), RTNConfig(disable_opt_rtn=True)],
     nsamples=256,
     seqlen=512,
 )
@@ -549,7 +549,7 @@ ar = AutoRound(model, tokenizer, algorithm="awq", scheme="W4A16")
 # INT8/W8A8 + AWQ + RTN
 ar = AutoRound(
     "meta-llama/Llama-3.1-8B-Instruct",
-    alg_configs=[AWQConfig(apply_clip=True), RTNConfig()],
+    alg_configs=[AWQConfig(apply_clip=True), RTNConfig(disable_opt_rtn=True)],
     scheme="INT8",
     nsamples=256,
     seqlen=512,
diff --git a/docs/step_by_step_CN.md b/docs/step_by_step_CN.md
index 0d346bbb7e..a3e76cc4fd 100644
--- a/docs/step_by_step_CN.md
+++ b/docs/step_by_step_CN.md
@@ -385,7 +385,7 @@ from auto_round import AWQConfig, AutoRound, RTNConfig
 ar = AutoRound(
     "meta-llama/Llama-3.1-8B-Instruct",
     scheme="INT8",
-    alg_configs=[AWQConfig(apply_clip=True), RTNConfig()],
+    alg_configs=[AWQConfig(apply_clip=True), RTNConfig(disable_opt_rtn=True)],
     nsamples=256,
     seqlen=512,
 )
@@ -547,7 +547,7 @@ ar = AutoRound(model, tokenizer, algorithm="awq", scheme="W4A16")
 # INT8/W8A8 + AWQ + RTN
 ar = AutoRound(
     "meta-llama/Llama-3.1-8B-Instruct",
-    alg_configs=[AWQConfig(apply_clip=True), RTNConfig()],
+    alg_configs=[AWQConfig(apply_clip=True), RTNConfig(disable_opt_rtn=True)],
     scheme="INT8",
     nsamples=256,
     seqlen=512,
diff --git a/test/test_cpu/algorithms/test_awq.py b/test/test_cpu/algorithms/test_awq.py
index 890e713c5c..477e4af277 100644
--- a/test/test_cpu/algorithms/test_awq.py
+++ b/test/test_cpu/algorithms/test_awq.py
@@ -297,8 +297,58 @@ def test_awq_moe_skip_moe(self, tiny_qwen_moe_model_path):
 
         del model
 
+    def test_explicit_awq_mapping_preserves_activation_hook_target(self):
+        """Custom AWQ mappings should keep activation_hook_target for non-standard balance inputs."""
+        import torch.nn as nn
+
+        from auto_round.algorithms.transforms.awq.mappings import resolve_mappings
+
+        class TinyBlock(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.smooth = nn.Linear(4, 4, bias=False)
+                self.hook = nn.Linear(4, 4, bias=False)
+                self.balance = nn.Linear(4, 4, bias=False)
+
+        class TinyModel(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.layers = nn.ModuleList([TinyBlock()])
+
+        resolved = resolve_mappings(
+            TinyModel(),
+            user_mappings=[
+                {
+                    "smooth_layer": "smooth$",
+                    "balance_layers": ["balance$"],
+                    "activation_hook_target": "hook",
+                }
+            ],
+        )
+
+        assert len(resolved) == 1
+        assert resolved[0].activation_hook_target == "hook"
+
+    def test_hybrid_attention_mapping_short_layer_types_falls_back(self):
+        """Malformed hybrid configs should fall back instead of raising IndexError."""
+        from types import SimpleNamespace
+
+        import torch.nn as nn
+
+        from auto_round.algorithms.transforms.awq.mappings import _build_hybrid_attention_mappings
+
+        class BadHybridModel(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.config = SimpleNamespace(
+                    layer_types=["full_attention", "linear_attention"],
+                    num_hidden_layers=3,
+                )
+
+        assert _build_hybrid_attention_mappings(BadHybridModel()) is None
+
     def test_awq_ignored_layer_skips_mapping(self):
-        """A mapping containing an ignore_layers / bits>=16 layer is skipped so it stays pure."""
+        """A mapping containing an ignore_layers / bits>=16 layer is skipped as one smooth-scale group."""
         import torch.nn as nn
 
         from auto_round.algorithms.transforms.awq.base import AWQTransform
@@ -330,6 +380,280 @@ def test_awq_ignored_layer_skips_mapping(self):
 
         transform._qdq_tool.layer_config = {q.global_name: {"bits": 4}, k.global_name: {"bits": 16}}
         assert transform._mapping_has_ignored_layer(mapping) is True
+        assert transform._mapping_is_smoothable(mapping) is False
+
+    def test_awq_mixed_balance_quant_params_skip_mapping(self, monkeypatch):
+        """Balance layers sharing one AWQ scale must share the same resolved quantization params."""
+        import torch.nn as nn
+
+        from auto_round.algorithms.transforms.awq import base as awq_base
+        from auto_round.algorithms.transforms.awq.base import AWQTransform
+        from auto_round.algorithms.transforms.awq.config import AWQConfig
+        from auto_round.algorithms.transforms.awq.mappings import ResolvedMapping
+
+        transform = AWQTransform(AWQConfig(bits=4, group_size=128, sym=True, data_type="int"))
+
+        ln = nn.LayerNorm(8)
+        q = nn.Linear(8, 8, bias=False)
+        q.global_name = "model.layers.0.self_attn.q_proj"
+        k = nn.Linear(8, 8, bias=False)
+        k.global_name = "model.layers.0.self_attn.k_proj"
+
+        mapping = ResolvedMapping(
+            smooth_name="model.layers.0.input_layernorm",
+            smooth_layer=ln,
+            balance_names=[q.global_name, k.global_name],
+            balance_layers=[q, k],
+            parent_name="model.layers.0.self_attn",
+            parent=nn.Module(),
+        )
+
+        warnings = []
+
+        def fake_warning(message, *args, **kwargs):
+            warnings.append(message % args if args else message)
+
+        monkeypatch.setattr(awq_base.logger, "warning", fake_warning)
+
+        transform._qdq_tool.layer_config = {
+            q.global_name: {"bits": 4, "group_size": 128, "sym": True, "data_type": "int"},
+            k.global_name: {"bits": 4, "group_size": 128, "sym": True, "data_type": "int"},
+        }
+        assert transform._mapping_has_mixed_quant_params(mapping) is False
+        assert transform._mapping_is_smoothable(mapping) is True
+
+        transform._qdq_tool.layer_config = {
+            q.global_name: {
+                "bits": 4,
+                "group_size": 128,
+                "sym": True,
+                "data_type": "int",
+                "disable_opt_rtn": False,
+            },
+            k.global_name: {
+                "bits": 4,
+                "group_size": 128,
+                "sym": True,
+                "data_type": "int",
+                "disable_opt_rtn": True,
+            },
+        }
+        assert transform._mapping_has_mixed_quant_params(mapping) is False
+        assert transform._mapping_is_smoothable(mapping) is True
+
+        transform._qdq_tool.layer_config = {
+            q.global_name: {"bits": 4, "group_size": 128, "sym": True, "data_type": "int"},
+            k.global_name: {"bits": 4, "group_size": 128, "sym": True, "data_type": "mx_fp"},
+        }
+        assert transform._mapping_has_mixed_quant_params(mapping) is True
+        assert transform._mapping_is_smoothable(mapping) is False
+        assert any("different quantization parameters" in warning for warning in warnings)
+
+    def test_awq_grid_search_uses_per_balance_layer_quant_func(self):
+        """Direct grid search should pass each balance layer's own resolved quant function."""
+        import torch.nn as nn
+
+        from auto_round.algorithms.transforms.awq.base import AWQTransform
+        from auto_round.algorithms.transforms.awq.config import AWQConfig
+        from auto_round.algorithms.transforms.awq.mappings import ResolvedMapping
+
+        transform = AWQTransform(AWQConfig(bits=4, group_size=8, sym=True, data_type="int", duo_scaling=False))
+
+        ln = nn.LayerNorm(8)
+        q = nn.Linear(8, 8, bias=False)
+        q.global_name = "model.layers.0.self_attn.q_proj"
+        k = nn.Linear(8, 8, bias=False)
+        k.global_name = "model.layers.0.self_attn.k_proj"
+        mapping = ResolvedMapping(
+            smooth_name="model.layers.0.input_layernorm",
+            smooth_layer=ln,
+            balance_names=[q.global_name, k.global_name],
+            balance_layers=[q, k],
+            parent_name="model.layers.0.self_attn",
+            parent=nn.Module(),
+        )
+
+        transform._qdq_tool.layer_config = {
+            q.global_name: {"bits": 4, "group_size": 8, "sym": True, "data_type": "int"},
+            k.global_name: {"bits": 8, "group_size": 8, "sym": True, "data_type": "mx_fp"},
+        }
+        records = []
+
+        def fake_resolve_quant_funcs(params):
+            return f"{params['data_type']}_{params['bits']}", None
+
+        def fake_qdq(weight, params, *, quant_func=None, opt_quant_func=None, imatrix=None):
+            records.append((params["data_type"], params["bits"], quant_func))
+            return weight
+
+        transform._qdq_tool.resolve_quant_funcs = fake_resolve_quant_funcs
+        transform._qdq_tool.qdq = fake_qdq
+
+        transform._grid_search_scales(mapping, torch.ones(8))
+
+        assert ("int", 4, "int_4") in records
+        assert ("mx_fp", 8, "mx_fp_8") in records
+
+    def test_awq_activation_stats_use_balance_layer_input_for_gated_mlp(self):
+        """up_proj -> down_proj stats must use down_proj input, not raw up_proj output."""
+        import torch.nn as nn
+
+        from auto_round.algorithms.transforms.awq.base import AWQTransform
+        from auto_round.algorithms.transforms.awq.config import AWQConfig
+        from auto_round.algorithms.transforms.awq.mappings import ResolvedMapping
+
+        class TinyGatedBlock(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.up_proj = nn.Linear(4, 6, bias=False)
+                self.gate_proj = nn.Linear(4, 6, bias=False)
+                self.down_proj = nn.Linear(6, 4, bias=False)
+
+            def forward(self, x):
+                return self.down_proj(torch.nn.functional.silu(self.gate_proj(x)) * self.up_proj(x))
+
+        class TinyModel(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.block = TinyGatedBlock()
+
+            def forward(self, x):
+                return self.block(x)
+
+        torch.manual_seed(0)
+        model = TinyModel()
+        transform = AWQTransform(
+            AWQConfig(bits=4, group_size=2, sym=True, data_type="int", apply_clip=True, clip_n_sample_token=32)
+        )
+        mapping = ResolvedMapping(
+            smooth_name="block.up_proj",
+            smooth_layer=model.block.up_proj,
+            balance_names=["block.down_proj"],
+            balance_layers=[model.block.down_proj],
+            parent_name="block",
+            parent=model.block,
+        )
+        transform._block_mappings = {"block": [mapping]}
+
+        x = torch.randn(2, 3, 4)
+        handles = transform._register_awq_hooks(model, model.block, "block")
+        try:
+            model(x)
+        finally:
+            for handle in handles:
+                handle.remove()
+
+        expected = torch.nn.functional.silu(model.block.gate_proj(x)) * model.block.up_proj(x)
+        expected_feat = expected.detach().flatten(0, -2)
+        raw_up_feat = model.block.up_proj(x).detach().flatten(0, -2)
+
+        act_sum, act_count = transform._activation_stats["block.up_proj"]
+        assert torch.allclose(act_sum, expected_feat.abs().sum(dim=0))
+        assert act_count == expected_feat.shape[0]
+        assert torch.allclose(transform._clip_input_feat["block.up_proj"], expected_feat.float().cpu())
+        assert not torch.allclose(act_sum, raw_up_feat.abs().sum(dim=0))
+
+    def test_awq_parent_cache_recursively_detaches_tensor_containers(self):
+        """Parent replay cache should not retain tensors inside list/dict containers."""
+        import torch.nn as nn
+
+        from auto_round.algorithms.transforms.awq.base import AWQTransform
+        from auto_round.algorithms.transforms.awq.config import AWQConfig
+        from auto_round.algorithms.transforms.awq.mappings import ResolvedMapping
+
+        class ContainerParent(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.smooth = nn.Linear(4, 4, bias=False)
+                self.balance = nn.Linear(4, 4, bias=False)
+
+            def forward(self, items=None, payload=None):
+                return self.balance(items[0]) + self.balance(payload["x"])
+
+        class TinyModel(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.block = ContainerParent()
+
+            def forward(self, items=None, payload=None):
+                return self.block(items=items, payload=payload)
+
+        model = TinyModel()
+        transform = AWQTransform(AWQConfig(bits=4, group_size=4, sym=True, data_type="int"))
+        mapping = ResolvedMapping(
+            smooth_name="block.smooth",
+            smooth_layer=model.block.smooth,
+            balance_names=["block.balance"],
+            balance_layers=[model.block.balance],
+            parent_name="block",
+            parent=model.block,
+        )
+        transform._block_mappings = {"block": [mapping]}
+
+        items = [torch.randn(1, 4, requires_grad=True)]
+        payload = {"x": torch.randn(1, 4, requires_grad=True)}
+        handles = transform._register_awq_hooks(model, model.block, "block")
+        try:
+            model(items=items, payload=payload)
+        finally:
+            for handle in handles:
+                handle.remove()
+
+        _, cached_kwargs = transform._parent_args_cache[model.block][0]
+        assert cached_kwargs["items"] is not items
+        assert cached_kwargs["payload"] is not payload
+        assert cached_kwargs["items"][0].device.type == "cpu"
+        assert cached_kwargs["payload"]["x"].device.type == "cpu"
+        assert cached_kwargs["items"][0].requires_grad is False
+        assert cached_kwargs["payload"]["x"].requires_grad is False
+
+    def test_awq_parent_replay_microbatch_matches_full_batch(self):
+        """AWQ parent replay microbatching should preserve exact parent-output loss."""
+        import torch.nn as nn
+
+        from auto_round.algorithms.transforms.awq.base import AWQTransform
+        from auto_round.algorithms.transforms.awq.config import AWQConfig
+
+        class RecordingParent(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.proj = nn.Linear(4, 4, bias=False)
+                self.batch_sizes = []
+
+            def forward(self, hidden_states, *, position_ids=None, cache_position=None, payload=None):
+                self.batch_sizes.append(hidden_states.shape[0])
+                assert position_ids is not None and position_ids.shape[0] == hidden_states.shape[0]
+                assert cache_position is not None and cache_position.shape == (3,)
+                add = payload["add"] if payload is not None else 0
+                return self.proj(hidden_states + add)
+
+        torch.manual_seed(0)
+        parent = RecordingParent()
+        hidden_states = torch.randn(5, 3, 4)
+        position_ids = torch.arange(15).reshape(5, 3)
+        cache_position = torch.arange(3)
+        payload = {"add": torch.randn(5, 3, 4)}
+        kwargs_list = [((hidden_states,), {"position_ids": position_ids, "cache_position": cache_position, "payload": payload})]
+
+        full = AWQTransform(AWQConfig(bits=4, group_size=4, sym=True, data_type="int"))
+        micro = AWQTransform(AWQConfig(bits=4, group_size=4, sym=True, data_type="int", smooth_batch_size=2))
+
+        full_outputs = full._run_parent_samples(parent, kwargs_list)
+        parent.batch_sizes.clear()
+        micro_outputs = micro._run_parent_samples(parent, kwargs_list)
+
+        assert parent.batch_sizes == [2, 2, 1]
+        assert len(full_outputs) == 1
+        assert len(micro_outputs) == 3
+        assert torch.allclose(torch.cat(micro_outputs, dim=0), full_outputs[0])
+
+        ref_outputs = [out + 0.125 for out in micro_outputs]
+        parent.batch_sizes.clear()
+        streamed_loss = micro._compute_parent_loss(parent, kwargs_list, ref_outputs)
+        expected_loss = micro._compute_loss(ref_outputs, micro_outputs)
+
+        assert parent.batch_sizes == [2, 2, 1]
+        assert streamed_loss == expected_loss
 
     def test_awq_smooth_seqlen_truncates_parent_forward_inputs(self):
         """smooth_seqlen replay cache should truncate matching sequence dimensions consistently."""

From f4dd0e47dcaf64039521b3cef00c3131a883502e Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Fri, 7 Aug 2026 06:51:22 +0000
Subject: [PATCH 3/4] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 auto_round/algorithms/transforms/awq/base.py     | 5 +----
 auto_round/algorithms/transforms/awq/mappings.py | 3 +--
 test/unit/test_cpu/algorithms/test_awq.py        | 4 +++-
 3 files changed, 5 insertions(+), 7 deletions(-)

diff --git a/auto_round/algorithms/transforms/awq/base.py b/auto_round/algorithms/transforms/awq/base.py
index dcb747d3e0..619b919127 100644
--- a/auto_round/algorithms/transforms/awq/base.py
+++ b/auto_round/algorithms/transforms/awq/base.py
@@ -575,10 +575,7 @@ def _mapping_has_mixed_quant_params(self, mapping: ResolvedMapping) -> bool:
         if all(signature == first for signature in signatures[1:]):
             return False
 
-        details = {
-            name: dict(signature)
-            for name, signature in zip(mapping.balance_names, signatures)
-        }
+        details = {name: dict(signature) for name, signature in zip(mapping.balance_names, signatures)}
         logger.warning(
             "AWQ: skipping smoothing for '%s' because balance layers in the same mapping "
             "have different quantization parameters: %s.",
diff --git a/auto_round/algorithms/transforms/awq/mappings.py b/auto_round/algorithms/transforms/awq/mappings.py
index c04c79417f..e122c2bea9 100644
--- a/auto_round/algorithms/transforms/awq/mappings.py
+++ b/auto_round/algorithms/transforms/awq/mappings.py
@@ -483,8 +483,7 @@ def resolve_mappings(
     """
     if user_mappings is not None:
         mapping_defs = [
-            AWQMapping(m["smooth_layer"], m["balance_layers"], m.get("activation_hook_target"))
-            for m in user_mappings
+            AWQMapping(m["smooth_layer"], m["balance_layers"], m.get("activation_hook_target")) for m in user_mappings
         ]
     else:
         mapping_defs = _get_mappings_for_model(model)
diff --git a/test/unit/test_cpu/algorithms/test_awq.py b/test/unit/test_cpu/algorithms/test_awq.py
index 03a1abd110..e3565ddf25 100644
--- a/test/unit/test_cpu/algorithms/test_awq.py
+++ b/test/unit/test_cpu/algorithms/test_awq.py
@@ -633,7 +633,9 @@ def forward(self, hidden_states, *, position_ids=None, cache_position=None, payl
         position_ids = torch.arange(15).reshape(5, 3)
         cache_position = torch.arange(3)
         payload = {"add": torch.randn(5, 3, 4)}
-        kwargs_list = [((hidden_states,), {"position_ids": position_ids, "cache_position": cache_position, "payload": payload})]
+        kwargs_list = [
+            ((hidden_states,), {"position_ids": position_ids, "cache_position": cache_position, "payload": payload})
+        ]
 
         full = AWQTransform(AWQConfig(bits=4, group_size=4, sym=True, data_type="int"))
         micro = AWQTransform(AWQConfig(bits=4, group_size=4, sym=True, data_type="int", smooth_batch_size=2))

From 409874f8b0e551822522fb62c102aba1048a01e4 Mon Sep 17 00:00:00 2001
From: WeiweiZhang1 <weiwei1.zhang@intel.com>
Date: Fri, 7 Aug 2026 14:54:48 +0800
Subject: [PATCH 4/4] refactor: simplify AWQ block mapping lookup

Signed-off-by: WeiweiZhang1 <weiwei1.zhang@intel.com>
---
 auto_round/algorithms/transforms/awq/base.py | 17 +++++++++--------
 1 file changed, 9 insertions(+), 8 deletions(-)

diff --git a/auto_round/algorithms/transforms/awq/base.py b/auto_round/algorithms/transforms/awq/base.py
index dcb747d3e0..91cdab0fda 100644
--- a/auto_round/algorithms/transforms/awq/base.py
+++ b/auto_round/algorithms/transforms/awq/base.py
@@ -192,6 +192,14 @@ def _slice_batch_args_kwargs(args: tuple, kwargs: dict, actual_batch: int, start
     return new_args, new_kwargs
 
 
+def _iter_block_names_for_mapping(model: torch.nn.Module) -> list[str]:
+    """Return block names sorted from most-specific to least-specific for mapping lookup."""
+    from auto_round.utils.common import flatten_list
+    from auto_round.utils.model import get_block_names
+
+    return sorted((name for name in flatten_list(get_block_names(model)) if name), key=len, reverse=True)
+
+
 @register_pipeline_member(AWQConfig)
 class AWQTransform(BasePreprocessor):
     """AWQ transform: activation-aware weight smoothing pre-processor.
@@ -288,14 +296,7 @@ def prepare_run(self, composer: "AlgorithmComposer" = None) -> None:
                 cls_name,
             )
 
-        from auto_round.utils.common import flatten_list
-        from auto_round.utils.model import get_block_names
-
-        try:
-            iter_block_names = [b for b in flatten_list(get_block_names(model)) if b]
-        except Exception:  # noqa: BLE001 - fall back to prefix heuristic if block discovery fails
-            iter_block_names = []
-        iter_block_names.sort(key=len, reverse=True)
+        iter_block_names = _iter_block_names_for_mapping(model)
 
         self._block_mappings = {}
         for m in self._resolved_mappings: