feat(MAD-238): paged ml8-4 calibration pipeline + iter 5 PagedLinear

Layer-sequential ml8-4 calibration that runs the entire model through
the WeightPager VRAM ring instead of loading every weight resident.
Unlocks calibration of models larger than VRAM (35B-A3B target) on a
single 32 GiB card.

Pipeline (scripts/calibration/):
  calibrate_ml8_paged.py    main driver: load HF model, swap nn.Linear
                            with PagedLinear, run calibration data
                            layer-by-layer with each layer's weights
                            paged in on-demand and evicted after.
  awq.py                    AWQ rescaling (per-output-channel) layered
                            over the centroid quantizer.
  kronecker_rotation.py     Kronecker-product Hadamard rotation (small
                            factor on rows, large factor on columns)
                            for outlier mitigation.
  ml8_to_gguf.py            blob -> GGUF tensor format converter.
  ml8_calibration_sweep.sh  matrix-sweep harness (percdamp / Hadamard /
                            group_size cells).
  MAD_238_DESIGN.md         design doc.

Plus per-module test_*.py covering awq, centroid snap, kronecker
rotation, ml8_io round-trip, ml8_to_gguf format.

Python bindings (python_bindings/wp/):
  wp_bindings.cpp     pybind11 wrappers over the WeightPager runtime.
  paged_linear.py     PagedLinear nn.Module: torch tensor wrapped over
                      a pager-managed VRAM ptr; .weight getter calls
                      pager.ensure() and caches the torch view until
                      the slot pointer changes (zero-copy hot path,
                      eager VRAM release on cache invalidate).
  wp_torch.py         helpers: swap_linears_with_paged, page assignment
                      from HF layer names.
  test_*.py           unit + smoke tests for native and torch paths.

Parity gate (Qwen3.5-4B, Cell E recipe, MLP-only):
  HF-loaded baseline    : 8.318 PPL (100k tokens, wikitext-2)
  paged calibration     : 8.328 PPL
  delta                 : +0.0099 (within ROCm GEMM non-determinism
                          noise; root-cause confirmed via byte-identical
                          HF safetensors vs bf16 GGUF for layer 0)

Also lands miscellaneous in-flight artifacts:

  tests/wmma_rdna4_probe{,_v2}.cu  Standalone probes for the RDNA4
                                   WMMA layout investigation (parked
                                   until probe-driven mma.cuh fix).

  tests/perf-baseline/             Accumulated bench JSONs from
                                   May 19-23 paged-paths / NIAH-KV /
                                   PPL-KV runs.

  gguf_f16_to_bf16.py              F16->BF16 GGUF recast utility built
                                   during 0.048 PPL drift investigation,
                                   then superseded by direct byte-
                                   comparison; kept for reuse.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

Author: kmbandy

python_bindings/wp/paged_linear.py

@@ -0,0 +1,191 @@
+"""PagedLinear — nn.Linear whose .weight comes from a wp_native.WeightPager.
+
+Design: at model-load time, swap each target nn.Linear in an HF transformers
+model with a PagedLinear. HF's forward code uses `self.q_proj(x)` and accesses
+`.weight` transparently — the paged page-fault is invisible to the layer's
+arithmetic.
+
+Two weight paths:
+  1. weight_override (post-quant): a torch tensor set by the calibration loop
+     after quantizing this layer. While set, .weight returns the override and
+     the pager is bypassed. Used so subsequent layers' forward pass sees the
+     quantized output of this layer.
+  2. paged (pre-quant): ensure_as_torch() reads from the pager slot. Used
+     during forward passes that need the original weight (Hessian collection
+     pass for THIS layer).
+
+Lifetime: once weight_override is set, the pager slot is free to LRU-evict.
+
+Caching: the paged path caches the materialized tensor keyed by the slot's
+src_ptr. Because GGUF pages are read-only, identical src_ptr → identical bytes,
+so cache reuse is safe. Each .weight access calls pager.ensure() (O(1) hash
+lookup when resident); if the returned ptr matches the cached one, the cached
+tensor is returned without re-memcpy. When eviction reassigns the slot to a
+different page, src_ptr changes and the cache invalidates automatically.
+"""
+from __future__ import annotations
+
+from typing import Mapping, Optional
+
+import torch
+import torch.nn as nn
+
+
+class PagedLinear(nn.Linear):
+    """nn.Linear with weight backed by a WeightPager (or a post-quant override)."""
+
+    def __init__(self, pager, page_idx: int,
+                 weight_shape, weight_dtype: torch.dtype,
+                 bias: bool = False):
+        # weight_shape is (out_features, in_features) per HF convention.
+        out_features, in_features = int(weight_shape[0]), int(weight_shape[1])
+        super().__init__(in_features, out_features, bias=bias)
+
+        # Drop the auto-allocated weight Parameter — pager (or override) owns it.
+        if "weight" in self._parameters:
+            del self._parameters["weight"]
+
+        self.pager = pager
+        self.page_idx = page_idx
+        self.weight_shape = (out_features, in_features)
+        self.weight_dtype = weight_dtype
+        self.weight_override: Optional[torch.Tensor] = None
+        # Device index for the paged path (only used when override is unset).
+        self.device_idx = 0
+        # Cache for the page-faulted weight. _cached_src_ptr=0 is the
+        # "no cache" sentinel (real GPU pointers are always nonzero).
+        self._cached_weight: Optional[torch.Tensor] = None
+        self._cached_src_ptr: int = 0
+
+    def _materialize_weight(self, src_ptr: int) -> torch.Tensor:
+        """Allocate a fresh CUDA tensor and memcpy the slot bytes into it.
+
+        Extracted as a method so tests can override it without touching GPU.
+        Subclasses MUST return a torch.Tensor with shape=self.weight_shape and
+        dtype=self.weight_dtype.
+        """
+        import wp_native
+        out = torch.empty(self.weight_shape, dtype=self.weight_dtype,
+                          device=f"cuda:{self.device_idx}")
+        nbytes = out.element_size() * out.numel()
+        wp_native.device_memcpy(out.data_ptr(), src_ptr, nbytes)
+        return out
+
+    @property
+    def weight(self) -> torch.Tensor:
+        """Resolved weight tensor: override first, then cache, then materialize.
+
+        Caching: pager.ensure() returns the slot's src_ptr (O(1) when resident).
+        If src_ptr matches what we last materialized from, the cached tensor is
+        still valid (GGUF bytes are immutable) and we return it without copying.
+        When eviction reassigns the slot to a different page, src_ptr changes
+        and we re-materialize.
+        """
+        if self.weight_override is not None:
+            return self.weight_override
+        if self.pager is None:
+            raise RuntimeError(
+                "PagedLinear.weight: no pager attached AND no weight_override set. "
+                "Either wire a wp_native.WeightPager + page_idx, or assign weight_override."
+            )
+        src_ptr = self.pager.ensure(self.page_idx)
+        if src_ptr == 0:
+            raise RuntimeError(
+                f"PagedLinear.weight: pager.ensure(page_idx={self.page_idx}) "
+                f"returned null — page out of range or pager error."
+            )
+        if self._cached_weight is not None and self._cached_src_ptr == src_ptr:
+            return self._cached_weight
+        out = self._materialize_weight(src_ptr)
+        self._cached_weight = out
+        self._cached_src_ptr = src_ptr
+        return out
+
+    @weight.setter
+    def weight(self, value):
+        """Allow `layer.weight = some_tensor` to set the override.
+
+        Clears the paged-path cache. The override takes precedence in the getter,
+        but releasing the cached tensor frees its VRAM eagerly — useful in the
+        calibration loop where post-quant overrides are set per-layer."""
+        self.weight_override = value
+        self._cached_weight = None
+        self._cached_src_ptr = 0
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        """Compute in weight dtype, return in input dtype — fully transparent.
+
+        Page-loaded weights have fixed dtype (from the GGUF). The upstream
+        layer may produce a different dtype (notably: HF's dtype inference
+        sees the parameter-less PagedLinear and falls back to its model
+        default). Cast input to the weight's dtype, do the matmul, then
+        cast output back to the input's dtype so the next layer in the
+        model gets what it expects.
+        """
+        import torch.nn.functional as F
+        w = self.weight
+        in_dtype = input.dtype
+        if in_dtype != w.dtype:
+            input = input.to(w.dtype)
+        out = F.linear(input, w, self.bias)
+        if out.dtype != in_dtype:
+            out = out.to(in_dtype)
+        return out
+
+
+def swap_linears_with_paged(
+    root: nn.Module,
+    pager,
+    name_map: Mapping[str, str],
+    dtype: torch.dtype,
+    device_idx: int = 0,
+) -> int:
+    """Walk root, replace each nn.Linear whose module path is in name_map with PagedLinear.
+
+    name_map: {module_path_in_root → catalog_name_in_pager}.
+    e.g. {"model.layers.0.mlp.gate_proj": "blk.0.ffn_gate.weight"}.
+
+    Returns the number of modules replaced.
+    """
+    n_swapped = 0
+    # First pass: collect (parent_module, child_name, old_linear, catalog_name)
+    # tuples so we don't mutate while iterating.
+    to_swap = []
+    for module_path, catalog_name in name_map.items():
+        try:
+            parent_path, _, child_name = module_path.rpartition(".")
+            parent = root
+            if parent_path:
+                for part in parent_path.split("."):
+                    parent = getattr(parent, part) if not part.isdigit() else parent[int(part)]
+            child = getattr(parent, child_name) if not child_name.isdigit() else parent[int(child_name)]
+        except (AttributeError, IndexError):
+            continue
+        if not isinstance(child, nn.Linear):
+            continue
+        page_idx = pager.find_page(catalog_name)
+        if page_idx < 0:
+            continue
+        to_swap.append((parent, child_name, child, page_idx))
+
+    # Second pass: replace
+    for parent, child_name, old, page_idx in to_swap:
+        new_layer = PagedLinear(
+            pager=pager,
+            page_idx=page_idx,
+            weight_shape=old.weight.shape,  # (out, in) for nn.Linear
+            weight_dtype=dtype,
+            bias=old.bias is not None,
+        )
+        new_layer.device_idx = device_idx
+        if old.bias is not None:
+            new_layer.bias = old.bias  # keep existing bias parameter
+        if child_name.isdigit():
+            parent[int(child_name)] = new_layer
+        else:
+            setattr(parent, child_name, new_layer)
+        n_swapped += 1
+    return n_swapped
+
+
+__all__ = ["PagedLinear", "swap_linears_with_paged"]

python_bindings/wp/setup.py

@@ -0,0 +1,52 @@
+"""Build the wp_native pybind11 module.
+
+Usage:
+    cd python_bindings/wp
+    python3 setup.py build_ext --inplace
+
+Produces wp_native*.so in this directory, importable as `wp_native`.
+
+Links against the libllama.so built at <repo_root>/build-hip/bin/.
+"""
+import sys
+from pathlib import Path
+
+from pybind11.setup_helpers import Pybind11Extension, build_ext
+from setuptools import setup
+
+REPO_ROOT = Path(__file__).resolve().parent.parent.parent
+BUILD_HIP_BIN = REPO_ROOT / "build-hip" / "bin"
+
+assert REPO_ROOT.is_dir(), REPO_ROOT
+assert BUILD_HIP_BIN.is_dir(), f"missing build-hip/bin — did you cmake --build build-hip? ({BUILD_HIP_BIN})"
+assert (BUILD_HIP_BIN / "libllama.so").exists() or (BUILD_HIP_BIN / "libllama.so.0").exists(), \
+    f"missing libllama.so in {BUILD_HIP_BIN}"
+
+ROCM_INCLUDE = "/opt/rocm/include"
+ROCM_LIB = "/opt/rocm/lib"
+
+ext = Pybind11Extension(
+    "wp_native",
+    sources=["wp_bindings.cpp"],
+    include_dirs=[
+        str(REPO_ROOT),                       # for "src/weight-pager/..." headers
+        str(REPO_ROOT / "ggml" / "include"),  # for forward decls of ggml types
+        ROCM_INCLUDE,                          # hip/hip_runtime.h
+    ],
+    library_dirs=[str(BUILD_HIP_BIN), ROCM_LIB],
+    libraries=["llama", "ggml-hip", "ggml-base", "amdhip64"],
+    runtime_library_dirs=[str(BUILD_HIP_BIN), ROCM_LIB],
+    cxx_std=17,
+    extra_compile_args=[
+        "-O2", "-Wall",
+        "-D__HIP_PLATFORM_AMD__=1",  # required for hip_runtime.h
+    ],
+)
+
+setup(
+    name="wp_native",
+    version="0.1.0",
+    ext_modules=[ext],
+    cmdclass={"build_ext": build_ext},
+    zip_safe=False,
+)