"""Dequantize FP8 and NVFP4 checkpoints to plain float weights.

OBLITERATUS performs weight surgery in float space.  Checkpoints stored in
FP8 (DeepSeek block-wise or compressed-tensors per-channel) or NVFP4
(ModelOpt or compressed-tensors) are detected at load time and dequantized
to the requested float dtype before the model is handed to the pipeline.
Output is always saved as plain float (BF16 by default) — re-quantization
is deliberately out of scope.

Pure torch + safetensors; no new dependencies.

Layout conventions handled
--------------------------
FP8 block-wise (DeepSeek-style, ``quant_method: "fp8"``):
    ``<name>.weight``           float8_e4m3fn, shape (M, N)
    ``<name>.weight_scale_inv`` float32, shape (ceil(M/128), ceil(N/128))
    dequant: w * scale_inv broadcast over ``weight_block_size`` blocks.

FP8 per-channel (compressed-tensors, ``num_bits: 8, type: "float"``):
    ``<name>.weight``        float8_e4m3fn, shape (M, N)
    ``<name>.weight_scale``  float32, shape (M, 1) or scalar
    dequant: w * weight_scale.

NVFP4 (ModelOpt, ``quant_algo: "NVFP4"``):
    ``<name>.weight``        uint8, shape (M, N/2) — two E2M1 nibbles per
                             byte along the input dim, low nibble first
    ``<name>.weight_scale``  float8_e4m3fn, shape (M, N/16) — one scale
                             per 16-element group
    ``<name>.weight_scale_2`` float32 scalar — global scale (amax/2688)
    dequant: e2m1_values * weight_scale * weight_scale_2

NVFP4 (compressed-tensors, ``format: "nvfp4-pack-quantized"``):
    ``<name>.weight_packed``       uint8 packed E2M1 values
    ``<name>.weight_scale``        FP8 block scales (multiplicative)
    ``<name>.weight_global_scale`` FP32 reciprocal global scale
    dequant: e2m1_values * weight_scale / weight_global_scale.
"""

from __future__ import annotations

import json
import logging
import os
from dataclasses import dataclass, field
from enum import Enum
from typing import Dict, Optional, Tuple

import torch

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Dtypes / constants
# ---------------------------------------------------------------------------

def _fp8_dtypes() -> frozenset:
    """torch float8 dtypes available in this build (torch >= 2.1)."""
    out = set()
    for name in ("float8_e4m3fn", "float8_e5m2", "float8_e4m3fnuz", "float8_e5m2fnuz"):
        dt = getattr(torch, name, None)
        if dt is not None:
            out.add(dt)
    return frozenset(out)


FP8_DTYPES = _fp8_dtypes()

# E2M1 magnitude values indexed by the low 3 bits; bit 3 is the sign.
E2M1_POSITIVE = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0)
E2M1_LUT = torch.tensor(
    [v for v in E2M1_POSITIVE] + [-v for v in E2M1_POSITIVE],
    dtype=torch.float32,
)

NVFP4_GROUP_SIZE = 16
FP8_DEFAULT_BLOCK = (128, 128)


def is_fp8_dtype(dtype: torch.dtype) -> bool:
    return dtype in FP8_DTYPES


# ---------------------------------------------------------------------------
# Scheme detection
# ---------------------------------------------------------------------------

class QuantScheme(Enum):
    NONE = "none"
    FP8_BLOCKWISE = "fp8_blockwise"
    FP8_PER_CHANNEL_CT = "fp8_per_channel_compressed_tensors"
    NVFP4_MODELOPT = "nvfp4_modelopt"
    NVFP4_CT = "nvfp4_compressed_tensors"
    UNSUPPORTED = "unsupported"


@dataclass
class QuantDetection:
    scheme: QuantScheme
    reason: str = ""
    block_size: Tuple[int, int] = FP8_DEFAULT_BLOCK
    group_size: int = NVFP4_GROUP_SIZE
    scale_is_inverse: bool = False
    global_scale_is_inverse: bool = False
    raw_quant_config: dict = field(default_factory=dict)


def _load_json_from_checkpoint(
    model_name_or_path: str,
    filename: str,
    token: Optional[str] = None,
    revision: Optional[str] = None,
    local_files_only: bool = False,
) -> Optional[dict]:
    """Read a JSON metadata file from a local dir or the HF hub."""
    if os.path.isdir(model_name_or_path):
        path = os.path.join(model_name_or_path, filename)
        if not os.path.exists(path):
            return None
        with open(path, "r", encoding="utf-8") as fh:
            return json.load(fh)
    try:
        from huggingface_hub import hf_hub_download
        from huggingface_hub.utils import EntryNotFoundError
    except Exception:  # pragma: no cover - huggingface_hub always ships w/ transformers
        return None
    try:
        path = hf_hub_download(
            model_name_or_path,
            filename,
            token=token,
            revision=revision,
            local_files_only=local_files_only,
        )
    except EntryNotFoundError:
        return None
    except Exception as exc:
        logger.debug("could not fetch %s for %s: %s", filename, model_name_or_path, exc)
        return None
    with open(path, "r", encoding="utf-8") as fh:
        return json.load(fh)


def _safetensors_key_names(
    model_name_or_path: str,
    config_json: Optional[dict],
    token: Optional[str] = None,
    revision: Optional[str] = None,
    local_files_only: bool = False,
) -> set:
    """Collect tensor key names from the safetensors index (no weight loads)."""
    keys: set = set()
    index = _load_json_from_checkpoint(
        model_name_or_path,
        "model.safetensors.index.json",
        token=token,
        revision=revision,
        local_files_only=local_files_only,
    )
    if index and "weight_map" in index:
        keys.update(index["weight_map"].keys())
        return keys
    # Single-file checkpoint: open headers only.
    if os.path.isdir(model_name_or_path):
        st_path = os.path.join(model_name_or_path, "model.safetensors")
    else:
        try:
            from huggingface_hub import hf_hub_download
            st_path = hf_hub_download(
                model_name_or_path,
                "model.safetensors",
                token=token,
                revision=revision,
                local_files_only=local_files_only,
            )
        except Exception:
            return keys
    if not os.path.exists(st_path):
        return keys
    try:
        from safetensors import safe_open
        with safe_open(st_path, framework="pt", device="cpu") as fh:
            keys.update(fh.keys())
    except Exception as exc:
        logger.debug("could not read safetensors header %s: %s", st_path, exc)
    return keys


def detect_quant_scheme(
    model_name_or_path: str,
    token: Optional[str] = None,
    revision: Optional[str] = None,
    local_files_only: bool = False,
) -> QuantDetection:
    """Classify a checkpoint's quantization without loading any weights.

    Peeks at ``config.json``'s ``quantization_config`` plus safetensors
    key names.  Anything quantized that we do not explicitly support is
    reported as UNSUPPORTED with a human-readable reason — the loader
    turns that into a loud error rather than silent corruption.
    """
    config_json = _load_json_from_checkpoint(
        model_name_or_path,
        "config.json",
        token=token,
        revision=revision,
        local_files_only=local_files_only,
    ) or {}
    qcfg = config_json.get("quantization_config")
    if not qcfg:
        return QuantDetection(QuantScheme.NONE)

    quant_method = str(qcfg.get("quant_method", "")).lower()
    raw = dict(qcfg)

    if quant_method == "fp8":
        block = qcfg.get("weight_block_size")
        if block:
            return QuantDetection(
                QuantScheme.FP8_BLOCKWISE,
                block_size=(int(block[0]), int(block[1])),
                raw_quant_config=raw,
            )
        keys = _safetensors_key_names(
            model_name_or_path,
            config_json,
            token=token,
            revision=revision,
            local_files_only=local_files_only,
        )
        if any(k.endswith("weight_scale_inv") for k in keys):
            return QuantDetection(QuantScheme.FP8_BLOCKWISE, raw_quant_config=raw)
        return QuantDetection(
            QuantScheme.UNSUPPORTED,
            reason=(
                f"quant_method 'fp8' without weight_block_size/weight_scale_inv "
                f"(activation scheme {qcfg.get('activation_scheme')!r}); only "
                f"DeepSeek-style block-wise FP8 and compressed-tensors FP8 are "
                f"supported — use a BF16 checkpoint"
            ),
            raw_quant_config=raw,
        )

    if quant_method == "modelopt":
        algo = str(qcfg.get("quant_algo", "")).upper()
        kv = str(qcfg.get("kv_cache_quant_algo", "") or "").upper()
        if "MIXED" in algo:
            # Mixed-precision ModelOpt (e.g. FP8 mixer + NVFP4 experts):
            # inspect config_groups — any 4-bit group means NVFP4 tensors are
            # present; dequantization is per-tensor so FP8 tensors still take
            # the FP8 path automatically.
            bits = {
                (grp or {}).get("weights", {}).get("num_bits")
                for grp in (qcfg.get("config_groups") or {}).values()
            }
            if 4 in bits:
                return QuantDetection(QuantScheme.NVFP4_MODELOPT, raw_quant_config=raw)
            return QuantDetection(QuantScheme.FP8_PER_CHANNEL_CT, raw_quant_config=raw)
        if "NVFP4" in algo or "FP4" in algo:
            return QuantDetection(QuantScheme.NVFP4_MODELOPT, raw_quant_config=raw)
        if "FP8" in algo:
            # ModelOpt FP8 checkpoints are per-tensor; treat like the
            # per-channel path with scalar scales.
            return QuantDetection(QuantScheme.FP8_PER_CHANNEL_CT, raw_quant_config=raw)
        return QuantDetection(
            QuantScheme.UNSUPPORTED,
            reason=f"modelopt quant_algo {algo or kv or 'unknown'!r} not supported",
            raw_quant_config=raw,
        )

    if quant_method == "compressed-tensors":
        groups = qcfg.get("config_groups") or {}
        wcfg = {}
        for grp in groups.values():
            w = (grp or {}).get("weights") or {}
            if w.get("num_bits") is not None:
                wcfg = w
                break
        num_bits = wcfg.get("num_bits")
        wtype = str(wcfg.get("type", "")).lower()
        if num_bits == 8 and wtype == "float":
            return QuantDetection(QuantScheme.FP8_PER_CHANNEL_CT, raw_quant_config=raw)
        if (
            num_bits == 4
            and wtype == "float"
            and qcfg.get("format") == "nvfp4-pack-quantized"
        ):
            gs = int(wcfg.get("group_size") or NVFP4_GROUP_SIZE)
            return QuantDetection(
                QuantScheme.NVFP4_CT,
                group_size=gs,
                global_scale_is_inverse=True,
                raw_quant_config=raw,
            )
        return QuantDetection(
            QuantScheme.UNSUPPORTED,
            reason=(
                f"compressed-tensors weights num_bits={num_bits} type={wtype!r} "
                f"format={qcfg.get('format')!r} not supported (need float 8-bit "
                f"or NVFP4 4-bit with format 'nvfp4-pack-quantized')"
            ),
            raw_quant_config=raw,
        )

    if quant_method in ("gptq", "awq", "bitsandbytes", "bitsandbytes_4bit",
                        "bitsandbytes_8bit", "mxfp4"):
        # Handled elsewhere in the loader / surgery layer.
        return QuantDetection(QuantScheme.NONE, raw_quant_config=raw)

    return QuantDetection(
        QuantScheme.UNSUPPORTED,
        reason=f"quant_method {quant_method!r} not supported — use a BF16 checkpoint",
        raw_quant_config=raw,
    )


# ---------------------------------------------------------------------------
# FP8 dequantization
# ---------------------------------------------------------------------------

def _upcast_fp8(t: torch.Tensor) -> torch.Tensor:
    """float8 (or uint8-viewed-as-float8) tensor → float32."""
    if t.dtype in FP8_DTYPES:
        return t.to(torch.float32)
    if t.dtype == torch.uint8 and getattr(torch, "float8_e4m3fn", None) is not None:
        return t.view(torch.float8_e4m3fn).to(torch.float32)
    raise RuntimeError(f"cannot upcast dtype {t.dtype} as FP8")


def dequantize_fp8_blockwise(
    w_fp8: torch.Tensor,
    scale_inv: torch.Tensor,
    block_size: Tuple[int, int] = FP8_DEFAULT_BLOCK,
) -> torch.Tensor:
    """DeepSeek-style block-wise FP8 → float32.

    ``w_fp8``: (M, N) float8. ``scale_inv``: (ceil(M/bm), ceil(N/bn)) — one
    scale per ``block_size`` tile, applied multiplicatively.
    """
    w = _upcast_fp8(w_fp8)
    bm, bn = int(block_size[0]), int(block_size[1])
    M, N = w.shape[-2], w.shape[-1]
    s = scale_inv.to(torch.float32)
    if s.shape[-2] * bm < M or s.shape[-1] * bn < N:
        raise RuntimeError(
            f"weight_scale_inv shape {tuple(s.shape)} incompatible with "
            f"weight shape {(M, N)} and block {block_size}"
        )
    s = s.repeat_interleave(bm, dim=-2).repeat_interleave(bn, dim=-1)
    s = s[..., :M, :N]
    w.mul_(s)  # in-place: float32 weights are already 2x the BF16 size
    return w


def dequantize_fp8_per_channel(
    w_fp8: torch.Tensor,
    scale: torch.Tensor,
    scale_is_inverse: bool = False,
) -> torch.Tensor:
    """Per-channel (or per-tensor) FP8 → float32."""
    w = _upcast_fp8(w_fp8)
    s = scale.to(torch.float32)
    while s.ndim < w.ndim:
        s = s.unsqueeze(-1)
    if scale_is_inverse:
        w.div_(s)
    else:
        w.mul_(s)
    return w


# ---------------------------------------------------------------------------
# NVFP4 dequantization
# ---------------------------------------------------------------------------

_NATIVE_FP4_OK: Optional[bool] = None
_NATIVE_FP4_LOW_FIRST: bool = True


def _native_fp4_upcast_works() -> bool:
    """Probe: does this torch build upcast float4_e2m1fn_x2 correctly?

    Op support for the packed-FP4 dtype is spotty across torch versions,
    so we verify with a known byte pattern instead of trusting a version
    check.  Byte 0x1B = low nibble 0xB (-1.5), high nibble 0x1 (+0.5).
    Also detects the nibble order the native path uses.
    """
    global _NATIVE_FP4_OK, _NATIVE_FP4_LOW_FIRST
    if _NATIVE_FP4_OK is not None:
        return _NATIVE_FP4_OK
    _NATIVE_FP4_OK = False
    dt = getattr(torch, "float4_e2m1fn_x2", None)
    if dt is None:
        return False
    try:
        packed = torch.tensor([0x1B], dtype=torch.uint8)
        vals = packed.view(dt).to(torch.float32).flatten()
        if vals.numel() != 2:
            return False
        got = vals.tolist()
        # ModelOpt convention is low-nibble-first: byte 0x1B → [-1.5, +0.5].
        if got == [-1.5, 0.5]:
            _NATIVE_FP4_OK, _NATIVE_FP4_LOW_FIRST = True, True
        elif got == [0.5, -1.5]:
            _NATIVE_FP4_OK, _NATIVE_FP4_LOW_FIRST = True, False
    except Exception:
        _NATIVE_FP4_OK = False
    return _NATIVE_FP4_OK


def _unpack_e2m1_native(packed: torch.Tensor) -> torch.Tensor:
    """uint8 (…, K) → float32 (…, 2K) via the native FP4 dtype."""
    dt = torch.float4_e2m1fn_x2
    vals = packed.view(dt).to(torch.float32)
    vals = vals.reshape(*packed.shape[:-1], packed.shape[-1] * 2)
    if not _NATIVE_FP4_LOW_FIRST:
        # Native path emitted high nibble first: swap adjacent pairs.
        vals = vals.reshape(*vals.shape[:-1], -1, 2).flip(-1).reshape(vals.shape)
    return vals


def _unpack_e2m1_manual(packed: torch.Tensor) -> torch.Tensor:
    """uint8 (…, K) → float32 (…, 2K), low nibble first, via LUT.

    Memory-conscious: LUT indexing requires int64 indices (8 B/value), so
    the leading dim is processed in chunks to bound the transient index
    tensors — indexing a whole 30B-model shard naively peaks at >100 GB.
    """
    lut = E2M1_LUT.to(device=packed.device)
    out = torch.empty(
        *packed.shape[:-1], packed.shape[-1] * 2,
        dtype=torch.float32, device=packed.device,
    )
    flat = packed.reshape(-1, packed.shape[-1])
    out_flat = out.reshape(-1, packed.shape[-1] * 2)
    # ~64M nibbles per chunk → ≤0.5 GB transient int64
    chunk = max(1, (64 * 1024 * 1024) // (packed.shape[-1] * 2))
    for i in range(0, flat.shape[0], chunk):
        pk = flat[i:i + chunk]
        dst = out_flat[i:i + chunk]
        dst[:, 0::2] = lut[(pk & 0x0F).long()]
        dst[:, 1::2] = lut[(pk >> 4).long()]
    return out


def unpack_e2m1(packed: torch.Tensor, force_manual: bool = False) -> torch.Tensor:
    """Unpack NVFP4 nibbles to float32, preferring the native dtype."""
    if not force_manual and _native_fp4_upcast_works():
        try:
            return _unpack_e2m1_native(packed)
        except Exception:
            logger.debug("native FP4 upcast failed at runtime; using manual LUT")
    return _unpack_e2m1_manual(packed)


def _upcast_fp8_scales(t: torch.Tensor) -> torch.Tensor:
    """Block scales may be stored as float8 or as raw uint8 bytes."""
    if t.dtype in FP8_DTYPES or t.dtype == torch.uint8:
        return _upcast_fp8(t)
    return t.to(torch.float32)


def dequantize_nvfp4(
    packed_uint8: torch.Tensor,
    block_scale: torch.Tensor,
    global_scale: Optional[torch.Tensor],
    out_shape: Optional[Tuple[int, ...]] = None,
    scale_is_inverse: bool = False,
    global_scale_is_inverse: bool = False,
    group_size: int = NVFP4_GROUP_SIZE,
    force_manual: bool = False,
) -> torch.Tensor:
    """NVFP4 → float32.

    ``packed_uint8``: (M, N/2) uint8, two E2M1 values per byte, low nibble
    first along the input dim.  ``block_scale``: one FP8-E4M3 scale per
    ``group_size`` elements.  ``global_scale``: FP32 scalar (ModelOpt
    ``weight_scale_2`` = amax/2688); None means 1.0.  ModelOpt stores both
    scales multiplicatively.  Compressed-tensors keeps the block scale but
    stores the global scale as its reciprocal, selected with
    ``global_scale_is_inverse``.  ``scale_is_inverse`` is retained for
    explicitly reciprocal block-scale layouts.
    """
    vals = unpack_e2m1(packed_uint8, force_manual=force_manual)
    *lead, N = vals.shape
    if N % group_size != 0:
        raise RuntimeError(
            f"unpacked NVFP4 dim {N} not divisible by group_size {group_size}"
        )
    vals = vals.reshape(*lead, N // group_size, group_size)
    bs = _upcast_fp8_scales(block_scale).reshape(*lead, N // group_size, 1)
    # In-place scaling — full-model float32 copies are 2x the BF16 size each.
    if scale_is_inverse:
        vals.div_(bs)
    else:
        vals.mul_(bs)
    vals = vals.reshape(*lead, N)
    if global_scale is not None:
        gs = global_scale.to(torch.float32)
        if global_scale_is_inverse:
            vals.div_(gs)
        else:
            vals.mul_(gs)
    if out_shape is not None and tuple(vals.shape) != tuple(out_shape):
        vals = vals.reshape(out_shape)
    return vals


# ---------------------------------------------------------------------------
# State-dict level dequantization (used by the loader)
# ---------------------------------------------------------------------------

_SCALE_SUFFIXES = (
    # Dense-linears (``foo.weight`` + ``foo.weight_scale``)
    ".weight_scale_inv",
    ".weight_scale_2",
    ".weight_global_scale",
    ".weight_scale",
    ".input_scale",
    ".input_global_scale",
    ".activation_scale",
    # Fused MoE experts (``experts.gate_up_proj`` + ``experts.gate_up_proj_scale``)
    "_scale_inv",
    "_scale_2",
    "_global_scale",
    "_input_scale",
    "_scale",
)


def _strip_known_suffix(key: str) -> Optional[str]:
    for suf in _SCALE_SUFFIXES:
        if key.endswith(suf):
            return key[: -len(suf)]
    return None


def _scale_sibling_names(weight_key: str) -> tuple[tuple[str, str], ...]:
    """Return candidate ``(key, role)`` pairs for one quantized weight."""
    if weight_key.endswith(".weight_packed"):
        base = weight_key[: -len(".weight_packed")]
        dotted = True
    elif weight_key.endswith(".weight"):
        base = weight_key[: -len(".weight")]
        dotted = True
    else:
        base = weight_key
        dotted = False
    suffixes = (
        (
            (".weight_scale_inv", "scale_inv"),
            (".weight_scale_2", "global_scale"),
            (".weight_global_scale", "global_scale"),
            (".weight_scale", "block_scale"),
        )
        if dotted
        else (
            # Fused-MoE naming:
            # experts.gate_up_proj{,_scale,_scale_2,_scale_inv}
            ("_scale_inv", "scale_inv"),
            ("_scale_2", "global_scale"),
            ("_global_scale", "global_scale"),
            ("_scale", "block_scale"),
        )
    )
    return tuple((base + suffix, role) for suffix, role in suffixes)


def find_scale_siblings(state_dict: Dict[str, torch.Tensor], weight_key: str) -> dict:
    """Locate the scale tensors belonging to ``weight_key`` (``foo.weight``)."""
    out = {}
    for key, role in _scale_sibling_names(weight_key):
        if key in state_dict and role not in out:
            out[role] = state_dict[key]
    return out


def is_scale_key(
    key: str,
    state_dict: Optional[Dict[str, torch.Tensor]] = None,
    known_keys: Optional[set[str]] = None,
) -> bool:
    """True for keys holding quantization scales, not real parameters.

    With ``state_dict`` given, the key only counts as a scale if its base
    tensor exists in the shard — otherwise a legitimate float parameter
    that merely *ends* in ``_scale`` (e.g. ``logit_scale``) would be
    silently dropped from the dequantized checkpoint.
    """
    base = _strip_known_suffix(key)
    if base is None:
        return False
    keys = known_keys if known_keys is not None else state_dict
    if keys is not None:
        # Dotted scales pair with dense or packed weights; fused-MoE scales
        # pair directly with ``base``.
        if (
            base not in keys
            and (base + ".weight") not in keys
            and (base + ".weight_packed") not in keys
        ):
            return False
    return True


def dequantize_state_dict(
    state_dict: Dict[str, torch.Tensor],
    detection: QuantDetection,
    out_dtype: torch.dtype = torch.bfloat16,
    sibling_tensors: Optional[Dict[str, torch.Tensor]] = None,
    known_keys: Optional[set[str]] = None,
) -> Dict[str, torch.Tensor]:
    """Dequantize every quantized weight group in one checkpoint shard.

    Returns a new dict with float tensors only: scale/aux tensors are
    dropped, quantized weights become ``out_dtype``.  Works generically
    over naming conventions (dense linears and fused MoE experts alike)
    because grouping is by key suffix, not by module type.
    """
    scale_lookup = state_dict
    if sibling_tensors:
        scale_lookup = {**state_dict, **sibling_tensors}
    out: Dict[str, torch.Tensor] = {}
    for key, tensor in state_dict.items():
        if is_scale_key(key, state_dict, known_keys):
            # Scales are consumed together with their weight (or dropped if
            # the weight is absent — e.g. quantization status tensors).
            continue
        sib = find_scale_siblings(scale_lookup, key)
        if is_fp8_dtype(tensor.dtype):
            if "scale_inv" in sib:
                w = dequantize_fp8_blockwise(tensor, sib["scale_inv"], detection.block_size)
            elif "block_scale" in sib:
                w = dequantize_fp8_per_channel(
                    tensor, sib["block_scale"], detection.scale_is_inverse
                )
            else:
                raise RuntimeError(
                    f"FP8 tensor {key!r} has no weight_scale/weight_scale_inv "
                    f"sibling — unsupported layout; report a bug"
                )
            out[key] = w.to(out_dtype)
        elif tensor.dtype == torch.uint8 and ("block_scale" in sib or "scale_inv" in sib):
            gs = sib.get("global_scale")
            scale = sib.get("block_scale", sib.get("scale_inv"))
            if detection.scheme in (QuantScheme.NVFP4_MODELOPT, QuantScheme.NVFP4_CT):
                w = dequantize_nvfp4(
                    tensor, scale, gs,
                    scale_is_inverse=detection.scale_is_inverse,
                    global_scale_is_inverse=detection.global_scale_is_inverse,
                    group_size=detection.group_size,
                )
            else:
                raise RuntimeError(
                    f"packed uint8 tensor {key!r} in a {detection.scheme.value} "
                    f"checkpoint — unsupported layout; report a bug"
                )
            output_key = key[: -len("_packed")] if key.endswith(".weight_packed") else key
            out[output_key] = w.to(out_dtype)
        elif tensor.dtype == torch.uint8 and key.endswith((".weight", ".weight_packed")):
            raise RuntimeError(
                f"uint8 weight {key!r} has no recognizable scale siblings "
                f"(keys present: {[k for k in state_dict if k.startswith(key[:-7])]}) "
                f"— unsupported packed layout; use a BF16 checkpoint"
            )
        else:
            out[key] = tensor
    return out



# ---------------------------------------------------------------------------
# Checkpoint materialization (dequantize an on-disk checkpoint to float)
# ---------------------------------------------------------------------------

def materialize_dequantized_checkpoint(
    model_name_or_path: str,
    detection: QuantDetection,
    out_dtype: torch.dtype = torch.bfloat16,
    token: Optional[str] = None,
    revision: Optional[str] = None,
    local_files_only: bool = False,
) -> Tuple[str, str]:
    """Write a dequantized float copy of a checkpoint to a temp dir.

    Returns ``(tmp_dir, source_dir)``.  The tmp dir is a complete
    standalone checkpoint (config sans ``quantization_config``, tokenizer,
    safetensors shards + index) that ``from_pretrained`` can load through
    the normal float path — no custom module placement needed.  Shards are
    processed one at a time so peak RAM stays near one shard.
    """
    import shutil
    import tempfile

    from safetensors import safe_open
    from safetensors.torch import load_file, save_file

    if os.path.isdir(model_name_or_path):
        src = model_name_or_path
    else:
        from huggingface_hub import snapshot_download

        src = snapshot_download(
            model_name_or_path,
            token=token,
            revision=revision,
            local_files_only=local_files_only,
            allow_patterns=[
                "*.json", "*.safetensors", "*.py", "tokenizer*", "*.model",
                "*.txt", "chat_template*", "special_tokens_map.json",
                "generation_config.json",
            ],
        )

    tmp = tempfile.mkdtemp(prefix="obliteratus_dequant_")
    logger.warning(
        "Dequantizing %s checkpoint %s -> %s (temporary dir %s)",
        detection.scheme.value, model_name_or_path, out_dtype, tmp,
    )

    try:
        # Copy non-weight files; strip quantization_config from config.json.
        for name in os.listdir(src):
            s = os.path.join(src, name)
            if not os.path.isfile(s) or name.endswith(".safetensors"):
                continue
            d = os.path.join(tmp, name)
            if name == "config.json":
                with open(s, "r", encoding="utf-8") as fh:
                    cfg = json.load(fh)
                cfg.pop("quantization_config", None)
                with open(d, "w", encoding="utf-8") as fh:
                    json.dump(cfg, fh, indent=2)
            else:
                shutil.copy2(s, d)

        # Locate shards.
        index_path = os.path.join(src, "model.safetensors.index.json")
        old_index = None
        if os.path.exists(index_path):
            with open(index_path, "r", encoding="utf-8") as fh:
                old_index = json.load(fh)
            shards = sorted(set(old_index["weight_map"].values()))
        elif os.path.exists(os.path.join(src, "model.safetensors")):
            shards = ["model.safetensors"]
        else:
            raise RuntimeError(
                f"Quantized checkpoint '{model_name_or_path}' has no safetensors "
                f"weights (pytorch_model.bin quantized checkpoints are not "
                f"supported) — use a BF16 checkpoint"
            )

        weight_map = dict(old_index.get("weight_map") or {}) if old_index else {}
        known_keys = set(weight_map) if weight_map else None
        new_weight_map: Dict[str, str] = {}
        new_total_size = 0
        for i, shard in enumerate(shards, 1):
            sd = load_file(os.path.join(src, shard), device="cpu")
            external_scales: Dict[str, torch.Tensor] = {}
            external_by_shard: Dict[str, list[str]] = {}
            for weight_key in sd:
                for scale_key, _role in _scale_sibling_names(weight_key):
                    scale_shard = weight_map.get(scale_key)
                    if scale_shard and scale_shard != shard:
                        external_by_shard.setdefault(scale_shard, []).append(scale_key)
            for scale_shard, scale_keys in external_by_shard.items():
                with safe_open(
                    os.path.join(src, scale_shard), framework="pt", device="cpu",
                ) as handle:
                    for scale_key in scale_keys:
                        external_scales[scale_key] = handle.get_tensor(scale_key)

            out_sd = dequantize_state_dict(
                sd,
                detection,
                out_dtype=out_dtype,
                sibling_tensors=external_scales,
                known_keys=known_keys,
            )
            if out_sd:
                save_file(out_sd, os.path.join(tmp, shard), metadata={"format": "pt"})
            for k, tensor in out_sd.items():
                new_weight_map[k] = shard
                new_total_size += tensor.numel() * tensor.element_size()
            logger.info(
                "dequantized shard %d/%d (%s): %d tensors -> %d float tensors",
                i, len(shards), shard, len(sd), len(out_sd),
            )

        if old_index is not None:
            metadata = dict(old_index.get("metadata") or {})
            metadata["total_size"] = new_total_size
            with open(
                os.path.join(tmp, "model.safetensors.index.json"),
                "w",
                encoding="utf-8",
            ) as fh:
                json.dump({"metadata": metadata, "weight_map": new_weight_map}, fh, indent=2)
    except BaseException:
        shutil.rmtree(tmp, ignore_errors=True)
        raise

    return tmp, src
