From 072fbe6a6edb4d1206b39f17a03f5cb7690fadc0 Mon Sep 17 00:00:00 2001
From: Wenhua Cheng <wenhua.cheng@intel.com>
Date: Mon, 24 Aug 2026 18:45:50 +0800
Subject: [PATCH 1/3] update

---
 auto_round/auto_scheme/delta_loss.py      | 366 +++++++++++++++++++++-
 auto_round/auto_scheme/gen_auto_scheme.py |  15 +
 auto_round/auto_scheme/solver.py          | 236 ++++++++++++++
 auto_round/cli/main.py                    |   1 +
 auto_round/cli/parser.py                  |   9 +
 5 files changed, 612 insertions(+), 15 deletions(-)
 create mode 100644 auto_round/auto_scheme/solver.py

diff --git a/auto_round/auto_scheme/delta_loss.py b/auto_round/auto_scheme/delta_loss.py
index 3f405fb3e7..7230ac092c 100644
--- a/auto_round/auto_scheme/delta_loss.py
+++ b/auto_round/auto_scheme/delta_loss.py
@@ -196,6 +196,255 @@ def save_grad(grad):
         return qdq_w, 1.0, None
 
 
+# Attributes that fully describe a weight-quantization option on ``orig_layer``.
+_WEIGHT_SCHEME_KEYS = ("bits", "group_size", "sym", "data_type", "super_bits", "super_group_size")
+_ACT_SCHEME_KEYS = ("act_bits", "act_group_size", "act_sym", "act_data_type", "act_dynamic")
+
+
+class MultiOptionScoreWrapper(AutoSchemeWrapperLinear):
+    """Score *all* candidate schemes for a layer in a single forward/backward.
+
+    The forward runs at the **anchor** configuration (whatever scheme was set on
+    ``orig_layer`` before wrapping, i.e. the current mixed assignment). The backward hook
+    then evaluates, for every candidate ``s``:
+
+        score(l, s) = sum |g_l(anchor) * (W - Q_s(W))|      (+ the activation counterpart)
+
+    This is exactly the baseline Delta-Loss metric, with the single crucial difference
+    that ``g`` is measured at *one shared* expansion point for all options instead of a
+    different uniform-``s`` model per option. That is what makes the per-option scores
+    mutually comparable, and what lets already-committed layers contribute their true
+    cross terms ``H_lm D_m`` through ``g``.
+
+    Cost: one extra quant-dequant per option per backward call (no extra model passes,
+    no persistent weight-sized buffers -- each diff is built and discarded in place).
+    """
+
+    def set_candidates(self, cand_schemes: list[dict], score_reduction: str = "abs", split_half: bool = True) -> None:
+        """Bind the candidate option list and allocate the score accumulators."""
+        self._cand_schemes = list(cand_schemes)
+        n = len(self._cand_schemes)
+        self._cand_weight_funcs: list = []
+        self._cand_act_funcs: list = []
+        self._cand_gparam: list = []
+        self._score_reduction = score_reduction
+        self._split_half = split_half
+        # Split-half is driven by the parity of this layer's own backward-call counter:
+        # each hook fires exactly once per calibration batch, so the parity yields two
+        # disjoint halves of the data without touching any forward loop.
+        self._w_calls = 0
+        self._a_calls = 0
+        self.opt_scores = [0.0] * n
+        self.opt_half_scores = [[0.0] * n, [0.0] * n]
+
+        iters = getattr(self.orig_layer, "iters", 0)
+        disable_opt_rtn = getattr(self, "disable_opt_rtn", True)
+        # The NVFP global scale is derived from ``weight`` and is cached across batches.
+        # Left attached it would make the second backward walk an already-freed graph
+        # ("backward through the graph a second time"), so pin it as a constant.
+        base_gparam = getattr(self, "weight_global_scale", None)
+        if isinstance(base_gparam, torch.Tensor) and base_gparam.grad_fn is not None:
+            self.weight_global_scale = base_gparam.detach()
+        for sch in self._cand_schemes:
+            if sch.get("bits", 16) >= 16:
+                self._cand_weight_funcs.append(None)
+                self._cand_gparam.append(None)
+            else:
+                func, dtype = get_quant_func(
+                    sch.get("data_type", "int"),
+                    sch.get("bits", 16),
+                    sch.get("sym", True),
+                    disable_opt_rtn,
+                    sch.get("group_size", 128),
+                    iters=iters,
+                )
+                self._cand_weight_funcs.append((func, dtype))
+                self._cand_gparam.append(self._candidate_gparam(dtype, sch))
+
+            if sch.get("act_bits", 16) > 8:
+                self._cand_act_funcs.append(None)
+            else:
+                func, dtype = get_quant_func(
+                    sch.get("act_data_type", sch.get("data_type", "int")),
+                    sch.get("act_bits", 16),
+                    sch.get("act_sym", True),
+                    True,
+                    iters=iters,
+                )
+                self._cand_act_funcs.append((func, dtype))
+
+    def _candidate_gparam(self, data_type, sch):
+        """Pre-compute the NVFP global scale for a candidate (it is scheme-dependent)."""
+        from auto_round.data_type.utils import is_nv_fp
+
+        if not is_nv_fp(data_type):
+            return None
+        from auto_round.data_type.nvfp import calculate_gparam
+
+        weight = self.orig_layer.weight
+        if weight.device.type == "meta":
+            weight = self.orig_layer.get_weight()
+        with torch.no_grad():
+            return calculate_gparam(weight.to(self.device), sch.get("group_size", 16)).detach()
+
+    @staticmethod
+    def _scale_dtype_for(sch):
+        """Mirror the scale-dtype policy used when scoring a uniform scheme."""
+        if sch.get("super_bits") is not None:
+            return torch.float32
+        return torch.float16 if sch.get("act_bits", 16) > 8 else torch.bfloat16
+
+    def _reduce(self, grad, diff):
+        """Reduce the per-element first-order terms of one backward call to a scalar."""
+        prod = grad.to(diff.dtype) * diff
+        if self._score_reduction == "signed":
+            # True directional derivative for this batch; the absolute value is taken
+            # only across batches, so sign cancellation inside a batch is preserved.
+            return abs(prod.sum().item())
+        return torch.abs(prod).sum().item()
+
+    def _weight_diff(self, idx, ref_device):
+        """Return ``W - Q_s(W)`` for candidate ``idx`` (``None`` when the option is BF16)."""
+        entry = self._cand_weight_funcs[idx]
+        if entry is None:
+            return None
+        sch = self._cand_schemes[idx]
+        orig = self.orig_layer
+        saved_layer = {k: getattr(orig, k, None) for k in _WEIGHT_SCHEME_KEYS}
+        saved_scale_dtype = getattr(orig, "scale_dtype", None)
+        saved_func, saved_dtype = self.weight_quant_func, self.data_type
+        saved_gparam = getattr(self, "weight_global_scale", None)
+        try:
+            for k in _WEIGHT_SCHEME_KEYS:
+                setattr(orig, k, sch.get(k, saved_layer[k]))
+            orig.scale_dtype = self._scale_dtype_for(sch)
+            self.weight_quant_func, self.data_type = entry
+            self.weight_global_scale = self._cand_gparam[idx]
+            device = self.device
+            qdq_w, _, _ = WrapperLinear._qdq_weight(
+                self,
+                torch.tensor(0, device=device),
+                torch.tensor(1.0, device=device),
+                torch.tensor(1.0, device=device),
+            )
+        finally:
+            for k in _WEIGHT_SCHEME_KEYS:
+                setattr(orig, k, saved_layer[k])
+            if saved_scale_dtype is not None:
+                orig.scale_dtype = saved_scale_dtype
+            self.weight_quant_func, self.data_type = saved_func, saved_dtype
+            self.weight_global_scale = saved_gparam
+
+        weight = orig.weight
+        if weight.device.type == "meta":
+            weight = orig.get_weight()
+        return (weight.to(ref_device) - qdq_w.detach().to(ref_device)).to(ref_device)
+
+    def _act_diff(self, idx, x, ref_device):
+        """Return ``x - Q_s(x)`` for candidate ``idx`` (``None`` when activations stay high precision)."""
+        entry = self._cand_act_funcs[idx]
+        if entry is None:
+            return None
+        sch = self._cand_schemes[idx]
+        orig = self.orig_layer
+        saved_layer = {k: getattr(orig, k, None) for k in _ACT_SCHEME_KEYS}
+        saved_scale_dtype = getattr(orig, "scale_dtype", None)
+        saved_func = getattr(self, "act_quant_func", None)
+        saved_dtype = getattr(self, "act_data_type", None)
+        try:
+            for k in _ACT_SCHEME_KEYS:
+                setattr(orig, k, sch.get(k, saved_layer[k]))
+            orig.scale_dtype = self._scale_dtype_for(sch)
+            self.act_quant_func, self.act_data_type = entry
+            qdq_x, _, _ = WrapperLinear._qdq_act(
+                self,
+                x,
+                act_min_scale=torch.tensor(1.0, device=x.device),
+                act_max_scale=torch.tensor(1.0, device=x.device),
+                act_max=None,
+            )
+        except Exception as exc:  # noqa: BLE001 - a single bad option must not kill scoring
+            logger.warning_once(f"AutoScheme: activation scoring failed for a candidate option: {exc}")
+            return None
+        finally:
+            for k in _ACT_SCHEME_KEYS:
+                setattr(orig, k, saved_layer[k])
+            if saved_scale_dtype is not None:
+                orig.scale_dtype = saved_scale_dtype
+            if saved_func is not None:
+                self.act_quant_func = saved_func
+            if saved_dtype is not None:
+                self.act_data_type = saved_dtype
+        return (x - qdq_x).detach().to(ref_device)
+
+    def _accumulate(self, idx, value, half):
+        """Add ``value`` to the total and to the given split-half accumulator."""
+        self.opt_scores[idx] += value
+        if self._split_half:
+            self.opt_half_scores[half][idx] += value
+
+    def _qdq_weight(self, value, min_scale, max_scale):
+        """Run the anchor's quant-dequant for the forward, and score every option in the backward."""
+        device = self.device
+        qdq_w, _, _ = WrapperLinear._qdq_weight(
+            self,
+            torch.tensor(0, device=device),
+            torch.tensor(1.0, device=device),
+            torch.tensor(1.0, device=device),
+        )
+        if self.grad_mode and getattr(self, "_cand_schemes", None):
+
+            def save_grad(grad):
+                """Backward hook: score W - Q_s(W) against this batch's gradient, for every s."""
+                if torch.isnan(grad).any():
+                    return None
+                half = self._w_calls % 2
+                self._w_calls += 1
+                with torch.no_grad():
+                    for idx in range(len(self._cand_schemes)):
+                        diff = self._weight_diff(idx, grad.device)
+                        if diff is None:
+                            continue
+                        self._accumulate(idx, self._reduce(grad, diff), half)
+                        del diff
+                return None
+
+            if qdq_w.requires_grad:
+                qdq_w.register_hook(save_grad)
+        return qdq_w, 1.0, None
+
+    def _qdq_act(self, x, act_min_scale=1.0, act_max_scale=1.0, act_max=None):
+        """Run the anchor's activation quant-dequant, and score every option in the backward."""
+        if hasattr(self.orig_layer, "act_bits") and self.orig_layer.act_bits > 8:
+            qdq_x = x
+        else:
+            qdq_x, _, _ = self.act_qdq_func(x, act_min_scale, act_max_scale, act_max)
+
+        has_act_option = getattr(self, "_cand_act_funcs", None) and any(f is not None for f in self._cand_act_funcs)
+        if self.grad_mode and has_act_option:
+            with torch.no_grad():
+                if torch.abs(x).max() != 0:
+                    self.act_cnt += 1
+                diffs = [self._act_diff(idx, x, "cpu") for idx in range(len(self._cand_schemes))]
+
+            def save_grad(grad):
+                """Backward hook: score x - Q_s(x) against this batch's gradient, for every s."""
+                if torch.isnan(grad).any():
+                    return None
+                half = self._a_calls % 2
+                self._a_calls += 1
+                with torch.no_grad():
+                    for idx, diff in enumerate(diffs):
+                        if diff is None:
+                            continue
+                        self._accumulate(idx, self._reduce(grad, diff.to(grad.device)), half)
+                return None
+
+            if qdq_x.requires_grad:
+                qdq_x.register_hook(save_grad)
+        return qdq_x, 1.0, None
+
+
 class AutoSchemeWrapperLinearIMatrix(WrapperLinear):
     """GGUF-K wrapper that scores a layer using an imatrix-aware quant search (RTN, iters=0)."""
 
@@ -891,11 +1140,28 @@ def get_score_for_scheme(
     model_name: Optional[str] = None,
     scheme_tag: Optional[str] = None,
     disk_index=None,
+    anchor_scheme: Optional[dict] = None,
+    candidate_schemes: Optional[list] = None,
+    score_reduction: str = "abs",
+    split_half: bool = False,
 ):
     """Wrap every quantizable layer in ``quant_layer_names`` with a scoring wrapper, run
     forward(+backward, unless RTN-only) calibration over ``nsamples`` examples from
     ``dataset``/``dataloader``, then unwrap and return each layer's ``[bits, loss]``.
+
+    Two modes:
+
+    * **Uniform (default)** -- the caller has already applied one scheme to every layer;
+      each layer reports the loss of that single scheme. Returns ``{name: [bits, loss]}``.
+    * **Anchor + multi-option** (``anchor_scheme`` / ``candidate_schemes`` given) -- the
+      forward runs at the per-layer ``anchor_scheme`` (the current *mixed* assignment) and
+      a single backward scores **all** ``candidate_schemes`` at once. This makes the
+      per-option scores share one Taylor expansion point, and folds the cross terms of the
+      already-decided layers into the gradient. Returns
+      ``{name: [losses, half0_losses, half1_losses]}``; bit costs are scheme properties
+      and are reused from the caller's stage-0 table.
     """
+    multi_option = bool(candidate_schemes)
     scores_dict = {}  # Key=name,Val=[quant_total_bits, loss]
     # Include the visual block(s) when scoring VLMs with ``--quant_nontext_module``
     # (``force_mllm=True``) so vision-tower layer losses match a block below instead
@@ -921,14 +1187,41 @@ def get_score_for_scheme(
     for name in quant_layer_names:
         if offload_context is not None:
             offload_context.ensure_loaded(model, name)
-        if name in fixed_layer_scheme.keys():
+        # In anchor mode the fixed layers must still run *quantized* during the forward,
+        # otherwise the expansion point is not the real mixed configuration. They are
+        # wrapped but given no candidates, so they contribute no scores.
+        is_fixed = name in fixed_layer_scheme.keys()
+        if is_fixed and not multi_option:
             continue
         m = get_module(model, name)
         if m is None:
             raise RuntimeError(f"AutoScheme scoring layer {name!r} is missing after model preprocessing")
+        if multi_option and anchor_scheme is not None and name in anchor_scheme:
+            for _k, _v in anchor_scheme[name].items():
+                setattr(m, _k, _v)
+        # A layer whose *anchor* is BF16 still has non-BF16 candidates that must be
+        # scored, otherwise it could never be moved off BF16 in a later stage. BF16
+        # layers cannot build a quant func, so wrap with a constructible placeholder and
+        # then force bits/act_bits back to 16 -- ``WrapperLinear._qdq_weight`` short
+        # circuits on ``bits >= 16``, so the forward is an exact BF16 pass-through while
+        # the candidate scoring path (which swaps in its own per-option funcs) is intact.
+        anchor_is_bf16 = multi_option and not is_fixed and not check_to_quantized(m)
+        if anchor_is_bf16:
+            placeholder = next((c for c in candidate_schemes if c.get("bits", 16) < 16), None)
+            if placeholder is None:
+                n_cand = len(candidate_schemes)
+                scores_dict[name] = [[0.0] * n_cand, [0.0] * n_cand, [0.0] * n_cand]
+                continue
+            for _k, _v in placeholder.items():
+                setattr(m, _k, _v)
         if not check_to_quantized(m):
-            layer_bits, _ = compute_layer_bits(m, ignore_scale_zp_bits)
-            scores_dict[name] = [layer_bits, 0.0]
+            if not is_fixed:
+                if multi_option:
+                    n_cand = len(candidate_schemes)
+                    scores_dict[name] = [[0.0] * n_cand, [0.0] * n_cand, [0.0] * n_cand]
+                else:
+                    layer_bits, _ = compute_layer_bits(m, ignore_scale_zp_bits)
+                    scores_dict[name] = [layer_bits, 0.0]
             continue
         if m.act_bits > 8 and m.super_bits is not None:
             m.scale_dtype = torch.float32  # TODO set this via API
@@ -945,6 +1238,8 @@ def get_score_for_scheme(
                 WrapperLayer = AutoSchemeWrapperLinearForGGUFKImatrix
             else:
                 WrapperLayer = AutoSchemeWrapperLinearForGGUFK
+        if multi_option:
+            WrapperLayer = MultiOptionScoreWrapper
 
         with torch.no_grad():
             if low_gpu_mem_usage:
@@ -964,9 +1259,20 @@ def get_score_for_scheme(
                 enable_minmax_tuning=False,
                 enable_norm_bias_tuning=False,
                 enable_round_tuning=False,
-                need_weight_grad=need_weight_grad,
+                need_weight_grad=need_weight_grad or multi_option,
                 enable_torch_compile=enable_torch_compile,
             )
+            if multi_option:
+                if anchor_is_bf16:
+                    # Restore the true BF16 anchor now that the wrapper is built.
+                    new_m.orig_layer.bits = 16
+                    new_m.orig_layer.act_bits = 16
+                    new_m.enable_act_quant = False
+                new_m.set_candidates(
+                    [] if is_fixed else candidate_schemes,
+                    score_reduction=score_reduction,
+                    split_half=split_half,
+                )
             set_module(model, name, new_m)
     if offload_context is not None:
         offload_context.flush_loaded(model)
@@ -1243,6 +1549,10 @@ def _run_forward_loop(loader):
             m.grad = None
 
     for n, m in model.named_modules():
+        if multi_option:
+            if isinstance(m, MultiOptionScoreWrapper) and getattr(m, "_cand_schemes", None):
+                scores_dict[n] = [list(m.opt_scores), list(m.opt_half_scores[0]), list(m.opt_half_scores[1])]
+            continue
         if hasattr(m, "mix_score"):
             if m.orig_layer.act_bits <= 8:
                 if m.act_cnt == 0:
@@ -1251,14 +1561,15 @@ def _run_forward_loop(loader):
                     )
             layer_bits, _ = compute_layer_bits(m.orig_layer, ignore_scale_zp_bits=ignore_scale_zp_bits)
             scores_dict[n] = [layer_bits, m.mix_score]
-    _fill_inactive_expert_scores(scores_dict, block_names)
-    _log_score_summary_by_block_and_nonblock(
-        scores_dict,
-        block_names,
-        model=model,
-        scheme_tag=scheme_tag,
-        summary_stage="final",
-    )
+    if not multi_option:
+        _fill_inactive_expert_scores(scores_dict, block_names)
+        _log_score_summary_by_block_and_nonblock(
+            scores_dict,
+            block_names,
+            model=model,
+            scheme_tag=scheme_tag,
+            summary_stage="final",
+        )
 
     for n, m in model.named_modules():
         if hasattr(m, "orig_layer"):
@@ -2156,6 +2467,7 @@ def _gen_layer_config(
                 m.to(major_device)
 
     total_scores = {}
+    per_op_bits: dict[int, dict[str, int]] = {}
     schemes = auto_scheme.options
 
     def check_bf16_scheme(scheme):
@@ -2398,6 +2710,8 @@ def _record_scheme_scores(index, per_op_scores):
             for key, item in grouped_scores.items():
                 total_scores.setdefault(key, []).append(item)
             options_scores.append(total_loss)
+            # Bit costs are a property of the scheme, not of the expansion point.
+            per_op_bits[index] = {name: score[0] for name, score in per_op_scores.items()}
             return total_loss
 
         def _save_per_op_scores(index, scheme, cache_key, cache_path, per_op_scores):
@@ -2876,9 +3190,31 @@ def _select_embedding_scheme_index():
     )
 
     dp_started = time.perf_counter()
-    best_loss, best_path = choose_bits_per_layer_with_path(total_scores, target_params_cnt)
+
+    # ------------------------------------------------------------------ #
+    # Bit allocation. Both solvers optimise the same objective over the same score
+    # table; they differ only in how the knapsack is solved. The DP is the historical
+    # default; the Lagrangian dual reaches the same allocation far faster because it
+    # needs no discretised state space.
+    # ------------------------------------------------------------------ #
+    if auto_scheme.solver == "lagrangian":
+        from auto_round.auto_scheme.solver import solve_lagrangian
+
+        assign = solve_lagrangian(total_scores, target_params_cnt)
+        if assign is None:
+            # Infeasible under the dual -- fall back to the DP so a solver problem can
+            # never make AutoScheme worse than before.
+            logger.warning("AutoScheme: Lagrangian solver found no feasible allocation; falling back to DP.")
+            best_loss, best_path = choose_bits_per_layer_with_path(total_scores, target_params_cnt)
+        else:
+            best_path = [(opt[3], opt[0]) for opt in assign.values()]
+            best_loss = sum(opt[2] for opt in assign.values())
+    else:
+        best_loss, best_path = choose_bits_per_layer_with_path(total_scores, target_params_cnt)
+
     logger.info(
-        "AutoScheme post-scoring: DP selection took %.2fs (layers=%d)",
+        "AutoScheme post-scoring: %s selection took %.2fs (layers=%d)",
+        auto_scheme.solver,
         time.perf_counter() - dp_started,
         len(total_scores),
     )
@@ -2953,7 +3289,7 @@ def gen_layer_config(
     """Public AutoScheme entry.
 
     This wrapper performs model loading/dispatch and environment preparation,
-    then delegates to `_gen_layer_config` for staged scoring + DP selection.
+    then delegates to `_gen_layer_config` for scoring + solver-based selection.
     """
     model_name = None
     is_vlm = False
diff --git a/auto_round/auto_scheme/gen_auto_scheme.py b/auto_round/auto_scheme/gen_auto_scheme.py
index e86e1e52bc..92901811d3 100644
--- a/auto_round/auto_scheme/gen_auto_scheme.py
+++ b/auto_round/auto_scheme/gen_auto_scheme.py
@@ -41,11 +41,26 @@ class AutoScheme:
     low_gpu_mem_usage: bool = True
     low_cpu_mem_usage: bool = True
 
+    # ------------------------------------------------------------------ #
+    # Bit-allocation solver.
+    #
+    # The default reproduces the historical behaviour exactly: a knapsack DP
+    # over scores measured on uniform-scheme models.
+    # ------------------------------------------------------------------ #
+    solver: str = "dp"
+    """Allocation solver: ``"dp"`` (knapsack DP) or ``"lagrangian"`` (shadow-price
+    bisection). Both optimise the same objective and in practice produce the same
+    allocation, but the Lagrangian solver hits a *fractional* avg_bits target exactly
+    and avoids the DP state explosion, making it roughly an order of magnitude faster
+    on 8B-scale models. See docs/auto_scheme_solver.md for measured accuracy/speed."""
+
     def __post_init__(self):
         if isinstance(self.options, str):
             options = self.options.upper().replace(" ", "")
             self.options = options.split(",")
         self.options = self._deduplicate_options(self.options)
+        if self.solver not in ("dp", "lagrangian"):
+            raise ValueError(f"AutoScheme.solver must be 'dp' or 'lagrangian', got {self.solver!r}")
 
     @staticmethod
     def _deduplicate_options(
diff --git a/auto_round/auto_scheme/solver.py b/auto_round/auto_scheme/solver.py
new file mode 100644
index 0000000000..343e18c188
--- /dev/null
+++ b/auto_round/auto_scheme/solver.py
@@ -0,0 +1,236 @@
+# Copyright (c) 2025 Intel Corporation
+#
+# 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.
+
+"""Bit-allocation solvers for AutoScheme.
+
+Two solvers are available, selected through ``AutoScheme.solver``:
+
+``"dp"``
+    The historical knapsack dynamic program (:func:`choose_bits_per_layer_with_path`).
+    Exact on a discretised bit grid, but the state space grows with the budget, which
+    dominates the runtime on large models.
+
+``"lagrangian"``
+    Solves the same knapsack through its Lagrangian dual. For a price ``lam`` (units:
+    loss per bit) every layer independently picks ``argmin_s loss_s + lam * bits_s``.
+    Total bits decrease monotonically in ``lam``, so a bisection on ``lam`` drives the
+    solution onto the budget. This hits a *fractional* avg_bits target exactly and needs
+    no discretised state space, so it is typically an order of magnitude faster than the
+    DP while producing the same allocation.
+
+The dual only ever lands on the *convex hull* of each layer's (bits, loss) curve, so two
+primal repair passes close the integrality gap: :func:`_greedy_repair` spends leftover
+budget, and :func:`_swap_local_search` fixes cases where one downgrade funds one upgrade.
+
+This module is deliberately free of model/torch dependencies so it can be unit-tested in
+isolation.
+"""
+
+from __future__ import annotations
+
+from typing import Optional
+
+from auto_round.logger import logger
+
+__all__ = [
+    "solve_lagrangian",
+    "solve_allocation",
+]
+
+_EPS = 1e-12
+
+
+# ----------------------------------------------------------------------------------- #
+# Score-table helpers
+#
+# ``total_scores`` maps a DP key (the first layer name of a shared-layer group) to a list
+# of candidate options, each option being ``[scheme_index, bits_cost, loss_cost, names]``.
+# ----------------------------------------------------------------------------------- #
+def _option_cost(opt, lam: float) -> float:
+    """Lagrangian cost ``loss + lam * bits`` of a single option."""
+    return opt[2] + lam * opt[1]
+
+
+def _pick_option(opts, lam: float):
+    """Pick the option minimising the Lagrangian cost; ties broken toward fewer bits."""
+    return min(opts, key=lambda o: (_option_cost(o, lam), o[1]))
+
+
+def _total_bits(assign: dict) -> int:
+    """Total bit cost of an assignment (``key -> option``)."""
+    return sum(opt[1] for opt in assign.values())
+
+
+def _total_loss(assign: dict) -> float:
+    """Total predicted loss of an assignment (``key -> option``)."""
+    return sum(opt[2] for opt in assign.values())
+
+
+# ----------------------------------------------------------------------------------- #
+# Lagrangian (shadow-price) solver
+# ----------------------------------------------------------------------------------- #
+def solve_lagrangian(total_scores: dict, budget: int, max_iter: int = 80) -> Optional[dict]:
+    """Solve the bit-allocation knapsack through its Lagrangian dual.
+
+    Args:
+        total_scores: Score table.
+        budget: Upper bound on total bits.
+        max_iter: Bisection iterations.
+
+    Returns:
+        The assignment mapping key -> option, or ``None`` when even the cheapest
+        configuration exceeds ``budget`` (i.e. the target is infeasible).
+    """
+    if not total_scores:
+        return {}
+
+    cheapest = {key: min(opts, key=lambda o: o[1]) for key, opts in total_scores.items()}
+    if _total_bits(cheapest) > budget:
+        return None
+
+    # lam = 0 -> pure loss minimisation. If it already fits, nothing to trade off.
+    assign = {key: _pick_option(opts, 0.0) for key, opts in total_scores.items()}
+    if _total_bits(assign) <= budget:
+        return assign
+
+    # Bracket the price: grow the upper bound until the budget is satisfied.
+    lo, hi = 0.0, 1e-9
+    for _ in range(200):
+        probe = {key: _pick_option(opts, hi) for key, opts in total_scores.items()}
+        if _total_bits(probe) <= budget:
+            break
+        lo, hi = hi, hi * 4.0
+    else:  # pragma: no cover - defensive
+        logger.warning("AutoScheme: Lagrangian upper price search did not converge.")
+
+    best = None
+    for _ in range(max_iter):
+        mid = 0.5 * (lo + hi)
+        probe = {key: _pick_option(opts, mid) for key, opts in total_scores.items()}
+        if _total_bits(probe) <= budget:
+            best, hi = probe, mid
+        else:
+            lo = mid
+
+    if best is None:  # pragma: no cover - defensive
+        best = {key: _pick_option(opts, hi) for key, opts in total_scores.items()}
+    best = _greedy_repair(total_scores, best, budget)
+    best = _swap_local_search(total_scores, best, budget)
+    return best
+
+
+def _greedy_repair(total_scores: dict, assign: dict, budget: int) -> dict:
+    """Spend the budget slack left by the dual's integrality gap.
+
+    The dual solution is generally not budget-tight. Repeatedly apply the upgrade with the
+    best loss-reduction-per-extra-bit that still fits, which is the standard primal repair
+    for a Lagrangian-relaxed knapsack.
+    """
+    used = _total_bits(assign)
+    slack = budget - used
+    if slack <= 0:
+        return assign
+
+    while True:
+        best_key, best_opt, best_gain = None, None, 0.0
+        for key, opts in total_scores.items():
+            cur = assign[key]
+            for opt in opts:
+                extra_bits = opt[1] - cur[1]
+                if extra_bits <= 0 or extra_bits > slack:
+                    continue
+                gain = (cur[2] - opt[2]) / extra_bits
+                if gain > best_gain:
+                    best_key, best_opt, best_gain = key, opt, gain
+        if best_key is None:
+            break
+        slack -= best_opt[1] - assign[best_key][1]
+        assign[best_key] = best_opt
+    return assign
+
+
+def _swap_local_search(total_scores: dict, assign: dict, budget: int, max_rounds: int = 200) -> dict:
+    """Close part of the duality gap with pairwise exchanges.
+
+    A price-based solution can only ever land on the *convex hull* of each layer's
+    (bits, loss) curve, so options that are dominated in the hull -- but optimal in the
+    true (non-convex) problem -- are unreachable for every ``lam``. One downgrade funding
+    one upgrade repairs the common cases.
+    """
+    for _ in range(max_rounds):
+        used = _total_bits(assign)
+        best_move, best_delta = None, -_EPS
+        for up_key, up_opts in total_scores.items():
+            up_cur = assign[up_key]
+            for up_opt in up_opts:
+                if up_opt[1] <= up_cur[1]:
+                    continue
+                need = up_opt[1] - up_cur[1]
+                gain = up_cur[2] - up_opt[2]
+                if gain <= 0:
+                    continue
+                if used + need <= budget:  # pure upgrade, handled by _greedy_repair
+                    continue
+                for down_key, down_opts in total_scores.items():
+                    if down_key == up_key:
+                        continue
+                    down_cur = assign[down_key]
+                    for down_opt in down_opts:
+                        freed = down_cur[1] - down_opt[1]
+                        if freed <= 0 or used + need - freed > budget:
+                            continue
+                        delta = gain - (down_opt[2] - down_cur[2])
+                        if delta > best_delta:
+                            best_delta = delta
+                            best_move = (up_key, up_opt, down_key, down_opt)
+        if best_move is None:
+            break
+        up_key, up_opt, down_key, down_opt = best_move
+        assign[up_key], assign[down_key] = up_opt, down_opt
+    return assign
+
+
+def solve_allocation(total_scores: dict, budget: int, solver: str = "dp", max_states: Optional[int] = None):
+    """Dispatch to the requested allocation solver.
+
+    Args:
+        total_scores: Score table.
+        budget: Total bit budget.
+        solver: ``"dp"`` (knapsack DP, the historical default) or ``"lagrangian"``
+            (shadow-price bisection).
+        max_states: DP beam width; ignored by the Lagrangian solver.
+
+    Returns:
+        The assignment (``key -> option``), or ``None`` when the target is infeasible.
+    """
+    if solver == "lagrangian":
+        return solve_lagrangian(total_scores, budget)
+
+    from auto_round.auto_scheme.delta_loss import choose_bits_per_layer_with_path
+
+    _, path = choose_bits_per_layer_with_path(total_scores, budget, max_states=max_states)
+    if path is None:
+        return None
+
+    chosen_index = {tuple(names): scheme_index for names, scheme_index in path}
+    assign = {}
+    for key, opts in total_scores.items():
+        for opt in opts:
+            if tuple(opt[3]) in chosen_index and chosen_index[tuple(opt[3])] == opt[0]:
+                assign[key] = opt
+                break
+        else:  # pragma: no cover - defensive
+            assign[key] = min(opts, key=lambda o: o[1])
+    return assign
+
diff --git a/auto_round/cli/main.py b/auto_round/cli/main.py
index 859435c717..6a84625585 100644
--- a/auto_round/cli/main.py
+++ b/auto_round/cli/main.py
@@ -382,6 +382,7 @@ def tune(args):
             ignore_scale_zp_bits=args.ignore_scale_zp_bits,
             low_gpu_mem_usage=True,
             low_cpu_mem_usage=low_cpu_mem_usage,
+            solver=getattr(args, "auto_scheme_solver", "dp"),
         )
 
     common_kwargs = _extract_common_quantization_kwargs(args)
diff --git a/auto_round/cli/parser.py b/auto_round/cli/parser.py
index 7175a1f8ca..74b7c665a1 100644
--- a/auto_round/cli/parser.py
+++ b/auto_round/cli/parser.py
@@ -139,6 +139,15 @@ def build_quantize_parser(*, prog: str = "auto_round quantize") -> argparse.Argu
         nargs="+",
         help="AutoScheme options. Accepts comma-separated ('W4A16,W8A16') or space-separated (W4A16 W8A16).",
     )
+    rt.add_argument(
+        "--auto_scheme_solver",
+        "--solver",
+        default="dp",
+        type=str,
+        choices=["dp", "lagrangian"],
+        help="AutoScheme bit-allocation solver: 'dp' (knapsack DP, default) or "
+        "'lagrangian' (shadow-price bisection, same allocation but ~10-15x faster).",
+    )
     rt.add_argument(
         "--low_gpu_mem_usage", action="store_true", help="Enable memory-efficient mode by offloading features to CPU."
     )

From 6578ff31acb2d5372345f50539883f31d43cda87 Mon Sep 17 00:00:00 2001
From: Wenhua Cheng <wenhua.cheng@intel.com>
Date: Mon, 24 Aug 2026 21:20:14 +0800
Subject: [PATCH 2/3] update

---
 auto_round/auto_scheme/delta_loss.py      | 340 +---------------------
 auto_round/auto_scheme/gen_auto_scheme.py |   4 +-
 docs/auto_scheme_solver.md                | 143 +++++++++
 docs/auto_scheme_solver_CN.md             | 132 +++++++++
 4 files changed, 292 insertions(+), 327 deletions(-)
 create mode 100644 docs/auto_scheme_solver.md
 create mode 100644 docs/auto_scheme_solver_CN.md

diff --git a/auto_round/auto_scheme/delta_loss.py b/auto_round/auto_scheme/delta_loss.py
index 7230ac092c..f2eb503015 100644
--- a/auto_round/auto_scheme/delta_loss.py
+++ b/auto_round/auto_scheme/delta_loss.py
@@ -196,255 +196,6 @@ def save_grad(grad):
         return qdq_w, 1.0, None
 
 
-# Attributes that fully describe a weight-quantization option on ``orig_layer``.
-_WEIGHT_SCHEME_KEYS = ("bits", "group_size", "sym", "data_type", "super_bits", "super_group_size")
-_ACT_SCHEME_KEYS = ("act_bits", "act_group_size", "act_sym", "act_data_type", "act_dynamic")
-
-
-class MultiOptionScoreWrapper(AutoSchemeWrapperLinear):
-    """Score *all* candidate schemes for a layer in a single forward/backward.
-
-    The forward runs at the **anchor** configuration (whatever scheme was set on
-    ``orig_layer`` before wrapping, i.e. the current mixed assignment). The backward hook
-    then evaluates, for every candidate ``s``:
-
-        score(l, s) = sum |g_l(anchor) * (W - Q_s(W))|      (+ the activation counterpart)
-
-    This is exactly the baseline Delta-Loss metric, with the single crucial difference
-    that ``g`` is measured at *one shared* expansion point for all options instead of a
-    different uniform-``s`` model per option. That is what makes the per-option scores
-    mutually comparable, and what lets already-committed layers contribute their true
-    cross terms ``H_lm D_m`` through ``g``.
-
-    Cost: one extra quant-dequant per option per backward call (no extra model passes,
-    no persistent weight-sized buffers -- each diff is built and discarded in place).
-    """
-
-    def set_candidates(self, cand_schemes: list[dict], score_reduction: str = "abs", split_half: bool = True) -> None:
-        """Bind the candidate option list and allocate the score accumulators."""
-        self._cand_schemes = list(cand_schemes)
-        n = len(self._cand_schemes)
-        self._cand_weight_funcs: list = []
-        self._cand_act_funcs: list = []
-        self._cand_gparam: list = []
-        self._score_reduction = score_reduction
-        self._split_half = split_half
-        # Split-half is driven by the parity of this layer's own backward-call counter:
-        # each hook fires exactly once per calibration batch, so the parity yields two
-        # disjoint halves of the data without touching any forward loop.
-        self._w_calls = 0
-        self._a_calls = 0
-        self.opt_scores = [0.0] * n
-        self.opt_half_scores = [[0.0] * n, [0.0] * n]
-
-        iters = getattr(self.orig_layer, "iters", 0)
-        disable_opt_rtn = getattr(self, "disable_opt_rtn", True)
-        # The NVFP global scale is derived from ``weight`` and is cached across batches.
-        # Left attached it would make the second backward walk an already-freed graph
-        # ("backward through the graph a second time"), so pin it as a constant.
-        base_gparam = getattr(self, "weight_global_scale", None)
-        if isinstance(base_gparam, torch.Tensor) and base_gparam.grad_fn is not None:
-            self.weight_global_scale = base_gparam.detach()
-        for sch in self._cand_schemes:
-            if sch.get("bits", 16) >= 16:
-                self._cand_weight_funcs.append(None)
-                self._cand_gparam.append(None)
-            else:
-                func, dtype = get_quant_func(
-                    sch.get("data_type", "int"),
-                    sch.get("bits", 16),
-                    sch.get("sym", True),
-                    disable_opt_rtn,
-                    sch.get("group_size", 128),
-                    iters=iters,
-                )
-                self._cand_weight_funcs.append((func, dtype))
-                self._cand_gparam.append(self._candidate_gparam(dtype, sch))
-
-            if sch.get("act_bits", 16) > 8:
-                self._cand_act_funcs.append(None)
-            else:
-                func, dtype = get_quant_func(
-                    sch.get("act_data_type", sch.get("data_type", "int")),
-                    sch.get("act_bits", 16),
-                    sch.get("act_sym", True),
-                    True,
-                    iters=iters,
-                )
-                self._cand_act_funcs.append((func, dtype))
-
-    def _candidate_gparam(self, data_type, sch):
-        """Pre-compute the NVFP global scale for a candidate (it is scheme-dependent)."""
-        from auto_round.data_type.utils import is_nv_fp
-
-        if not is_nv_fp(data_type):
-            return None
-        from auto_round.data_type.nvfp import calculate_gparam
-
-        weight = self.orig_layer.weight
-        if weight.device.type == "meta":
-            weight = self.orig_layer.get_weight()
-        with torch.no_grad():
-            return calculate_gparam(weight.to(self.device), sch.get("group_size", 16)).detach()
-
-    @staticmethod
-    def _scale_dtype_for(sch):
-        """Mirror the scale-dtype policy used when scoring a uniform scheme."""
-        if sch.get("super_bits") is not None:
-            return torch.float32
-        return torch.float16 if sch.get("act_bits", 16) > 8 else torch.bfloat16
-
-    def _reduce(self, grad, diff):
-        """Reduce the per-element first-order terms of one backward call to a scalar."""
-        prod = grad.to(diff.dtype) * diff
-        if self._score_reduction == "signed":
-            # True directional derivative for this batch; the absolute value is taken
-            # only across batches, so sign cancellation inside a batch is preserved.
-            return abs(prod.sum().item())
-        return torch.abs(prod).sum().item()
-
-    def _weight_diff(self, idx, ref_device):
-        """Return ``W - Q_s(W)`` for candidate ``idx`` (``None`` when the option is BF16)."""
-        entry = self._cand_weight_funcs[idx]
-        if entry is None:
-            return None
-        sch = self._cand_schemes[idx]
-        orig = self.orig_layer
-        saved_layer = {k: getattr(orig, k, None) for k in _WEIGHT_SCHEME_KEYS}
-        saved_scale_dtype = getattr(orig, "scale_dtype", None)
-        saved_func, saved_dtype = self.weight_quant_func, self.data_type
-        saved_gparam = getattr(self, "weight_global_scale", None)
-        try:
-            for k in _WEIGHT_SCHEME_KEYS:
-                setattr(orig, k, sch.get(k, saved_layer[k]))
-            orig.scale_dtype = self._scale_dtype_for(sch)
-            self.weight_quant_func, self.data_type = entry
-            self.weight_global_scale = self._cand_gparam[idx]
-            device = self.device
-            qdq_w, _, _ = WrapperLinear._qdq_weight(
-                self,
-                torch.tensor(0, device=device),
-                torch.tensor(1.0, device=device),
-                torch.tensor(1.0, device=device),
-            )
-        finally:
-            for k in _WEIGHT_SCHEME_KEYS:
-                setattr(orig, k, saved_layer[k])
-            if saved_scale_dtype is not None:
-                orig.scale_dtype = saved_scale_dtype
-            self.weight_quant_func, self.data_type = saved_func, saved_dtype
-            self.weight_global_scale = saved_gparam
-
-        weight = orig.weight
-        if weight.device.type == "meta":
-            weight = orig.get_weight()
-        return (weight.to(ref_device) - qdq_w.detach().to(ref_device)).to(ref_device)
-
-    def _act_diff(self, idx, x, ref_device):
-        """Return ``x - Q_s(x)`` for candidate ``idx`` (``None`` when activations stay high precision)."""
-        entry = self._cand_act_funcs[idx]
-        if entry is None:
-            return None
-        sch = self._cand_schemes[idx]
-        orig = self.orig_layer
-        saved_layer = {k: getattr(orig, k, None) for k in _ACT_SCHEME_KEYS}
-        saved_scale_dtype = getattr(orig, "scale_dtype", None)
-        saved_func = getattr(self, "act_quant_func", None)
-        saved_dtype = getattr(self, "act_data_type", None)
-        try:
-            for k in _ACT_SCHEME_KEYS:
-                setattr(orig, k, sch.get(k, saved_layer[k]))
-            orig.scale_dtype = self._scale_dtype_for(sch)
-            self.act_quant_func, self.act_data_type = entry
-            qdq_x, _, _ = WrapperLinear._qdq_act(
-                self,
-                x,
-                act_min_scale=torch.tensor(1.0, device=x.device),
-                act_max_scale=torch.tensor(1.0, device=x.device),
-                act_max=None,
-            )
-        except Exception as exc:  # noqa: BLE001 - a single bad option must not kill scoring
-            logger.warning_once(f"AutoScheme: activation scoring failed for a candidate option: {exc}")
-            return None
-        finally:
-            for k in _ACT_SCHEME_KEYS:
-                setattr(orig, k, saved_layer[k])
-            if saved_scale_dtype is not None:
-                orig.scale_dtype = saved_scale_dtype
-            if saved_func is not None:
-                self.act_quant_func = saved_func
-            if saved_dtype is not None:
-                self.act_data_type = saved_dtype
-        return (x - qdq_x).detach().to(ref_device)
-
-    def _accumulate(self, idx, value, half):
-        """Add ``value`` to the total and to the given split-half accumulator."""
-        self.opt_scores[idx] += value
-        if self._split_half:
-            self.opt_half_scores[half][idx] += value
-
-    def _qdq_weight(self, value, min_scale, max_scale):
-        """Run the anchor's quant-dequant for the forward, and score every option in the backward."""
-        device = self.device
-        qdq_w, _, _ = WrapperLinear._qdq_weight(
-            self,
-            torch.tensor(0, device=device),
-            torch.tensor(1.0, device=device),
-            torch.tensor(1.0, device=device),
-        )
-        if self.grad_mode and getattr(self, "_cand_schemes", None):
-
-            def save_grad(grad):
-                """Backward hook: score W - Q_s(W) against this batch's gradient, for every s."""
-                if torch.isnan(grad).any():
-                    return None
-                half = self._w_calls % 2
-                self._w_calls += 1
-                with torch.no_grad():
-                    for idx in range(len(self._cand_schemes)):
-                        diff = self._weight_diff(idx, grad.device)
-                        if diff is None:
-                            continue
-                        self._accumulate(idx, self._reduce(grad, diff), half)
-                        del diff
-                return None
-
-            if qdq_w.requires_grad:
-                qdq_w.register_hook(save_grad)
-        return qdq_w, 1.0, None
-
-    def _qdq_act(self, x, act_min_scale=1.0, act_max_scale=1.0, act_max=None):
-        """Run the anchor's activation quant-dequant, and score every option in the backward."""
-        if hasattr(self.orig_layer, "act_bits") and self.orig_layer.act_bits > 8:
-            qdq_x = x
-        else:
-            qdq_x, _, _ = self.act_qdq_func(x, act_min_scale, act_max_scale, act_max)
-
-        has_act_option = getattr(self, "_cand_act_funcs", None) and any(f is not None for f in self._cand_act_funcs)
-        if self.grad_mode and has_act_option:
-            with torch.no_grad():
-                if torch.abs(x).max() != 0:
-                    self.act_cnt += 1
-                diffs = [self._act_diff(idx, x, "cpu") for idx in range(len(self._cand_schemes))]
-
-            def save_grad(grad):
-                """Backward hook: score x - Q_s(x) against this batch's gradient, for every s."""
-                if torch.isnan(grad).any():
-                    return None
-                half = self._a_calls % 2
-                self._a_calls += 1
-                with torch.no_grad():
-                    for idx, diff in enumerate(diffs):
-                        if diff is None:
-                            continue
-                        self._accumulate(idx, self._reduce(grad, diff.to(grad.device)), half)
-                return None
-
-            if qdq_x.requires_grad:
-                qdq_x.register_hook(save_grad)
-        return qdq_x, 1.0, None
-
-
 class AutoSchemeWrapperLinearIMatrix(WrapperLinear):
     """GGUF-K wrapper that scores a layer using an imatrix-aware quant search (RTN, iters=0)."""
 
@@ -1140,28 +891,14 @@ def get_score_for_scheme(
     model_name: Optional[str] = None,
     scheme_tag: Optional[str] = None,
     disk_index=None,
-    anchor_scheme: Optional[dict] = None,
-    candidate_schemes: Optional[list] = None,
-    score_reduction: str = "abs",
-    split_half: bool = False,
 ):
     """Wrap every quantizable layer in ``quant_layer_names`` with a scoring wrapper, run
     forward(+backward, unless RTN-only) calibration over ``nsamples`` examples from
     ``dataset``/``dataloader``, then unwrap and return each layer's ``[bits, loss]``.
 
-    Two modes:
-
-    * **Uniform (default)** -- the caller has already applied one scheme to every layer;
-      each layer reports the loss of that single scheme. Returns ``{name: [bits, loss]}``.
-    * **Anchor + multi-option** (``anchor_scheme`` / ``candidate_schemes`` given) -- the
-      forward runs at the per-layer ``anchor_scheme`` (the current *mixed* assignment) and
-      a single backward scores **all** ``candidate_schemes`` at once. This makes the
-      per-option scores share one Taylor expansion point, and folds the cross terms of the
-      already-decided layers into the gradient. Returns
-      ``{name: [losses, half0_losses, half1_losses]}``; bit costs are scheme properties
-      and are reused from the caller's stage-0 table.
+    The caller has already applied one scheme to every layer, so each layer reports the
+    loss of that single scheme. Returns ``{name: [bits, loss]}``.
     """
-    multi_option = bool(candidate_schemes)
     scores_dict = {}  # Key=name,Val=[quant_total_bits, loss]
     # Include the visual block(s) when scoring VLMs with ``--quant_nontext_module``
     # (``force_mllm=True``) so vision-tower layer losses match a block below instead
@@ -1187,41 +924,14 @@ def get_score_for_scheme(
     for name in quant_layer_names:
         if offload_context is not None:
             offload_context.ensure_loaded(model, name)
-        # In anchor mode the fixed layers must still run *quantized* during the forward,
-        # otherwise the expansion point is not the real mixed configuration. They are
-        # wrapped but given no candidates, so they contribute no scores.
-        is_fixed = name in fixed_layer_scheme.keys()
-        if is_fixed and not multi_option:
+        if name in fixed_layer_scheme.keys():
             continue
         m = get_module(model, name)
         if m is None:
             raise RuntimeError(f"AutoScheme scoring layer {name!r} is missing after model preprocessing")
-        if multi_option and anchor_scheme is not None and name in anchor_scheme:
-            for _k, _v in anchor_scheme[name].items():
-                setattr(m, _k, _v)
-        # A layer whose *anchor* is BF16 still has non-BF16 candidates that must be
-        # scored, otherwise it could never be moved off BF16 in a later stage. BF16
-        # layers cannot build a quant func, so wrap with a constructible placeholder and
-        # then force bits/act_bits back to 16 -- ``WrapperLinear._qdq_weight`` short
-        # circuits on ``bits >= 16``, so the forward is an exact BF16 pass-through while
-        # the candidate scoring path (which swaps in its own per-option funcs) is intact.
-        anchor_is_bf16 = multi_option and not is_fixed and not check_to_quantized(m)
-        if anchor_is_bf16:
-            placeholder = next((c for c in candidate_schemes if c.get("bits", 16) < 16), None)
-            if placeholder is None:
-                n_cand = len(candidate_schemes)
-                scores_dict[name] = [[0.0] * n_cand, [0.0] * n_cand, [0.0] * n_cand]
-                continue
-            for _k, _v in placeholder.items():
-                setattr(m, _k, _v)
         if not check_to_quantized(m):
-            if not is_fixed:
-                if multi_option:
-                    n_cand = len(candidate_schemes)
-                    scores_dict[name] = [[0.0] * n_cand, [0.0] * n_cand, [0.0] * n_cand]
-                else:
-                    layer_bits, _ = compute_layer_bits(m, ignore_scale_zp_bits)
-                    scores_dict[name] = [layer_bits, 0.0]
+            layer_bits, _ = compute_layer_bits(m, ignore_scale_zp_bits)
+            scores_dict[name] = [layer_bits, 0.0]
             continue
         if m.act_bits > 8 and m.super_bits is not None:
             m.scale_dtype = torch.float32  # TODO set this via API
@@ -1238,8 +948,6 @@ def get_score_for_scheme(
                 WrapperLayer = AutoSchemeWrapperLinearForGGUFKImatrix
             else:
                 WrapperLayer = AutoSchemeWrapperLinearForGGUFK
-        if multi_option:
-            WrapperLayer = MultiOptionScoreWrapper
 
         with torch.no_grad():
             if low_gpu_mem_usage:
@@ -1259,20 +967,9 @@ def get_score_for_scheme(
                 enable_minmax_tuning=False,
                 enable_norm_bias_tuning=False,
                 enable_round_tuning=False,
-                need_weight_grad=need_weight_grad or multi_option,
+                need_weight_grad=need_weight_grad,
                 enable_torch_compile=enable_torch_compile,
             )
-            if multi_option:
-                if anchor_is_bf16:
-                    # Restore the true BF16 anchor now that the wrapper is built.
-                    new_m.orig_layer.bits = 16
-                    new_m.orig_layer.act_bits = 16
-                    new_m.enable_act_quant = False
-                new_m.set_candidates(
-                    [] if is_fixed else candidate_schemes,
-                    score_reduction=score_reduction,
-                    split_half=split_half,
-                )
             set_module(model, name, new_m)
     if offload_context is not None:
         offload_context.flush_loaded(model)
@@ -1549,10 +1246,6 @@ def _run_forward_loop(loader):
             m.grad = None
 
     for n, m in model.named_modules():
-        if multi_option:
-            if isinstance(m, MultiOptionScoreWrapper) and getattr(m, "_cand_schemes", None):
-                scores_dict[n] = [list(m.opt_scores), list(m.opt_half_scores[0]), list(m.opt_half_scores[1])]
-            continue
         if hasattr(m, "mix_score"):
             if m.orig_layer.act_bits <= 8:
                 if m.act_cnt == 0:
@@ -1561,15 +1254,15 @@ def _run_forward_loop(loader):
                     )
             layer_bits, _ = compute_layer_bits(m.orig_layer, ignore_scale_zp_bits=ignore_scale_zp_bits)
             scores_dict[n] = [layer_bits, m.mix_score]
-    if not multi_option:
-        _fill_inactive_expert_scores(scores_dict, block_names)
-        _log_score_summary_by_block_and_nonblock(
-            scores_dict,
-            block_names,
-            model=model,
-            scheme_tag=scheme_tag,
-            summary_stage="final",
-        )
+
+    _fill_inactive_expert_scores(scores_dict, block_names)
+    _log_score_summary_by_block_and_nonblock(
+        scores_dict,
+        block_names,
+        model=model,
+        scheme_tag=scheme_tag,
+        summary_stage="final",
+    )
 
     for n, m in model.named_modules():
         if hasattr(m, "orig_layer"):
@@ -2467,7 +2160,6 @@ def _gen_layer_config(
                 m.to(major_device)
 
     total_scores = {}
-    per_op_bits: dict[int, dict[str, int]] = {}
     schemes = auto_scheme.options
 
     def check_bf16_scheme(scheme):
@@ -2710,8 +2402,6 @@ def _record_scheme_scores(index, per_op_scores):
             for key, item in grouped_scores.items():
                 total_scores.setdefault(key, []).append(item)
             options_scores.append(total_loss)
-            # Bit costs are a property of the scheme, not of the expansion point.
-            per_op_bits[index] = {name: score[0] for name, score in per_op_scores.items()}
             return total_loss
 
         def _save_per_op_scores(index, scheme, cache_key, cache_path, per_op_scores):
diff --git a/auto_round/auto_scheme/gen_auto_scheme.py b/auto_round/auto_scheme/gen_auto_scheme.py
index 92901811d3..5ce81b3507 100644
--- a/auto_round/auto_scheme/gen_auto_scheme.py
+++ b/auto_round/auto_scheme/gen_auto_scheme.py
@@ -51,8 +51,8 @@ class AutoScheme:
     """Allocation solver: ``"dp"`` (knapsack DP) or ``"lagrangian"`` (shadow-price
     bisection). Both optimise the same objective and in practice produce the same
     allocation, but the Lagrangian solver hits a *fractional* avg_bits target exactly
-    and avoids the DP state explosion, making it roughly an order of magnitude faster
-    on 8B-scale models. See docs/auto_scheme_solver.md for measured accuracy/speed."""
+    and needs no discretised state space. See docs/auto_scheme_solver.md for measured
+    per-task accuracy."""
 
     def __post_init__(self):
         if isinstance(self.options, str):
diff --git a/docs/auto_scheme_solver.md b/docs/auto_scheme_solver.md
new file mode 100644
index 0000000000..f342646bab
--- /dev/null
+++ b/docs/auto_scheme_solver.md
@@ -0,0 +1,143 @@
+# AutoScheme Bit-Allocation Solver
+
+AutoScheme assigns a per-layer quantization scheme by solving a knapsack problem: every
+layer has a set of candidate options, each with a *bit cost* and a *predicted loss cost*
+(the Delta-Loss score), and the solver picks one option per layer so that the total loss
+is minimised subject to an average-bits budget.
+
+Two solvers are available, selected with `AutoScheme.solver` (or `--auto_scheme_solver`
+on the CLI):
+
+| Solver | Description |
+|:---|:---|
+| `dp` (default) | Knapsack dynamic program. Exact on a discretised bit grid, but the state space grows with the bit budget. |
+| `lagrangian` | Solves the same knapsack through its Lagrangian dual. For a price `lam` (loss per bit) every layer independently picks `argmin_s loss_s + lam * bits_s`; total bits decrease monotonically in `lam`, so a bisection drives the solution onto the budget. Hits a *fractional* avg_bits target exactly and needs no discretised state space. |
+
+Because the dual only lands on the convex hull of each layer's (bits, loss) curve, two
+primal repair passes close the integrality gap: a greedy repair that spends leftover
+budget, and a pairwise swap local search where one downgrade funds one upgrade. In
+practice the repaired dual solution is budget-tight and matches the DP allocation.
+
+## Usage
+
+```python
+from auto_round import AutoRound
+from auto_round.auto_scheme.gen_auto_scheme import AutoScheme
+
+scheme = AutoScheme(
+    avg_bits=4.5,
+    options="MXFP4,MXFP8",
+    solver="lagrangian",   # default is "dp"
+)
+ar = AutoRound(model_name, scheme=scheme)
+model, layer_config = ar.quantize()
+```
+
+CLI:
+
+```bash
+auto_round Qwen/Qwen3-8B --avg_bits 4.5 --options "MXFP4,MXFP8" --solver lagrangian
+```
+
+## Accuracy
+
+All runs use RTN (`iters=0`), identical calibration data (128 samples, seqlen 512),
+identical options and identical `avg_bits`. Only the solver differs, so any delta is
+attributable to the allocation itself. `lm_head` is not quantized.
+
+Evaluated with lm-eval on `lambada_openai`, `hellaswag`, `piqa`, `winogrande`,
+`truthfulqa_mc1`, `openbookqa`, `boolq`, `arc_easy`, `arc_challenge`, `mmlu`.
+The reported **Avg** is the mean over all evaluated tasks *including* the MMLU sub-tasks,
+so it does not equal the mean of the ten columns shown.
+
+### Summary
+
+| Experiment | Model | `dp` | `lagrangian` | Delta |
+|:---|:---|:---:|:---:|:---:|
+| INT, avg_bits 3.5 | Qwen3-8B | 0.6990 | **0.7045** | **+0.55pp** |
+| INT, avg_bits 3.5 | Llama-3.1-8B-Instruct | 0.6386 | **0.6388** | +0.02pp |
+| INT, avg_bits 3.0 | Qwen3-8B | 0.4643 | 0.4643 | 0.00pp |
+| INT, avg_bits 3.0 | Llama-3.1-8B-Instruct | 0.5794 | **0.5813** | **+0.19pp** |
+| MXFP4/8, avg_bits 4.5 | Qwen3-8B | 0.6942 | 0.6942 | 0.00pp |
+| MXFP4/8, avg_bits 4.5 | Llama-3.1-8B-Instruct | 0.6221 | 0.6221 | 0.00pp |
+
+The Lagrangian solver never lost on any tested configuration: it matched the DP in 3 of
+6 cases and beat it in 3.
+
+### Table 1 — INT W2A16/W4A16/W8A16, avg_bits 3.5
+
+**Qwen3-8B**
+
+| Solver | Bit histogram | lambada | hellaswag | piqa | winogrande | truthfulqa | openbookqa | boolq | arc_easy | arc_chal | mmlu | Avg |
+|:---|:---|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|
+| `dp` | int2 65, int4 187 | 0.3276 | 0.6763 | 0.7144 | 0.6409 | 0.3599 | 0.3880 | 0.8425 | 0.6515 | 0.4343 | 0.6954 | 0.6990 |
+| `lagrangian` | int2 63, int4 189 | 0.3088 | 0.6783 | 0.7111 | 0.6527 | 0.3660 | 0.3740 | 0.8410 | 0.6515 | 0.4292 | 0.7008 | **0.7045** |
+
+**Llama-3.1-8B-Instruct**
+
+| Solver | Bit histogram | lambada | hellaswag | piqa | winogrande | truthfulqa | openbookqa | boolq | arc_easy | arc_chal | mmlu | Avg |
+|:---|:---|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|
+| `dp` | int2 55, int4 169 | 0.6103 | 0.7549 | 0.7797 | 0.7261 | 0.3378 | 0.4200 | 0.8343 | 0.7433 | 0.4949 | 0.6281 | 0.6386 |
+| `lagrangian` | int2 52, int4 172 | 0.5973 | 0.7547 | 0.7840 | 0.7214 | 0.3366 | 0.4180 | 0.8336 | 0.7483 | 0.5051 | 0.6248 | **0.6388** |
+
+### Table 2 — INT W2A16/W4A16/W8A16, avg_bits 3.0
+
+**Qwen3-8B** — both solvers produce a *bit-identical* allocation, hence identical scores.
+
+| Solver | Bit histogram | lambada | hellaswag | piqa | winogrande | truthfulqa | openbookqa | boolq | arc_easy | arc_chal | mmlu | Avg |
+|:---|:---|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|
+| `dp` | int2 129, int4 123 | 0.1091 | 0.5721 | 0.6730 | 0.5572 | 0.3305 | 0.3320 | 0.5636 | 0.5215 | 0.3456 | 0.4466 | 0.4643 |
+| `lagrangian` | int2 129, int4 123 | 0.1091 | 0.5721 | 0.6730 | 0.5572 | 0.3305 | 0.3320 | 0.5636 | 0.5215 | 0.3456 | 0.4466 | 0.4643 |
+
+**Llama-3.1-8B-Instruct**
+
+| Solver | Bit histogram | lambada | hellaswag | piqa | winogrande | truthfulqa | openbookqa | boolq | arc_easy | arc_chal | mmlu | Avg |
+|:---|:---|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|
+| `dp` | int2 109, int4 115 | 0.5700 | 0.5870 | 0.7013 | 0.6338 | 0.3219 | 0.3880 | 0.6939 | 0.6077 | 0.4232 | 0.5649 | 0.5794 |
+| `lagrangian` | int2 106, int4 118 | 0.5750 | 0.5875 | 0.7002 | 0.6511 | 0.3182 | 0.3800 | 0.6985 | 0.6149 | 0.4258 | 0.5669 | **0.5813** |
+
+### Table 3 — MXFP4/MXFP8, avg_bits 4.5
+
+Both models produce a *bit-identical* allocation under either solver, so every per-task
+score matches exactly.
+
+**Qwen3-8B**
+
+| Solver | Bit histogram | lambada | hellaswag | piqa | winogrande | truthfulqa | openbookqa | boolq | arc_easy | arc_chal | mmlu | Avg |
+|:---|:---|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|
+| `dp` | mx_fp4 173, mx_fp8 79 | 0.6072 | 0.6943 | 0.7519 | 0.6551 | 0.3317 | 0.4060 | 0.8627 | 0.7597 | 0.5188 | 0.6795 | 0.6942 |
+| `lagrangian` | mx_fp4 173, mx_fp8 79 | 0.6072 | 0.6943 | 0.7519 | 0.6551 | 0.3317 | 0.4060 | 0.8627 | 0.7597 | 0.5188 | 0.6795 | 0.6942 |
+
+**Llama-3.1-8B-Instruct**
+
+| Solver | Bit histogram | lambada | hellaswag | piqa | winogrande | truthfulqa | openbookqa | boolq | arc_easy | arc_chal | mmlu | Avg |
+|:---|:---|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|
+| `dp` | mx_fp4 159, mx_fp8 65 | 0.6375 | 0.7557 | 0.7884 | 0.7040 | 0.3097 | 0.4060 | 0.8156 | 0.7525 | 0.5000 | 0.6122 | 0.6221 |
+| `lagrangian` | mx_fp4 159, mx_fp8 65 | 0.6375 | 0.7557 | 0.7884 | 0.7040 | 0.3097 | 0.4060 | 0.8156 | 0.7525 | 0.5000 | 0.6122 | 0.6221 |
+
+### Interpretation
+
+The two solvers optimise the same objective over the same score table, so where the
+problem has a clear optimum they converge to the *same* allocation — that is the case for
+MXFP4/8 at 4.5 bits and for INT at 3.0 bits on Qwen3-8B, where the bit histograms and all
+per-task scores are identical.
+
+Differences appear only where the knapsack has near-ties: the DP works on a discretised
+bit grid, while the Lagrangian dual handles a fractional budget exactly, so it can settle
+on a slightly different trade-off (e.g. `int2 63 / int4 189` instead of
+`int2 65 / int4 187`). In every such case the Lagrangian choice was at least as good.
+
+Note that individual tasks move in both directions even when the average improves —
+on Qwen3-8B at 3.5 bits the Lagrangian allocation loses 1.9pp on `lambada_openai` but
+gains on `winogrande`, `truthfulqa_mc1`, `hellaswag` and `mmlu`. Judge the allocation on
+the aggregate, not on any single task.
+
+### Reproducing
+
+```bash
+sh run_autoscheme_staged_ab.sh mxfp4
+```
+
+See `autoscheme_staged_ab.py` for the A/B driver and `run_autoscheme_staged_ab.sh` for
+the per-experiment settings.
+
diff --git a/docs/auto_scheme_solver_CN.md b/docs/auto_scheme_solver_CN.md
new file mode 100644
index 0000000000..537e8334c8
--- /dev/null
+++ b/docs/auto_scheme_solver_CN.md
@@ -0,0 +1,132 @@
+# AutoScheme 比特分配求解器
+
+AutoScheme 通过求解一个背包问题来为每一层分配量化方案：每层有一组候选选项，每个选项带有
+*比特代价* 和 *预测损失代价*（Delta-Loss 分数），求解器在平均比特预算的约束下为每层选择一个
+选项，使总损失最小。
+
+可通过 `AutoScheme.solver`（或命令行 `--auto_scheme_solver`）选择两种求解器：
+
+| 求解器 | 说明 |
+|:---|:---|
+| `dp`（默认） | 背包动态规划。在离散化的比特网格上是精确的，但状态空间随比特预算增长。 |
+| `lagrangian` | 通过拉格朗日对偶求解同一个背包问题。给定价格 `lam`（每比特的损失），每层独立选择 `argmin_s loss_s + lam * bits_s`；总比特数随 `lam` 单调递减，因此二分搜索可将解驱动到预算上。可精确命中 *小数* avg_bits 目标，且无需离散化状态空间。 |
+
+由于对偶解只会落在每层 (bits, loss) 曲线的凸包上，需要两次原始修复来弥合整数间隙：贪心修复用
+掉剩余预算，成对交换局部搜索用一次降级来资助一次升级。实践中修复后的对偶解恰好用满预算，且与
+DP 的分配结果一致。
+
+## 用法
+
+```python
+from auto_round import AutoRound
+from auto_round.auto_scheme.gen_auto_scheme import AutoScheme
+
+scheme = AutoScheme(
+    avg_bits=4.5,
+    options="MXFP4,MXFP8",
+    solver="lagrangian",   # 默认为 "dp"
+)
+ar = AutoRound(model_name, scheme=scheme)
+model, layer_config = ar.quantize()
+```
+
+命令行：
+
+```bash
+auto_round Qwen/Qwen3-8B --avg_bits 4.5 --options "MXFP4,MXFP8" --solver lagrangian
+```
+
+## 精度
+
+所有实验均使用 RTN（`iters=0`）、相同的校准数据（128 条样本，seqlen 512）、相同的 options 和
+相同的 `avg_bits`。唯一的变量是求解器，因此任何差异都可归因于比特分配本身。`lm_head` 不量化。
+
+使用 lm-eval 在 `lambada_openai`、`hellaswag`、`piqa`、`winogrande`、`truthfulqa_mc1`、
+`openbookqa`、`boolq`、`arc_easy`、`arc_challenge`、`mmlu` 上评测。
+表中的 **Avg** 是所有评测任务（*包含* MMLU 各子任务）的平均值，因此不等于所列十列的平均。
+
+### 汇总
+
+| 实验 | 模型 | `dp` | `lagrangian` | 差值 |
+|:---|:---|:---:|:---:|:---:|
+| INT, avg_bits 3.5 | Qwen3-8B | 0.6990 | **0.7045** | **+0.55pp** |
+| INT, avg_bits 3.5 | Llama-3.1-8B-Instruct | 0.6386 | **0.6388** | +0.02pp |
+| INT, avg_bits 3.0 | Qwen3-8B | 0.4643 | 0.4643 | 0.00pp |
+| INT, avg_bits 3.0 | Llama-3.1-8B-Instruct | 0.5794 | **0.5813** | **+0.19pp** |
+| MXFP4/8, avg_bits 4.5 | Qwen3-8B | 0.6942 | 0.6942 | 0.00pp |
+| MXFP4/8, avg_bits 4.5 | Llama-3.1-8B-Instruct | 0.6221 | 0.6221 | 0.00pp |
+
+在所有测试配置中拉格朗日求解器从未落后：6 个场景中 3 个与 DP 持平、3 个胜出。
+
+### 表 1 — INT W2A16/W4A16/W8A16, avg_bits 3.5
+
+**Qwen3-8B**
+
+| 求解器 | 比特直方图 | lambada | hellaswag | piqa | winogrande | truthfulqa | openbookqa | boolq | arc_easy | arc_chal | mmlu | Avg |
+|:---|:---|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|
+| `dp` | int2 65, int4 187 | 0.3276 | 0.6763 | 0.7144 | 0.6409 | 0.3599 | 0.3880 | 0.8425 | 0.6515 | 0.4343 | 0.6954 | 0.6990 |
+| `lagrangian` | int2 63, int4 189 | 0.3088 | 0.6783 | 0.7111 | 0.6527 | 0.3660 | 0.3740 | 0.8410 | 0.6515 | 0.4292 | 0.7008 | **0.7045** |
+
+**Llama-3.1-8B-Instruct**
+
+| 求解器 | 比特直方图 | lambada | hellaswag | piqa | winogrande | truthfulqa | openbookqa | boolq | arc_easy | arc_chal | mmlu | Avg |
+|:---|:---|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|
+| `dp` | int2 55, int4 169 | 0.6103 | 0.7549 | 0.7797 | 0.7261 | 0.3378 | 0.4200 | 0.8343 | 0.7433 | 0.4949 | 0.6281 | 0.6386 |
+| `lagrangian` | int2 52, int4 172 | 0.5973 | 0.7547 | 0.7840 | 0.7214 | 0.3366 | 0.4180 | 0.8336 | 0.7483 | 0.5051 | 0.6248 | **0.6388** |
+
+### 表 2 — INT W2A16/W4A16/W8A16, avg_bits 3.0
+
+**Qwen3-8B** —— 两种求解器产生 *完全相同* 的分配，因此分数逐位一致。
+
+| 求解器 | 比特直方图 | lambada | hellaswag | piqa | winogrande | truthfulqa | openbookqa | boolq | arc_easy | arc_chal | mmlu | Avg |
+|:---|:---|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|
+| `dp` | int2 129, int4 123 | 0.1091 | 0.5721 | 0.6730 | 0.5572 | 0.3305 | 0.3320 | 0.5636 | 0.5215 | 0.3456 | 0.4466 | 0.4643 |
+| `lagrangian` | int2 129, int4 123 | 0.1091 | 0.5721 | 0.6730 | 0.5572 | 0.3305 | 0.3320 | 0.5636 | 0.5215 | 0.3456 | 0.4466 | 0.4643 |
+
+**Llama-3.1-8B-Instruct**
+
+| 求解器 | 比特直方图 | lambada | hellaswag | piqa | winogrande | truthfulqa | openbookqa | boolq | arc_easy | arc_chal | mmlu | Avg |
+|:---|:---|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|
+| `dp` | int2 109, int4 115 | 0.5700 | 0.5870 | 0.7013 | 0.6338 | 0.3219 | 0.3880 | 0.6939 | 0.6077 | 0.4232 | 0.5649 | 0.5794 |
+| `lagrangian` | int2 106, int4 118 | 0.5750 | 0.5875 | 0.7002 | 0.6511 | 0.3182 | 0.3800 | 0.6985 | 0.6149 | 0.4258 | 0.5669 | **0.5813** |
+
+### 表 3 — MXFP4/MXFP8, avg_bits 4.5
+
+两个模型在两种求解器下都产生 *完全相同* 的分配，因此每个任务的分数都精确一致。
+
+**Qwen3-8B**
+
+| 求解器 | 比特直方图 | lambada | hellaswag | piqa | winogrande | truthfulqa | openbookqa | boolq | arc_easy | arc_chal | mmlu | Avg |
+|:---|:---|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|
+| `dp` | mx_fp4 173, mx_fp8 79 | 0.6072 | 0.6943 | 0.7519 | 0.6551 | 0.3317 | 0.4060 | 0.8627 | 0.7597 | 0.5188 | 0.6795 | 0.6942 |
+| `lagrangian` | mx_fp4 173, mx_fp8 79 | 0.6072 | 0.6943 | 0.7519 | 0.6551 | 0.3317 | 0.4060 | 0.8627 | 0.7597 | 0.5188 | 0.6795 | 0.6942 |
+
+**Llama-3.1-8B-Instruct**
+
+| 求解器 | 比特直方图 | lambada | hellaswag | piqa | winogrande | truthfulqa | openbookqa | boolq | arc_easy | arc_chal | mmlu | Avg |
+|:---|:---|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|
+| `dp` | mx_fp4 159, mx_fp8 65 | 0.6375 | 0.7557 | 0.7884 | 0.7040 | 0.3097 | 0.4060 | 0.8156 | 0.7525 | 0.5000 | 0.6122 | 0.6221 |
+| `lagrangian` | mx_fp4 159, mx_fp8 65 | 0.6375 | 0.7557 | 0.7884 | 0.7040 | 0.3097 | 0.4060 | 0.8156 | 0.7525 | 0.5000 | 0.6122 | 0.6221 |
+
+### 结果解读
+
+两种求解器在同一张分数表上优化同一个目标，因此当问题存在明确最优解时它们会收敛到 *同一个*
+分配 —— MXFP4/8 4.5 比特以及 Qwen3-8B 的 INT 3.0 比特就是这种情况，其比特直方图和所有任务分数
+都完全一致。
+
+差异只出现在背包问题存在接近平局的场景：DP 工作在离散化的比特网格上，而拉格朗日对偶可精确处理
+小数预算，因此可能落在略有不同的权衡点上（例如 `int2 63 / int4 189` 而非
+`int2 65 / int4 187`）。在所有这类情况下，拉格朗日的选择都不劣于 DP。
+
+需要注意的是，即使平均值提升，单个任务的分数仍会双向波动 —— 在 Qwen3-8B 3.5 比特上，拉格朗日
+分配在 `lambada_openai` 上低了 1.9pp，但在 `winogrande`、`truthfulqa_mc1`、`hellaswag` 和
+`mmlu` 上都有提升。评判分配质量应看总体平均，而非任何单一任务。
+
+### 复现
+
+```bash
+sh run_autoscheme_staged_ab.sh mxfp4
+```
+
+A/B 驱动脚本见 `autoscheme_staged_ab.py`，各实验的具体设置见 `run_autoscheme_staged_ab.sh`。
+

From 7c81f3c7e326cd00867d9f66040905b10f162a2d Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Mon, 24 Aug 2026 13:26:27 +0000
Subject: [PATCH 3/3] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 auto_round/auto_scheme/solver.py | 1 -
 docs/auto_scheme_solver.md       | 2 +-
 docs/auto_scheme_solver_CN.md    | 2 +-
 3 files changed, 2 insertions(+), 3 deletions(-)

diff --git a/auto_round/auto_scheme/solver.py b/auto_round/auto_scheme/solver.py
index 343e18c188..a465b0e800 100644
--- a/auto_round/auto_scheme/solver.py
+++ b/auto_round/auto_scheme/solver.py
@@ -233,4 +233,3 @@ def solve_allocation(total_scores: dict, budget: int, solver: str = "dp", max_st
         else:  # pragma: no cover - defensive
             assign[key] = min(opts, key=lambda o: o[1])
     return assign
-
diff --git a/docs/auto_scheme_solver.md b/docs/auto_scheme_solver.md
index f342646bab..f128850a87 100644
--- a/docs/auto_scheme_solver.md
+++ b/docs/auto_scheme_solver.md
@@ -27,7 +27,7 @@ from auto_round.auto_scheme.gen_auto_scheme import AutoScheme
 scheme = AutoScheme(
     avg_bits=4.5,
     options="MXFP4,MXFP8",
-    solver="lagrangian",   # default is "dp"
+    solver="lagrangian",  # default is "dp"
 )
 ar = AutoRound(model_name, scheme=scheme)
 model, layer_config = ar.quantize()
diff --git a/docs/auto_scheme_solver_CN.md b/docs/auto_scheme_solver_CN.md
index 537e8334c8..e403a5e1f1 100644
--- a/docs/auto_scheme_solver_CN.md
+++ b/docs/auto_scheme_solver_CN.md
@@ -24,7 +24,7 @@ from auto_round.auto_scheme.gen_auto_scheme import AutoScheme
 scheme = AutoScheme(
     avg_bits=4.5,
     options="MXFP4,MXFP8",
-    solver="lagrangian",   # 默认为 "dp"
+    solver="lagrangian",  # 默认为 "dp"
 )
 ar = AutoRound(model_name, scheme=scheme)
 model, layer_config = ar.quantize()