Publish K96-tail source patches

reproduction/closure/patches/kquant/changed-files/kquant/candidate_hessian.py

@@ -0,0 +1,292 @@
+"""Candidate-specific covariance helpers for offline expert encoding studies."""
+
+from __future__ import annotations
+
+import hashlib
+import math
+from collections.abc import Mapping, Sequence
+
+import torch
+
+
+def partition_documents(
+    documents: Sequence[str], *, modulus: int, confirmation_index: int
+) -> tuple[list[str], list[str]]:
+    """Deterministically split document hashes into fit and confirmation sets."""
+
+    if modulus < 2:
+        raise ValueError("document partition modulus must be at least 2")
+    if not 0 <= confirmation_index < modulus:
+        raise ValueError("confirmation index must be in the partition modulus")
+    if not documents or len(set(documents)) != len(documents):
+        raise ValueError("documents must be a nonempty unique sequence")
+    fit: list[str] = []
+    confirmation: list[str] = []
+    for document in documents:
+        digest = hashlib.blake2b(document.encode(), digest_size=8).digest()
+        bucket = int.from_bytes(digest, "little") % modulus
+        (confirmation if bucket == confirmation_index else fit).append(document)
+    if not fit or not confirmation:
+        raise ValueError("document partition produced an empty side")
+    return fit, confirmation
+
+
+def subset_request_documents(
+    request_documents: Mapping[int, str], documents: Sequence[str]
+) -> dict[int, str]:
+    """Return request/document entries belonging to ``documents`` exactly."""
+
+    selected = set(documents)
+    if not selected:
+        raise ValueError("document subset must be nonempty")
+    result = {
+        int(request): document
+        for request, document in request_documents.items()
+        if document in selected
+    }
+    covered = set(result.values())
+    if covered != selected:
+        missing = sorted(selected - covered)
+        raise ValueError(f"documents lack request epochs: {missing[:8]}")
+    return result
+
+
+def request_mask(
+    observations: torch.Tensor, request_documents: Mapping[int, str]
+) -> torch.Tensor:
+    """Select capture rows whose encoded request epoch is in a corpus subset."""
+
+    if observations.ndim != 1:
+        raise ValueError("observations must be one-dimensional")
+    if not request_documents:
+        raise ValueError("request document mapping must be nonempty")
+    steps = torch.bitwise_right_shift(observations, 32)
+    allowed = torch.tensor(
+        sorted(request_documents), dtype=steps.dtype, device=steps.device
+    )
+    return torch.isin(steps, allowed)
+
+
+def weighted_covariance(
+    rows: torch.Tensor,
+    weights: torch.Tensor,
+    *,
+    device: torch.device,
+    chunk_rows: int = 256,
+    return_cpu: bool = True,
+) -> tuple[torch.Tensor, float]:
+    """Return ``sum(w * x xT) / sum(w)`` as symmetric FP32."""
+
+    if rows.ndim != 2 or weights.ndim != 1 or rows.shape[0] != weights.shape[0]:
+        raise ValueError("invalid covariance sample shapes")
+    if not rows.shape[0] or not rows.shape[1]:
+        raise ValueError("covariance samples must be nonempty")
+    if chunk_rows <= 0:
+        raise ValueError("chunk_rows must be positive")
+    if not torch.all(torch.isfinite(rows)) or not torch.all(torch.isfinite(weights)):
+        raise ValueError("covariance samples must be finite")
+    if torch.any(weights < 0):
+        raise ValueError("covariance weights must be nonnegative")
+
+    dimension = int(rows.shape[1])
+    accumulator = torch.zeros(
+        (dimension, dimension), dtype=torch.float32, device=device
+    )
+    denominator_tensor = weights.to(device=device, dtype=torch.float64).sum()
+    if not bool(torch.isfinite(denominator_tensor)) or not bool(
+        denominator_tensor > 0
+    ):
+        raise ValueError("covariance weights must have a positive finite sum")
+    for begin in range(0, int(rows.shape[0]), chunk_rows):
+        values = rows[begin : begin + chunk_rows].to(
+            device=device, dtype=torch.float32
+        )
+        local_weights = weights[begin : begin + chunk_rows].to(
+            device=device, dtype=torch.float32
+        )
+        accumulator.addmm_(values.T, values * local_weights[:, None])
+    covariance = accumulator.div_(denominator_tensor.to(torch.float32))
+    covariance = ((covariance + covariance.T) * 0.5).contiguous()
+    denominator = float(denominator_tensor)
+    if return_cpu:
+        covariance = covariance.cpu()
+    if not torch.all(torch.isfinite(covariance)):
+        raise ValueError("weighted covariance is non-finite")
+    return covariance, denominator
+
+
+def shrink_covariance(
+    global_covariance: torch.Tensor,
+    local_covariance: torch.Tensor,
+    alpha: float | torch.Tensor,
+) -> torch.Tensor:
+    """Convexly shrink a low-support local covariance toward a global one."""
+
+    if isinstance(alpha, torch.Tensor):
+        if alpha.numel() != 1 or not bool(torch.isfinite(alpha)) or not bool(
+            (alpha >= 0) & (alpha <= 1)
+        ):
+            raise ValueError("alpha must be a finite scalar in [0, 1]")
+        blend_weight: float | torch.Tensor = alpha.to(
+            device=global_covariance.device, dtype=torch.float32
+        )
+    else:
+        alpha_value = float(alpha)
+        if not math.isfinite(alpha_value) or not 0.0 <= alpha_value <= 1.0:
+            raise ValueError("alpha must be finite and in [0, 1]")
+        blend_weight = alpha_value
+    if (
+        global_covariance.ndim != 2
+        or global_covariance.shape[0] != global_covariance.shape[1]
+        or local_covariance.shape != global_covariance.shape
+    ):
+        raise ValueError("global and local covariances must be equal-size squares")
+    if not torch.all(torch.isfinite(global_covariance)) or not torch.all(
+        torch.isfinite(local_covariance)
+    ):
+        raise ValueError("covariances must be finite")
+    blended = torch.lerp(
+        global_covariance.float(), local_covariance.float(), blend_weight
+    )
+    return ((blended + blended.T) * 0.5).contiguous()
+
+
+def weighted_effective_sample_size(weights: torch.Tensor) -> float:
+    """Return Kish effective sample size for nonnegative row weights."""
+
+    if weights.ndim != 1 or not weights.numel():
+        raise ValueError("weights must be a nonempty vector")
+    if not torch.all(torch.isfinite(weights)) or torch.any(weights < 0):
+        raise ValueError("weights must be finite and nonnegative")
+    value = _weighted_effective_sample_size_tensor(weights)
+    return float(value)
+
+
+def _weighted_effective_sample_size_tensor(weights: torch.Tensor) -> torch.Tensor:
+    """Kish support as a scalar tensor on the input numerical device."""
+
+    if weights.ndim != 1 or not weights.numel():
+        raise ValueError("weights must be a nonempty vector")
+    if not torch.all(torch.isfinite(weights)) or torch.any(weights < 0):
+        raise ValueError("weights must be finite and nonnegative")
+    weights64 = weights.double()
+    total = weights64.sum()
+    squared_total = weights64.square().sum()
+    if total <= 0 or squared_total <= 0:
+        raise ValueError("weights must have positive effective support")
+    return total.square() / squared_total
+
+
+def adaptive_identity_shrinkage(
+    local_covariance: torch.Tensor,
+    weights: torch.Tensor,
+    *,
+    max_local_alpha: float = 0.75,
+) -> tuple[torch.Tensor, dict[str, float]]:
+    """Shrink a weighted local covariance toward its scaled identity.
+
+    The Oracle Approximating Shrinkage expression is evaluated with Kish
+    effective sample size in place of an unweighted row count.  The target is
+    always ``trace(local) / dimension * I``.  This is important for expert
+    ``H2`` matrices: intermediate coordinates have no common meaning across
+    experts, so a pooled post-SiTU covariance is not a valid prior.
+    ``max_local_alpha`` preserves a nonzero identity contribution even for
+    very well-supported experts.
+    """
+
+    max_local_alpha = float(max_local_alpha)
+    if not math.isfinite(max_local_alpha) or not 0.0 <= max_local_alpha <= 1.0:
+        raise ValueError("max_local_alpha must be finite and in [0, 1]")
+    if local_covariance.ndim != 2 or (
+        local_covariance.shape[0] != local_covariance.shape[1]
+    ):
+        raise ValueError("local covariance must be square")
+    if not torch.all(torch.isfinite(local_covariance)):
+        raise ValueError("local covariance must be finite")
+
+    effective_rows_tensor = _weighted_effective_sample_size_tensor(weights)
+    dimension = int(local_covariance.shape[0])
+    local64 = local_covariance.double()
+    local_trace = torch.trace(local64)
+    if local_trace <= 0:
+        raise ValueError("local covariance trace must be positive")
+
+    mu = local_trace / dimension
+    mean_square = local64.square().mean()
+    denominator = (effective_rows_tensor + 1.0) * (
+        mean_square - mu.square() / dimension
+    )
+    raw_oas = (mean_square + mu.square()) / denominator.clamp_min(
+        torch.finfo(local64.dtype).tiny
+    )
+    oas_shrinkage_tensor = torch.where(
+        denominator <= 0,
+        torch.ones_like(raw_oas),
+        raw_oas.clamp(0.0, 1.0),
+    )
+    local_alpha_tensor = torch.minimum(
+        torch.as_tensor(max_local_alpha, device=local64.device, dtype=local64.dtype),
+        1.0 - oas_shrinkage_tensor,
+    )
+    scaled_identity = torch.eye(
+        dimension,
+        dtype=torch.float32,
+        device=local_covariance.device,
+    ).mul_(mu.to(dtype=torch.float32))
+    blended = shrink_covariance(
+        scaled_identity, local_covariance, local_alpha_tensor.to(torch.float32)
+    )
+    # Evidence scalars are detached only after the numerical result is frozen.
+    effective_rows = float(effective_rows_tensor)
+    oas_shrinkage = float(oas_shrinkage_tensor)
+    local_alpha = float(local_alpha_tensor)
+    identity_scale = float(mu)
+    return blended, {
+        "effective_sample_size": effective_rows,
+        "oas_shrinkage": oas_shrinkage,
+        "local_alpha": local_alpha,
+        "max_local_alpha": max_local_alpha,
+        "identity_scale": identity_scale,
+    }
+
+
+def covariance_comparison(
+    reference: torch.Tensor, candidate: torch.Tensor
+) -> dict[str, float]:
+    """Compact scale/shape diagnostics for two covariance matrices."""
+
+    if reference.ndim != 2 or reference.shape[0] != reference.shape[1]:
+        raise ValueError("reference covariance must be square")
+    if candidate.shape != reference.shape:
+        raise ValueError("candidate covariance shape differs")
+    reference = reference.float()
+    candidate = candidate.float()
+    if not torch.all(torch.isfinite(reference)) or not torch.all(torch.isfinite(candidate)):
+        raise ValueError("covariances must be finite")
+    reference_norm = torch.linalg.vector_norm(reference.double())
+    reference_diagonal = torch.diagonal(reference).double()
+    candidate_diagonal = torch.diagonal(candidate).double()
+    diagonal_denominator = torch.linalg.vector_norm(
+        reference_diagonal
+    ) * torch.linalg.vector_norm(candidate_diagonal)
+    trace = torch.trace(reference.double())
+    return {
+        "relative_frobenius_difference": float(
+            torch.linalg.vector_norm((candidate - reference).double())
+            / reference_norm
+        )
+        if reference_norm > 0
+        else math.inf,
+        "trace_ratio": float(torch.trace(candidate.double()) / trace)
+        if trace != 0
+        else math.inf,
+        "diagonal_cosine_similarity": float(
+            torch.dot(reference_diagonal, candidate_diagonal)
+            / diagonal_denominator
+        )
+        if diagonal_denominator > 0
+        else 0.0,
+        "candidate_symmetry_max_abs": float(
+            (candidate - candidate.T).abs().max()
+        ),
+    }

reproduction/closure/patches/kquant/changed-files/kquant/csrc/qsrt_quantize_tiles_kernel.cuh

@@ -0,0 +1,492 @@
+#pragma once
+
+// Templated trellis-quantization (Viterbi) kernel, instantiated per (K, cb) in
+// comp_units/quantize_tiles_inst_k*.cu
+
+#include <cuda_fp16.h>
+#include <cuda_bf16.h>
+#include <cublas_v2.h>
+#include <cstdio>
+#include "exl3_compat/util.h"
+#include "exl3_compat/util.cuh"
+#include "exl3_compat/quant/codebook.cuh"
+
+#define QUANTIZE_TILES_K2_NUM_THREADS 1024
+#define QUANTIZE_TILES_K3_NUM_THREADS 640
+#define QUANTIZE_TILES_K4_NUM_THREADS 704
+#ifndef H_INF
+#define H_INF __ushort_as_half(0x7c00)
+#endif
+#ifndef H_MAX_FINITE
+#define H_MAX_FINITE __ushort_as_half(0x7bff)
+#endif
+
+template <int K, int cb>
+__global__ __launch_bounds__(
+    K == 2 ? QUANTIZE_TILES_K2_NUM_THREADS :
+        (K == 3 ? QUANTIZE_TILES_K3_NUM_THREADS : QUANTIZE_TILES_K4_NUM_THREADS),
+    K == 2 ? 1 : 2)
+void quantize_tiles_kernel
+(
+    const float* __restrict__ input_tiles_ptr,
+    float* __restrict__ output_tiles_ptr,
+    uint16_t* __restrict__ output_indices_ptr,
+    half* __restrict__ temp_costs_ptr,
+    uint8_t* __restrict__ temp_edges_ptr,
+    int tailbite_context
+)
+{
+    extern __shared__ uint8_t shbuf[];
+    uint8_t* sh = shbuf;
+
+    constexpr int Kr = 16 - K;
+    constexpr int max_q = 1 << K;
+    constexpr int edges = 65536 >> K;
+    constexpr int decisions_per_row =
+        K == 2 ? edges / 4 : (K <= 4 ? edges / 2 : edges);
+    constexpr int decision_shift = 16 - 2 * K;
+    constexpr int predecessor_step = 1 << decision_shift;
+    constexpr int num_threads =
+        K == 2 ? QUANTIZE_TILES_K2_NUM_THREADS :
+            (K == 3 ? QUANTIZE_TILES_K3_NUM_THREADS : QUANTIZE_TILES_K4_NUM_THREADS);
+
+    const int tile_idx = blockIdx.x;
+    const int thread = threadIdx.x;
+    const float* input_tile = input_tiles_ptr + 256 * tile_idx;
+    float* output_tile = output_tiles_ptr + 256 * tile_idx;
+    uint16_t* output_indices = output_indices_ptr + 256 * tile_idx;
+    uint8_t* temp_edges = temp_edges_ptr + 256 * decisions_per_row * tile_idx;
+
+    half* sh_input_tile = (half*) sh; sh += 256 * sizeof(half);
+    half* sh_min = (half*) sh; sh += 32 * sizeof(half);
+    int* sh_idx = (int*) sh; sh += 32 * sizeof(int);
+
+    half* sh_temp_costs = (half*) sh;
+    half* temp_costs = K >= 2 ? sh_temp_costs : temp_costs_ptr + 2 * edges * tile_idx;
+    half* temp_costs_inc = temp_costs + edges;
+
+    if (thread < 256) sh_input_tile[thread] = __float2half_rn(input_tile[thread]);
+    __syncthreads();
+
+    auto forward = [&](int roll, int steps, int pre_state, int trace_start)
+    {
+        int ri = roll & 255;
+        half* t = temp_costs;
+        temp_costs = temp_costs_inc;
+        temp_costs_inc = t;
+
+        for (int out_edge_idx = 2 * thread; out_edge_idx < edges; out_edge_idx += 2 * num_threads)
+        {
+            const half2 w2 = __half2half2(sh_input_tile[ri]);
+            int in_edge_idx = out_edge_idx >> K;
+            uint32_t product0 = 0;
+            uint32_t product1 = 0;
+            uint2 sqg_k2_labels{};
+            uint4 sqg_k3_labels{};
+            uint4 sqg_k4_labels0{};
+            uint4 sqg_k4_labels1{};
+            half2 decoded2;
+            if constexpr (cb == 1)
+            {
+                product0 = mul_const_u32<0xCBAC1FEDu>(out_edge_idx);
+                product1 = product0 + 0xCBAC1FEDu;
+                decoded2 = decode_mcg_product_2(product0, product1);
+            }
+            else if constexpr (cb == 2)
+            {
+                product0 = out_edge_idx * 0x83DCD12Du;
+                product1 = product0 + 0x83DCD12Du;
+                decoded2 = decode_mul1_product_2(product0, product1);
+            }
+            else if constexpr (cb == 4 && K == 2)
+            {
+                sqg_k2_labels =
+                    reinterpret_cast<const uint2*>(sqg_e4m3_lut)[out_edge_idx >> 1];
+                decoded2 = decode_sqg_k2_pair(sqg_k2_labels, 0);
+            }
+            else if constexpr (cb == 4 && K == 3)
+            {
+                sqg_k3_labels =
+                    reinterpret_cast<const uint4*>(sqg_e4m3_lut)[out_edge_idx >> 1];
+                decoded2 = decode_sqg_k3_pair(sqg_k3_labels, 0);
+            }
+            else if constexpr (cb == 4 && K == 4)
+            {
+                const int packed_index = out_edge_idx;
+                sqg_k4_labels0 =
+                    reinterpret_cast<const uint4*>(sqg_e4m3_lut)[packed_index];
+                sqg_k4_labels1 =
+                    reinterpret_cast<const uint4*>(sqg_e4m3_lut)[packed_index + 1];
+                decoded2 = decode_sqg_k4_pair(sqg_k4_labels0, sqg_k4_labels1, 0);
+            }
+            else
+            {
+                decoded2 = decode_3inst_2<cb>(out_edge_idx, out_edge_idx + 1);
+            }
+            half2 dh2 = __hsub2(decoded2, w2);
+            // Positive infinity denotes an unreachable state. Reachable path
+            // costs must remain finite when their FP16 sum overflows, or
+            // traceback can confuse a saturated path with an unreachable one
+            // and violate the constrained cyclic boundary.
+            half2 min_err2 = __hmin2(
+                __hmul2(dh2, dh2), __half2half2(H_MAX_FINITE));
+            if (pre_state >= 0 && in_edge_idx != pre_state) min_err2 = __half2half2(H_INF);
+            int min_k0 = 0;
+            int min_k1 = 0;
+
+            #pragma unroll
+            for (int k = 1; k < max_q; ++k)
+            {
+                in_edge_idx += predecessor_step;
+                if constexpr (cb == 1)
+                {
+                    // MCG multiplication is linear modulo 2^32 across successive branch states.
+                    constexpr uint32_t product_step = 0xCBAC1FEDu << Kr;
+                    product0 += product_step;
+                    product1 += product_step;
+                    decoded2 = decode_mcg_product_2(product0, product1);
+                }
+                else if constexpr (cb == 2)
+                {
+                    // The mul1 multiplication is equally linear modulo 2^32.
+                    constexpr uint32_t product_step = 0x83DCD12Du << Kr;
+                    product0 += product_step;
+                    product1 += product_step;
+                    decoded2 = decode_mul1_product_2(product0, product1);
+                }
+                else if constexpr (cb == 4 && K == 2)
+                {
+                    decoded2 = decode_sqg_k2_pair(sqg_k2_labels, k);
+                }
+                else if constexpr (cb == 4 && K == 3)
+                {
+                    decoded2 = decode_sqg_k3_pair(sqg_k3_labels, k);
+                }
+                else if constexpr (cb == 4 && K == 4)
+                {
+                    decoded2 = decode_sqg_k4_pair(
+                        sqg_k4_labels0, sqg_k4_labels1, k);
+                }
+                else
+                {
+                    const int state0 = (k << Kr) | out_edge_idx;
+                    decoded2 = decode_3inst_2<cb>(state0, state0 + 1);
+                }
+                dh2 = __hsub2(decoded2, w2);
+                half2 err2 = __hmin2(
+                    __hmul2(dh2, dh2), __half2half2(H_MAX_FINITE));
+                if (pre_state >= 0 && in_edge_idx != pre_state) err2 = __half2half2(H_INF);
+                const unsigned less = __hlt2_mask(err2, min_err2);
+                min_err2 = __hmin2(err2, min_err2);
+                if (less & 0xffffu) min_k0 = k;
+                if (less >> 16) min_k1 = k;
+            }
+
+            reinterpret_cast<half2*>(temp_costs)[out_edge_idx >> 1] = min_err2;
+            if (trace_start == 0)
+            {
+                if constexpr (K == 2)
+                {
+                    const unsigned pair = min_k0 | (min_k1 << 2);
+                    const unsigned next = __shfl_down_sync(0xffffffff, pair, 1);
+                    if ((thread & 1) == 0)
+                        temp_edges[decisions_per_row * ri + (out_edge_idx >> 2)] =
+                            (uint8_t) (pair | (next << 4));
+                }
+                else if constexpr (K <= 4)
+                    temp_edges[decisions_per_row * ri + (out_edge_idx >> 1)] =
+                        (uint8_t) (min_k0 | (min_k1 << 4));
+                else
+                {
+                    temp_edges[decisions_per_row * ri + out_edge_idx] = (uint8_t) min_k0;
+                    temp_edges[decisions_per_row * ri + out_edge_idx + 1] = (uint8_t) min_k1;
+                }
+            }
+        }
+        __syncthreads();
+
+        for (int i = 1; i < steps; ++i)
+        {
+            ri = (i + roll) & 255;
+            t = temp_costs;
+            temp_costs = temp_costs_inc;
+            temp_costs_inc = t;
+
+            for (int out_edge_idx = 2 * thread; out_edge_idx < edges; out_edge_idx += 2 * num_threads)
+            {
+                const half2 w2 = __half2half2(sh_input_tile[ri]);
+                int in_edge_idx = out_edge_idx >> K;
+                uint32_t product0 = 0;
+                uint32_t product1 = 0;
+                uint2 sqg_k2_labels{};
+                uint4 sqg_k3_labels{};
+                uint4 sqg_k4_labels0{};
+                uint4 sqg_k4_labels1{};
+                half2 decoded2;
+                if constexpr (cb == 1)
+                {
+                    product0 = mul_const_u32<0xCBAC1FEDu>(out_edge_idx);
+                    product1 = product0 + 0xCBAC1FEDu;
+                    decoded2 = decode_mcg_product_2(product0, product1);
+                }
+                else if constexpr (cb == 2)
+                {
+                    product0 = out_edge_idx * 0x83DCD12Du;
+                    product1 = product0 + 0x83DCD12Du;
+                    decoded2 = decode_mul1_product_2(product0, product1);
+                }
+                else if constexpr (cb == 4 && K == 2)
+                {
+                    sqg_k2_labels =
+                        reinterpret_cast<const uint2*>(sqg_e4m3_lut)[out_edge_idx >> 1];
+                    decoded2 = decode_sqg_k2_pair(sqg_k2_labels, 0);
+                }
+                else if constexpr (cb == 4 && K == 3)
+                {
+                    sqg_k3_labels =
+                        reinterpret_cast<const uint4*>(sqg_e4m3_lut)[out_edge_idx >> 1];
+                    decoded2 = decode_sqg_k3_pair(sqg_k3_labels, 0);
+                }
+                else if constexpr (cb == 4 && K == 4)
+                {
+                    const int packed_index = out_edge_idx;
+                    sqg_k4_labels0 =
+                        reinterpret_cast<const uint4*>(sqg_e4m3_lut)[packed_index];
+                    sqg_k4_labels1 =
+                        reinterpret_cast<const uint4*>(sqg_e4m3_lut)[packed_index + 1];
+                    decoded2 = decode_sqg_k4_pair(
+                        sqg_k4_labels0, sqg_k4_labels1, 0);
+                }
+                else
+                {
+                    decoded2 = decode_3inst_2<cb>(out_edge_idx, out_edge_idx + 1);
+                }
+                half2 dh2 = __hsub2(decoded2, w2);
+                half predecessor_cost = temp_costs_inc[in_edge_idx];
+                half2 min_err2 = __heq(predecessor_cost, H_INF)
+                    ? __half2half2(H_INF)
+                    : __hmin2(
+                        __hfma2(
+                            dh2, dh2, __half2half2(predecessor_cost)),
+                        __half2half2(H_MAX_FINITE));
+                int min_k0 = 0;
+                int min_k1 = 0;
+
+                #pragma unroll
+                for (int k = 1; k < max_q; ++k)
+                {
+                    in_edge_idx += predecessor_step;
+                    if constexpr (cb == 1)
+                    {
+                        // MCG multiplication is linear modulo 2^32 across successive branch states.
+                        constexpr uint32_t product_step = 0xCBAC1FEDu << Kr;
+                        product0 += product_step;
+                        product1 += product_step;
+                        decoded2 = decode_mcg_product_2(product0, product1);
+                    }
+                    else if constexpr (cb == 2)
+                    {
+                        // The mul1 multiplication is equally linear modulo 2^32.
+                        constexpr uint32_t product_step = 0x83DCD12Du << Kr;
+                        product0 += product_step;
+                        product1 += product_step;
+                        decoded2 = decode_mul1_product_2(product0, product1);
+                    }
+                    else if constexpr (cb == 4 && K == 2)
+                    {
+                        decoded2 = decode_sqg_k2_pair(sqg_k2_labels, k);
+                    }
+                    else if constexpr (cb == 4 && K == 3)
+                    {
+                        decoded2 = decode_sqg_k3_pair(sqg_k3_labels, k);
+                    }
+                    else if constexpr (cb == 4 && K == 4)
+                    {
+                        decoded2 = decode_sqg_k4_pair(
+                            sqg_k4_labels0, sqg_k4_labels1, k);
+                    }
+                    else
+                    {
+                        const int state0 = (k << Kr) | out_edge_idx;
+                        decoded2 = decode_3inst_2<cb>(state0, state0 + 1);
+                    }
+                    dh2 = __hsub2(decoded2, w2);
+                    predecessor_cost = temp_costs_inc[in_edge_idx];
+                    half2 err2 = __heq(predecessor_cost, H_INF)
+                        ? __half2half2(H_INF)
+                        : __hmin2(
+                            __hfma2(
+                                dh2, dh2, __half2half2(predecessor_cost)),
+                            __half2half2(H_MAX_FINITE));
+                    const unsigned less = __hlt2_mask(err2, min_err2);
+                    min_err2 = __hmin2(err2, min_err2);
+                    if (less & 0xffffu) min_k0 = k;
+                    if (less >> 16) min_k1 = k;
+                }
+
+                reinterpret_cast<half2*>(temp_costs)[out_edge_idx >> 1] = min_err2;
+                if (i >= trace_start)
+                {
+                    if constexpr (K == 2)
+                    {
+                        const unsigned pair = min_k0 | (min_k1 << 2);
+                        const unsigned next = __shfl_down_sync(0xffffffff, pair, 1);
+                        if ((thread & 1) == 0)
+                            temp_edges[decisions_per_row * ri + (out_edge_idx >> 2)] =
+                                (uint8_t) (pair | (next << 4));
+                    }
+                    else if constexpr (K <= 4)
+                        temp_edges[decisions_per_row * ri + (out_edge_idx >> 1)] =
+                            (uint8_t) (min_k0 | (min_k1 << 4));
+                    else
+                    {
+                        temp_edges[decisions_per_row * ri + out_edge_idx] = (uint8_t) min_k0;
+                        temp_edges[decisions_per_row * ri + out_edge_idx + 1] = (uint8_t) min_k1;
+                    }
+                }
+            }
+            __syncthreads();
+        }
+    };
+
+    auto argmin_cost = [&]()
+    {
+        // Preserve the historical 1024-thread tie-breaking order.
+        const int lane_id = thread & 31;
+        const int warp_id = thread >> 5;
+        if constexpr (K == 2)
+        {
+            // With 1024 K2 threads, each warp represents exactly one of the
+            // 32 reduction groups used by the historical 512-thread kernel's
+            // two local streams.  The strict-less reduction therefore keeps
+            // the same deterministic tie order.
+            half local_min = H_INF;
+            int local_idx = -1;
+            for (int e = thread; e < edges; e += num_threads)
+            {
+                half v = temp_costs_inc[e];
+                if (__hlt(v, local_min)) { local_min = v; local_idx = e; }
+            }
+            #pragma unroll
+            for (int offset = 16; offset > 0; offset >>= 1)
+            {
+                half other_min = __shfl_down_sync(0xffffffff, local_min, offset);
+                int other_idx = __shfl_down_sync(0xffffffff, local_idx, offset);
+                if (__hlt(other_min, local_min)) { local_min = other_min; local_idx = other_idx; }
+            }
+            sh_min[warp_id] = local_min;
+            sh_idx[warp_id] = local_idx;
+        }
+        else
+        {
+            // Preserve the established 512-thread reduction tree even when
+            // the forward pass uses a wider block.  This keeps exact-tie
+            // path selection stable across encoder-kernel revisions.
+            if (thread < 512)
+            {
+                half local_min0 = H_INF;
+                half local_min1 = H_INF;
+                int local_idx0 = -1;
+                int local_idx1 = -1;
+                for (int e = thread; e < edges; e += 1024)
+                {
+                    half v = temp_costs_inc[e];
+                    if (__hlt(v, local_min0)) { local_min0 = v; local_idx0 = e; }
+                }
+                for (int e = thread + 512; e < edges; e += 1024)
+                {
+                    half v = temp_costs_inc[e];
+                    if (__hlt(v, local_min1)) { local_min1 = v; local_idx1 = e; }
+                }
+                #pragma unroll
+                for (int offset = 16; offset > 0; offset >>= 1)
+                {
+                    half other_min0 = __shfl_down_sync(0xffffffff, local_min0, offset);
+                    int other_idx0 = __shfl_down_sync(0xffffffff, local_idx0, offset);
+                    if (__hlt(other_min0, local_min0)) { local_min0 = other_min0; local_idx0 = other_idx0; }
+                    half other_min1 = __shfl_down_sync(0xffffffff, local_min1, offset);
+                    int other_idx1 = __shfl_down_sync(0xffffffff, local_idx1, offset);
+                    if (__hlt(other_min1, local_min1)) { local_min1 = other_min1; local_idx1 = other_idx1; }
+                }
+                sh_min[warp_id] = local_min0;
+                sh_idx[warp_id] = local_idx0;
+                sh_min[16 + warp_id] = local_min1;
+                sh_idx[16 + warp_id] = local_idx1;
+            }
+        }
+        __syncthreads();
+
+        int local_idx = 0;
+        if (warp_id == 0)
+        {
+            half local_min = sh_min[lane_id];
+            local_idx = sh_idx[lane_id];
+            #pragma unroll
+            for (int offset = 16; offset > 0; offset >>= 1)
+            {
+                half other_min = __shfl_down_sync(0xffffffff, local_min, offset);
+                int other_idx = __shfl_down_sync(0xffffffff, local_idx, offset);
+                if (__hlt(other_min, local_min)) { local_min = other_min; local_idx = other_idx; }
+            }
+        }
+        return local_idx;
+    };
+
+    auto backward = [&](int roll, int steps, bool write, int edge)
+    {
+        if (thread == 0)
+        {
+            for (int i = steps - 1; i >= 0; --i)
+            {
+                const int ri = (i + roll) & 255;
+                int decision;
+                if constexpr (K == 2)
+                {
+                    const uint8_t packed =
+                        temp_edges[decisions_per_row * ri + (edge >> 2)];
+                    decision = (packed >> ((edge & 3) * 2)) & 0x03;
+                }
+                else if constexpr (K <= 4)
+                {
+                    const uint8_t packed =
+                        temp_edges[decisions_per_row * ri + (edge >> 1)];
+                    decision = (packed >> ((edge & 1) * 4)) & 0x0f;
+                }
+                else
+                    decision = (int) temp_edges[decisions_per_row * ri + edge];
+                const int prev_edge =
+                    (decision << decision_shift) | (edge >> K);
+                const int encoded = (prev_edge << K) | edge;
+                edge = prev_edge;
+                if (write)
+                {
+                    output_indices[ri] = (uint16_t) encoded;
+                    if constexpr (cb == 4 && K == 2)
+                        output_tile[ri] = __half2float(decode_sqg_k2_scalar(encoded));
+                    else if constexpr (cb == 4 && K == 3)
+                        output_tile[ri] = __half2float(decode_sqg_k3_scalar(encoded));
+                    else if constexpr (cb == 4 && K == 4)
+                        output_tile[ri] = __half2float(decode_sqg_k4_scalar(encoded));
+                    else
+                        output_tile[ri] = __half2float(decode_3inst<cb>(encoded));
+                }
+                else if (ri == 0) break;
+            }
+        }
+        if (thread == 0) sh_idx[0] = edge;
+        __syncthreads();
+        return sh_idx[0];
+    };
+
+    // Infer a cyclic boundary state from symmetric context around element 0,
+    // then run the ordinary complete constrained pass.  The historical path
+    // uses 128 symbols on either side (a full 256-step primer).  Smaller
+    // contexts are an encoder-only search approximation: the final pass is
+    // still a valid closed trellis and the bitstream/decoder are unchanged.
+    const int primer_roll = 256 - tailbite_context;
+    const int primer_steps = 2 * tailbite_context;
+    forward(primer_roll, primer_steps, -1, tailbite_context);
+    int end_state = backward(primer_roll, primer_steps, false, argmin_cost());
+    forward(0, 256, end_state, 0);
+    backward(0, 256, true, end_state);
+}

reproduction/closure/patches/kquant/changed-files/kquant/csrc/sqg_quantize.cu

@@ -0,0 +1,358 @@
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <cuda_fp16.h>
+#include <cuda_fp8.h>
+#include <torch/extension.h>
+
+#include "exl3_compat/util.cuh"
+#include "exl3_compat/quant/codebook.cuh"
+
+// Keep only the table pointer in broadcast-friendly constant memory.  The
+// labels themselves remain in ordinary read-only global memory: a warp's
+// Viterbi lanes normally request different states, which would serialize a
+// 64-KiB constant-memory table despite its attractive nominal footprint.
+__device__ __constant__ const uint8_t* sqg_e4m3_lut;
+
+template <>
+__device__ inline half decode_3inst<4>(uint32_t state)
+{
+    const __nv_fp8_storage_t fp8 = sqg_e4m3_lut[state & 0xffffu];
+    const __half_raw raw = __nv_cvt_fp8_to_halfraw(fp8, __NV_E4M3);
+    return __ushort_as_half(raw.x);
+}
+
+template <>
+__device__ inline half2 decode_3inst_2<4>(uint32_t state0, uint32_t state1)
+{
+    // The trellis kernel always asks for an aligned consecutive pair.  Load
+    // both E4M3 labels together so a warp consumes every byte in each cache
+    // sector instead of issuing separate even- and odd-byte gathers.
+    const __nv_fp8x2_storage_t fp8x2 =
+        reinterpret_cast<const __nv_fp8x2_storage_t*>(sqg_e4m3_lut)[state0 >> 1];
+    const __half2_raw raw = __nv_cvt_fp8x2_to_halfraw2(fp8x2, __NV_E4M3);
+    return __halves2half2(__ushort_as_half(raw.x), __ushort_as_half(raw.y));
+}
+
+__device__ inline half2 decode_sqg_fp8x2(__nv_fp8x2_storage_t fp8x2)
+{
+    const __half2_raw raw = __nv_cvt_fp8x2_to_halfraw2(fp8x2, __NV_E4M3);
+    return __halves2half2(__ushort_as_half(raw.x), __ushort_as_half(raw.y));
+}
+
+__device__ inline half2 decode_sqg_k2_pair(const uint2& packed, int predecessor)
+{
+    const uint32_t word = predecessor < 2 ? packed.x : packed.y;
+    return decode_sqg_fp8x2(
+        static_cast<__nv_fp8x2_storage_t>(word >> (16 * (predecessor & 1))));
+}
+
+__device__ inline half decode_sqg_k2_scalar(uint32_t state)
+{
+    constexpr int edges = 65536 >> 2;
+    const int predecessor = state >> 14;
+    const int out_edge = state & (edges - 1);
+    const __nv_fp8x2_storage_t pair =
+        reinterpret_cast<const __nv_fp8x2_storage_t*>(sqg_e4m3_lut)[
+            (out_edge >> 1) * 4 + predecessor];
+    const __nv_fp8_storage_t fp8 =
+        static_cast<__nv_fp8_storage_t>(pair >> (8 * (out_edge & 1)));
+    const __half_raw raw = __nv_cvt_fp8_to_halfraw(fp8, __NV_E4M3);
+    return __ushort_as_half(raw.x);
+}
+
+__device__ inline half2 decode_sqg_k3_pair(const uint4& packed, int predecessor)
+{
+    uint32_t word;
+    switch (predecessor >> 1)
+    {
+        case 0: word = packed.x; break;
+        case 1: word = packed.y; break;
+        case 2: word = packed.z; break;
+        default: word = packed.w; break;
+    }
+    return decode_sqg_fp8x2(
+        static_cast<__nv_fp8x2_storage_t>(word >> (16 * (predecessor & 1))));
+}
+
+__device__ inline half2 decode_sqg_k4_pair(
+    const uint4& packed0, const uint4& packed1, int predecessor)
+{
+    const uint4& packed = predecessor < 8 ? packed0 : packed1;
+    const int local = predecessor & 7;
+    uint32_t word;
+    switch (local >> 1)
+    {
+        case 0: word = packed.x; break;
+        case 1: word = packed.y; break;
+        case 2: word = packed.z; break;
+        default: word = packed.w; break;
+    }
+    return decode_sqg_fp8x2(
+        static_cast<__nv_fp8x2_storage_t>(word >> (16 * (local & 1))));
+}
+
+__device__ inline half decode_sqg_k4_scalar(uint32_t state)
+{
+    constexpr int edges = 65536 >> 4;
+    const int predecessor = state >> 12;
+    const int out_edge = state & (edges - 1);
+    const __nv_fp8x2_storage_t pair =
+        reinterpret_cast<const __nv_fp8x2_storage_t*>(sqg_e4m3_lut)[
+            (out_edge >> 1) * 16 + predecessor];
+    const __nv_fp8_storage_t fp8 =
+        static_cast<__nv_fp8_storage_t>(pair >> (8 * (out_edge & 1)));
+    const __half_raw raw = __nv_cvt_fp8_to_halfraw(fp8, __NV_E4M3);
+    return __ushort_as_half(raw.x);
+}
+
+__device__ inline half decode_sqg_k3_scalar(uint32_t state)
+{
+    constexpr int edges = 65536 >> 3;
+    const int predecessor = state >> 13;
+    const int out_edge = state & (edges - 1);
+    const __nv_fp8x2_storage_t pair =
+        reinterpret_cast<const __nv_fp8x2_storage_t*>(sqg_e4m3_lut)[
+            (out_edge >> 1) * 8 + predecessor];
+    const __nv_fp8_storage_t fp8 =
+        static_cast<__nv_fp8_storage_t>(pair >> (8 * (out_edge & 1)));
+    const __half_raw raw = __nv_cvt_fp8_to_halfraw(fp8, __NV_E4M3);
+    return __ushort_as_half(raw.x);
+}
+
+#include "qsrt_quantize_tiles_kernel.cuh"
+
+namespace
+{
+
+template <int K, int codebook>
+void launch_trellis(
+    const float* input,
+    float* output,
+    uint16_t* indices,
+    half* costs,
+    uint8_t* edges,
+    int tiles,
+    int max_batch,
+    int shared_memory,
+    int tailbite_context,
+    cudaStream_t stream)
+{
+    const int threads = K == 2 ? 1024 : (K == 3 ? 640 : 704);
+    cudaFuncSetAttribute(
+        quantize_tiles_kernel<K, codebook>,
+        cudaFuncAttributeMaxDynamicSharedMemorySize,
+        shared_memory);
+    cudaFuncSetAttribute(
+        quantize_tiles_kernel<K, codebook>,
+        cudaFuncAttributePreferredSharedMemoryCarveout,
+        K == 4 ? cudaSharedmemCarveoutMaxL1 : cudaSharedmemCarveoutMaxShared);
+    for (int start = 0; start < tiles; start += max_batch)
+    {
+        const int batch = min(max_batch, tiles - start);
+        quantize_tiles_kernel<K, codebook><<<batch, threads, shared_memory, stream>>>(
+            input + 256 * start,
+            output + 256 * start,
+            indices + 256 * start,
+            costs,
+            edges,
+            tailbite_context);
+        C10_CUDA_KERNEL_LAUNCH_CHECK();
+    }
+}
+
+}  // namespace
+
+void quantize_tiles_sqg_cuda(
+    at::Tensor input_tiles,
+    at::Tensor output_tiles,
+    at::Tensor output_indices,
+    at::Tensor temp_costs,
+    at::Tensor temp_edges,
+    at::Tensor codebook,
+    int64_t K,
+    int64_t tailbite_context)
+{
+    const at::cuda::OptionalCUDAGuard guard(input_tiles.device());
+    const cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();
+
+    TORCH_CHECK(input_tiles.is_cuda(), "input_tiles must be CUDA");
+    TORCH_CHECK(input_tiles.scalar_type() == at::kFloat, "input_tiles must be FP32");
+    TORCH_CHECK(input_tiles.dim() == 2 && input_tiles.size(1) == 256,
+                "input_tiles must have shape [tiles, 256]");
+    TORCH_CHECK(output_tiles.sizes() == input_tiles.sizes(),
+                "output_tiles shape mismatch");
+    TORCH_CHECK(output_tiles.scalar_type() == at::kFloat,
+                "output_tiles must be FP32");
+    TORCH_CHECK(output_indices.sizes() == input_tiles.sizes(),
+                "output_indices shape mismatch");
+    TORCH_CHECK(output_indices.scalar_type() == at::kShort,
+                "output_indices must be int16");
+    TORCH_CHECK(K >= 2 && K <= 6, "SQG validation kernel supports K2--K6");
+    TORCH_CHECK(tailbite_context >= 1 && tailbite_context <= 128,
+                "tailbite_context must be in 1..128");
+    TORCH_CHECK(codebook.is_cuda() && codebook.device() == input_tiles.device(),
+                "codebook must be on the input CUDA device");
+    TORCH_CHECK(codebook.scalar_type() == at::kByte && codebook.numel() == 65536,
+                "codebook must contain 65,536 uint8 E4M3 values");
+    TORCH_CHECK(codebook.is_contiguous(), "codebook must be contiguous");
+
+    const int edges_per_state = 65536 >> K;
+    const int decision_entries =
+        K == 2 ? edges_per_state / 4 : (K <= 4 ? edges_per_state / 2 : edges_per_state);
+    TORCH_CHECK(temp_costs.scalar_type() == at::kHalf && temp_costs.dim() == 3,
+                "temp_costs must be rank-3 FP16");
+    TORCH_CHECK(temp_costs.size(1) == 2 && temp_costs.size(2) == edges_per_state,
+                "temp_costs shape mismatch");
+    TORCH_CHECK(temp_edges.scalar_type() == at::kByte && temp_edges.dim() == 3,
+                "temp_edges must be rank-3 uint8");
+    TORCH_CHECK(temp_edges.size(1) == 256 && temp_edges.size(2) == decision_entries,
+                "temp_edges shape mismatch");
+
+    const uint8_t* codebook_pointer = codebook.data_ptr<uint8_t>();
+    C10_CUDA_CHECK(cudaMemcpyToSymbolAsync(
+        sqg_e4m3_lut,
+        &codebook_pointer,
+        sizeof(codebook_pointer),
+        0,
+        cudaMemcpyHostToDevice,
+        stream));
+
+    int device = -1;
+    C10_CUDA_CHECK(cudaGetDevice(&device));
+    int multiprocessors = 0;
+    C10_CUDA_CHECK(cudaDeviceGetAttribute(
+        &multiprocessors, cudaDevAttrMultiProcessorCount, device));
+    const int max_batch = min(
+        static_cast<int>(temp_costs.size(0)), 3 * multiprocessors);
+    const int shared_memory =
+        2 * edges_per_state * static_cast<int>(sizeof(half)) + 512 + 64 + 128;
+    const int tiles = static_cast<int>(input_tiles.size(0));
+    if (!tiles) return;
+
+    switch (K)
+    {
+        case 2:
+            launch_trellis<2, 4>(
+                input_tiles.data_ptr<float>(), output_tiles.data_ptr<float>(),
+                reinterpret_cast<uint16_t*>(output_indices.data_ptr<int16_t>()),
+                reinterpret_cast<half*>(temp_costs.data_ptr<at::Half>()),
+                temp_edges.data_ptr<uint8_t>(), tiles, max_batch, shared_memory,
+                static_cast<int>(tailbite_context), stream);
+            break;
+        case 3:
+            launch_trellis<3, 4>(
+                input_tiles.data_ptr<float>(), output_tiles.data_ptr<float>(),
+                reinterpret_cast<uint16_t*>(output_indices.data_ptr<int16_t>()),
+                reinterpret_cast<half*>(temp_costs.data_ptr<at::Half>()),
+                temp_edges.data_ptr<uint8_t>(), tiles, max_batch, shared_memory,
+                static_cast<int>(tailbite_context), stream);
+            break;
+        case 4:
+            launch_trellis<4, 4>(
+                input_tiles.data_ptr<float>(), output_tiles.data_ptr<float>(),
+                reinterpret_cast<uint16_t*>(output_indices.data_ptr<int16_t>()),
+                reinterpret_cast<half*>(temp_costs.data_ptr<at::Half>()),
+                temp_edges.data_ptr<uint8_t>(), tiles, max_batch, shared_memory,
+                static_cast<int>(tailbite_context), stream);
+            break;
+        case 5:
+            launch_trellis<5, 4>(
+                input_tiles.data_ptr<float>(), output_tiles.data_ptr<float>(),
+                reinterpret_cast<uint16_t*>(output_indices.data_ptr<int16_t>()),
+                reinterpret_cast<half*>(temp_costs.data_ptr<at::Half>()),
+                temp_edges.data_ptr<uint8_t>(), tiles, max_batch, shared_memory,
+                static_cast<int>(tailbite_context), stream);
+            break;
+        case 6:
+            launch_trellis<6, 4>(
+                input_tiles.data_ptr<float>(), output_tiles.data_ptr<float>(),
+                reinterpret_cast<uint16_t*>(output_indices.data_ptr<int16_t>()),
+                reinterpret_cast<half*>(temp_costs.data_ptr<at::Half>()),
+                temp_edges.data_ptr<uint8_t>(), tiles, max_batch, shared_memory,
+                static_cast<int>(tailbite_context), stream);
+            break;
+    }
+}
+
+void quantize_tiles_procedural_cuda(
+    at::Tensor input_tiles,
+    at::Tensor output_tiles,
+    at::Tensor output_indices,
+    at::Tensor temp_costs,
+    at::Tensor temp_edges,
+    int64_t K,
+    int64_t codebook,
+    int64_t tailbite_context)
+{
+    const at::cuda::OptionalCUDAGuard guard(input_tiles.device());
+    const cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();
+
+    TORCH_CHECK(input_tiles.is_cuda(), "input_tiles must be CUDA");
+    TORCH_CHECK(input_tiles.scalar_type() == at::kFloat, "input_tiles must be FP32");
+    TORCH_CHECK(input_tiles.dim() == 2 && input_tiles.size(1) == 256,
+                "input_tiles must have shape [tiles, 256]");
+    TORCH_CHECK(output_tiles.sizes() == input_tiles.sizes(),
+                "output_tiles shape mismatch");
+    TORCH_CHECK(output_tiles.scalar_type() == at::kFloat,
+                "output_tiles must be FP32");
+    TORCH_CHECK(output_indices.sizes() == input_tiles.sizes(),
+                "output_indices shape mismatch");
+    TORCH_CHECK(output_indices.scalar_type() == at::kShort,
+                "output_indices must be int16");
+    TORCH_CHECK(K >= 2 && K <= 4, "procedural control supports K2--K4");
+    TORCH_CHECK(codebook == 1 || codebook == 2,
+                "procedural codebook must be 1 (MCG) or 2 (MUL1)");
+    TORCH_CHECK(tailbite_context >= 1 && tailbite_context <= 128,
+                "tailbite_context must be in 1..128");
+
+    const int edges_per_state = 65536 >> K;
+    const int decision_entries = K == 2 ? edges_per_state / 4 : edges_per_state / 2;
+    TORCH_CHECK(temp_costs.scalar_type() == at::kHalf && temp_costs.dim() == 3,
+                "temp_costs must be rank-3 FP16");
+    TORCH_CHECK(temp_costs.size(1) == 2 && temp_costs.size(2) == edges_per_state,
+                "temp_costs shape mismatch");
+    TORCH_CHECK(temp_edges.scalar_type() == at::kByte && temp_edges.dim() == 3,
+                "temp_edges must be rank-3 uint8");
+    TORCH_CHECK(temp_edges.size(1) == 256 && temp_edges.size(2) == decision_entries,
+                "temp_edges shape mismatch");
+
+    int device = -1;
+    C10_CUDA_CHECK(cudaGetDevice(&device));
+    int multiprocessors = 0;
+    C10_CUDA_CHECK(cudaDeviceGetAttribute(
+        &multiprocessors, cudaDevAttrMultiProcessorCount, device));
+    const int max_batch = min(
+        static_cast<int>(temp_costs.size(0)), 3 * multiprocessors);
+    const int shared_memory =
+        2 * edges_per_state * static_cast<int>(sizeof(half)) + 512 + 64 + 128;
+    const int tiles = static_cast<int>(input_tiles.size(0));
+    if (!tiles) return;
+
+#define LAUNCH_PROCEDURAL(K_VALUE, CODEBOOK_VALUE) \
+    launch_trellis<K_VALUE, CODEBOOK_VALUE>( \
+        input_tiles.data_ptr<float>(), output_tiles.data_ptr<float>(), \
+        reinterpret_cast<uint16_t*>(output_indices.data_ptr<int16_t>()), \
+        reinterpret_cast<half*>(temp_costs.data_ptr<at::Half>()), \
+        temp_edges.data_ptr<uint8_t>(), tiles, max_batch, shared_memory, \
+        static_cast<int>(tailbite_context), stream)
+
+    if (codebook == 1)
+    {
+        switch (K)
+        {
+            case 2: LAUNCH_PROCEDURAL(2, 1); break;
+            case 3: LAUNCH_PROCEDURAL(3, 1); break;
+            case 4: LAUNCH_PROCEDURAL(4, 1); break;
+        }
+    }
+    else
+    {
+        switch (K)
+        {
+            case 2: LAUNCH_PROCEDURAL(2, 2); break;
+            case 3: LAUNCH_PROCEDURAL(3, 2); break;
+            case 4: LAUNCH_PROCEDURAL(4, 2); break;
+        }
+    }
+#undef LAUNCH_PROCEDURAL
+}

reproduction/closure/patches/kquant/changed-files/kquant/exl3_reference.py

@@ -0,0 +1,281 @@
+"""Pure-PyTorch reference operations for EXL3 procedural codebooks."""
+
+from __future__ import annotations
+
+import math
+
+import torch
+
+from kquant.sqg_e4m3 import (
+    SQG_XOR_CHEB_T12,
+    SQG_CHEB,
+    SQG_CHEB_NORMAL_E4M3,
+    SQG_NORMAL_E4M3,
+    sqg_codebook,
+)
+
+
+TILE_VALUES = 256
+HADAMARD_BLOCK = 128
+CODEBOOK_SQG_NORMAL_E4M3 = SQG_NORMAL_E4M3
+CODEBOOK_SQG_XOR_CHEB_T12 = SQG_XOR_CHEB_T12
+CODEBOOK_SQG_CHEB_NORMAL_E4M3 = SQG_CHEB_NORMAL_E4M3
+CODEBOOK_SQG_CHEB = SQG_CHEB
+QSRT_CODEBOOKS = (
+    CODEBOOK_SQG_XOR_CHEB_T12,
+    CODEBOOK_SQG_NORMAL_E4M3,
+    CODEBOOK_SQG_CHEB_NORMAL_E4M3,
+    CODEBOOK_SQG_CHEB,
+)
+
+
+def _validate_bits(bits: int) -> None:
+    if isinstance(bits, bool) or not isinstance(bits, int) or bits not in range(2, 7):
+        raise ValueError("the SQG reference supports K=2 through K=6")
+
+
+def reconstruct_trellis_states(
+    edge_symbols: torch.Tensor, bits: int
+) -> torch.Tensor:
+    """Recover cyclic 16-bit trellis states from their stored K-bit edges.
+
+    A state is the low 16 bits of ``(previous_state << K) | edge``.  EXL's
+    256-symbol tile is cyclic, so state zero also incorporates the final edge
+    symbols in the tile.
+    """
+
+    _validate_bits(bits)
+    if edge_symbols.ndim < 1 or edge_symbols.shape[-1] != TILE_VALUES:
+        raise ValueError("edge_symbols must have a final dimension of 256")
+    if edge_symbols.dtype == torch.bool or edge_symbols.is_floating_point():
+        raise TypeError("edge_symbols must use an integer dtype")
+    edges = edge_symbols.to(dtype=torch.int64) & ((1 << bits) - 1)
+    # Unroll the cyclic recurrence in state space.  At position i, edge i is
+    # in the low K bits, edge i-1 is above it, and so on.  Only ceil(16 / K)
+    # edges can survive the low-16-bit mask.  Expressing that closed form with
+    # 4--8 rolls avoids launching one tensor operation for every one of the
+    # 256 symbols when this reference check runs on CUDA.
+    states = torch.zeros_like(edges)
+    for lag in range(math.ceil(16 / bits)):
+        states |= torch.roll(edges, shifts=lag, dims=-1) << (lag * bits)
+    return (states & 0xFFFF).to(dtype=torch.int16).contiguous()
+
+
+def _decode_codebook_states(
+    states: torch.Tensor,
+    codebook: str,
+    *,
+    bits: int | None = None,
+    rate_axis: str | None = None,
+) -> torch.Tensor:
+    if codebook in QSRT_CODEBOOKS:
+        if bits is None:
+            raise ValueError("SQG reconstruction requires an explicit trellis rate")
+        _validate_bits(bits)
+        indices = (states.to(dtype=torch.int64) & 0xFFFF).long()
+        values = sqg_codebook(
+            bits,
+            codebook,
+            rate_axis=rate_axis,
+            device=states.device,
+            dtype=torch.float16,
+        )
+        return values.index_select(0, indices.flatten()).reshape_as(states)
+    raise ValueError(
+        f"unsupported QSRT codebook {codebook!r}; expected one of {QSRT_CODEBOOKS}"
+    )
+
+
+def tensor_core_permutation(device: torch.device | str = "cpu") -> torch.Tensor:
+    """Return the native EXL 16x16 row-major -> tensor-core element order."""
+
+    permutation = [0] * TILE_VALUES
+    for thread in range(32):
+        row0 = (thread % 4) * 2
+        row1 = row0 + 1
+        row2 = row0 + 8
+        row3 = row0 + 9
+        column0 = thread // 4
+        column1 = column0 + 8
+        permutation[thread * 8 + 0] = row0 * 16 + column0
+        permutation[thread * 8 + 1] = row1 * 16 + column0
+        permutation[thread * 8 + 2] = row2 * 16 + column0
+        permutation[thread * 8 + 3] = row3 * 16 + column0
+        permutation[thread * 8 + 4] = row0 * 16 + column1
+        permutation[thread * 8 + 5] = row1 * 16 + column1
+        permutation[thread * 8 + 6] = row2 * 16 + column1
+        permutation[thread * 8 + 7] = row3 * 16 + column1
+    return torch.tensor(permutation, dtype=torch.long, device=device)
+
+
+def decode_regularized_weight(
+    states: torch.Tensor,
+    *,
+    codebook: str = CODEBOOK_SQG_XOR_CHEB_T12,
+    codebook_values: torch.Tensor | None = None,
+    bits: int | None = None,
+) -> torch.Tensor:
+    """Decode ``[K/16, N/16, 256]`` states to transformed ``[K, N]``."""
+
+    if states.ndim != 3 or states.shape[-1] != TILE_VALUES:
+        raise ValueError("states must have shape [K/16, N/16, 256]")
+    if codebook_values is None:
+        decoded = _decode_codebook_states(states, codebook, bits=bits).float()
+    else:
+        if codebook_values.ndim != 1 or codebook_values.numel() != 1 << 16:
+            raise ValueError("custom codebook must contain exactly 65,536 values")
+        if not codebook_values.is_floating_point():
+            raise TypeError("custom codebook must be floating point")
+        if not bool(torch.isfinite(codebook_values).all()):
+            raise ValueError("custom codebook must be finite")
+        indices = (states.to(dtype=torch.int64) & 0xFFFF).long()
+        decoded = codebook_values.to(device=states.device).index_select(
+            0, indices.flatten()
+        ).reshape_as(states).float()
+    inverse = torch.argsort(tensor_core_permutation(states.device))
+    decoded = decoded.index_select(-1, inverse)
+    k_tiles, n_tiles, _ = decoded.shape
+    return (
+        decoded.reshape(k_tiles, n_tiles, 16, 16)
+        .permute(0, 2, 1, 3)
+        .reshape(k_tiles * 16, n_tiles * 16)
+        .contiguous()
+    )
+
+
+def decode_qsrt_regularized_weight(
+    states: torch.Tensor,
+    *,
+    rate_axis: str,
+    tile_bits: tuple[int, ...],
+    codebook: str,
+) -> torch.Tensor:
+    """Decode a K2--K6 tile map to transformed ``[K, N]`` weights.
+
+    SQG uses a distinct state labelling for every supported K, so the physical
+    tile rate is part of the numerical decode contract.
+    """
+
+    if states.ndim != 3 or states.shape[-1] != TILE_VALUES:
+        raise ValueError("states must have shape [K/16, N/16, 256]")
+    if rate_axis not in ("k", "n"):
+        raise ValueError("rate_axis must be 'k' or 'n'")
+    axis = 0 if rate_axis == "k" else 1
+    if len(tile_bits) != states.shape[axis]:
+        raise ValueError("tile_bits length does not match the selected rate axis")
+    if any(
+        isinstance(value, bool) or not isinstance(value, int) or value not in range(2, 7)
+        for value in tile_bits
+    ):
+        raise ValueError("tile_bits must contain only integer K2 through K6")
+    if codebook not in QSRT_CODEBOOKS:
+        raise ValueError(f"unsupported QSRT codebook: {codebook}")
+
+    decoded_tiles = torch.empty_like(states, dtype=torch.float32)
+    bits_tensor = torch.tensor(tile_bits, dtype=torch.long, device=states.device)
+    for bits in sorted(set(tile_bits)):
+        positions = torch.nonzero(bits_tensor == bits).flatten()
+        selected = states.index_select(axis, positions)
+        values = _decode_codebook_states(
+            selected, codebook, bits=bits, rate_axis=rate_axis
+        ).float()
+        decoded_tiles.index_copy_(axis, positions, values)
+
+    inverse = torch.argsort(tensor_core_permutation(states.device))
+    decoded_tiles = decoded_tiles.index_select(-1, inverse)
+    k_tiles, n_tiles, _ = decoded_tiles.shape
+    return (
+        decoded_tiles.reshape(k_tiles, n_tiles, 16, 16)
+        .permute(0, 2, 1, 3)
+        .reshape(k_tiles * 16, n_tiles * 16)
+        .contiguous()
+    )
+
+
+def normalized_hadamard(
+    *,
+    device: torch.device | str,
+    dtype: torch.dtype = torch.float32,
+    size: int = HADAMARD_BLOCK,
+) -> torch.Tensor:
+    if size <= 0 or size & (size - 1):
+        raise ValueError("Hadamard size must be a positive power of two")
+    matrix = torch.ones((1, 1), dtype=dtype, device=device)
+    while matrix.shape[0] < size:
+        matrix = torch.cat(
+            (
+                torch.cat((matrix, matrix), dim=1),
+                torch.cat((matrix, -matrix), dim=1),
+            ),
+            dim=0,
+        )
+    return matrix * (1.0 / math.sqrt(size))
+
+
+def _blockwise_hadamard_left(weight: torch.Tensor) -> torch.Tensor:
+    if weight.shape[0] % HADAMARD_BLOCK:
+        raise ValueError("EXL K dimension must be divisible by 128")
+    hadamard = normalized_hadamard(device=weight.device, dtype=weight.dtype)
+    blocks = weight.reshape(-1, HADAMARD_BLOCK, weight.shape[1])
+    return torch.matmul(hadamard, blocks).reshape_as(weight).contiguous()
+
+
+def _blockwise_hadamard_right(weight: torch.Tensor) -> torch.Tensor:
+    if weight.shape[1] % HADAMARD_BLOCK:
+        raise ValueError("EXL N dimension must be divisible by 128")
+    hadamard = normalized_hadamard(device=weight.device, dtype=weight.dtype)
+    blocks = weight.reshape(weight.shape[0], -1, HADAMARD_BLOCK)
+    return torch.matmul(blocks, hadamard).reshape_as(weight).contiguous()
+
+
+def decode_exl3_weight(
+    states: torch.Tensor,
+    suh: torch.Tensor,
+    svh: torch.Tensor,
+    *,
+    codebook: str = CODEBOOK_SQG_XOR_CHEB_T12,
+    codebook_values: torch.Tensor | None = None,
+    bits: int | None = None,
+) -> torch.Tensor:
+    """Decode stored states/scales to an EXL-oriented ``[K, N]`` weight."""
+
+    weight = decode_regularized_weight(
+        states, codebook=codebook, codebook_values=codebook_values, bits=bits
+    )
+    if suh.ndim != 1 or suh.numel() != weight.shape[0]:
+        raise ValueError("suh length does not match the EXL K dimension")
+    if svh.ndim != 1 or svh.numel() != weight.shape[1]:
+        raise ValueError("svh length does not match the EXL N dimension")
+    weight = _blockwise_hadamard_left(weight)
+    weight *= suh.to(device=weight.device, dtype=weight.dtype).unsqueeze(1)
+    weight = _blockwise_hadamard_right(weight)
+    weight *= svh.to(device=weight.device, dtype=weight.dtype).unsqueeze(0)
+    return weight.contiguous()
+
+
+def decode_qsrt_weight(
+    states: torch.Tensor,
+    suh: torch.Tensor,
+    svh: torch.Tensor,
+    *,
+    rate_axis: str,
+    tile_bits: tuple[int, ...],
+    codebook: str = CODEBOOK_SQG_XOR_CHEB_T12,
+) -> torch.Tensor:
+    """Decode stored QSRT states and scales to an EXL ``[K, N]`` weight."""
+
+    weight = decode_qsrt_regularized_weight(
+        states,
+        rate_axis=rate_axis,
+        tile_bits=tile_bits,
+        codebook=codebook,
+    )
+    if suh.ndim != 1 or suh.numel() != weight.shape[0]:
+        raise ValueError("suh length does not match the EXL K dimension")
+    if svh.ndim != 1 or svh.numel() != weight.shape[1]:
+        raise ValueError("svh length does not match the EXL N dimension")
+    weight = _blockwise_hadamard_left(weight)
+    weight *= suh.to(device=weight.device, dtype=weight.dtype).unsqueeze(1)
+    weight = _blockwise_hadamard_right(weight)
+    weight *= svh.to(device=weight.device, dtype=weight.dtype).unsqueeze(0)
+    return weight.contiguous()

reproduction/closure/patches/kquant/changed-files/kquant/qsrt.py

@@ -0,0 +1,1141 @@
+"""Reference codec geometry for Kimi-K3 QSRT candidates.
+
+Candidate payloads describe complete experts and contain no tensor-parallel
+ownership.  TP sharding is a property of the atom-major storage/runtime layer.
+The ordinary expert weights use ``[out_features, in_features]`` orientation,
+while the EXL encoder consumes their transpose, ``[K, N]``.  Consequently the
+shared 3072-channel intermediate axis is EXL ``N`` for ``w1``/``w3`` and EXL
+``K`` for ``w2``.
+
+The reference packer preserves the native EXL bit ordering for each 16x16
+trellis tile.  Its inverse reconstructs both the stored K-bit edge symbols and
+the full 16-bit cyclic trellis states consumed by the SQG decoder.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import dataclass
+from typing import Literal
+
+import torch
+
+from kquant.exl3_reference import (
+    CODEBOOK_SQG_XOR_CHEB_T12,
+    CODEBOOK_SQG_CHEB_NORMAL_E4M3,
+    CODEBOOK_SQG_CHEB,
+    CODEBOOK_SQG_NORMAL_E4M3,
+    QSRT_CODEBOOKS,
+    decode_qsrt_weight,
+    reconstruct_trellis_states,
+)
+
+
+EXPERTS_PER_LAYER = 896
+INTERMEDIATE_CHANNELS = 3072
+LATENT_CHANNELS = 3584
+CONTEXT_GROUP_CHANNELS = 4
+TILE_CHANNELS = 16
+TILE_VALUES = TILE_CHANNELS * TILE_CHANNELS
+RECORD_CHANNELS = 128
+PAIR_CHANNELS = 2 * RECORD_CHANNELS
+RECORDS_PER_EXPERT = INTERMEDIATE_CHANNELS // RECORD_CHANNELS
+PAIRS_PER_EXPERT = RECORDS_PER_EXPERT // 2
+TILES_PER_RECORD_AXIS = RECORD_CHANNELS // TILE_CHANNELS
+SCHEMA = "kquant_kimi_k3_qsrt_candidate_v1"
+LOGICAL_CANDIDATE_SCHEMAS = (SCHEMA,)
+STORAGE_ALIGNMENT = 4096
+PAIR_ROTATION_EXPERT_MULTIPLIER = 5
+
+# The phase-1 table uses one byte per expert.  Compressed entries contain
+# two four-bit transfer counts: high nibble r13 for fused w1/w3, low nibble r2
+# for w2.  Each value is in 0..2.  0xff denotes the exact X4T tier.  The nibbles
+# remain deliberately roomy so the fixed table is cheap to inspect and decode.
+FORMAT_BITS = 8
+FORMAT_X4T = 0xFF
+FORMAT_NIBBLE_MASK = 0x0F
+
+LAYER_HEADER_BYTES = STORAGE_ALIGNMENT
+FORMAT_TABLE_BYTES = (EXPERTS_PER_LAYER * FORMAT_BITS + 7) // 8
+FORMAT_SECTION_BYTES = STORAGE_ALIGNMENT
+HOT_FORMAT_BYTES = EXPERTS_PER_LAYER
+HOT_SLOT_BYTES = EXPERTS_PER_LAYER * torch.int16.itemsize
+HOT_METADATA_BYTES = HOT_FORMAT_BYTES + HOT_SLOT_BYTES
+
+# Shared hidden-side vectors are w1.suh, w3.suh, and w2.svh.  The three
+# corresponding intermediate-side vectors remain expert-local.  Their
+# atom-major storage is defined independently from a serving TP size.
+SHARED_SCALE_VECTORS = 3
+LOCAL_SCALE_VECTORS = 3
+SHARED_SCALE_ORDER = ("w1.suh", "w3.suh", "w2.svh")
+LOCAL_SCALE_ORDER = ("w1.svh", "w3.svh", "w2.suh")
+SCALE_BYTES = torch.float16.itemsize
+SHARED_SCALE_BYTES = SHARED_SCALE_VECTORS * LATENT_CHANNELS * SCALE_BYTES
+SHARED_SCALE_SECTION_BYTES = 6 * STORAGE_ALIGNMENT
+
+# P24 and P33 both carry a six-bit sum over 256 intermediate channels.
+PAIR_BYTES = PAIR_CHANNELS * LATENT_CHANNELS * 3 // 8
+MATRIX_TRELLIS_BYTES = PAIRS_PER_EXPERT * PAIR_BYTES
+EXPERT_TRELLIS_BYTES = 3 * MATRIX_TRELLIS_BYTES
+RateAxis = Literal["k", "n"]
+PairKind = Literal["P24", "P33"]
+
+MATRIX_RATE_AXIS: dict[str, RateAxis] = {
+    "w1": "n",
+    "w3": "n",
+    "w2": "k",
+}
+
+
+def pair_rotation(layer: int, expert: int) -> int:
+    """Cyclic logical-pair placement used to balance physical atom slots."""
+
+    if (
+        isinstance(layer, bool)
+        or not isinstance(layer, int)
+        or layer not in range(1, 93)
+    ):
+        raise ValueError("K3 MoE layer must be in 1..92")
+    if (
+        isinstance(expert, bool)
+        or not isinstance(expert, int)
+        or not 0 <= expert < EXPERTS_PER_LAYER
+    ):
+        raise ValueError(f"expert must be in 0..{EXPERTS_PER_LAYER - 1}")
+    return (PAIR_ROTATION_EXPERT_MULTIPLIER * expert + layer) % PAIRS_PER_EXPERT
+
+
+def logical_pair_for_slot(layer: int, expert: int, pair_slot: int) -> int:
+    """Logical candidate pair stored in one physical pair slot."""
+
+    if (
+        isinstance(pair_slot, bool)
+        or not isinstance(pair_slot, int)
+        or not 0 <= pair_slot < PAIRS_PER_EXPERT
+    ):
+        raise ValueError(f"pair_slot must be in 0..{PAIRS_PER_EXPERT - 1}")
+    return (pair_slot - pair_rotation(layer, expert)) % PAIRS_PER_EXPERT
+
+
+def physical_pair_slot(layer: int, expert: int, logical_pair: int) -> int:
+    """Physical slot containing one logical candidate pair."""
+
+    if (
+        isinstance(logical_pair, bool)
+        or not isinstance(logical_pair, int)
+        or not 0 <= logical_pair < PAIRS_PER_EXPERT
+    ):
+        raise ValueError(f"logical_pair must be in 0..{PAIRS_PER_EXPERT - 1}")
+    return (logical_pair + pair_rotation(layer, expert)) % PAIRS_PER_EXPERT
+
+
+@dataclass(frozen=True)
+class ModeSpec:
+    """One rate-transfer schedule, ordered from low to high importance.
+
+    ``context_bits`` contains one entry per final 128-channel record.
+    ``mode_id`` is the number of K2/K4 record exchanges.
+    """
+
+    mode_id: int
+    name: str
+    context_bits: tuple[int, ...]
+    def __post_init__(self) -> None:
+        if isinstance(self.mode_id, bool) or not isinstance(self.mode_id, int):
+            raise TypeError("mode_id must be an integer")
+        if not 0 <= self.mode_id <= 2:
+            raise ValueError("QSRT mode_id must be R0, R1, or R2")
+        if not self.name:
+            raise ValueError("mode name must not be empty")
+        if not self.context_bits or len(self.context_bits) % 2:
+            raise ValueError("a mode must contain an even, nonzero context count")
+        if any(bits not in (2, 3, 4) for bits in self.context_bits):
+            raise ValueError("Kimi-K3 QSRT supports only K2, K3, and K4")
+        if sum(self.context_bits) != 3 * len(self.context_bits):
+            raise ValueError("a mode must average exactly three trellis bits")
+        if INTERMEDIATE_CHANNELS % len(self.context_bits):
+            raise ValueError("contexts must divide the 3072-channel axis")
+        context_channels = INTERMEDIATE_CHANNELS // len(self.context_bits)
+        if context_channels % RECORD_CHANNELS:
+            raise ValueError("each context must contain whole 128-channel records")
+        if tuple(sorted(self.context_bits)) != self.context_bits:
+            raise ValueError("context rates must be monotone from K2 through K4")
+        records_per_context = RECORDS_PER_EXPERT // len(self.context_bits)
+        k2_records = self.context_bits.count(2) * records_per_context
+        k4_records = self.context_bits.count(4) * records_per_context
+        if k2_records != k4_records:
+            raise ValueError("K2 and K4 record counts must match")
+        if self.mode_id != k2_records:
+            raise ValueError(
+                f"mode_id must equal the rate-transfer count {k2_records}"
+            )
+
+    @property
+    def context_count(self) -> int:
+        return len(self.context_bits)
+
+    @property
+    def records_per_context(self) -> int:
+        return RECORDS_PER_EXPERT // self.context_count
+
+
+def _rate_transfer_mode(r: int) -> ModeSpec:
+    return ModeSpec(
+        r,
+        f"R{r}",
+        (2,) * r + (3,) * (RECORDS_PER_EXPERT - 2 * r) + (4,) * r,
+    )
+
+
+RATE_TRANSFER_MODES = tuple(_rate_transfer_mode(r) for r in range(3))
+R0 = RATE_TRANSFER_MODES[0]
+R1 = RATE_TRANSFER_MODES[1]
+R2 = RATE_TRANSFER_MODES[2]
+REPRESENTATIVE_MODE_CANDIDATES = RATE_TRANSFER_MODES
+EXPERIMENTAL_MODE_CANDIDATES = RATE_TRANSFER_MODES
+
+# Frozen phase-1 encoder policy.  The payload schema still parses the complete
+# ladder for controlled studies, but the initial all-expert run searches R0,
+# R1, and R2 independently for the fused w1/w3 family and for w2.  The
+# covariance alphas are offline encoder controls and consume no format bits.
+PHASE1_MODE_IDS = (0, 1, 2)
+PHASE1_MODE_CANDIDATES = tuple(RATE_TRANSFER_MODES[r] for r in PHASE1_MODE_IDS)
+PHASE1_SHARED_RATE_MODE = False
+PHASE1_H13_EXPERT_LOCAL_ALPHA = 0.0
+# Maximum local weight for adaptive, scaled-identity OAS H2 shrinkage.  The
+# actual alpha is support-dependent and can be substantially smaller.
+PHASE1_H2_EXPERT_LOCAL_ALPHA = 0.75
+PHASE1_H2_LOCAL_BASIS = "decoded_candidate_post_situ"
+PHASE1_H2_SHRINKAGE_POLICY = "weighted_oas_scaled_identity"
+_MODES_BY_NAME = {mode.name: mode for mode in RATE_TRANSFER_MODES}
+_MODES_BY_ID = {mode.mode_id: mode for mode in RATE_TRANSFER_MODES}
+
+
+def mode_from_context_bits(context_bits: Sequence[int]) -> ModeSpec:
+    """Resolve any legal equal-context ladder allocation to canonical ``R_r``."""
+
+    bits = tuple(context_bits)
+    if not bits or RECORDS_PER_EXPERT % len(bits):
+        raise ValueError("allocation contexts must divide the 24 records")
+    records_per_context = RECORDS_PER_EXPERT // len(bits)
+    r = bits.count(2) * records_per_context
+    # ModeSpec performs monotonicity, rate, and equal-donor validation.
+    ModeSpec(r, "candidate", bits)
+    return RATE_TRANSFER_MODES[r]
+
+
+@dataclass(frozen=True)
+class ExpertFormatSpec:
+    """Provisional per-expert tier and matrix-family rate schedules."""
+
+    r13: int | None
+    r2: int | None
+
+    def __post_init__(self) -> None:
+        if (self.r13 is None) != (self.r2 is None):
+            raise ValueError("r13 and r2 must both be absent or both be present")
+        for name, value in (("r13", self.r13), ("r2", self.r2)):
+            if value is None:
+                continue
+            if isinstance(value, bool) or not isinstance(value, int):
+                raise TypeError(f"{name} must be an integer")
+            if not 0 <= value <= 2:
+                raise ValueError(f"{name} must be R0, R1, or R2")
+
+    @classmethod
+    def x4t(cls) -> "ExpertFormatSpec":
+        return cls(None, None)
+
+    @classmethod
+    def compressed(cls, r13: int, r2: int | None = None) -> "ExpertFormatSpec":
+        return cls(r13, r13 if r2 is None else r2)
+
+    @property
+    def is_x4t(self) -> bool:
+        return self.r13 is None
+
+    @property
+    def name(self) -> str:
+        if self.is_x4t:
+            return "X4T"
+        assert self.r13 is not None and self.r2 is not None
+        return f"R{self.r13}" if self.r13 == self.r2 else f"R{self.r13}/R{self.r2}"
+
+    @property
+    def code(self) -> int:
+        if self.is_x4t:
+            return FORMAT_X4T
+        assert self.r13 is not None and self.r2 is not None
+        return (self.r13 << 4) | self.r2
+
+    @classmethod
+    def from_code(cls, code: int) -> "ExpertFormatSpec":
+        if code == FORMAT_X4T:
+            return cls.x4t()
+        r13 = code >> 4
+        r2 = code & FORMAT_NIBBLE_MASK
+        if r13 > 2 or r2 > 2:
+            raise ValueError(f"format table contains invalid rate code 0x{code:02x}")
+        return cls.compressed(r13, r2)
+
+
+ExpertFormatLike = str | ModeSpec | ExpertFormatSpec | tuple[int, int]
+
+
+def align_up(value: int, alignment: int = STORAGE_ALIGNMENT) -> int:
+    if isinstance(value, bool) or not isinstance(value, int) or value < 0:
+        raise ValueError("value must be a non-negative integer")
+    if (
+        isinstance(alignment, bool)
+        or not isinstance(alignment, int)
+        or alignment <= 0
+        or alignment & (alignment - 1)
+    ):
+        raise ValueError("alignment must be a positive power of two")
+    return (value + alignment - 1) & -alignment
+
+
+def resolve_expert_format(value: ExpertFormatLike) -> ExpertFormatSpec:
+    if isinstance(value, ExpertFormatSpec):
+        return value
+    if isinstance(value, ModeSpec):
+        return ExpertFormatSpec.compressed(value.mode_id)
+    if isinstance(value, tuple):
+        if len(value) != 2:
+            raise ValueError("expert format tuple must be (r13, r2)")
+        return ExpertFormatSpec.compressed(value[0], value[1])
+    if not isinstance(value, str):
+        raise TypeError(
+            "expert format must be a name, ModeSpec, ExpertFormatSpec, or pair"
+        )
+    normalized = value.upper()
+    if normalized == "X4T":
+        return ExpertFormatSpec.x4t()
+    if "/" in normalized:
+        left, right = normalized.split("/", 1)
+        return ExpertFormatSpec.compressed(
+            resolve_mode(left).mode_id, resolve_mode(right).mode_id
+        )
+    return ExpertFormatSpec.compressed(resolve_mode(normalized).mode_id)
+
+
+def pack_qsrt_format_table(formats: Sequence[ExpertFormatLike]) -> torch.Tensor:
+    """Pack 896 one-byte expert format codes."""
+
+    if len(formats) != EXPERTS_PER_LAYER:
+        raise ValueError(f"format table must contain {EXPERTS_PER_LAYER} experts")
+    return torch.tensor(
+        [resolve_expert_format(value).code for value in formats],
+        dtype=torch.uint8,
+    ).contiguous()
+
+
+def unpack_qsrt_format_table(payload: torch.Tensor) -> tuple[str, ...]:
+    """Decode the format table and reject invalid rate nibbles."""
+
+    if payload.dtype != torch.uint8:
+        raise TypeError("format table must use torch.uint8")
+    if payload.ndim != 1 or payload.numel() != FORMAT_TABLE_BYTES:
+        raise ValueError(
+            f"format table must contain exactly {FORMAT_TABLE_BYTES} bytes"
+        )
+    return tuple(
+        ExpertFormatSpec.from_code(code).name for code in payload.cpu().tolist()
+    )
+
+
+def pack_qsrt_format_section(formats: Sequence[ExpertFormatLike]) -> torch.Tensor:
+    """Build the canonical 4-KiB on-disk format-table section."""
+
+    table = pack_qsrt_format_table(formats)
+    section = torch.zeros(FORMAT_SECTION_BYTES, dtype=torch.uint8)
+    section[:FORMAT_TABLE_BYTES].copy_(table)
+    return section
+
+
+def unpack_qsrt_format_section(payload: torch.Tensor) -> tuple[str, ...]:
+    if payload.dtype != torch.uint8:
+        raise TypeError("format section must use torch.uint8")
+    if payload.ndim != 1 or payload.numel() != FORMAT_SECTION_BYTES:
+        raise ValueError(
+            f"format section must contain exactly {FORMAT_SECTION_BYTES} bytes"
+        )
+    if bool(torch.any(payload[FORMAT_TABLE_BYTES:] != 0)):
+        raise ValueError("format section has nonzero alignment padding")
+    return unpack_qsrt_format_table(payload[:FORMAT_TABLE_BYTES])
+
+
+def expand_qsrt_hot_metadata(
+    formats: Sequence[ExpertFormatLike],
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Expand disk codes to uint8 schedules and uint16 tier-local slots."""
+
+    codes = pack_qsrt_format_table(formats)
+    # Validate before deriving slots, including codes supplied through strings.
+    unpack_qsrt_format_table(codes)
+    slots = torch.empty(EXPERTS_PER_LAYER, dtype=torch.int16)
+    compressed_slot = 0
+    kept_slot = 0
+    for expert, code in enumerate(codes.tolist()):
+        if code == FORMAT_X4T:
+            slots[expert] = kept_slot
+            kept_slot += 1
+        else:
+            slots[expert] = compressed_slot
+            compressed_slot += 1
+    return codes, slots
+
+
+def _validate_scale_vector(
+    scale: torch.Tensor, length: int, name: str, device: torch.device | None = None
+) -> None:
+    if scale.dtype != torch.float16:
+        raise TypeError(f"{name} must use torch.float16")
+    if scale.ndim != 1 or scale.numel() != length:
+        raise ValueError(f"{name} must contain exactly {length} values")
+    if not scale.is_contiguous():
+        raise ValueError(f"{name} must be contiguous")
+    if device is not None and scale.device != device:
+        raise ValueError(f"{name} must be on {device}")
+    if not bool(torch.all(torch.isfinite(scale))):
+        raise ValueError(f"{name} contains non-finite values")
+
+
+def pack_qsrt_shared_scale_section(
+    w1_suh: torch.Tensor, w3_suh: torch.Tensor, w2_svh: torch.Tensor
+) -> torch.Tensor:
+    """Pack the three layer-shared hidden-side FP16 vectors and zero padding."""
+
+    scales = (w1_suh, w3_suh, w2_svh)
+    device = scales[0].device
+    for name, scale in zip(SHARED_SCALE_ORDER, scales, strict=True):
+        _validate_scale_vector(scale, LATENT_CHANNELS, name, device)
+    section = torch.zeros(
+        SHARED_SCALE_SECTION_BYTES, dtype=torch.uint8, device=device
+    )
+    cursor = 0
+    for scale in scales:
+        raw = scale.view(torch.uint8)
+        section[cursor : cursor + raw.numel()].copy_(raw)
+        cursor += raw.numel()
+    if cursor != SHARED_SCALE_BYTES:
+        raise AssertionError("shared scale byte accounting drifted")
+    return section
+
+
+def unpack_qsrt_shared_scale_section(
+    payload: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    if payload.dtype != torch.uint8:
+        raise TypeError("shared scale section must use torch.uint8")
+    if payload.ndim != 1 or payload.numel() != SHARED_SCALE_SECTION_BYTES:
+        raise ValueError(
+            "shared scale section must contain exactly "
+            f"{SHARED_SCALE_SECTION_BYTES} bytes"
+        )
+    if bool(torch.any(payload[SHARED_SCALE_BYTES:] != 0)):
+        raise ValueError("shared scale section has nonzero alignment padding")
+    values = []
+    cursor = 0
+    vector_bytes = LATENT_CHANNELS * SCALE_BYTES
+    for name in SHARED_SCALE_ORDER:
+        vector = (
+            payload[cursor : cursor + vector_bytes]
+            .clone()
+            .view(torch.float16)
+            .contiguous()
+        )
+        _validate_scale_vector(vector, LATENT_CHANNELS, name)
+        values.append(vector)
+        cursor += vector_bytes
+    return tuple(values)  # type: ignore[return-value]
+
+
+def resolve_mode(mode: ModeSpec | str | int) -> ModeSpec:
+    if isinstance(mode, ModeSpec):
+        return mode
+    if isinstance(mode, str):
+        try:
+            return _MODES_BY_NAME[mode.upper()]
+        except KeyError as exc:
+            raise ValueError(f"unknown QSRT mode: {mode}") from exc
+    if isinstance(mode, int) and not isinstance(mode, bool):
+        try:
+            return _MODES_BY_ID[mode]
+        except KeyError as exc:
+            raise ValueError(f"unknown QSRT mode ID: {mode}") from exc
+    raise TypeError("mode must be a ModeSpec, name, or integer ID")
+
+
+def matrix_rate_axis(matrix: str) -> RateAxis:
+    try:
+        return MATRIX_RATE_AXIS[matrix]
+    except KeyError as exc:
+        raise ValueError(f"unknown expert matrix: {matrix}") from exc
+
+
+def record_contexts(mode: ModeSpec | str | int) -> tuple[int, ...]:
+    """Return physical record contexts, paired low/high then moving inward."""
+
+    spec = resolve_mode(mode)
+    contexts: list[int] = []
+    for low_context in range(spec.context_count // 2):
+        high_context = spec.context_count - 1 - low_context
+        for _ in range(spec.records_per_context):
+            contexts.extend((low_context, high_context))
+    if len(contexts) != RECORDS_PER_EXPERT:
+        raise AssertionError("mode construction did not cover 24 records")
+    return tuple(contexts)
+
+
+def record_bits(mode: ModeSpec | str | int) -> tuple[int, ...]:
+    spec = resolve_mode(mode)
+    return tuple(spec.context_bits[context] for context in record_contexts(spec))
+
+
+def pair_bits(mode: ModeSpec | str | int) -> tuple[tuple[int, int], ...]:
+    records = record_bits(mode)
+    pairs = tuple(zip(records[0::2], records[1::2], strict=True))
+    if len(pairs) != PAIRS_PER_EXPERT or any(sum(pair) != 6 for pair in pairs):
+        raise ValueError("every QSRT record pair must have a six-bit rate sum")
+    return pairs
+
+
+def pair_kinds(mode: ModeSpec | str | int) -> tuple[PairKind, ...]:
+    kinds: list[PairKind] = []
+    for pair in pair_bits(mode):
+        if pair == (2, 4):
+            kinds.append("P24")
+        elif pair == (3, 3):
+            kinds.append("P33")
+        else:
+            raise ValueError(f"unsupported QSRT pair schedule: {pair}")
+    return tuple(kinds)
+
+
+def _validate_block_contexts(
+    block_contexts: torch.Tensor, mode: ModeSpec
+) -> None:
+    expected_groups = INTERMEDIATE_CHANNELS // CONTEXT_GROUP_CHANNELS
+    if block_contexts.ndim != 1 or block_contexts.numel() != expected_groups:
+        raise ValueError(
+            f"block_contexts must contain {expected_groups} four-channel groups"
+        )
+    if block_contexts.dtype == torch.bool or block_contexts.is_floating_point():
+        raise TypeError("block_contexts must use an integer dtype")
+    if block_contexts.numel() == 0:
+        raise ValueError("block_contexts must not be empty")
+    if int(block_contexts.min()) < 0 or int(block_contexts.max()) >= mode.context_count:
+        raise ValueError("block_contexts contains an out-of-range context ID")
+    counts = torch.bincount(
+        block_contexts.to(dtype=torch.long), minlength=mode.context_count
+    )
+    expected = expected_groups // mode.context_count
+    if counts.tolist() != [expected] * mode.context_count:
+        raise ValueError("every context must have equal population")
+
+
+def importance_group_order(
+    block_contexts: torch.Tensor, mode: ModeSpec | str | int
+) -> torch.Tensor:
+    """Return stable low-to-high context order for the LDLQ recursion."""
+
+    spec = resolve_mode(mode)
+    _validate_block_contexts(block_contexts, spec)
+    return torch.argsort(block_contexts, stable=True)
+
+
+def storage_group_order(
+    block_contexts: torch.Tensor, mode: ModeSpec | str | int
+) -> torch.Tensor:
+    """Return the physical group order forming canonical P24/P33 pairs."""
+
+    spec = resolve_mode(mode)
+    _validate_block_contexts(block_contexts, spec)
+    groups_per_record = RECORD_CHANNELS // CONTEXT_GROUP_CHANNELS
+    by_context = []
+    for context in range(spec.context_count):
+        indices = torch.nonzero(block_contexts == context).flatten()
+        by_context.append(indices.reshape(-1, groups_per_record))
+
+    chunks = []
+    for low_context in range(spec.context_count // 2):
+        high_context = spec.context_count - 1 - low_context
+        low_records = by_context[low_context]
+        high_records = by_context[high_context]
+        if low_records.shape != high_records.shape:
+            raise ValueError("paired contexts have different record populations")
+        for record in range(low_records.shape[0]):
+            chunks.extend((low_records[record], high_records[record]))
+    return torch.cat(chunks)
+
+
+def rate_transfer_record_order(
+    donor_records: Sequence[int],
+    recipient_records: Sequence[int],
+    *,
+    middle_records: Sequence[int] | None = None,
+) -> tuple[int, ...]:
+    """Build a low-to-high record order for an arbitrary equal-rate exchange.
+
+    Record IDs refer to an existing 24-record partition, usually the
+    activation-ranked partition produced by calibration.  Donors occupy the
+    K2 end of an ``R_r`` schedule and recipients occupy the K4 end.  The
+    recipient suffix is reversed so ``donor_records[i]`` and
+    ``recipient_records[i]`` become one physical P24 pair after the canonical
+    balanced-atom repack.
+
+    This order can be baked into the common w1/w3-row and w2-column
+    permutation.  Therefore an arbitrary *shared* donor/recipient map does
+    not require a serving-time per-record map or a different payload format;
+    only the offline selector changes.
+    """
+
+    donors = tuple(donor_records)
+    recipients = tuple(recipient_records)
+    if len(donors) != len(recipients):
+        raise ValueError("donor and recipient record counts must match")
+    if len(donors) > PAIRS_PER_EXPERT:
+        raise ValueError(f"at most {PAIRS_PER_EXPERT} rate exchanges are allowed")
+
+    selected = donors + recipients
+    if any(
+        isinstance(record, bool)
+        or not isinstance(record, int)
+        or not 0 <= record < RECORDS_PER_EXPERT
+        for record in selected
+    ):
+        raise ValueError(
+            f"record IDs must be integers in 0..{RECORDS_PER_EXPERT - 1}"
+        )
+    if len(set(selected)) != len(selected):
+        raise ValueError("donor and recipient records must be unique and disjoint")
+
+    remaining = tuple(
+        record for record in range(RECORDS_PER_EXPERT) if record not in set(selected)
+    )
+    if middle_records is None:
+        middle = remaining
+    else:
+        middle = tuple(middle_records)
+        if len(middle) != len(remaining) or set(middle) != set(remaining):
+            raise ValueError(
+                "middle_records must be a permutation of all unselected records"
+            )
+    return donors + middle + tuple(reversed(recipients))
+
+
+def relabel_block_contexts_for_record_order(
+    block_contexts: torch.Tensor,
+    record_order: Sequence[int],
+) -> torch.Tensor:
+    """Relabel a 24-record partition to follow ``record_order``.
+
+    ``record_order[new_record]`` is the old record ID.  Sorting the returned
+    labels with :func:`importance_group_order` yields the requested record
+    order while preserving every complete 128-channel record.
+    """
+
+    _validate_block_contexts(block_contexts, R0)
+    order = tuple(record_order)
+    if len(order) != RECORDS_PER_EXPERT:
+        raise ValueError(f"record_order must contain {RECORDS_PER_EXPERT} records")
+    if any(
+        isinstance(record, bool)
+        or not isinstance(record, int)
+        or not 0 <= record < RECORDS_PER_EXPERT
+        for record in order
+    ) or len(set(order)) != RECORDS_PER_EXPERT:
+        raise ValueError("record_order must be a permutation of all record IDs")
+    inverse = torch.empty(
+        RECORDS_PER_EXPERT, dtype=torch.long, device=block_contexts.device
+    )
+    inverse[torch.tensor(order, dtype=torch.long, device=block_contexts.device)] = (
+        torch.arange(RECORDS_PER_EXPERT, device=block_contexts.device)
+    )
+    return inverse[block_contexts.to(dtype=torch.long)]
+
+
+def expand_group_order(group_order: torch.Tensor) -> torch.Tensor:
+    """Expand four-channel group IDs into a new-index -> old-index permutation."""
+
+    expected_groups = INTERMEDIATE_CHANNELS // CONTEXT_GROUP_CHANNELS
+    if group_order.ndim != 1 or group_order.numel() != expected_groups:
+        raise ValueError(f"group_order must contain {expected_groups} groups")
+    _validate_permutation(group_order, expected_groups, "group_order")
+    offsets = torch.arange(CONTEXT_GROUP_CHANNELS, device=group_order.device)
+    return (
+        group_order.to(dtype=torch.long)[:, None] * CONTEXT_GROUP_CHANNELS
+        + offsets[None]
+    ).flatten()
+
+
+def _validate_permutation(
+    permutation: torch.Tensor, length: int, name: str = "permutation"
+) -> None:
+    if permutation.ndim != 1 or permutation.numel() != length:
+        raise ValueError(f"{name} must contain exactly {length} entries")
+    if permutation.dtype == torch.bool or permutation.is_floating_point():
+        raise TypeError(f"{name} must use an integer dtype")
+    values = permutation.to(dtype=torch.long)
+    expected = torch.arange(length, device=values.device)
+    if not torch.equal(torch.sort(values).values, expected):
+        raise ValueError(f"{name} must be a bijection over 0..{length - 1}")
+
+
+def invert_permutation(permutation: torch.Tensor) -> torch.Tensor:
+    _validate_permutation(permutation, permutation.numel())
+    return torch.argsort(permutation.to(dtype=torch.long))
+
+
+def coupled_intermediate_permutation(
+    w1: torch.Tensor,
+    w3: torch.Tensor,
+    w2: torch.Tensor,
+    permutation: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Apply one exact SiTU-neuron permutation to ordinary weight matrices.
+
+    ``permutation[new_channel]`` is the corresponding old channel.  The
+    returned matrices satisfy ``w1' = w1[P]``, ``w3' = w3[P]``, and
+    ``w2' = w2[:, P]``.
+    """
+
+    if w1.ndim != 2 or w3.ndim != 2 or w2.ndim != 2:
+        raise ValueError("expert weights must all be two-dimensional")
+    if w1.shape != w3.shape:
+        raise ValueError("w1 and w3 must have identical shapes")
+    intermediate, hidden = w1.shape
+    if tuple(w2.shape) != (hidden, intermediate):
+        raise ValueError("w2 must have shape [hidden, intermediate]")
+    _validate_permutation(permutation, intermediate)
+    row_index = permutation.to(device=w1.device, dtype=torch.long)
+    w3_index = permutation.to(device=w3.device, dtype=torch.long)
+    column_index = permutation.to(device=w2.device, dtype=torch.long)
+    return (
+        w1.index_select(0, row_index).contiguous(),
+        w3.index_select(0, w3_index).contiguous(),
+        w2.index_select(1, column_index).contiguous(),
+    )
+
+
+def record_repack_order(
+    source_group_order: torch.Tensor, destination_group_order: torch.Tensor
+) -> torch.Tensor:
+    """Map destination records to source records without splitting a record."""
+
+    expected_groups = INTERMEDIATE_CHANNELS // CONTEXT_GROUP_CHANNELS
+    _validate_permutation(source_group_order, expected_groups, "source_group_order")
+    _validate_permutation(
+        destination_group_order, expected_groups, "destination_group_order"
+    )
+    groups_per_record = RECORD_CHANNELS // CONTEXT_GROUP_CHANNELS
+    source_records = source_group_order.cpu().reshape(-1, groups_per_record)
+    destination_records = destination_group_order.cpu().reshape(-1, groups_per_record)
+    source_lookup = {
+        tuple(record.tolist()): index for index, record in enumerate(source_records)
+    }
+    if len(source_lookup) != RECORDS_PER_EXPERT:
+        raise ValueError("source order contains duplicate or malformed records")
+    try:
+        order = [
+            source_lookup[tuple(record.tolist())]
+            for record in destination_records
+        ]
+    except KeyError as exc:
+        raise ValueError("destination order splits or changes a source record") from exc
+    return torch.tensor(order, dtype=torch.long, device=source_group_order.device)
+
+
+def expand_record_order(record_order: torch.Tensor) -> torch.Tensor:
+    """Expand destination-record -> source-record mapping to channel indices."""
+
+    _validate_permutation(record_order, RECORDS_PER_EXPERT, "record_order")
+    offsets = torch.arange(RECORD_CHANNELS, device=record_order.device)
+    return (
+        record_order.to(dtype=torch.long)[:, None] * RECORD_CHANNELS + offsets[None]
+    ).flatten()
+
+
+def repack_rate_axis(
+    tensor: torch.Tensor, record_order: torch.Tensor, axis: int
+) -> torch.Tensor:
+    """Move complete 128-channel blocks on a tensor's logical rate axis."""
+
+    if tensor.ndim == 0:
+        raise ValueError("cannot repack a scalar tensor")
+    normalized_axis = axis % tensor.ndim
+    if tensor.shape[normalized_axis] != INTERMEDIATE_CHANNELS:
+        raise ValueError("rate axis must contain exactly 3072 channels")
+    channels = expand_record_order(record_order).to(device=tensor.device)
+    return tensor.index_select(normalized_axis, channels).contiguous()
+
+
+def repack_encoded_records(
+    encoded: torch.Tensor, record_order: torch.Tensor, rate_axis: RateAxis
+) -> torch.Tensor:
+    """Reorder complete 128-channel records in ``[K/16, N/16, 256]`` states."""
+
+    _validate_encoded_shape(encoded, rate_axis)
+    _validate_permutation(record_order, RECORDS_PER_EXPERT, "record_order")
+    order = record_order.to(device=encoded.device, dtype=torch.long)
+    if rate_axis == "k":
+        records = encoded.reshape(
+            RECORDS_PER_EXPERT,
+            TILES_PER_RECORD_AXIS,
+            encoded.shape[1],
+            TILE_VALUES,
+        )
+        return records.index_select(0, order).reshape_as(encoded).contiguous()
+    records = (
+        encoded.reshape(
+            encoded.shape[0],
+            RECORDS_PER_EXPERT,
+            TILES_PER_RECORD_AXIS,
+            TILE_VALUES,
+        )
+        .permute(1, 0, 2, 3)
+        .contiguous()
+    )
+    reordered = records.index_select(0, order)
+    return (
+        reordered.permute(1, 0, 2, 3).contiguous().reshape_as(encoded).contiguous()
+    )
+
+
+def _validate_k(bits: int) -> None:
+    if isinstance(bits, bool) or not isinstance(bits, int) or bits not in range(2, 7):
+        raise ValueError("SQG trellis packing supports K=2 through K=6")
+
+
+def pack_trellis_edges(indices: torch.Tensor, bits: int) -> torch.Tensor:
+    """Pack low K bits with the native EXL ``pack_trellis`` word ordering."""
+
+    _validate_k(bits)
+    if indices.ndim < 1 or indices.shape[-1] != TILE_VALUES:
+        raise ValueError("indices must have a final dimension of 256")
+    if indices.dtype == torch.bool or indices.is_floating_point():
+        raise TypeError("indices must use an integer dtype")
+    values = indices.to(dtype=torch.int64) & ((1 << bits) - 1)
+    spans = values.reshape(*values.shape[:-1], TILE_CHANNELS, TILE_CHANNELS)
+    symbol_shifts = torch.arange(bits - 1, -1, -1, device=values.device)
+    bitstream = ((spans[..., None] >> symbol_shifts) & 1).reshape(
+        *values.shape[:-1], TILE_CHANNELS, bits * TILE_CHANNELS
+    )
+    word_shifts = torch.arange(15, -1, -1, device=values.device)
+    words = (
+        bitstream.reshape(*values.shape[:-1], TILE_CHANNELS, bits, 16)
+        << word_shifts
+    ).sum(dim=-1)
+    flat = words.reshape(*values.shape[:-1], TILE_CHANNELS * bits)
+    # Native pack.cu applies SWAP16: adjacent uint16 words exchange places.
+    swapped = flat.reshape(*flat.shape[:-1], -1, 2).flip(-1).reshape(flat.shape)
+    return swapped.to(dtype=torch.int16).contiguous()
+
+
+def unpack_trellis_edges(packed: torch.Tensor, bits: int) -> torch.Tensor:
+    """Invert :func:`pack_trellis_edges` to the stored K-bit edge symbols."""
+
+    _validate_k(bits)
+    expected_words = TILE_CHANNELS * bits
+    if packed.ndim < 1 or packed.shape[-1] != expected_words:
+        raise ValueError(f"packed tiles must end in {expected_words} int16 words")
+    if packed.dtype != torch.int16:
+        raise TypeError("packed trellis words must use torch.int16")
+    words = packed.to(dtype=torch.int64) & 0xFFFF
+    words = words.reshape(*words.shape[:-1], -1, 2).flip(-1).reshape(words.shape)
+    words = words.reshape(*words.shape[:-1], TILE_CHANNELS, bits)
+    word_shifts = torch.arange(15, -1, -1, device=words.device)
+    bitstream = ((words[..., None] >> word_shifts) & 1).reshape(
+        *words.shape[:-2], TILE_CHANNELS, bits * TILE_CHANNELS
+    )
+    symbol_bits = bitstream.reshape(
+        *words.shape[:-2], TILE_CHANNELS, TILE_CHANNELS, bits
+    )
+    symbol_shifts = torch.arange(bits - 1, -1, -1, device=words.device)
+    values = (symbol_bits << symbol_shifts).sum(dim=-1)
+    return values.reshape(*words.shape[:-2], TILE_VALUES).to(torch.int16).contiguous()
+
+
+def unpack_trellis_states(packed: torch.Tensor, bits: int) -> torch.Tensor:
+    """Unpack native edge words and reconstruct every cyclic 16-bit state."""
+
+    return reconstruct_trellis_states(unpack_trellis_edges(packed, bits), bits)
+
+
+@dataclass(frozen=True)
+class QSRTTrellisDescriptor:
+    """TP-independent description of one complete packed expert matrix."""
+
+    mode_id: int
+    rate_axis: RateAxis
+    k_tiles: int
+    n_tiles: int
+    schema: str = SCHEMA
+    intermediate_channels: int = INTERMEDIATE_CHANNELS
+    record_channels: int = RECORD_CHANNELS
+
+    def __post_init__(self) -> None:
+        # Candidate payloads are allocation-independent and stay in logical
+        # pair order. Physical ownership is applied only by the atom sharder.
+        if self.schema not in LOGICAL_CANDIDATE_SCHEMAS:
+            raise ValueError(f"unsupported QSRT schema: {self.schema}")
+        if self.intermediate_channels != INTERMEDIATE_CHANNELS:
+            raise ValueError("Kimi-K3 QSRT requires a 3072-channel intermediate axis")
+        if self.record_channels != RECORD_CHANNELS:
+            raise ValueError("Kimi-K3 QSRT requires 128-channel records")
+        resolve_mode(self.mode_id)
+        if self.rate_axis not in ("k", "n"):
+            raise ValueError("rate_axis must be 'k' or 'n'")
+        if self.k_tiles <= 0 or self.n_tiles <= 0:
+            raise ValueError("tile dimensions must be positive")
+        rate_tiles = self.k_tiles if self.rate_axis == "k" else self.n_tiles
+        if rate_tiles * TILE_CHANNELS != INTERMEDIATE_CHANNELS:
+            raise ValueError("rate axis must contain exactly 3072 channels")
+        orthogonal_tiles = self.n_tiles if self.rate_axis == "k" else self.k_tiles
+        if orthogonal_tiles * TILE_CHANNELS != LATENT_CHANNELS:
+            raise ValueError("orthogonal axis must contain exactly 3584 channels")
+
+    @property
+    def mode(self) -> ModeSpec:
+        return resolve_mode(self.mode_id)
+
+    @property
+    def orthogonal_tiles(self) -> int:
+        return self.n_tiles if self.rate_axis == "k" else self.k_tiles
+
+    @property
+    def tiles_per_record(self) -> int:
+        return TILES_PER_RECORD_AXIS * self.orthogonal_tiles
+
+    @property
+    def words_per_pair(self) -> int:
+        return self.tiles_per_record * TILE_CHANNELS * 6
+
+    @property
+    def payload_words(self) -> int:
+        return PAIRS_PER_EXPERT * self.words_per_pair
+
+    @property
+    def payload_bytes(self) -> int:
+        return self.payload_words * torch.int16.itemsize
+
+    @property
+    def weight_count(self) -> int:
+        return self.k_tiles * self.n_tiles * TILE_VALUES
+
+    def to_manifest(self) -> dict[str, object]:
+        return {
+            "schema": self.schema,
+            "mode_id": self.mode_id,
+            "mode": self.mode.name,
+            "rate_axis": self.rate_axis,
+            "k_tiles": self.k_tiles,
+            "n_tiles": self.n_tiles,
+            "intermediate_channels": self.intermediate_channels,
+            "record_channels": self.record_channels,
+            "pair_channels": PAIR_CHANNELS,
+            "record_bits": list(record_bits(self.mode)),
+            "pair_kinds": list(pair_kinds(self.mode)),
+            "words_per_pair": self.words_per_pair,
+            "payload_words": self.payload_words,
+            "payload_bytes": self.payload_bytes,
+        }
+
+
+@dataclass(frozen=True)
+class PackedQSRTTrellis:
+    descriptor: QSRTTrellisDescriptor
+    payload: torch.Tensor
+
+    def validate(self) -> None:
+        if self.payload.dtype != torch.int16:
+            raise TypeError("QSRT trellis payload must use torch.int16")
+        if self.payload.ndim != 1:
+            raise ValueError("QSRT trellis payload must be one-dimensional")
+        if self.payload.numel() != self.descriptor.payload_words:
+            raise ValueError(
+                "QSRT trellis payload length does not match its descriptor"
+            )
+
+
+def _validate_encoded_shape(encoded: torch.Tensor, rate_axis: RateAxis) -> None:
+    if encoded.ndim != 3 or encoded.shape[-1] != TILE_VALUES:
+        raise ValueError("encoded states must have shape [K/16, N/16, 256]")
+    if encoded.dtype == torch.bool or encoded.is_floating_point():
+        raise TypeError("encoded states must use an integer dtype")
+    if rate_axis not in ("k", "n"):
+        raise ValueError("rate_axis must be 'k' or 'n'")
+    rate_tiles = encoded.shape[0] if rate_axis == "k" else encoded.shape[1]
+    if rate_tiles * TILE_CHANNELS != INTERMEDIATE_CHANNELS:
+        raise ValueError("encoded rate axis must contain exactly 3072 channels")
+
+
+def _record_tiles(
+    encoded: torch.Tensor, record: int, rate_axis: RateAxis
+) -> torch.Tensor:
+    begin = record * TILES_PER_RECORD_AXIS
+    if rate_axis == "k":
+        return encoded.narrow(0, begin, TILES_PER_RECORD_AXIS).reshape(-1, TILE_VALUES)
+    return encoded.narrow(1, begin, TILES_PER_RECORD_AXIS).reshape(-1, TILE_VALUES)
+
+
+def pack_qsrt_trellis(
+    encoded: torch.Tensor,
+    mode: ModeSpec | str | int,
+    *,
+    rate_axis: RateAxis,
+    schema: str = SCHEMA,
+    validate_states: bool = True,
+) -> PackedQSRTTrellis:
+    """Pack a complete matrix into fixed-size balanced record pairs."""
+
+    _validate_encoded_shape(encoded, rate_axis)
+    spec = resolve_mode(mode)
+    descriptor = QSRTTrellisDescriptor(
+        mode_id=spec.mode_id,
+        rate_axis=rate_axis,
+        k_tiles=encoded.shape[0],
+        n_tiles=encoded.shape[1],
+        schema=schema,
+    )
+    pieces = []
+    for record, bits in enumerate(record_bits(spec)):
+        record_states = _record_tiles(encoded, record, rate_axis)
+        record_payload = pack_trellis_edges(record_states, bits)
+        if validate_states and not torch.equal(
+            unpack_trellis_states(record_payload, bits), record_states.to(torch.int16)
+        ):
+            raise ValueError("encoded input is not a closed cyclic trellis path")
+        pieces.append(record_payload)
+    payload = torch.cat([piece.reshape(-1) for piece in pieces]).contiguous()
+    packed = PackedQSRTTrellis(descriptor, payload)
+    packed.validate()
+    return packed
+
+
+def unpack_qsrt_trellis_edges(packed: PackedQSRTTrellis) -> torch.Tensor:
+    """Unpack a complete QSRT payload to physical-order K-bit edge symbols."""
+
+    packed.validate()
+    descriptor = packed.descriptor
+    output = torch.empty(
+        (descriptor.k_tiles, descriptor.n_tiles, TILE_VALUES),
+        dtype=torch.int16,
+        device=packed.payload.device,
+    )
+    cursor = 0
+    for record, bits in enumerate(record_bits(descriptor.mode)):
+        word_count = descriptor.tiles_per_record * TILE_CHANNELS * bits
+        words = packed.payload.narrow(0, cursor, word_count).reshape(
+            descriptor.tiles_per_record, TILE_CHANNELS * bits
+        )
+        tiles = unpack_trellis_edges(words, bits)
+        begin = record * TILES_PER_RECORD_AXIS
+        if descriptor.rate_axis == "k":
+            output.narrow(0, begin, TILES_PER_RECORD_AXIS).copy_(
+                tiles.reshape(
+                    TILES_PER_RECORD_AXIS, descriptor.n_tiles, TILE_VALUES
+                )
+            )
+        else:
+            output.narrow(1, begin, TILES_PER_RECORD_AXIS).copy_(
+                tiles.reshape(
+                    descriptor.k_tiles, TILES_PER_RECORD_AXIS, TILE_VALUES
+                )
+            )
+        cursor += word_count
+    if cursor != packed.payload.numel():
+        raise ValueError("QSRT trellis payload has trailing data")
+    return output
+
+
+def unpack_qsrt_trellis_states(packed: PackedQSRTTrellis) -> torch.Tensor:
+    """Unpack a complete QSRT payload to its physical-order 16-bit states."""
+
+    edges = unpack_qsrt_trellis_edges(packed)
+    descriptor = packed.descriptor
+    output = torch.empty_like(edges)
+    rates = record_bits(descriptor.mode)
+    # The recurrence is elementwise over every dimension except the trailing
+    # 256-symbol path.  Group equal-K records so each matrix performs at most
+    # three Python-level recurrences instead of one per 128-channel record.
+    for bits in sorted(set(rates)):
+        records = [
+            record for record, record_bits_value in enumerate(rates)
+            if record_bits_value == bits
+        ]
+        if descriptor.rate_axis == "k":
+            grouped_edges = torch.stack(
+                [
+                    edges.narrow(
+                        0,
+                        record * TILES_PER_RECORD_AXIS,
+                        TILES_PER_RECORD_AXIS,
+                    )
+                    for record in records
+                ]
+            )
+            grouped_states = reconstruct_trellis_states(grouped_edges, bits)
+            for index, record in enumerate(records):
+                output.narrow(
+                    0,
+                    record * TILES_PER_RECORD_AXIS,
+                    TILES_PER_RECORD_AXIS,
+                ).copy_(grouped_states[index])
+        else:
+            grouped_edges = torch.stack(
+                [
+                    edges.narrow(
+                        1,
+                        record * TILES_PER_RECORD_AXIS,
+                        TILES_PER_RECORD_AXIS,
+                    )
+                    for record in records
+                ]
+            )
+            grouped_states = reconstruct_trellis_states(grouped_edges, bits)
+            for index, record in enumerate(records):
+                output.narrow(
+                    1,
+                    record * TILES_PER_RECORD_AXIS,
+                    TILES_PER_RECORD_AXIS,
+                ).copy_(grouped_states[index])
+    return output
+
+
+def decode_qsrt_exl3_weight(
+    packed: PackedQSRTTrellis,
+    suh: torch.Tensor,
+    svh: torch.Tensor,
+    *,
+    codebook: str = CODEBOOK_SQG_XOR_CHEB_T12,
+) -> torch.Tensor:
+    """Decode a complete QSRT payload to its physical EXL-oriented weight."""
+
+    packed.validate()
+    _validate_scale_vector(
+        suh,
+        packed.descriptor.k_tiles * TILE_CHANNELS,
+        "suh",
+        packed.payload.device,
+    )
+    _validate_scale_vector(
+        svh,
+        packed.descriptor.n_tiles * TILE_CHANNELS,
+        "svh",
+        packed.payload.device,
+    )
+    tile_bits = tuple(
+        bits
+        for rate in record_bits(packed.descriptor.mode)
+        for bits in (rate,) * TILES_PER_RECORD_AXIS
+    )
+    return decode_qsrt_weight(
+        unpack_qsrt_trellis_states(packed),
+        suh,
+        svh,
+        rate_axis=packed.descriptor.rate_axis,
+        tile_bits=tile_bits,
+        codebook=codebook,
+    )

reproduction/closure/patches/kquant/changed-files/kquant/sqg_e4m3.py

@@ -0,0 +1,438 @@
+"""QSRT and research-control SQG E4M3 reconstruction tables.
+
+This module implements the independently proposed L16 SQG labelling.  It
+deliberately separates the rolling trellis graph from its numerical labels:
+the encoder still stores K branch bits per coefficient, while a deterministic
+K-specific mapping turns each 16-bit transition state into an E4M3 value.
+
+The primary ``sqg_xor_cheb_t12`` codebook composes the carry-mixed bijective SQG rank
+map with a modal 12-bit compression of the Chebyshev-derived finite-E4M3
+staircase.  The older exact-graph and R44 mappings remain explicit research
+controls.  All paths return raw E4M3 bytes so the offline encoder and serving
+kernel can share one byte-exact numerical contract.
+"""
+
+from __future__ import annotations
+
+from functools import lru_cache
+
+import torch
+
+
+SQG_NORMAL_E4M3 = "sqg-normal-e4m3"
+SQG_CHEB_NORMAL_E4M3 = "sqg-cheb-normal-e4m3"
+SQG_CHEB = "sqg_cheb"
+SQG_XOR_CHEB_T12 = "sqg_xor_cheb_t12"
+
+_TRANSITIONS = 1 << 16
+_CLIP = 1.0 / 2048.0
+_P = (1.25667142, 2.87422731, -9.02398882, 5.36810336, -0.46703015)
+_Q = (1.0, 2.07630930, -8.08332684, 6.32135736, -1.31208298)
+
+
+def _validate(bits: int, mode: str) -> None:
+    if isinstance(bits, bool) or not isinstance(bits, int) or bits not in range(2, 7):
+        raise ValueError("SQG supports integer K2 through K6")
+    if mode != "normal":
+        raise ValueError("the supported R44 mode is 'normal'")
+
+
+def _mix_width(
+    values: torch.Tensor,
+    *,
+    width: int,
+    multiplier_a: int,
+    multiplier_b: int,
+    shift_a: int,
+    shift_b: int,
+    shift_c: int,
+) -> torch.Tensor:
+    """Apply the report's bijective XOR/multiply mixer modulo ``2**width``."""
+
+    mask = (1 << width) - 1
+    result = values & mask
+    for multiplier, shift in (
+        (multiplier_a | 1, shift_a),
+        (multiplier_b | 1, shift_b),
+    ):
+        result ^= result >> min(max(shift, 1), width - 1)
+        result = (result * multiplier) & mask
+    result ^= result >> min(max(shift_c, 1), width - 1)
+    return result & mask
+
+
+def _reverse_low_bits(values: torch.Tensor, bits: int) -> torch.Tensor:
+    result = torch.zeros_like(values)
+    for index in range(bits):
+        result |= ((values >> index) & 1) << (bits - 1 - index)
+    return result
+
+
+def _r44_inverse_normal(probability: torch.Tensor) -> torch.Tensor:
+    a = 2.0 * probability - 1.0
+    x = a.square()
+    numerator = torch.full_like(x, _P[-1])
+    denominator = torch.full_like(x, _Q[-1])
+    for coefficient in reversed(_P[:-1]):
+        numerator = numerator * x + coefficient
+    for coefficient in reversed(_Q[:-1]):
+        denominator = denominator * x + coefficient
+    return a * numerator / denominator
+
+
+@lru_cache(maxsize=None)
+def _sqg_cpu_bytes(bits: int, mode: str) -> torch.Tensor:
+    _validate(bits, mode)
+    width = 16 - bits
+    branches = 1 << bits
+    states = torch.arange(_TRANSITIONS, dtype=torch.int64)
+    history = states >> bits
+    branch = states & (branches - 1)
+
+    phase = _mix_width(
+        history,
+        width=width,
+        multiplier_a=0x65AF,
+        multiplier_b=0x16BF,
+        shift_a=6,
+        shift_b=4,
+        shift_c=5,
+    )
+    syndrome_hash = _mix_width(
+        history ^ 0x5105,
+        width=width,
+        multiplier_a=0x8693,
+        multiplier_b=0x2A21,
+        shift_a=2,
+        shift_b=4,
+        shift_c=4,
+    )
+    syndrome = syndrome_hash & (branches - 1)
+    stratum = (
+        7 * (_reverse_low_bits(branch, bits) ^ syndrome)
+    ) & (branches - 1)
+    rank = (stratum << width) | phase
+    probability = ((rank.double() + 0.5) / _TRANSITIONS).clamp(
+        _CLIP, 1.0 - _CLIP
+    )
+    gaussian = _r44_inverse_normal(probability)
+    values = (1.5 * gaussian).float()
+    if not bool(torch.isfinite(values).all()):
+        raise RuntimeError("SQG compander produced a non-finite value")
+    return values.to(torch.float8_e4m3fn).view(torch.uint8).contiguous()
+
+
+def sqg_e4m3_bytes(
+    bits: int,
+    mode: str,
+    *,
+    device: torch.device | str = "cpu",
+) -> torch.Tensor:
+    """Return the 65,536 raw finite-E4M3 labels for an L16 SQG graph."""
+
+    _validate(bits, mode)
+    return _sqg_cpu_bytes(bits, mode).to(device=device).contiguous()
+
+
+def sqg_e4m3_codebook(
+    bits: int,
+    mode: str,
+    *,
+    device: torch.device | str = "cpu",
+    dtype: torch.dtype = torch.float16,
+) -> torch.Tensor:
+    """Return the SQG table widened from its exact E4M3 wire values."""
+
+    if not dtype.is_floating_point:
+        raise TypeError("SQG codebook dtype must be floating point")
+    raw = sqg_e4m3_bytes(bits, mode, device=device)
+    return raw.view(torch.float8_e4m3fn).to(dtype=dtype).contiguous()
+
+
+def sqg_e4m3_bytes_from_rank_lut(
+    bits: int,
+    rank_lut: torch.Tensor,
+    *,
+    device: torch.device | str = "cpu",
+) -> torch.Tensor:
+    """Map a shared rank-indexed E4M3 law onto the unchanged SQG graph.
+
+    ``rank_lut`` is deliberately independent of ``bits``.  K2 through K6 may
+    expose different outgoing menus because their stratum widths differ, but
+    they retain the same history mixer, syndrome mixer, branch permutation,
+    and scalar reconstruction law.  This is the controlled interface used to
+    evaluate alternative companders without introducing rate-specific graph
+    edits.
+    """
+
+    _validate(bits, "normal")
+    if rank_lut.dtype != torch.uint8:
+        raise TypeError("SQG rank LUT must contain raw uint8 E4M3 labels")
+    if rank_lut.ndim != 1 or rank_lut.numel() != _TRANSITIONS:
+        raise ValueError("SQG rank LUT must contain exactly 65,536 labels")
+    rank_lut = rank_lut.detach().to(device="cpu").contiguous()
+    values = rank_lut.view(torch.float8_e4m3fn).float()
+    if not bool(torch.isfinite(values).all()):
+        raise ValueError("SQG rank LUT must contain only finite E4M3 labels")
+    return rank_lut.index_select(0, sqg_rank_permutation(bits)).to(
+        device=device
+    ).contiguous()
+
+
+def sqg_e4m3_codebook_from_rank_lut(
+    bits: int,
+    rank_lut: torch.Tensor,
+    *,
+    device: torch.device | str = "cpu",
+    dtype: torch.dtype = torch.float16,
+) -> torch.Tensor:
+    """Widen a shared rank-indexed E4M3 law after SQG graph labelling."""
+
+    if not dtype.is_floating_point:
+        raise TypeError("SQG codebook dtype must be floating point")
+    raw = sqg_e4m3_bytes_from_rank_lut(bits, rank_lut, device=device)
+    return raw.view(torch.float8_e4m3fn).to(dtype=dtype).contiguous()
+
+
+def _finite_e4m3_rank_lut(values: torch.Tensor) -> torch.Tensor:
+    """Round a monotone real rank law to canonical finite E4M3 bytes."""
+
+    if values.ndim != 1 or values.numel() != _TRANSITIONS:
+        raise ValueError("SQG rank-law values must contain exactly 65,536 entries")
+    if not values.is_floating_point() or not bool(torch.isfinite(values).all()):
+        raise ValueError("SQG rank-law values must be finite floating point")
+    raw = values.float().to(torch.float8_e4m3fn).view(torch.uint8).contiguous()
+    # E4M3 has signed zero.  The numerical law uses one canonical zero byte so
+    # its exact representation is stable across CPU/GPU conversion paths.
+    raw = raw.clone()
+    raw[(raw & 0x7F) == 0] = 0
+    decoded = raw.view(torch.float8_e4m3fn).float()
+    if not bool(torch.isfinite(decoded).all()):
+        raise RuntimeError("SQG rank law rounded to a non-finite E4M3 value")
+    if not bool(torch.all(decoded[1:] >= decoded[:-1])):
+        raise RuntimeError("SQG rank law is not monotone after E4M3 rounding")
+    return raw
+
+
+@lru_cache(maxsize=1)
+def sqg_cheb_normal_rank_e4m3_bytes() -> torch.Tensor:
+    """Return the exact full-tail normal target used by SQG-Cheb.
+
+    The Chebyshev evaluator is an implementation of this finite staircase,
+    not a different reconstruction law.  Computing the midpoint quantiles
+    here gives a compact scientific reference for candidate generation and
+    reproduces the exhaustively validated L16 SQG-Cheb table byte-for-byte.
+    """
+
+    ranks = torch.arange(_TRANSITIONS, dtype=torch.float64)
+    probability = (ranks + 0.5) / _TRANSITIONS
+    values = 1.5 * torch.special.ndtri(probability)
+    return _finite_e4m3_rank_lut(values)
+
+
+@lru_cache(maxsize=None)
+def sqg_cheb_normal_e4m3_bytes(bits: int) -> torch.Tensor:
+    """Map the exact SQG-Cheb normal staircase onto the frozen SQG graph."""
+
+    return sqg_e4m3_bytes_from_rank_lut(
+        bits, sqg_cheb_normal_rank_e4m3_bytes()
+    ).contiguous()
+
+
+@lru_cache(maxsize=1)
+def sqg_xor_cheb_t12_rank_lut_bytes() -> torch.Tensor:
+    """Return the frozen 4 KiB modal T12 reconstruction staircase.
+
+    Each byte represents sixteen consecutive ranks of the exact normal
+    staircase.  The modal byte is selected per block; ties select the lower
+    raw byte.  This is the same immutable construction used by the B12X
+    ``sqg_xor_cheb_t12`` decoder.
+    """
+
+    exact = sqg_cheb_normal_rank_e4m3_bytes().reshape(1 << 12, 16)
+    result = torch.empty(1 << 12, dtype=torch.uint8)
+    for index, block in enumerate(exact):
+        labels, counts = torch.unique(block, return_counts=True)
+        result[index] = labels[counts == counts.max()].min()
+    return result.contiguous()
+
+
+@lru_cache(maxsize=None)
+def sqg_xor_rank_permutation(bits: int) -> torch.Tensor:
+    """Return the primary carry-mixed SQG rank for every L16 codeword.
+
+    The two triangular xorshifts are bijections over the retained history,
+    the multiplier is odd, and bit reversal followed by XOR is a branch
+    permutation.  The resulting map is therefore bijective over all 65,536
+    ``(stratum, phase)`` coordinates while preserving one outgoing edge per
+    stratum in every state.
+    """
+
+    _validate(bits, "normal")
+    width = 16 - bits
+    history_mask = (1 << width) - 1
+    branch_mask = (1 << bits) - 1
+    codeword = torch.arange(_TRANSITIONS, dtype=torch.int64)
+    history = codeword >> bits
+    branch = codeword & branch_mask
+
+    mixed = history ^ (history >> 11)
+    mixed ^= (mixed << 11) & history_mask
+    product = (0x3FA7D929 * mixed + 0xC928FD8E) & 0xFFFFFFFF
+    phase = product & history_mask
+    syndrome = product >> (32 - bits)
+    stratum = _reverse_low_bits(branch, bits) ^ syndrome
+    return ((stratum << width) | phase).contiguous()
+
+
+@lru_cache(maxsize=None)
+def _sqg_xor_cheb_t12_cpu_bytes(bits: int) -> torch.Tensor:
+    ranks = sqg_xor_rank_permutation(bits)
+    return sqg_xor_cheb_t12_rank_lut_bytes().index_select(0, ranks >> 4).contiguous()
+
+
+def sqg_xor_cheb_t12_bytes(
+    bits: int,
+    *,
+    device: torch.device | str = "cpu",
+) -> torch.Tensor:
+    """Return the primary QSRT-E4M3 labels for all L16 codewords."""
+
+    _validate(bits, "normal")
+    return _sqg_xor_cheb_t12_cpu_bytes(bits).to(device=device).contiguous()
+
+
+def sqg_codebook_bytes(
+    bits: int,
+    codebook: str,
+    *,
+    rate_axis: str | None = None,
+    device: torch.device | str = "cpu",
+) -> torch.Tensor:
+    """Return exact state-indexed E4M3 labels for a named SQG codebook."""
+
+    if codebook == SQG_XOR_CHEB_T12:
+        return sqg_xor_cheb_t12_bytes(bits, device=device)
+    if codebook == SQG_NORMAL_E4M3:
+        return sqg_e4m3_bytes(bits, "normal", device=device)
+    if codebook == SQG_CHEB_NORMAL_E4M3:
+        return sqg_cheb_normal_e4m3_bytes(bits).to(device=device).contiguous()
+    if codebook == SQG_CHEB:
+        if rate_axis not in ("k", "n"):
+            raise ValueError(
+                "the w2-only K2-Q8H4 profile requires rate_axis 'k' or 'n'"
+            )
+        rank_lut = sqg_cheb_normal_rank_e4m3_bytes()
+        if bits == 2 and rate_axis == "k":
+            return sqg_k2_eight_stratum_e4m3_bytes_from_rank_lut(
+                rank_lut, history_bit=4, device=device
+            )
+        return sqg_e4m3_bytes_from_rank_lut(
+            bits, rank_lut, device=device
+        ).contiguous()
+    raise ValueError(f"unsupported SQG codebook: {codebook!r}")
+
+
+def sqg_codebook(
+    bits: int,
+    codebook: str,
+    *,
+    rate_axis: str | None = None,
+    device: torch.device | str = "cpu",
+    dtype: torch.dtype = torch.float16,
+) -> torch.Tensor:
+    """Return a named SQG codebook widened from its exact E4M3 labels."""
+
+    if not dtype.is_floating_point:
+        raise TypeError("SQG codebook dtype must be floating point")
+    return (
+        sqg_codebook_bytes(bits, codebook, rate_axis=rate_axis, device=device)
+        .view(torch.float8_e4m3fn)
+        .to(dtype=dtype)
+        .contiguous()
+    )
+
+
+@lru_cache(maxsize=None)
+def sqg_k2_eight_stratum_rank_permutation(
+    history_bit: int = 2,
+) -> torch.Tensor:
+    """Return a state-conditioned eight-stratum labelling for a K2 trellis.
+
+    The physical graph still appends two branch bits and retains fourteen
+    state bits.  Numerical labelling reinterprets the newest retained history
+    bit plus the two new branch bits as a virtual three-bit stratum selector;
+    the remaining thirteen bits select phase.  A K2 state therefore exposes
+    four of eight strata, while its retained history chooses which four.
+    Globally this is the same bijective rank law used by K3.  ``history_bit``
+    selects which retained physical codeword bit becomes the most-significant
+    virtual stratum bit.  Bit 2 is the newest retained bit and reproduces the
+    exact K3 transition-label function over the unmodified codeword.
+    """
+
+    if isinstance(history_bit, bool) or not isinstance(history_bit, int):
+        raise TypeError("K2 eight-stratum history bit must be an integer")
+    if not 2 <= history_bit < 16:
+        raise ValueError("K2 eight-stratum history bit must be in [2, 15]")
+    k3_ranks = sqg_rank_permutation(3)
+    if history_bit == 2:
+        return k3_ranks
+
+    transitions = torch.arange(_TRANSITIONS, dtype=torch.int64)
+    bit_2 = (transitions >> 2) & 1
+    bit_h = (transitions >> history_bit) & 1
+    swap = bit_2 ^ bit_h
+    permuted = transitions ^ (swap << 2) ^ (swap << history_bit)
+    return k3_ranks.index_select(0, permuted).contiguous()
+
+
+def sqg_k2_eight_stratum_e4m3_bytes_from_rank_lut(
+    rank_lut: torch.Tensor,
+    *,
+    history_bit: int = 2,
+    device: torch.device | str = "cpu",
+) -> torch.Tensor:
+    """Map one shared rank staircase onto state-conditioned K2 octiles."""
+
+    if rank_lut.dtype != torch.uint8:
+        raise TypeError("SQG rank LUT must contain raw uint8 E4M3 labels")
+    if rank_lut.ndim != 1 or rank_lut.numel() != _TRANSITIONS:
+        raise ValueError("SQG rank LUT must contain exactly 65,536 labels")
+    rank_lut = rank_lut.detach().to(device="cpu").contiguous()
+    values = rank_lut.view(torch.float8_e4m3fn).float()
+    if not bool(torch.isfinite(values).all()):
+        raise ValueError("SQG rank LUT must contain only finite E4M3 labels")
+    return rank_lut.index_select(
+        0, sqg_k2_eight_stratum_rank_permutation(history_bit)
+    ).to(device=device).contiguous()
+
+
+def sqg_rank_permutation(bits: int) -> torch.Tensor:
+    """Return the pre-projection quantile rank for structural validation."""
+
+    _validate(bits, "normal")
+    width = 16 - bits
+    branches = 1 << bits
+    transitions = torch.arange(_TRANSITIONS, dtype=torch.int64)
+    history = transitions >> bits
+    branch = transitions & (branches - 1)
+    phase = _mix_width(
+        history,
+        width=width,
+        multiplier_a=0x65AF,
+        multiplier_b=0x16BF,
+        shift_a=6,
+        shift_b=4,
+        shift_c=5,
+    )
+    syndrome = _mix_width(
+        history ^ 0x5105,
+        width=width,
+        multiplier_a=0x8693,
+        multiplier_b=0x2A21,
+        shift_a=2,
+        shift_b=4,
+        shift_c=4,
+    ) & (branches - 1)
+    stratum = (7 * (_reverse_low_bits(branch, bits) ^ syndrome)) & (branches - 1)
+    return ((stratum << width) | phase).contiguous()

reproduction/closure/patches/kquant/changed-files/kquant/sqg_quantizer.py

@@ -0,0 +1,270 @@
+"""Research-only bridge from EXL's production encoder to SQG E4M3 tables."""
+
+from __future__ import annotations
+
+import hashlib
+import importlib.machinery
+import importlib.util
+import os
+import re
+import sys
+from collections.abc import Mapping
+from functools import lru_cache
+from pathlib import Path
+
+import torch
+from torch.utils.cpp_extension import load
+
+from kquant.sqg_e4m3 import sqg_e4m3_bytes
+
+
+_EXTENSION_NAME = "kquant_sqg_quantize_ext_v22"
+_EXTENSION_PATH_ENV = "KQUANT_SQG_EXTENSION_PATH"
+_EXTENSION_SHA256_ENV = "KQUANT_SQG_EXTENSION_SHA256"
+_REQUIRE_PREBUILT_ENV = "KQUANT_SQG_REQUIRE_PREBUILT"
+_SHA256_RE = re.compile(r"[0-9a-f]{64}")
+
+
+def _sha256_file(path: Path) -> str:
+    digest = hashlib.sha256()
+    with path.open("rb") as handle:
+        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def _import_prebuilt_extension(path: Path):
+    """Import one exact pybind extension without invoking the JIT builder."""
+
+    spec = importlib.util.spec_from_file_location(_EXTENSION_NAME, path)
+    if spec is None or spec.loader is None:
+        raise ImportError(f"cannot construct an import spec for SQG extension: {path}")
+    previous = sys.modules.get(_EXTENSION_NAME)
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[_EXTENSION_NAME] = module
+    try:
+        spec.loader.exec_module(module)
+    except BaseException:
+        if previous is None:
+            sys.modules.pop(_EXTENSION_NAME, None)
+        else:
+            sys.modules[_EXTENSION_NAME] = previous
+        raise
+    return module
+
+
+def _load_prebuilt_extension(raw_path: str, expected_sha256: str):
+    candidate = Path(raw_path)
+    if not candidate.is_absolute():
+        raise ValueError(f"{_EXTENSION_PATH_ENV} must be an absolute path")
+    if candidate.is_symlink():
+        raise ValueError(f"{_EXTENSION_PATH_ENV} may not name a symlink")
+    try:
+        path = candidate.resolve(strict=True)
+    except FileNotFoundError as exc:
+        raise FileNotFoundError(
+            f"sealed SQG extension does not exist: {candidate}"
+        ) from exc
+    if not path.is_file():
+        raise ValueError(f"sealed SQG extension is not a regular file: {path}")
+    allowed_names = {
+        f"{_EXTENSION_NAME}{suffix}"
+        for suffix in importlib.machinery.EXTENSION_SUFFIXES
+    }
+    if path.name not in allowed_names:
+        raise ValueError(
+            "sealed SQG extension filename must identify "
+            f"{_EXTENSION_NAME}: {path.name}"
+        )
+    if _SHA256_RE.fullmatch(expected_sha256) is None:
+        raise ValueError(f"{_EXTENSION_SHA256_ENV} must be 64 lowercase hex digits")
+    observed_sha256 = _sha256_file(path)
+    if observed_sha256 != expected_sha256:
+        raise RuntimeError(
+            "sealed SQG extension SHA256 mismatch: "
+            f"{observed_sha256} != {expected_sha256}"
+        )
+    module = _import_prebuilt_extension(path)
+    module_path = Path(str(getattr(module, "__file__", ""))).resolve()
+    if module_path != path:
+        raise RuntimeError(
+            f"loaded SQG extension origin differs: {module_path} != {path}"
+        )
+    if not callable(getattr(module, "quantize_tiles_procedural", None)) or not callable(
+        getattr(module, "quantize_tiles_sqg", None)
+    ):
+        raise RuntimeError(
+            "sealed SQG extension lacks required quantization entry points"
+        )
+    if _sha256_file(path) != expected_sha256:
+        raise RuntimeError("sealed SQG extension changed while it was being imported")
+    return module
+
+
+@lru_cache(maxsize=1)
+def _extension():
+    require_value = os.environ.get(_REQUIRE_PREBUILT_ENV, "")
+    if require_value not in ("", "0", "1"):
+        raise ValueError(f"{_REQUIRE_PREBUILT_ENV} must be unset, 0, or 1")
+    raw_path = os.environ.get(_EXTENSION_PATH_ENV, "")
+    expected_sha256 = os.environ.get(_EXTENSION_SHA256_ENV, "")
+    if require_value == "1" or raw_path or expected_sha256:
+        if not raw_path or not expected_sha256:
+            raise RuntimeError(
+                f"{_EXTENSION_PATH_ENV} and {_EXTENSION_SHA256_ENV} are both required "
+                "for sealed prebuilt SQG loading"
+            )
+        return _load_prebuilt_extension(raw_path, expected_sha256)
+
+    project = Path(__file__).resolve().parents[1]
+    return load(
+        name=_EXTENSION_NAME,
+        sources=[
+            str(project / "kquant/csrc/sqg_quantize.cpp"),
+            str(project / "kquant/csrc/sqg_quantize.cu"),
+        ],
+        extra_include_paths=[str(project / "kquant/csrc")],
+        extra_cflags=["-O3"],
+        extra_cuda_cflags=[
+            "-O3",
+            "--use_fast_math",
+            "-lineinfo",
+            "-Xcudafe",
+            "--diag_suppress=177",
+            "-Xcudafe",
+            "--diag_suppress=20012",
+        ],
+        verbose=False,
+    )
+
+
+@lru_cache(maxsize=None)
+def _sqg_temp_buffers(device: torch.device, bits: int):
+    """Allocate the packed traceback buffers consumed by kquant's kernel."""
+
+    multiprocessors = torch.cuda.get_device_properties(device).multi_processor_count
+    edges = 65536 >> bits
+    decisions = edges // 4 if bits == 2 else (edges // 2 if bits <= 4 else edges)
+    free_bytes, _ = torch.cuda.mem_get_info(device)
+    decision_bytes_per_tile = 256 * decisions
+    affordable = max(256, int(free_bytes * 0.5) // decision_bytes_per_tile)
+    max_batch = min(max(256, 3 * multiprocessors), affordable)
+    costs = torch.zeros(
+        (max_batch, 2, edges), dtype=torch.float16, device=device
+    )
+    traceback = torch.empty(
+        (max_batch, 256, decisions), dtype=torch.uint8, device=device
+    )
+    return costs, traceback
+
+
+def install_sqg_quantizer(quantizer_module) -> None:
+    """Teach a loaded EXL encoder module to consume ``sqg_e4m3_lut``.
+
+    The patch is process-local. SQG and explicit MCG/MUL1 controls use
+    kquant's CUDA extension with one Viterbi/tail-biting implementation. A
+    ``None`` entry in a rate-specific mapping explicitly selects MCG,
+    permitting controlled hybrid rate-curve studies. Calls without any
+    kquant codebook argument retain the unmodified upstream EXL behavior.
+    """
+
+    if getattr(quantizer_module, "_kquant_sqg_installed", False):
+        return
+    original = quantizer_module.quantize_tiles
+
+    device_luts: dict[tuple[str, int, str], torch.Tensor] = {}
+    transposed_sqg_luts: dict[
+        tuple[str, int, int], tuple[torch.Tensor, torch.Tensor]
+    ] = {}
+
+    def quantize_tiles(tiles: torch.Tensor, quant_args: dict):
+        codebook = quant_args.get("sqg_e4m3_lut")
+        rate_codebooks = quant_args.get("sqg_e4m3_luts_by_bits")
+        mode = quant_args.get("sqg_e4m3_mode")
+        if codebook is None and rate_codebooks is None and mode is None:
+            return original(tiles, quant_args)
+        if len(quant_args["devices"]) != 1:
+            raise ValueError("the SQG validation hook currently requires one CUDA device")
+        tiles = tiles.contiguous()
+        if tiles.dtype != torch.float32 or tiles.ndim != 2 or tiles.shape[1] != 256:
+            raise ValueError("SQG tiles must be contiguous FP32 [N, 256]")
+        bits = int(quant_args["K"])
+        if rate_codebooks is not None:
+            if codebook is not None or mode is not None:
+                raise ValueError(
+                    "rate-specific SQG LUTs cannot be combined with another SQG law"
+                )
+            if not isinstance(rate_codebooks, Mapping):
+                raise TypeError("sqg_e4m3_luts_by_bits must be a mapping")
+            if set(rate_codebooks) - set(range(2, 7)):
+                raise ValueError("rate-specific SQG LUT keys must be K2 through K6")
+            try:
+                codebook = rate_codebooks[bits]
+            except KeyError as exc:
+                raise ValueError(f"missing rate-specific SQG K{bits} LUT") from exc
+            if codebook is None:
+                if bits not in (2, 3, 4):
+                    raise ValueError(
+                        "non-SQG procedural controls remain limited to K2 through K4"
+                    )
+                output = torch.empty_like(tiles)
+                indices = torch.empty_like(tiles, dtype=torch.int16)
+                costs, edges = _sqg_temp_buffers(tiles.device, bits)
+                _extension().quantize_tiles_procedural(
+                    tiles,
+                    output,
+                    indices,
+                    costs,
+                    edges,
+                    bits,
+                    1,
+                    int(quant_args.get("tailbite_context", 128)),
+                )
+                return output, indices
+        elif codebook is None:
+            if mode != "normal":
+                raise ValueError("the supported R44 mode is 'normal'")
+            key = (str(tiles.device), bits, mode)
+            codebook = device_luts.get(key)
+            if codebook is None:
+                codebook = sqg_e4m3_bytes(bits, mode, device=tiles.device)
+                device_luts[key] = codebook
+        output = torch.empty_like(tiles)
+        indices = torch.empty_like(tiles, dtype=torch.int16)
+        costs, edges = _sqg_temp_buffers(tiles.device, bits)
+        lut = codebook.to(device=tiles.device, dtype=torch.uint8).contiguous()
+        if bits in (2, 3, 4):
+            source_key = codebook.data_ptr() if codebook.is_cuda else id(codebook)
+            cache_key = (str(tiles.device), bits, source_key)
+            cached = transposed_sqg_luts.get(cache_key)
+            if cached is None or cached[0] is not codebook:
+                # [predecessor, out-edge-pair, pair-byte] ->
+                # [out-edge-pair, predecessor, pair-byte].  Each CUDA thread
+                # can then fetch all predecessor labels in one uint2 (K2),
+                # one uint4 (K3), or two uint4s (K4), rather than issuing one
+                # gather per predecessor.
+                predecessors = 1 << bits
+                out_edge_pairs = (65536 >> bits) // 2
+                lut = (
+                    lut.reshape(predecessors, out_edge_pairs, 2)
+                    .permute(1, 0, 2)
+                    .contiguous()
+                    .reshape(-1)
+                )
+                transposed_sqg_luts[cache_key] = (codebook, lut)
+            else:
+                lut = cached[1]
+        _extension().quantize_tiles_sqg(
+            tiles,
+            output,
+            indices,
+            costs,
+            edges,
+            lut,
+            bits,
+            int(quant_args.get("tailbite_context", 128)),
+        )
+        return output, indices
+
+    quantizer_module.quantize_tiles = quantize_tiles
+    quantizer_module._kquant_sqg_installed = True

reproduction/closure/patches/kquant/changed-files/tests/test_qsrt.py

@@ -0,0 +1,486 @@
+from __future__ import annotations
+
+from dataclasses import replace
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from kquant.exl3_reference import reconstruct_trellis_states
+from kquant.qsrt import (
+    EXPERT_TRELLIS_BYTES,
+    EXPERTS_PER_LAYER,
+    FORMAT_SECTION_BYTES,
+    FORMAT_TABLE_BYTES,
+    HOT_METADATA_BYTES,
+    INTERMEDIATE_CHANNELS,
+    LAYER_HEADER_BYTES,
+    MATRIX_TRELLIS_BYTES,
+    PAIRS_PER_EXPERT,
+    PAIR_BYTES,
+    PHASE1_H13_EXPERT_LOCAL_ALPHA,
+    PHASE1_H2_EXPERT_LOCAL_ALPHA,
+    PHASE1_H2_LOCAL_BASIS,
+    PHASE1_H2_SHRINKAGE_POLICY,
+    PHASE1_MODE_CANDIDATES,
+    PHASE1_MODE_IDS,
+    PHASE1_SHARED_RATE_MODE,
+    RECORDS_PER_EXPERT,
+    RATE_TRANSFER_MODES,
+    R0,
+    R1,
+    R2,
+    SCHEMA,
+    SHARED_SCALE_BYTES,
+    SHARED_SCALE_SECTION_BYTES,
+    STORAGE_ALIGNMENT,
+    ExpertFormatSpec,
+    ModeSpec,
+    PackedQSRTTrellis,
+    QSRTTrellisDescriptor,
+    coupled_intermediate_permutation,
+    decode_qsrt_exl3_weight,
+    expand_group_order,
+    expand_qsrt_hot_metadata,
+    expand_record_order,
+    importance_group_order,
+    matrix_rate_axis,
+    mode_from_context_bits,
+    pack_qsrt_format_table,
+    pack_qsrt_format_section,
+    pack_qsrt_trellis,
+    pack_qsrt_shared_scale_section,
+    pack_trellis_edges,
+    rate_transfer_record_order,
+    record_repack_order,
+    relabel_block_contexts_for_record_order,
+    repack_encoded_records,
+    repack_rate_axis,
+    pair_bits,
+    pair_kinds,
+    logical_pair_for_slot,
+    physical_pair_slot,
+    record_bits,
+    record_contexts,
+    storage_group_order,
+    unpack_qsrt_trellis_edges,
+    unpack_qsrt_format_section,
+    unpack_qsrt_shared_scale_section,
+    unpack_qsrt_format_table,
+    unpack_qsrt_trellis_states,
+    unpack_trellis_edges,
+    unpack_trellis_states,
+)
+from kquant.tp_simulator import situ
+
+
+def test_mode_schedules_are_fixed_rate_pairs() -> None:
+    assert pair_kinds(R0) == ("P33",) * PAIRS_PER_EXPERT
+    assert pair_kinds(R1) == ("P24",) + ("P33",) * 11
+    assert pair_kinds(R2) == ("P24",) * 2 + ("P33",) * 10
+
+
+def test_pair_rotation_is_a_per_expert_bijection() -> None:
+    for layer, expert in ((1, 0), (1, 17), (40, 285), (92, 895)):
+        physical = [
+            physical_pair_slot(layer, expert, logical_pair)
+            for logical_pair in range(PAIRS_PER_EXPERT)
+        ]
+        assert sorted(physical) == list(range(PAIRS_PER_EXPERT))
+        assert [
+            logical_pair_for_slot(layer, expert, rank)
+            for rank in physical
+        ] == list(range(PAIRS_PER_EXPERT))
+
+
+def test_pair_rotation_balances_p24_ownership_across_experts() -> None:
+    counts = [0] * PAIRS_PER_EXPERT
+    r = 2
+    for expert in range(EXPERTS_PER_LAYER):
+        for rank in range(PAIRS_PER_EXPERT):
+            counts[rank] += logical_pair_for_slot(24, expert, rank) < r
+
+    assert max(counts) - min(counts) <= 1
+    for r, mode in enumerate(RATE_TRANSFER_MODES):
+        assert mode.mode_id == r
+        assert mode.name == f"R{r}"
+        assert len(record_bits(mode)) == RECORDS_PER_EXPERT
+        assert len(pair_bits(mode)) == PAIRS_PER_EXPERT
+        assert all(sum(pair) == 6 for pair in pair_bits(mode))
+        assert sum(record_bits(mode)) / RECORDS_PER_EXPERT == 3
+        assert pair_kinds(mode) == ("P24",) * r + ("P33",) * (12 - r)
+
+
+def test_qsrt_encoder_policy_is_frozen_to_r012() -> None:
+    assert PHASE1_MODE_IDS == (0, 1, 2)
+    assert tuple(mode.mode_id for mode in PHASE1_MODE_CANDIDATES) == PHASE1_MODE_IDS
+    assert not PHASE1_SHARED_RATE_MODE
+    assert PHASE1_H13_EXPERT_LOCAL_ALPHA == 0.0
+    assert PHASE1_H2_EXPERT_LOCAL_ALPHA == 0.75
+    assert PHASE1_H2_LOCAL_BASIS == "decoded_candidate_post_situ"
+    assert PHASE1_H2_SHRINKAGE_POLICY == "weighted_oas_scaled_identity"
+    assert len(RATE_TRANSFER_MODES) == 3
+
+
+def test_context_allocations_map_to_the_rate_transfer_ladder() -> None:
+    assert mode_from_context_bits((3, 3, 3, 3)) is R0
+    assert mode_from_context_bits((2,) + (3,) * 22 + (4,)) is R1
+    assert mode_from_context_bits((2, 2) + (3,) * 20 + (4, 4)) is R2
+    with pytest.raises(ValueError, match="monotone"):
+        mode_from_context_bits((2, 4) + (3,) * 20 + (2, 4))
+    with pytest.raises(ValueError, match="rate-transfer count 6"):
+        ModeSpec(1, "bad", (2, 3, 3, 4))
+
+
+def test_candidate_descriptor_contains_no_tensor_parallel_geometry() -> None:
+    descriptor = QSRTTrellisDescriptor(
+        mode_id=R2.mode_id,
+        rate_axis="k",
+        k_tiles=192,
+        n_tiles=224,
+    )
+    assert "tp_size" not in descriptor.to_manifest()
+
+
+def test_provisional_format_table_supports_separate_r13_r2() -> None:
+    formats = ["X4T", R0, R1, (1, 2)] * (EXPERTS_PER_LAYER // 4)
+    packed = pack_qsrt_format_table(formats)
+
+    assert packed.dtype == torch.uint8
+    assert packed.shape == (FORMAT_TABLE_BYTES,)
+    assert packed[:4].tolist() == [0xFF, 0x00, 0x11, 0x12]
+    assert unpack_qsrt_format_table(packed) == tuple(
+        ["X4T", "R0", "R1", "R1/R2"] * (EXPERTS_PER_LAYER // 4)
+    )
+
+    malformed = packed.clone()
+    malformed[0] = 0xD0
+    with pytest.raises(ValueError, match="invalid rate code"):
+        unpack_qsrt_format_table(malformed)
+    assert ExpertFormatSpec.compressed(1, 2).code == 0x12
+    assert ExpertFormatSpec.from_code(0x12).name == "R1/R2"
+    assert ExpertFormatSpec.from_code(0xFF).is_x4t
+    with pytest.raises(ValueError, match="896 experts"):
+        pack_qsrt_format_table([R0])
+    with pytest.raises(TypeError, match="torch.uint8"):
+        unpack_qsrt_format_table(packed.int())
+
+
+def test_format_section_padding_and_hot_slot_expansion_are_canonical() -> None:
+    formats = ["X4T", R0, R1, (1, 2)] * (EXPERTS_PER_LAYER // 4)
+    section = pack_qsrt_format_section(formats)
+    codes, slots = expand_qsrt_hot_metadata(formats)
+
+    assert section.shape == (FORMAT_SECTION_BYTES,)
+    assert torch.count_nonzero(section[FORMAT_TABLE_BYTES:]) == 0
+    assert unpack_qsrt_format_section(section) == tuple(
+        ["X4T", "R0", "R1", "R1/R2"] * (EXPERTS_PER_LAYER // 4)
+    )
+    assert codes[:8].tolist() == [255, 0, 17, 18, 255, 0, 17, 18]
+    assert slots[:8].tolist() == [0, 0, 1, 2, 1, 3, 4, 5]
+    assert codes.numel() + slots.numel() * slots.element_size() == HOT_METADATA_BYTES
+
+    malformed = section.clone()
+    malformed[-1] = 1
+    with pytest.raises(ValueError, match="nonzero alignment padding"):
+        unpack_qsrt_format_section(malformed)
+
+
+def test_shared_scale_section_round_trips_exactly() -> None:
+    shared = tuple(
+        torch.linspace(-1 + index, 1 + index, 3584, dtype=torch.float16)
+        for index in range(3)
+    )
+    section = pack_qsrt_shared_scale_section(*shared)
+    actual = unpack_qsrt_shared_scale_section(section)
+    assert all(
+        torch.equal(expected, observed)
+        for expected, observed in zip(shared, actual, strict=True)
+    )
+
+    malformed = section.clone()
+    malformed[-1] = 1
+    with pytest.raises(ValueError, match="nonzero alignment padding"):
+        unpack_qsrt_shared_scale_section(malformed)
+
+
+def test_matrix_rate_axis_tracks_shared_intermediate_dimension() -> None:
+    assert matrix_rate_axis("w1") == "n"
+    assert matrix_rate_axis("w3") == "n"
+    assert matrix_rate_axis("w2") == "k"
+    with pytest.raises(ValueError, match="unknown expert matrix"):
+        matrix_rate_axis("router")
+
+
+@pytest.mark.parametrize(
+    ("bits", "expected_prefix"),
+    [
+        (2, [6939] * 6),
+        (3, [30469, 1337, 1337, 14711, 14711, 30469]),
+        (4, [17767, 291, -12817, -30293, 17767, 291]),
+        (5, [12820, 68, 21686, -14526, -31643, 13775]),
+        (6, [-31984, 16, 8338, 20871, -11377, -29904]),
+    ],
+)
+def test_trellis_edge_pack_matches_native_word_order(
+    bits: int, expected_prefix: list[int]
+) -> None:
+    indices = (torch.arange(256, dtype=torch.int16) % (1 << bits))[None]
+    packed = pack_trellis_edges(indices, bits)
+
+    assert packed.shape == (1, 16 * bits)
+    assert packed[0, : len(expected_prefix)].tolist() == expected_prefix
+    torch.testing.assert_close(unpack_trellis_edges(packed, bits), indices)
+
+
+@pytest.mark.parametrize("bits", range(2, 7))
+def test_trellis_edge_pack_round_trips_only_stored_low_bits(bits: int) -> None:
+    generator = torch.Generator().manual_seed(100 + bits)
+    indices = torch.randint(-32768, 32768, (7, 256), generator=generator)
+    packed = pack_trellis_edges(indices, bits)
+    unpacked = unpack_trellis_edges(packed, bits)
+
+    torch.testing.assert_close(unpacked.int(), indices.int() & ((1 << bits) - 1))
+
+
+@pytest.mark.parametrize("bits", range(2, 7))
+def test_cyclic_trellis_states_round_trip_exactly(bits: int) -> None:
+    generator = torch.Generator().manual_seed(150 + bits)
+    edges = torch.randint(0, 1 << bits, (5, 256), generator=generator)
+    states = reconstruct_trellis_states(edges, bits)
+    packed = pack_trellis_edges(states, bits)
+
+    torch.testing.assert_close(unpack_trellis_states(packed, bits), states)
+
+
+@pytest.mark.parametrize("bits", range(2, 7))
+def test_cyclic_trellis_state_closed_form_matches_scalar_recurrence(bits: int) -> None:
+    generator = torch.Generator().manual_seed(175 + bits)
+    edges = torch.randint(0, 1 << bits, (3, 5, 256), generator=generator)
+
+    state = torch.zeros(edges.shape[:-1], dtype=torch.int64)
+    for edge in edges[..., -((16 + bits - 1) // bits) :].unbind(dim=-1):
+        state = ((state << bits) | edge) & 0xFFFF
+    expected = []
+    for edge in edges.unbind(dim=-1):
+        state = ((state << bits) | edge) & 0xFFFF
+        expected.append(state)
+    expected_states = torch.stack(expected, dim=-1).to(torch.int16)
+
+    assert torch.equal(reconstruct_trellis_states(edges, bits), expected_states)
+
+
+def _expected_edge_symbols(
+    encoded: torch.Tensor, record_bits: tuple[int, ...], rate_axis: str
+) -> torch.Tensor:
+    tile_bits = torch.tensor(record_bits).repeat_interleave(8)
+    masks = (1 << tile_bits) - 1
+    if rate_axis == "k":
+        masks = masks[:, None, None]
+    else:
+        masks = masks[None, :, None]
+    return (encoded.to(torch.int64) & masks).to(torch.int16)
+
+
+def _legal_qsrt_states(
+    shape: tuple[int, int, int], record_bits: tuple[int, ...], rate_axis: str, seed: int
+) -> torch.Tensor:
+    generator = torch.Generator().manual_seed(seed)
+    edges = torch.empty(shape, dtype=torch.int16)
+    states = torch.empty_like(edges)
+    for record, bits in enumerate(record_bits):
+        begin = record * 8
+        if rate_axis == "k":
+            record_edges = torch.randint(
+                0, 1 << bits, (8, shape[1], 256), generator=generator
+            ).to(torch.int16)
+            edges[begin : begin + 8] = record_edges
+            states[begin : begin + 8] = reconstruct_trellis_states(
+                record_edges, bits
+            )
+        else:
+            record_edges = torch.randint(
+                0, 1 << bits, (shape[0], 8, 256), generator=generator
+            ).to(torch.int16)
+            edges[:, begin : begin + 8] = record_edges
+            states[:, begin : begin + 8] = reconstruct_trellis_states(
+                record_edges, bits
+            )
+    return states
+
+
+@pytest.mark.parametrize(
+    ("rate_axis", "shape"),
+    [("k", (192, 224, 256)), ("n", (224, 192, 256))],
+)
+@pytest.mark.parametrize("mode", [R0, R1, R2])
+def test_pair_payload_closes_on_both_logical_axes(
+    rate_axis: str, shape: tuple[int, int, int], mode
+) -> None:
+    encoded = _legal_qsrt_states(
+        shape, record_bits(mode), rate_axis, 200 + mode.mode_id
+    )
+    packed = pack_qsrt_trellis(encoded, mode, rate_axis=rate_axis)
+    unpacked = unpack_qsrt_trellis_edges(packed)
+
+    assert packed.descriptor.schema == SCHEMA
+    assert packed.payload.numel() == encoded.numel() * 3 // 16
+    assert packed.descriptor.payload_bytes == packed.payload.numel() * 2
+    assert (
+        packed.descriptor.words_per_pair * PAIRS_PER_EXPERT
+        == packed.payload.numel()
+    )
+    assert packed.descriptor.to_manifest()["pair_kinds"] == list(
+        pair_kinds(mode)
+    )
+    torch.testing.assert_close(
+        unpacked, _expected_edge_symbols(encoded, record_bits(mode), rate_axis)
+    )
+    torch.testing.assert_close(unpack_qsrt_trellis_states(packed), encoded)
+
+
+def test_packer_rejects_noncyclic_states() -> None:
+    malformed = torch.zeros((192, 224, 256), dtype=torch.int16)
+    malformed[..., 0] = 1
+
+    with pytest.raises(ValueError, match="closed cyclic trellis path"):
+        pack_qsrt_trellis(malformed, R2, rate_axis="k")
+
+
+def test_payload_rejects_wrong_shape_dtype_and_length() -> None:
+    encoded = torch.zeros((192, 224, 256), dtype=torch.int16)
+    packed = pack_qsrt_trellis(encoded, R2, rate_axis="k")
+    suh = torch.ones(3072, dtype=torch.float16)
+    svh = torch.ones(3584, dtype=torch.float16)
+
+    with pytest.raises(ValueError, match="length"):
+        PackedQSRTTrellis(packed.descriptor, packed.payload[:-1]).validate()
+    with pytest.raises(TypeError, match="torch.int16"):
+        PackedQSRTTrellis(packed.descriptor, packed.payload.int()).validate()
+    with pytest.raises(ValueError, match="3072"):
+        pack_qsrt_trellis(encoded[:184], R2, rate_axis="k")
+    with pytest.raises(ValueError, match="rate_axis"):
+        replace(packed.descriptor, rate_axis="x")
+    with pytest.raises(TypeError, match="suh must use torch.float16"):
+        decode_qsrt_exl3_weight(packed, suh.float(), svh)
+    with pytest.raises(ValueError, match="exactly 3072"):
+        decode_qsrt_exl3_weight(packed, suh[:-1], svh)
+    nonfinite = svh.clone()
+    nonfinite[0] = torch.inf
+    with pytest.raises(ValueError, match="non-finite"):
+        decode_qsrt_exl3_weight(packed, suh, nonfinite)
+
+
+def test_context_orders_form_homogeneous_records() -> None:
+    block_contexts = torch.arange(24).repeat_interleave(32)
+    importance = importance_group_order(block_contexts, R2)
+    storage = storage_group_order(block_contexts, R2)
+    storage_channels = block_contexts.index_select(0, storage).repeat_interleave(4)
+
+    assert expand_group_order(storage).shape == (INTERMEDIATE_CHANNELS,)
+    assert torch.all(storage_channels.reshape(RECORDS_PER_EXPERT, 128).diff(dim=1) == 0)
+    assert (
+        storage_channels.reshape(RECORDS_PER_EXPERT, 128)[:, 0].tolist()
+        == list(record_contexts(R2))
+    )
+
+    order = record_repack_order(importance, storage)
+    encoded = (
+        torch.arange(RECORDS_PER_EXPERT)
+        .repeat_interleave(8 * 256)
+        .reshape(192, 1, 256)
+    )
+    repacked = repack_encoded_records(encoded, order, "k")
+    assert repacked.reshape(RECORDS_PER_EXPERT, -1)[:, 0].tolist() == order.tolist()
+
+
+def test_record_repack_rejects_a_split_record() -> None:
+    block_contexts = torch.arange(24).repeat_interleave(32)
+    source = importance_group_order(block_contexts, R2)
+    malformed = source.clone()
+    malformed[31], malformed[32] = malformed[32].clone(), malformed[31].clone()
+
+    with pytest.raises(ValueError, match="splits or changes"):
+        record_repack_order(source, malformed)
+
+
+def test_arbitrary_shared_record_map_folds_into_common_permutation() -> None:
+    # Start from a deliberately non-contiguous semantic partition so the test
+    # exercises relabeling rather than relying on identity group order.
+    generator = torch.Generator().manual_seed(47)
+    base_contexts = torch.arange(RECORDS_PER_EXPERT).repeat_interleave(32)
+    base_contexts = base_contexts[torch.randperm(base_contexts.numel(), generator=generator)]
+    donors = (7, 2)
+    recipients = (19, 11)
+    order = rate_transfer_record_order(donors, recipients)
+    relabeled = relabel_block_contexts_for_record_order(base_contexts, order)
+    importance = importance_group_order(relabeled, R2)
+    ordered_semantic_records = base_contexts.index_select(0, importance).reshape(
+        RECORDS_PER_EXPERT, 32
+    )[:, 0]
+
+    assert order[:2] == donors
+    assert order[-2:] == tuple(reversed(recipients))
+    assert ordered_semantic_records.tolist() == list(order)
+
+    storage = storage_group_order(relabeled, R2)
+    physical_records = base_contexts.index_select(0, storage).reshape(
+        RECORDS_PER_EXPERT, 32
+    )[:, 0]
+    assert list(zip(physical_records[0:4:2].tolist(), physical_records[1:4:2].tolist())) == list(
+        zip(donors, recipients, strict=True)
+    )
+
+
+@pytest.mark.parametrize(
+    "donors,recipients,middle,error",
+    [
+        ((1,), (), None, "counts must match"),
+        ((1,), (1,), None, "unique and disjoint"),
+        ((24,), (1,), None, "record IDs"),
+        ((1,), (2,), tuple(range(21)), "middle_records"),
+    ],
+)
+def test_arbitrary_record_map_rejects_malformed_assignments(
+    donors, recipients, middle, error
+) -> None:
+    with pytest.raises(ValueError, match=error):
+        rate_transfer_record_order(
+            donors, recipients, middle_records=middle
+        )
+
+
+def test_record_repack_moves_matching_scale_or_weight_axis() -> None:
+    order = torch.arange(RECORDS_PER_EXPERT - 1, -1, -1)
+    channels = expand_record_order(order)
+    vector = torch.arange(INTERMEDIATE_CHANNELS)
+    matrix = vector.repeat(3, 1)
+
+    torch.testing.assert_close(repack_rate_axis(vector, order, 0), vector[channels])
+    torch.testing.assert_close(
+        repack_rate_axis(matrix, order, 1), matrix.index_select(1, channels)
+    )
+
+
+def test_coupled_permutation_preserves_situ_expert_function() -> None:
+    generator = torch.Generator().manual_seed(300)
+    hidden = 7
+    intermediate = 64
+    inputs = torch.randn(5, hidden, dtype=torch.float64, generator=generator)
+    w1 = torch.randn(intermediate, hidden, dtype=torch.float64, generator=generator)
+    w3 = torch.randn(intermediate, hidden, dtype=torch.float64, generator=generator)
+    w2 = torch.randn(hidden, intermediate, dtype=torch.float64, generator=generator)
+    permutation = torch.randperm(intermediate, generator=generator)
+
+    reference_middle = situ(F.linear(inputs, w1), F.linear(inputs, w3))
+    reference = F.linear(reference_middle, w2)
+    permuted_w1, permuted_w3, permuted_w2 = coupled_intermediate_permutation(
+        w1, w3, w2, permutation
+    )
+    permuted_middle = situ(
+        F.linear(inputs, permuted_w1), F.linear(inputs, permuted_w3)
+    )
+    actual = F.linear(permuted_middle, permuted_w2)
+
+    torch.testing.assert_close(permuted_middle, reference_middle[:, permutation])
+    torch.testing.assert_close(actual, reference, atol=1e-12, rtol=1e-12)

reproduction/closure/patches/kquant/changed-files/tests/test_sqg_e4m3.py

@@ -0,0 +1,236 @@
+from __future__ import annotations
+
+import hashlib
+
+import pytest
+import torch
+
+from kquant.exl3_reference import (
+    CODEBOOK_SQG_XOR_CHEB_T12,
+    CODEBOOK_SQG_CHEB_NORMAL_E4M3,
+    CODEBOOK_SQG_CHEB,
+    CODEBOOK_SQG_NORMAL_E4M3,
+    decode_qsrt_regularized_weight,
+    decode_regularized_weight,
+)
+from kquant.sqg_e4m3 import (
+    sqg_xor_cheb_t12_bytes,
+    sqg_xor_cheb_t12_rank_lut_bytes,
+    sqg_xor_rank_permutation,
+    sqg_cheb_normal_e4m3_bytes,
+    sqg_cheb_normal_rank_e4m3_bytes,
+    sqg_codebook_bytes,
+    sqg_e4m3_bytes,
+    sqg_e4m3_bytes_from_rank_lut,
+    sqg_e4m3_codebook,
+    sqg_e4m3_codebook_from_rank_lut,
+    sqg_k2_eight_stratum_rank_permutation,
+    sqg_rank_permutation,
+)
+
+
+_SQG_XOR_CHEB_T12_SHA256 = {
+    "t12": "cca11fe5744c9c93a34f4217f342fbc0f74ecc8a007c076582424a505fc9da5e",
+    2: "62027916386245a84c86156a0a08b6cf07e41548871af0e56fb40780558f6293",
+    3: "afe7b3633e7d243b00b379b18ec4dca573722b3727cafef47fcb6470d7e7e6c9",
+    4: "5a9620f0c4d8f0a60d0b6fbea921dcecbd193e6febf31043aaa6403c20389c2f",
+    5: "a1caf5572c67d4face423478f423e1ed6149a8c920be5c28d4702946c93cae2f",
+    6: "8e0edc92e20eb92f5287dbdadc8e870fc500fc03c5eed996cae0781d197ebcb0",
+}
+
+
+def _sha256(tensor: torch.Tensor) -> str:
+    return hashlib.sha256(tensor.contiguous().numpy().tobytes()).hexdigest()
+
+
+def test_primary_sqg_xor_cheb_t12_matches_runtime_contract() -> None:
+    table = sqg_xor_cheb_t12_rank_lut_bytes()
+    assert table.dtype == torch.uint8
+    assert table.shape == (1 << 12,)
+    assert _sha256(table) == _SQG_XOR_CHEB_T12_SHA256["t12"]
+
+
+@pytest.mark.parametrize("bits", range(2, 7))
+def test_primary_sqg_xor_cheb_t12_graph_and_labels_match_runtime_contract(
+    bits: int,
+) -> None:
+    ranks = sqg_xor_rank_permutation(bits)
+    labels = sqg_xor_cheb_t12_bytes(bits)
+    width = 16 - bits
+    strata = (ranks >> width).reshape(1 << width, 1 << bits)
+
+    assert torch.equal(torch.sort(ranks).values, torch.arange(1 << 16))
+    assert torch.equal(
+        torch.sort(strata, dim=1).values,
+        torch.arange(1 << bits).expand_as(strata),
+    )
+    assert torch.equal(labels, sqg_xor_cheb_t12_rank_lut_bytes()[ranks >> 4])
+    assert torch.equal(labels, sqg_codebook_bytes(bits, CODEBOOK_SQG_XOR_CHEB_T12))
+    assert bool(torch.isfinite(labels.view(torch.float8_e4m3fn).float()).all())
+    assert not bool(((labels & 0x7F) == 0).logical_and(labels != 0).any())
+    assert _sha256(labels) == _SQG_XOR_CHEB_T12_SHA256[bits]
+
+
+@pytest.mark.parametrize("bits", range(2, 7))
+def test_sqg_preprojection_mapping_is_a_permutation(bits: int) -> None:
+    ranks = sqg_rank_permutation(bits)
+    assert ranks.dtype == torch.int64
+    assert ranks.shape == (1 << 16,)
+    assert torch.equal(torch.sort(ranks).values, torch.arange(1 << 16))
+
+
+@pytest.mark.parametrize("bits", range(2, 7))
+def test_sqg_codebook_round_trips_exact_e4m3(bits: int) -> None:
+    raw = sqg_e4m3_bytes(bits, "normal")
+    codebook = sqg_e4m3_codebook(bits, "normal")
+    assert raw.dtype == torch.uint8
+    assert codebook.dtype == torch.float16
+    assert raw.shape == codebook.shape == (1 << 16,)
+    assert torch.equal(codebook.to(torch.float8_e4m3fn).view(torch.uint8), raw)
+    assert bool(torch.isfinite(codebook).all())
+
+
+@pytest.mark.parametrize("bits", range(2, 7))
+def test_sqg_each_state_spans_all_coarse_strata(bits: int) -> None:
+    ranks = sqg_rank_permutation(bits)
+    width = 16 - bits
+    strata = (ranks >> width).reshape(1 << width, 1 << bits)
+    expected = torch.arange(1 << bits).expand_as(strata)
+    assert torch.equal(torch.sort(strata, dim=1).values, expected)
+
+
+@pytest.mark.parametrize("bits", range(2, 7))
+def test_shared_rank_lut_preserves_one_graph_for_every_rate(bits: int) -> None:
+    rank_lut = torch.arange(1 << 16, dtype=torch.int64).remainder(126).to(torch.uint8)
+    raw = sqg_e4m3_bytes_from_rank_lut(bits, rank_lut)
+    expected = rank_lut.index_select(0, sqg_rank_permutation(bits))
+    assert torch.equal(raw, expected)
+    assert torch.equal(
+        sqg_e4m3_codebook_from_rank_lut(bits, rank_lut)
+        .to(torch.float8_e4m3fn)
+        .view(torch.uint8),
+        expected,
+    )
+
+
+def test_shared_rank_lut_rejects_nonfinite_e4m3() -> None:
+    rank_lut = torch.zeros(1 << 16, dtype=torch.uint8)
+    rank_lut[0] = 0x7F
+    with pytest.raises(ValueError, match="finite E4M3"):
+        sqg_e4m3_bytes_from_rank_lut(2, rank_lut)
+
+
+@pytest.mark.parametrize("bits", range(2, 7))
+def test_sqg_cheb_normal_named_codebook_uses_exact_shared_rank_law(bits: int) -> None:
+    expected = sqg_cheb_normal_rank_e4m3_bytes().index_select(
+        0, sqg_rank_permutation(bits)
+    )
+    assert torch.equal(sqg_cheb_normal_e4m3_bytes(bits), expected)
+    assert torch.equal(
+        sqg_codebook_bytes(bits, CODEBOOK_SQG_CHEB_NORMAL_E4M3), expected
+    )
+    assert not torch.equal(
+        expected, sqg_codebook_bytes(bits, CODEBOOK_SQG_NORMAL_E4M3)
+    )
+
+
+@pytest.mark.parametrize("bits", range(2, 7))
+@pytest.mark.parametrize("rate_axis", ("k", "n"))
+def test_w2_k2_q8h4_profile_changes_only_k_axis_k2(
+    bits: int, rate_axis: str
+) -> None:
+    actual = sqg_codebook_bytes(
+        bits,
+        CODEBOOK_SQG_CHEB,
+        rate_axis=rate_axis,
+    )
+    rank_lut = sqg_cheb_normal_rank_e4m3_bytes()
+    ranks = (
+        sqg_k2_eight_stratum_rank_permutation(4)
+        if bits == 2 and rate_axis == "k"
+        else sqg_rank_permutation(bits)
+    )
+    assert torch.equal(actual, rank_lut.index_select(0, ranks))
+    if bits == 2 and rate_axis == "k":
+        assert not torch.equal(
+            actual, sqg_codebook_bytes(bits, CODEBOOK_SQG_CHEB_NORMAL_E4M3)
+        )
+
+
+def test_w2_k2_q8h4_profile_requires_rate_axis() -> None:
+    with pytest.raises(ValueError, match="requires rate_axis"):
+        sqg_codebook_bytes(2, CODEBOOK_SQG_CHEB)
+
+
+def test_reference_decoder_accepts_custom_codebook() -> None:
+    states = torch.arange(256, dtype=torch.int16).reshape(1, 1, 256)
+    codebook = sqg_e4m3_codebook(3, "normal")
+    decoded = decode_regularized_weight(
+        states, codebook_values=codebook
+    )
+    assert decoded.shape == (16, 16)
+    assert torch.equal(
+        torch.sort(decoded.flatten()).values,
+        torch.sort(codebook[:256].float()).values,
+    )
+
+
+@pytest.mark.parametrize("rate_axis", ("k", "n"))
+def test_reference_decoder_applies_rate_specific_sqg_tables(rate_axis: str) -> None:
+    states = torch.arange(3 * 3 * 256, dtype=torch.int32).to(torch.int16).reshape(
+        3, 3, 256
+    )
+    tile_bits = (2, 3, 4)
+    decoded = decode_qsrt_regularized_weight(
+        states,
+        rate_axis=rate_axis,
+        tile_bits=tile_bits,
+        codebook=CODEBOOK_SQG_NORMAL_E4M3,
+    )
+    axis = 0 if rate_axis == "k" else 1
+    pieces = []
+    for tile, bits in enumerate(tile_bits):
+        selected = states.narrow(axis, tile, 1)
+        pieces.append(
+            decode_regularized_weight(
+                selected,
+                codebook=CODEBOOK_SQG_NORMAL_E4M3,
+                bits=bits,
+            )
+        )
+    expected = torch.cat(pieces, dim=0 if rate_axis == "k" else 1)
+    assert torch.equal(decoded, expected)
+
+
+def test_reference_decoder_applies_w2_only_k2_q8h4_profile() -> None:
+    states = torch.arange(3 * 3 * 256, dtype=torch.int32).to(torch.int16).reshape(
+        3, 3, 256
+    )
+    tile_bits = (2, 3, 4)
+    down = decode_qsrt_regularized_weight(
+        states,
+        rate_axis="k",
+        tile_bits=tile_bits,
+        codebook=CODEBOOK_SQG_CHEB,
+    )
+    upstream = decode_qsrt_regularized_weight(
+        states,
+        rate_axis="n",
+        tile_bits=tile_bits,
+        codebook=CODEBOOK_SQG_CHEB,
+    )
+    native_down = decode_qsrt_regularized_weight(
+        states,
+        rate_axis="k",
+        tile_bits=tile_bits,
+        codebook=CODEBOOK_SQG_CHEB_NORMAL_E4M3,
+    )
+    native_upstream = decode_qsrt_regularized_weight(
+        states,
+        rate_axis="n",
+        tile_bits=tile_bits,
+        codebook=CODEBOOK_SQG_CHEB_NORMAL_E4M3,
+    )
+    assert not torch.equal(down[:16], native_down[:16])
+    assert torch.equal(down[16:], native_down[16:])
+    assert torch.equal(upstream, native_upstream)

reproduction/closure/patches/kquant/changed-files/tests/test_sqg_quantizer.py

@@ -0,0 +1,262 @@
+from __future__ import annotations
+
+import hashlib
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+import kquant.sqg_quantizer as sqg_quantizer
+from kquant.exl3_reference import reconstruct_trellis_states
+from kquant.sqg_quantizer import install_sqg_quantizer
+from kquant.sqg_e4m3 import sqg_xor_cheb_t12_bytes
+
+
+def _sealed_extension(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+    *,
+    filename: str = "kquant_sqg_quantize_ext_v22.so",
+) -> tuple[Path, SimpleNamespace]:
+    path = tmp_path / filename
+    payload = b"sealed-sqg-extension-test-double"
+    path.write_bytes(payload)
+    module = SimpleNamespace(
+        __file__=str(path),
+        quantize_tiles_procedural=lambda *_args: None,
+        quantize_tiles_sqg=lambda *_args: None,
+    )
+    monkeypatch.setenv("KQUANT_SQG_REQUIRE_PREBUILT", "1")
+    monkeypatch.setenv("KQUANT_SQG_EXTENSION_PATH", str(path))
+    monkeypatch.setenv(
+        "KQUANT_SQG_EXTENSION_SHA256", hashlib.sha256(payload).hexdigest()
+    )
+    return path, module
+
+
+def test_required_prebuilt_extension_uses_exact_path_and_never_jit_builds(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    path, module = _sealed_extension(tmp_path, monkeypatch)
+    imported: list[Path] = []
+
+    def fake_import(actual: Path):
+        imported.append(actual)
+        return module
+
+    monkeypatch.setattr(sqg_quantizer, "_import_prebuilt_extension", fake_import)
+    monkeypatch.setattr(
+        sqg_quantizer,
+        "load",
+        lambda **_kwargs: pytest.fail("JIT builder must not run in prebuilt mode"),
+    )
+    sqg_quantizer._extension.cache_clear()
+    try:
+        assert sqg_quantizer._extension() is module
+        assert imported == [path.resolve()]
+    finally:
+        sqg_quantizer._extension.cache_clear()
+
+
+def test_required_prebuilt_extension_rejects_wrong_filename_before_import_or_jit(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    _sealed_extension(tmp_path, monkeypatch, filename="foreign_extension.so")
+    monkeypatch.setattr(
+        sqg_quantizer,
+        "_import_prebuilt_extension",
+        lambda _path: pytest.fail("foreign extension must not be imported"),
+    )
+    monkeypatch.setattr(
+        sqg_quantizer,
+        "load",
+        lambda **_kwargs: pytest.fail("JIT builder must not run in prebuilt mode"),
+    )
+    sqg_quantizer._extension.cache_clear()
+    try:
+        with pytest.raises(ValueError, match="filename must identify"):
+            sqg_quantizer._extension()
+    finally:
+        sqg_quantizer._extension.cache_clear()
+
+
+def test_required_prebuilt_extension_rejects_hash_mismatch_before_import_or_jit(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    _sealed_extension(tmp_path, monkeypatch)
+    monkeypatch.setenv("KQUANT_SQG_EXTENSION_SHA256", "0" * 64)
+    monkeypatch.setattr(
+        sqg_quantizer,
+        "_import_prebuilt_extension",
+        lambda _path: pytest.fail("mismatched extension must not be imported"),
+    )
+    monkeypatch.setattr(
+        sqg_quantizer,
+        "load",
+        lambda **_kwargs: pytest.fail("JIT builder must not run in prebuilt mode"),
+    )
+    sqg_quantizer._extension.cache_clear()
+    try:
+        with pytest.raises(RuntimeError, match="SHA256 mismatch"):
+            sqg_quantizer._extension()
+    finally:
+        sqg_quantizer._extension.cache_clear()
+
+
+@pytest.mark.parametrize(
+    ("missing", "message"),
+    (
+        ("KQUANT_SQG_EXTENSION_PATH", "are both required"),
+        ("KQUANT_SQG_EXTENSION_SHA256", "are both required"),
+    ),
+)
+def test_required_prebuilt_extension_rejects_partial_binding_without_jit(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+    missing: str,
+    message: str,
+) -> None:
+    _sealed_extension(tmp_path, monkeypatch)
+    monkeypatch.delenv(missing)
+    monkeypatch.setattr(
+        sqg_quantizer,
+        "load",
+        lambda **_kwargs: pytest.fail("JIT builder must not run in prebuilt mode"),
+    )
+    sqg_quantizer._extension.cache_clear()
+    try:
+        with pytest.raises(RuntimeError, match=message):
+            sqg_quantizer._extension()
+    finally:
+        sqg_quantizer._extension.cache_clear()
+
+
+def test_rate_specific_none_dispatches_to_kquant_mcg(monkeypatch) -> None:
+    calls: list[tuple[torch.Tensor, dict]] = []
+    procedural_calls: list[tuple] = []
+
+    def original(tiles: torch.Tensor, args: dict):
+        calls.append((tiles, args))
+        return "upstream"
+
+    def procedural(*args):
+        procedural_calls.append(args)
+
+    monkeypatch.setattr(
+        sqg_quantizer,
+        "_sqg_temp_buffers",
+        lambda _device, _bits: (torch.empty(0), torch.empty(0)),
+    )
+    monkeypatch.setattr(
+        sqg_quantizer,
+        "_extension",
+        lambda: SimpleNamespace(quantize_tiles_procedural=procedural),
+    )
+
+    module = SimpleNamespace(quantize_tiles=original)
+    install_sqg_quantizer(module)
+    tiles = torch.zeros((1, 256), dtype=torch.float32)
+    result = module.quantize_tiles(
+        tiles,
+        {
+            "K": 2,
+            "devices": ["cuda:0"],
+            "sqg_e4m3_luts_by_bits": {2: None, 3: None, 4: None},
+        },
+    )
+
+    assert all(isinstance(tensor, torch.Tensor) for tensor in result)
+    assert not calls
+    assert len(procedural_calls) == 1
+    assert procedural_calls[0][5:] == (2, 1, 128)
+
+
+def test_sqg_luts_are_transposed_by_predecessor_for_vector_loads(
+    monkeypatch,
+) -> None:
+    calls: list[tuple] = []
+
+    monkeypatch.setattr(
+        sqg_quantizer,
+        "_sqg_temp_buffers",
+        lambda _device, _bits: (torch.empty(0), torch.empty(0)),
+    )
+    monkeypatch.setattr(
+        sqg_quantizer,
+        "_extension",
+        lambda: SimpleNamespace(quantize_tiles_sqg=lambda *args: calls.append(args)),
+    )
+
+    module = SimpleNamespace(quantize_tiles=lambda *_args: "upstream")
+    install_sqg_quantizer(module)
+    tiles = torch.zeros((1, 256), dtype=torch.float32)
+    for bits in (2, 3, 4):
+        raw = torch.arange(65536, dtype=torch.int64).to(torch.uint8)
+        module.quantize_tiles(
+            tiles,
+            {
+                "K": bits,
+                "devices": ["cuda:0"],
+                "sqg_e4m3_lut": raw,
+            },
+        )
+        actual = calls[-1][5]
+        predecessors = 1 << bits
+        out_edge_pairs = (65536 >> bits) // 2
+        expected = (
+            raw.reshape(predecessors, out_edge_pairs, 2)
+            .permute(1, 0, 2)
+            .contiguous()
+            .reshape(-1)
+        )
+        assert torch.equal(actual, expected)
+        assert calls[-1][6:] == (bits, 128)
+
+
+@pytest.mark.parametrize("bits", [5, 6])
+def test_k5_k6_sqg_luts_remain_in_direct_state_order(
+    monkeypatch, bits: int
+) -> None:
+    calls: list[tuple] = []
+    monkeypatch.setattr(
+        sqg_quantizer,
+        "_sqg_temp_buffers",
+        lambda _device, _bits: (torch.empty(0), torch.empty(0)),
+    )
+    monkeypatch.setattr(
+        sqg_quantizer,
+        "_extension",
+        lambda: SimpleNamespace(quantize_tiles_sqg=lambda *args: calls.append(args)),
+    )
+    module = SimpleNamespace(quantize_tiles=lambda *_args: "upstream")
+    install_sqg_quantizer(module)
+    raw = torch.arange(65536, dtype=torch.int64).to(torch.uint8)
+    module.quantize_tiles(
+        torch.zeros((1, 256), dtype=torch.float32),
+        {"K": bits, "devices": ["cuda:0"], "sqg_e4m3_lut": raw},
+    )
+    assert torch.equal(calls[-1][5], raw)
+    assert calls[-1][6:] == (bits, 128)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is unavailable")
+@pytest.mark.parametrize("bits", range(2, 7))
+def test_sqg_quantizer_keeps_saturated_cost_paths_closed(bits: int) -> None:
+    module = SimpleNamespace(quantize_tiles=lambda *_args: "upstream")
+    install_sqg_quantizer(module)
+    tiles = torch.full((2, 256), 1000.0, dtype=torch.float32, device="cuda")
+
+    _output, states = module.quantize_tiles(
+        tiles,
+        {
+            "K": bits,
+            "devices": ["cuda:0"],
+            "sqg_e4m3_lut": sqg_xor_cheb_t12_bytes(bits),
+            "tailbite_context": 128,
+        },
+    )
+
+    edges = states.to(torch.int64) & ((1 << bits) - 1)
+    expected = reconstruct_trellis_states(edges, bits)
+    assert torch.equal(states, expected)

reproduction/closure/patches/kquant/changes.patch

@@ -0,0 +1,1386 @@
+diff --git a/kquant/candidate_hessian.py b/kquant/candidate_hessian.py
+index 36f5f30..6603799 100644
+--- a/kquant/candidate_hessian.py
++++ b/kquant/candidate_hessian.py
+@@ -92,8 +92,10 @@ def weighted_covariance(
+     accumulator = torch.zeros(
+         (dimension, dimension), dtype=torch.float32, device=device
+     )
+-    denominator = float(weights.double().sum())
+-    if not math.isfinite(denominator) or denominator <= 0:
++    denominator_tensor = weights.to(device=device, dtype=torch.float64).sum()
++    if not bool(torch.isfinite(denominator_tensor)) or not bool(
++        denominator_tensor > 0
++    ):
+         raise ValueError("covariance weights must have a positive finite sum")
+     for begin in range(0, int(rows.shape[0]), chunk_rows):
+         values = rows[begin : begin + chunk_rows].to(
+@@ -103,8 +105,9 @@ def weighted_covariance(
+             device=device, dtype=torch.float32
+         )
+         accumulator.addmm_(values.T, values * local_weights[:, None])
+-    covariance = accumulator.div_(denominator)
++    covariance = accumulator.div_(denominator_tensor.to(torch.float32))
+     covariance = ((covariance + covariance.T) * 0.5).contiguous()
++    denominator = float(denominator_tensor)
+     if return_cpu:
+         covariance = covariance.cpu()
+     if not torch.all(torch.isfinite(covariance)):
+@@ -115,13 +118,23 @@ def weighted_covariance(
+ def shrink_covariance(
+     global_covariance: torch.Tensor,
+     local_covariance: torch.Tensor,
+-    alpha: float,
++    alpha: float | torch.Tensor,
+ ) -> torch.Tensor:
+     """Convexly shrink a low-support local covariance toward a global one."""
+ 
+-    alpha = float(alpha)
+-    if not math.isfinite(alpha) or not 0.0 <= alpha <= 1.0:
+-        raise ValueError("alpha must be finite and in [0, 1]")
++    if isinstance(alpha, torch.Tensor):
++        if alpha.numel() != 1 or not bool(torch.isfinite(alpha)) or not bool(
++            (alpha >= 0) & (alpha <= 1)
++        ):
++            raise ValueError("alpha must be a finite scalar in [0, 1]")
++        blend_weight: float | torch.Tensor = alpha.to(
++            device=global_covariance.device, dtype=torch.float32
++        )
++    else:
++        alpha_value = float(alpha)
++        if not math.isfinite(alpha_value) or not 0.0 <= alpha_value <= 1.0:
++            raise ValueError("alpha must be finite and in [0, 1]")
++        blend_weight = alpha_value
+     if (
+         global_covariance.ndim != 2
+         or global_covariance.shape[0] != global_covariance.shape[1]
+@@ -133,7 +146,7 @@ def shrink_covariance(
+     ):
+         raise ValueError("covariances must be finite")
+     blended = torch.lerp(
+-        global_covariance.float(), local_covariance.float(), alpha
++        global_covariance.float(), local_covariance.float(), blend_weight
+     )
+     return ((blended + blended.T) * 0.5).contiguous()
+ 
+@@ -141,6 +154,17 @@ def shrink_covariance(
+ def weighted_effective_sample_size(weights: torch.Tensor) -> float:
+     """Return Kish effective sample size for nonnegative row weights."""
+ 
++    if weights.ndim != 1 or not weights.numel():
++        raise ValueError("weights must be a nonempty vector")
++    if not torch.all(torch.isfinite(weights)) or torch.any(weights < 0):
++        raise ValueError("weights must be finite and nonnegative")
++    value = _weighted_effective_sample_size_tensor(weights)
++    return float(value)
++
++
++def _weighted_effective_sample_size_tensor(weights: torch.Tensor) -> torch.Tensor:
++    """Kish support as a scalar tensor on the input numerical device."""
++
+     if weights.ndim != 1 or not weights.numel():
+         raise ValueError("weights must be a nonempty vector")
+     if not torch.all(torch.isfinite(weights)) or torch.any(weights < 0):
+@@ -150,7 +174,7 @@ def weighted_effective_sample_size(weights: torch.Tensor) -> float:
+     squared_total = weights64.square().sum()
+     if total <= 0 or squared_total <= 0:
+         raise ValueError("weights must have positive effective support")
+-    return float(total.square() / squared_total)
++    return total.square() / squared_total
+ 
+ 
+ def adaptive_identity_shrinkage(
+@@ -180,7 +204,7 @@ def adaptive_identity_shrinkage(
+     if not torch.all(torch.isfinite(local_covariance)):
+         raise ValueError("local covariance must be finite")
+ 
+-    effective_rows = weighted_effective_sample_size(weights)
++    effective_rows_tensor = _weighted_effective_sample_size_tensor(weights)
+     dimension = int(local_covariance.shape[0])
+     local64 = local_covariance.double()
+     local_trace = torch.trace(local64)
+@@ -189,24 +213,34 @@ def adaptive_identity_shrinkage(
+ 
+     mu = local_trace / dimension
+     mean_square = local64.square().mean()
+-    denominator = (effective_rows + 1.0) * (
++    denominator = (effective_rows_tensor + 1.0) * (
+         mean_square - mu.square() / dimension
+     )
+-    if denominator <= 0:
+-        oas_shrinkage = 1.0
+-    else:
+-        oas_shrinkage = min(
+-            1.0,
+-            max(0.0, float((mean_square + mu.square()) / denominator)),
+-        )
+-    local_alpha = min(max_local_alpha, 1.0 - oas_shrinkage)
+-    identity_scale = float(mu)
++    raw_oas = (mean_square + mu.square()) / denominator.clamp_min(
++        torch.finfo(local64.dtype).tiny
++    )
++    oas_shrinkage_tensor = torch.where(
++        denominator <= 0,
++        torch.ones_like(raw_oas),
++        raw_oas.clamp(0.0, 1.0),
++    )
++    local_alpha_tensor = torch.minimum(
++        torch.as_tensor(max_local_alpha, device=local64.device, dtype=local64.dtype),
++        1.0 - oas_shrinkage_tensor,
++    )
+     scaled_identity = torch.eye(
+         dimension,
+         dtype=torch.float32,
+         device=local_covariance.device,
+-    ).mul_(identity_scale)
+-    blended = shrink_covariance(scaled_identity, local_covariance, local_alpha)
++    ).mul_(mu.to(dtype=torch.float32))
++    blended = shrink_covariance(
++        scaled_identity, local_covariance, local_alpha_tensor.to(torch.float32)
++    )
++    # Evidence scalars are detached only after the numerical result is frozen.
++    effective_rows = float(effective_rows_tensor)
++    oas_shrinkage = float(oas_shrinkage_tensor)
++    local_alpha = float(local_alpha_tensor)
++    identity_scale = float(mu)
+     return blended, {
+         "effective_sample_size": effective_rows,
+         "oas_shrinkage": oas_shrinkage,
+diff --git a/kquant/csrc/qsrt_quantize_tiles_kernel.cuh b/kquant/csrc/qsrt_quantize_tiles_kernel.cuh
+index 76febdd..f4299d1 100644
+--- a/kquant/csrc/qsrt_quantize_tiles_kernel.cuh
++++ b/kquant/csrc/qsrt_quantize_tiles_kernel.cuh
+@@ -17,6 +17,9 @@
+ #ifndef H_INF
+ #define H_INF __ushort_as_half(0x7c00)
+ #endif
++#ifndef H_MAX_FINITE
++#define H_MAX_FINITE __ushort_as_half(0x7bff)
++#endif
+ 
+ template <int K, int cb>
+ __global__ __launch_bounds__(
+@@ -121,7 +124,12 @@ void quantize_tiles_kernel
+                 decoded2 = decode_3inst_2<cb>(out_edge_idx, out_edge_idx + 1);
+             }
+             half2 dh2 = __hsub2(decoded2, w2);
+-            half2 min_err2 = __hmul2(dh2, dh2);
++            // Positive infinity denotes an unreachable state. Reachable path
++            // costs must remain finite when their FP16 sum overflows, or
++            // traceback can confuse a saturated path with an unreachable one
++            // and violate the constrained cyclic boundary.
++            half2 min_err2 = __hmin2(
++                __hmul2(dh2, dh2), __half2half2(H_MAX_FINITE));
+             if (pre_state >= 0 && in_edge_idx != pre_state) min_err2 = __half2half2(H_INF);
+             int min_k0 = 0;
+             int min_k1 = 0;
+@@ -165,7 +173,8 @@ void quantize_tiles_kernel
+                     decoded2 = decode_3inst_2<cb>(state0, state0 + 1);
+                 }
+                 dh2 = __hsub2(decoded2, w2);
+-                half2 err2 = __hmul2(dh2, dh2);
++                half2 err2 = __hmin2(
++                    __hmul2(dh2, dh2), __half2half2(H_MAX_FINITE));
+                 if (pre_state >= 0 && in_edge_idx != pre_state) err2 = __half2half2(H_INF);
+                 const unsigned less = __hlt2_mask(err2, min_err2);
+                 min_err2 = __hmin2(err2, min_err2);
+@@ -253,7 +262,13 @@ void quantize_tiles_kernel
+                     decoded2 = decode_3inst_2<cb>(out_edge_idx, out_edge_idx + 1);
+                 }
+                 half2 dh2 = __hsub2(decoded2, w2);
+-                half2 min_err2 = __hfma2(dh2, dh2, __half2half2(temp_costs_inc[in_edge_idx]));
++                half predecessor_cost = temp_costs_inc[in_edge_idx];
++                half2 min_err2 = __heq(predecessor_cost, H_INF)
++                    ? __half2half2(H_INF)
++                    : __hmin2(
++                        __hfma2(
++                            dh2, dh2, __half2half2(predecessor_cost)),
++                        __half2half2(H_MAX_FINITE));
+                 int min_k0 = 0;
+                 int min_k1 = 0;
+ 
+@@ -296,7 +311,13 @@ void quantize_tiles_kernel
+                         decoded2 = decode_3inst_2<cb>(state0, state0 + 1);
+                     }
+                     dh2 = __hsub2(decoded2, w2);
+-                    half2 err2 = __hfma2(dh2, dh2, __half2half2(temp_costs_inc[in_edge_idx]));
++                    predecessor_cost = temp_costs_inc[in_edge_idx];
++                    half2 err2 = __heq(predecessor_cost, H_INF)
++                        ? __half2half2(H_INF)
++                        : __hmin2(
++                            __hfma2(
++                                dh2, dh2, __half2half2(predecessor_cost)),
++                            __half2half2(H_MAX_FINITE));
+                     const unsigned less = __hlt2_mask(err2, min_err2);
+                     min_err2 = __hmin2(err2, min_err2);
+                     if (less & 0xffffu) min_k0 = k;
+diff --git a/kquant/csrc/sqg_quantize.cu b/kquant/csrc/sqg_quantize.cu
+index 42dcd88..e51210e 100644
+--- a/kquant/csrc/sqg_quantize.cu
++++ b/kquant/csrc/sqg_quantize.cu
+@@ -187,7 +187,7 @@ void quantize_tiles_sqg_cuda(
+                 "output_indices shape mismatch");
+     TORCH_CHECK(output_indices.scalar_type() == at::kShort,
+                 "output_indices must be int16");
+-    TORCH_CHECK(K >= 2 && K <= 4, "SQG validation kernel supports K2--K4");
++    TORCH_CHECK(K >= 2 && K <= 6, "SQG validation kernel supports K2--K6");
+     TORCH_CHECK(tailbite_context >= 1 && tailbite_context <= 128,
+                 "tailbite_context must be in 1..128");
+     TORCH_CHECK(codebook.is_cuda() && codebook.device() == input_tiles.device(),
+@@ -255,6 +255,22 @@ void quantize_tiles_sqg_cuda(
+                 temp_edges.data_ptr<uint8_t>(), tiles, max_batch, shared_memory,
+                 static_cast<int>(tailbite_context), stream);
+             break;
++        case 5:
++            launch_trellis<5, 4>(
++                input_tiles.data_ptr<float>(), output_tiles.data_ptr<float>(),
++                reinterpret_cast<uint16_t*>(output_indices.data_ptr<int16_t>()),
++                reinterpret_cast<half*>(temp_costs.data_ptr<at::Half>()),
++                temp_edges.data_ptr<uint8_t>(), tiles, max_batch, shared_memory,
++                static_cast<int>(tailbite_context), stream);
++            break;
++        case 6:
++            launch_trellis<6, 4>(
++                input_tiles.data_ptr<float>(), output_tiles.data_ptr<float>(),
++                reinterpret_cast<uint16_t*>(output_indices.data_ptr<int16_t>()),
++                reinterpret_cast<half*>(temp_costs.data_ptr<at::Half>()),
++                temp_edges.data_ptr<uint8_t>(), tiles, max_batch, shared_memory,
++                static_cast<int>(tailbite_context), stream);
++            break;
+     }
+ }
+ 
+diff --git a/kquant/exl3_encoder_backend.py b/kquant/exl3_encoder_backend.py
+index f810bf3..616b9e6 100644
+--- a/kquant/exl3_encoder_backend.py
++++ b/kquant/exl3_encoder_backend.py
+@@ -21,6 +21,18 @@ import threading
+ had_k, had_n = 128, 128
+ codebook_scale = 1.24371088
+ 
++# The original search used a 0.1..1.9 coarse bracket and could return 2.05
++# only because its five-point fine grid extended past the bracket.  GLM-5.2
++# contains healthy tensors with a stable minimum around 3.0, so keep the cheap
++# original bracket as stage one and expand only when its upper point wins.
++G_SCALE_SEARCH_INITIAL_MIN = 0.1
++G_SCALE_SEARCH_INITIAL_MAX = 1.9
++G_SCALE_SEARCH_COARSE_STEP = 0.2
++G_SCALE_SEARCH_EXPANDED_MAX = 3.9
++G_SCALE_SEARCH_FINE_STEP = 0.075
++G_SCALE_SEARCH_HARD_CEILING = 4.05
++G_SCALE_SEARCH_CEILING_TOLERANCE = 1e-6
++
+ codebook_mcg_mult = 0xCBAC1FED
+ codebook_mul1_mult = 0x83DCD12D
+ 
+@@ -319,6 +331,11 @@ def block_ldl(H: torch.Tensor, b: int, quant_args: dict, verbose: bool):
+     assert (n % b == 0)
+     m = n // b
+ 
++    keep_hessian_on_device = bool(quant_args.get("keep_hessian_on_device", False))
++    gpu_only_numerical_path = bool(quant_args.get("gpu_only_numerical_path", False))
++    if gpu_only_numerical_path and H.device.type != "cuda":
++        raise RuntimeError("GPU-only BlockLDLQ received a non-CUDA Hessian")
++
+     # Cholesky factorization: H = L @ L.T
+     # Try on GPU first
+     num_cholesky_retries = 0
+@@ -326,12 +343,14 @@ def block_ldl(H: torch.Tensor, b: int, quant_args: dict, verbose: bool):
+     while True:
+         try:
+             L = torch.linalg.cholesky(H)
+-            # H is not needed after this, move to CPU. Then overwrite H's GPU storage with L, since we can't otherwise
+-            # free up that VRAM as the tensor is referenced by the parent frame
+-            H_cpu = H.cpu()
+-            H.copy_(L)  # VRAM copy, tiny overhead
+-            L = H
+-            H = H_cpu
++            if not keep_hessian_on_device:
++                # Historical memory-saving path. Production retains both the
++                # caller-owned H and its factor on CUDA so neither numerical
++                # tensor crosses the host boundary.
++                H_cpu = H.cpu()
++                H.copy_(L)  # VRAM copy, tiny overhead
++                L = H
++                H = H_cpu
+             break
+ 
+         except torch._C._LinAlgError as e:
+@@ -346,6 +365,10 @@ def block_ldl(H: torch.Tensor, b: int, quant_args: dict, verbose: bool):
+         # Fall back on CPU factorization
+         except Exception as e:
+             if e.__class__.__name__ == "OutOfMemoryError" or "CUDA out of memory" in str(e) or "HIP out of memory" in str(e):
++                if gpu_only_numerical_path:
++                    raise RuntimeError(
++                        "GPU-only BlockLDLQ exhausted CUDA memory; CPU fallback is forbidden"
++                    ) from e
+                 retry_cpu = True
+                 break
+             else:
+@@ -380,7 +403,7 @@ def block_ldl(H: torch.Tensor, b: int, quant_args: dict, verbose: bool):
+     # Insert block identity matrices along the diagonal.
+     # TODO: Figure out if this is necessary. Diagonal blocks should already be identities after previous step
+     L_block = L.view(m, b, m, b).permute(0, 2, 1,3)
+-    dr = torch.arange(m)
++    dr = torch.arange(m, device=L.device)
+     L_block[dr, dr] = torch.stack([torch.eye(b, device = L.device, dtype = H.dtype)] * m)
+ 
+     return L, H  # , D.to(DL.device)
+@@ -1585,6 +1608,43 @@ def ldlq_batched(
+ 
+ finalize_capture_H_mutex = threading.Lock()
+ 
++
++def _explicit_vector(
++    value: torch.Tensor,
++    *,
++    length: int,
++    name: str,
++    device: torch.device,
++    orientation: str,
++    signs: bool = False,
++) -> torch.Tensor:
++    """Validate a caller-owned regularization vector without retaining it.
++
++    QSRT's original shared-input-scale hook stores the first tensor's profile
++    in module-global state.  Model ports that require a topology-neutral
++    layer-wide vector need a deterministic, caller-owned contract instead.
++    These optional vectors are cloned for each encode and never enter the
++    process-global cache.
++    """
++
++    if not isinstance(value, torch.Tensor):
++        raise TypeError(f"{name} must be a torch.Tensor")
++    if value.ndim != 1 or value.numel() != length:
++        raise ValueError(f"{name} must have shape ({length},)")
++    vector = value.detach().to(device=device, dtype=torch.float32, copy=True)
++    if not bool(torch.isfinite(vector).all()):
++        raise ValueError(f"{name} must contain only finite values")
++    if signs:
++        if not bool((vector.abs() == 1).all()):
++            raise ValueError(f"{name} must contain only -1 and +1")
++    elif not bool((vector > 0).all()):
++        raise ValueError(f"{name} must contain only positive values")
++    if orientation == "column":
++        return vector.unsqueeze(1).contiguous()
++    if orientation == "row":
++        return vector.unsqueeze(0).contiguous()
++    raise ValueError(f"unsupported explicit-vector orientation: {orientation}")
++
+ def finalize_capture_H(H_data: dict, quant_args: dict, verbose: bool):
+     with finalize_capture_H_mutex:
+ 
+@@ -1596,7 +1656,23 @@ def finalize_capture_H(H_data: dict, quant_args: dict, verbose: bool):
+ 
+             H = H_data["H"]
+             k = H.shape[0]
+-            su = (torch.randn(k, device = H_data["device"]).sign() + 1e-5).sign().to(torch.float).unsqueeze(1)
++            explicit_input_signs = quant_args.get("input_signs")
++            if explicit_input_signs is None:
++                su = (
++                    (torch.randn(k, device=H_data["device"]).sign() + 1e-5)
++                    .sign()
++                    .to(torch.float)
++                    .unsqueeze(1)
++                )
++            else:
++                su = _explicit_vector(
++                    explicit_input_signs,
++                    length=k,
++                    name="input_signs",
++                    device=H_data["device"],
++                    orientation="column",
++                    signs=True,
++                )
+             H_data["su"] = su
+ 
+             return True, None, None, su, None
+@@ -1608,7 +1684,28 @@ def finalize_capture_H(H_data: dict, quant_args: dict, verbose: bool):
+ 
+         H = H_data["H"]
+         if H_data["finalized"]:
+-            return H_data["q_fallback"], H, H_data["L"], H_data["su"], H_data["diag"]
++            explicit_input_signs = quant_args.get("input_signs")
++            if explicit_input_signs is not None:
++                expected = _explicit_vector(
++                    explicit_input_signs,
++                    length=H.shape[0],
++                    name="input_signs",
++                    device=H_data["su"].device,
++                    orientation="column",
++                    signs=True,
++                )
++                if not torch.equal(H_data["su"], expected):
++                    raise ValueError(
++                        "finalized Hessian input signs do not match the "
++                        "caller-supplied profile"
++                    )
++            return (
++                H_data["q_fallback"],
++                H,
++                H_data["L"],
++                H_data["su"],
++                H_data["diag"],
++            )
+ 
+         # Mean of samples summed up during forward pass
+         # Switch to uncalibrated fallback if no input activations or diagonal is too small (few activations)
+@@ -1634,7 +1731,23 @@ def finalize_capture_H(H_data: dict, quant_args: dict, verbose: bool):
+ 
+         # Random sign flips for input channel, fixed for the first linear layer to quantize with this H
+         k = H.shape[0]
+-        su = (torch.randn(k, device = H.device).sign() + 1e-5).sign().to(torch.float).unsqueeze(1)
++        explicit_input_signs = quant_args.get("input_signs")
++        if explicit_input_signs is None:
++            su = (
++                (torch.randn(k, device=H.device).sign() + 1e-5)
++                .sign()
++                .to(torch.float)
++                .unsqueeze(1)
++            )
++        else:
++            su = _explicit_vector(
++                explicit_input_signs,
++                length=k,
++                name="input_signs",
++                device=H.device,
++                orientation="column",
++                signs=True,
++            )
+         H_data["su"] = su
+ 
+         # Input had
+@@ -1648,14 +1761,21 @@ def finalize_capture_H(H_data: dict, quant_args: dict, verbose: bool):
+             L = None
+         else:
+             L, H = block_ldl(H, 16, quant_args, verbose)
+-            dr = torch.arange(k)
++            dr = torch.arange(k, device=L.device)
+             L[dr, dr] = 0
+ 
+         H_data["L"] = L
+ 
+-        # H is no longer needed except to compute proxy error so move to CPU
+-        H = H.cpu()
+-        H_data["H"] = H.cpu()
++        # Paid B300 production keeps caller-owned H/L on the quant device.
++        # CPU storage remains the default for historical/offline callers.
++        keep_hessian_on_device = bool(
++            quant_args.get("keep_hessian_on_device", False)
++        )
++        if keep_hessian_on_device:
++            H_data["H"] = H
++        else:
++            H = H.cpu()
++            H_data["H"] = H.cpu()
+ 
+         H_data["finalized"] = True
+         H_data["diag"] = diag
+@@ -1752,23 +1872,63 @@ def g_scale_search_batch(
+             i = j
+         return out
+ 
+-    # Stage 1: coarse grid over the full search range, on a subsample of each tensor's tiles
+-    coarse = [0.1 + 0.2 * i for i in range(10)]  # 0.1 .. 1.9, same range as the old golden-section bracket
++    # Stage 1: retain the historical cheap bracket on a tile subsample.  Only
++    # tensors whose minimum lands at the upper point pay for the expansion.
++    coarse = [
++        G_SCALE_SEARCH_INITIAL_MIN + G_SCALE_SEARCH_COARSE_STEP * i
++        for i in range(
++            round(
++                (G_SCALE_SEARCH_INITIAL_MAX - G_SCALE_SEARCH_INITIAL_MIN)
++                / G_SCALE_SEARCH_COARSE_STEP
++            )
++            + 1
++        )
++    ]
+     subs = [s[::3] for s in samples]
+     pairs1 = [(t, s) for t in range(n_t) for s in coarse]
+     mse1 = eval_pairs(pairs1, subs).view(n_t, len(coarse))
+-    centers = [coarse[c] for c in mse1.argmin(dim = 1).tolist()]
++    coarse_best = mse1.argmin(dim = 1).tolist()
++    centers = [coarse[c] for c in coarse_best]
++    expanded = [index == len(coarse) - 1 for index in coarse_best]
++
++    expansion = [
++        G_SCALE_SEARCH_INITIAL_MAX + G_SCALE_SEARCH_COARSE_STEP * (i + 1)
++        for i in range(
++            round(
++                (G_SCALE_SEARCH_EXPANDED_MAX - G_SCALE_SEARCH_INITIAL_MAX)
++                / G_SCALE_SEARCH_COARSE_STEP
++            )
++        )
++    ]
++    expanded_tensors = [t for t, needed in enumerate(expanded) if needed]
++    if expanded_tensors:
++        pairs_x = [(t, scale) for t in expanded_tensors for scale in expansion]
++        mse_x = eval_pairs(pairs_x, subs).view(len(expanded_tensors), len(expansion))
++        for row, tensor_index in enumerate(expanded_tensors):
++            values = torch.cat((mse1[tensor_index], mse_x[row]))
++            combined = coarse + expansion
++            centers[tensor_index] = combined[int(values.argmin().item())]
+ 
+     # Stage 2: fine grid around each tensor's coarse minimum, on its full sample
+-    step = 0.075
++    step = G_SCALE_SEARCH_FINE_STEP
+     fine = [[c + step * (i - 2) for i in range(5)] for c in centers]
+     pairs2 = [(t, s) for t in range(n_t) for s in fine[t]]
+     mse2 = eval_pairs(pairs2, samples).view(n_t, 5)
+     mse2_h = mse2.tolist()
+ 
+     results = []
++    evidence = []
+     for t in range(n_t):
+         best = min(range(5), key = lambda i: mse2_h[t][i])
++        if (
++            fine[t][best]
++            >= G_SCALE_SEARCH_HARD_CEILING - G_SCALE_SEARCH_CEILING_TOLERANCE
++        ):
++            raise RuntimeError(
++                "global scale search saturated at its hard upper ceiling: "
++                f"scale={fine[t][best]:.8g}, "
++                f"ceiling={G_SCALE_SEARCH_HARD_CEILING:.8g}"
++            )
+         if 0 < best < 4:
+             y0, y1, y2 = mse2_h[t][best - 1], mse2_h[t][best], mse2_h[t][best + 1]
+             denom = y0 - 2.0 * y1 + y2
+@@ -1777,7 +1937,31 @@ def g_scale_search_batch(
+         else:
+             offset = 0.0
+         best_scale = max(fine[t][best] + offset * step, 0.01)
++        if (
++            best_scale
++            >= G_SCALE_SEARCH_HARD_CEILING - G_SCALE_SEARCH_CEILING_TOLERANCE
++        ):
++            raise RuntimeError(
++                "global scale interpolation reached its hard upper ceiling: "
++                f"scale={best_scale:.8g}, "
++                f"ceiling={G_SCALE_SEARCH_HARD_CEILING:.8g}"
++            )
+         results.append((best_scale, mse2[t, best]))
++        evidence.append(
++            {
++                "policy": "raw_sample_adaptive_coarse_then_fine_v2",
++                "initial_min": G_SCALE_SEARCH_INITIAL_MIN,
++                "initial_max": G_SCALE_SEARCH_INITIAL_MAX,
++                "expanded": expanded[t],
++                "expanded_max": G_SCALE_SEARCH_EXPANDED_MAX,
++                "fine_min": fine[t][0],
++                "fine_max": fine[t][-1],
++                "hard_ceiling": G_SCALE_SEARCH_HARD_CEILING,
++                "selected_scale": best_scale,
++                "selected_is_below_hard_ceiling": True,
++            }
++        )
++    quant_args["g_scale_search_evidence"] = evidence
+     return results
+ 
+ 
+@@ -1867,6 +2051,16 @@ _SHARED_INPUT_SCALES: dict = {}
+ def output_signs(n: int, device: torch.device, quant_args: dict) -> torch.Tensor:
+     """Draw output signs from an optional transform-specific RNG stream."""
+ 
++    explicit_output_signs = quant_args.get("output_signs")
++    if explicit_output_signs is not None:
++        return _explicit_vector(
++            explicit_output_signs,
++            length=n,
++            name="output_signs",
++            device=device,
++            orientation="row",
++            signs=True,
++        )
+     if "sv_seed" in quant_args:
+         torch.manual_seed(quant_args["sv_seed"])
+     return (
+@@ -1920,35 +2114,81 @@ def regularize(
+             print(f"     - input state skew: {skew_factor.item():.6f}")
+ 
+         if force_out_scales is None:
+-            apply_out_scales = skew_factor.item() < 0.15
++            apply_out_scales_mask = skew_factor < 0.15
+         else:
+-            apply_out_scales = force_out_scales
++            apply_out_scales_mask = torch.tensor(
++                bool(force_out_scales), dtype=torch.bool, device=weight.device
++            )
+ 
+     else:
+-        apply_out_scales = True if force_out_scales is None else force_out_scales
++        apply_out_scales_mask = torch.tensor(
++            True if force_out_scales is None else bool(force_out_scales),
++            dtype=torch.bool,
++            device=weight.device,
++        )
+ 
+     if q_fallback:
+-        apply_out_scales = force_out_scales
+-
+-    # Apply output scales
+-    out_channel_scales = block_rms(weight, dim = 0, keepdim = True)
+-    mean = out_channel_scales.mean().item()
+-    if mean > 1e-30:
+-        out_channel_scales /= mean
+-        quant_args["zeros"] = False
++        # Preserve the historical truthiness of a missing fallback override.
++        apply_out_scales_mask = torch.tensor(
++            False if force_out_scales is None else bool(force_out_scales),
++            dtype=torch.bool,
++            device=weight.device,
++        )
++
++    # Apply output scales.  An explicit profile is caller-owned and bypasses
++    # first-tensor derivation, which is required for topology-neutral ports.
++    explicit_output_scales = quant_args.get("output_channel_scale_profile")
++    observed_out_channel_scales = block_rms(weight, dim = 0, keepdim = True)
++    if explicit_output_scales is not None:
++        if force_out_scales is False:
++            raise ValueError(
++                "output_channel_scale_profile requires apply_out_scales=True"
++            )
++        apply_out_scales_mask = torch.ones(
++            (), dtype=torch.bool, device=weight.device
++        )
++        out_channel_scales = _explicit_vector(
++            explicit_output_scales,
++            length=weight.shape[1],
++            name="output_channel_scale_profile",
++            device=weight.device,
++            orientation="row",
++        )
+     else:
+-        quant_args["zeros"] = True
+-        if force_out_scales is not None:
+-            apply_out_scales = True
+-    zero_out_scales = out_channel_scales.abs() < 1e-30
+-
+-    if apply_out_scales:
+-        out_channel_scales[zero_out_scales] = 0.1
+-        sv = (sv * out_channel_scales + 1e-10).float()
+-        if verbose:
+-            out_channel_std = out_channel_scales.std().item()
+-            out_channel_mean = out_channel_scales.mean().item()
+-            print(f"     - out ch scales std/mean: {out_channel_std:.6f}   {out_channel_mean:.6f}")
++        out_channel_scales = observed_out_channel_scales
++    mean = out_channel_scales.mean()
++    mean_positive = mean > 1e-30
++    if explicit_output_scales is None:
++        # Keep the data-dependent normalization decision on CUDA.  The
++        # clamped denominator is ignored by torch.where for an all-zero
++        # profile and therefore cannot alter that historical branch.
++        out_channel_scales = torch.where(
++            mean_positive,
++            out_channel_scales / mean.clamp_min(1e-30),
++            out_channel_scales,
++        )
++    if force_out_scales is not None:
++        # Historical behavior forces scaling for an all-zero tensor even when
++        # an explicit boolean override was supplied.
++        apply_out_scales_mask = torch.where(
++            mean_positive,
++            apply_out_scales_mask,
++            torch.ones_like(apply_out_scales_mask),
++        )
++    zero_out_scales = observed_out_channel_scales.abs() < 1e-30
++    if explicit_output_scales is not None and bool(zero_out_scales.any()):
++        raise ValueError(
++            "an explicit output scale profile cannot encode a zero output channel"
++        )
++
++    candidate_out_channel_scales = out_channel_scales.clone()
++    candidate_out_channel_scales[zero_out_scales] = 0.1
++    scaled_sv = (sv * candidate_out_channel_scales + 1e-10).float()
++    sv = torch.where(apply_out_scales_mask, scaled_sv, sv)
++    if verbose:
++        out_channel_std = candidate_out_channel_scales.std().item()
++        out_channel_mean = candidate_out_channel_scales.mean().item()
++        print(f"     - out ch scales std/mean: {out_channel_std:.6f}   {out_channel_mean:.6f}")
+ 
+     # Output sign flips (and scales)
+     weight /= sv
+@@ -1960,13 +2200,28 @@ def regularize(
+     blockwise_preapply_had_r_(weight, had_n)
+ 
+     # Input sign flips and scales
+-    in_channel_scales = block_rms(weight, dim = 1, keepdim = True)
+-    in_channel_scales[in_channel_scales.abs() < 1e-30] = 0.1
++    explicit_input_scales = quant_args.get("input_channel_scale_profile")
++    if explicit_input_scales is None:
++        in_channel_scales = block_rms(weight, dim = 1, keepdim = True)
++        in_channel_scales[in_channel_scales.abs() < 1e-30] = 0.1
++    else:
++        in_channel_scales = _explicit_vector(
++            explicit_input_scales,
++            length=weight.shape[0],
++            name="input_channel_scale_profile",
++            device=weight.device,
++            orientation="column",
++        )
+     # kquant shared-su mode: reuse the first tensor's channel-scale profile for
+     # every tensor with the same key (e.g. all experts of a MoE layer/matrix),
+     # so the stored su vector is identical across them and TP replicas can
+     # broadcast one copy. Trellis codes re-fit against the shared profile.
+     shared_key = quant_args.get("shared_input_scales_key")
++    if explicit_input_scales is not None and shared_key is not None:
++        raise ValueError(
++            "input_channel_scale_profile cannot be combined with the "
++            "module-global shared_input_scales_key"
++        )
+     if shared_key is not None:
+         in_channel_scales = _SHARED_INPUT_SCALES.setdefault(
+             shared_key, in_channel_scales)
+@@ -1976,7 +2231,49 @@ def regularize(
+ 
+     # Determine best scale for matrix by test quantizing a sample of tiles along a wrapped diagonal
+     if not skip_g_scale:
+-        g_scale, mse_scale = g_scale_gss(weight, False, quant_args, pb = pb)
++        override = quant_args.get("g_scale_override")
++        if override is None:
++            g_scale, mse_scale = g_scale_gss(weight, False, quant_args, pb = pb)
++        else:
++            if (
++                isinstance(override, bool)
++                or not isinstance(override, (int, float))
++                or not math.isfinite(float(override))
++                or float(override) <= 0.0
++                or float(override)
++                >= G_SCALE_SEARCH_HARD_CEILING - G_SCALE_SEARCH_CEILING_TOLERANCE
++            ):
++                raise ValueError(
++                    "g_scale_override must be positive, finite, and below the "
++                    "hard scale ceiling"
++                )
++            policy = quant_args.get("g_scale_override_policy")
++            evidence_id = quant_args.get("g_scale_override_evidence_id")
++            if policy != "activation_parametric_coupled_v1":
++                raise ValueError(
++                    "serial g_scale_override requires the activation-parametric "
++                    "coupled policy"
++                )
++            if (
++                not isinstance(evidence_id, str)
++                or len(evidence_id) != 64
++                or any(ch not in "0123456789abcdef" for ch in evidence_id)
++            ):
++                raise ValueError(
++                    "serial g_scale_override requires a lowercase SHA-256 evidence ID"
++                )
++            g_scale = float(override)
++            mse_scale = torch.tensor(float("nan"), device=weight.device)
++            quant_args["g_scale_search_evidence"] = [
++                {
++                    "policy": policy,
++                    "evidence_id": evidence_id,
++                    "raw_sample_search_performed": False,
++                    "hard_ceiling": G_SCALE_SEARCH_HARD_CEILING,
++                    "selected_scale": g_scale,
++                    "selected_is_below_hard_ceiling": True,
++                }
++            ]
+     else:
+         g_scale = 1.0
+     weight, su, sv = apply_g_scale(
+@@ -1987,6 +2284,11 @@ def regularize(
+         into_sv=bool(quant_args.get("g_scale_into_sv")),
+     )
+ 
++    # The transform is now numerically frozen.  Only manifest booleans cross
++    # to the host; neither value can feed back into weight construction.
++    apply_out_scales = bool(apply_out_scales_mask.item())
++    quant_args["zeros"] = bool((~mean_positive).item())
++
+     # ext.test_distribution(weight_os, dist_r, dist_ref, -3.8, 3.8)
+     # js_os = jsd(dist_r, dist_ref)
+ 
+@@ -2119,8 +2421,12 @@ def quantize_qsrt(
+             H_diag = H_diag.to(device)
+         sv = output_signs(n, device, quant_args)
+ 
+-        # Move stored L to CPU (if not already), move working L to device
+-        if H_data["L"] is not None:
++        # Historical callers store L on CPU.  The B300 path retains it on GPU
++        # so a shared H13 session does not round-trip a 6144-square factor for
++        # every gate/up candidate.
++        if H_data["L"] is not None and not quant_args.get(
++            "keep_hessian_on_device", False
++        ):
+             H_data["L"] = H_data["L"].cpu()
+         if L is not None:
+             L = L.to(device)
+@@ -2160,7 +2466,9 @@ def quantize_qsrt(
+ 
+         # Select device for work buffers (CPU is slower for small tensors but saves a lot of VRAM on big ones)
+         # TODO: Use pynvml or mem_get_info to predict whether CPU buffer is needed
+-        if weight_r.numel() > 5e8:
++        if weight_r.numel() > 5e8 and not quant_args.get(
++            "gpu_only_numerical_path", False
++        ):
+             weight_r = weight_r.cpu()
+ 
+         # Quantize
+@@ -2305,6 +2613,8 @@ def quantize_qsrt_batch(
+         "ldlq_tf32",
+         "tailbite_context",
+         "g_scale_override",
++        "keep_hessian_on_device",
++        "gpu_only_numerical_path",
+     )
+     for group in quant_args_groups:
+         reference = group[0]
+@@ -2338,7 +2648,9 @@ def quantize_qsrt_batch(
+         su = su.to(device)
+         if H_diag is not None:
+             H_diag = H_diag.to(device)
+-        if H_data["L"] is not None:
++        if H_data["L"] is not None and not args.get(
++            "keep_hessian_on_device", False
++        ):
+             H_data["L"] = H_data["L"].cpu()
+         sv = output_signs(weight.shape[1], device, args)
+         apply_out_scales, weight_r, _, su, sv = regularize(
+diff --git a/kquant/exl3_reference.py b/kquant/exl3_reference.py
+index b1a4e1e..01b4b5d 100644
+--- a/kquant/exl3_reference.py
++++ b/kquant/exl3_reference.py
+@@ -30,8 +30,8 @@ QSRT_CODEBOOKS = (
+ 
+ 
+ def _validate_bits(bits: int) -> None:
+-    if isinstance(bits, bool) or not isinstance(bits, int) or bits not in (2, 3, 4):
+-        raise ValueError("the QSRT reference supports K=2, K=3, or K=4")
++    if isinstance(bits, bool) or not isinstance(bits, int) or bits not in range(2, 7):
++        raise ValueError("the SQG reference supports K=2 through K=6")
+ 
+ 
+ def reconstruct_trellis_states(
+@@ -150,9 +150,9 @@ def decode_qsrt_regularized_weight(
+     tile_bits: tuple[int, ...],
+     codebook: str,
+ ) -> torch.Tensor:
+-    """Decode a K2/K3/K4 tile map to transformed ``[K, N]`` weights.
++    """Decode a K2--K6 tile map to transformed ``[K, N]`` weights.
+ 
+-    SQG uses a distinct state labelling for K2, K3, and K4, so the physical
++    SQG uses a distinct state labelling for every supported K, so the physical
+     tile rate is part of the numerical decode contract.
+     """
+ 
+@@ -164,10 +164,10 @@ def decode_qsrt_regularized_weight(
+     if len(tile_bits) != states.shape[axis]:
+         raise ValueError("tile_bits length does not match the selected rate axis")
+     if any(
+-        isinstance(value, bool) or not isinstance(value, int) or value not in (2, 3, 4)
++        isinstance(value, bool) or not isinstance(value, int) or value not in range(2, 7)
+         for value in tile_bits
+     ):
+-        raise ValueError("tile_bits must contain only integer K2, K3, or K4")
++        raise ValueError("tile_bits must contain only integer K2 through K6")
+     if codebook not in QSRT_CODEBOOKS:
+         raise ValueError(f"unsupported QSRT codebook: {codebook}")
+ 
+diff --git a/kquant/qsrt.py b/kquant/qsrt.py
+index 0424724..1bfd2a0 100644
+--- a/kquant/qsrt.py
++++ b/kquant/qsrt.py
+@@ -804,8 +804,8 @@ def repack_encoded_records(
+ 
+ 
+ def _validate_k(bits: int) -> None:
+-    if isinstance(bits, bool) or not isinstance(bits, int) or bits not in (2, 3, 4):
+-        raise ValueError("QSRT supports K=2, K=3, or K=4")
++    if isinstance(bits, bool) or not isinstance(bits, int) or bits not in range(2, 7):
++        raise ValueError("SQG trellis packing supports K=2 through K=6")
+ 
+ 
+ def pack_trellis_edges(indices: torch.Tensor, bits: int) -> torch.Tensor:
+diff --git a/kquant/sqg_e4m3.py b/kquant/sqg_e4m3.py
+index 2783b10..d7f831d 100644
+--- a/kquant/sqg_e4m3.py
++++ b/kquant/sqg_e4m3.py
+@@ -31,8 +31,8 @@ _Q = (1.0, 2.07630930, -8.08332684, 6.32135736, -1.31208298)
+ 
+ 
+ def _validate(bits: int, mode: str) -> None:
+-    if isinstance(bits, bool) or not isinstance(bits, int) or bits not in (2, 3, 4):
+-        raise ValueError("SQG supports integer K2, K3, or K4")
++    if isinstance(bits, bool) or not isinstance(bits, int) or bits not in range(2, 7):
++        raise ValueError("SQG supports integer K2 through K6")
+     if mode != "normal":
+         raise ValueError("the supported R44 mode is 'normal'")
+ 
+@@ -157,7 +157,7 @@ def sqg_e4m3_bytes_from_rank_lut(
+ ) -> torch.Tensor:
+     """Map a shared rank-indexed E4M3 law onto the unchanged SQG graph.
+ 
+-    ``rank_lut`` is deliberately independent of ``bits``.  K2/K3/K4 may
++    ``rank_lut`` is deliberately independent of ``bits``.  K2 through K6 may
+     expose different outgoing menus because their stratum widths differ, but
+     they retain the same history mixer, syndrome mixer, branch permutation,
+     and scalar reconstruction law.  This is the controlled interface used to
+diff --git a/kquant/sqg_quantizer.py b/kquant/sqg_quantizer.py
+index cc29e00..ade6f06 100644
+--- a/kquant/sqg_quantizer.py
++++ b/kquant/sqg_quantizer.py
+@@ -2,9 +2,15 @@
+ 
+ from __future__ import annotations
+ 
++import hashlib
++import importlib.machinery
++import importlib.util
++import os
++import re
++import sys
++from collections.abc import Mapping
+ from functools import lru_cache
+ from pathlib import Path
+-from collections.abc import Mapping
+ 
+ import torch
+ from torch.utils.cpp_extension import load
+@@ -12,11 +18,107 @@ from torch.utils.cpp_extension import load
+ from kquant.sqg_e4m3 import sqg_e4m3_bytes
+ 
+ 
++_EXTENSION_NAME = "kquant_sqg_quantize_ext_v22"
++_EXTENSION_PATH_ENV = "KQUANT_SQG_EXTENSION_PATH"
++_EXTENSION_SHA256_ENV = "KQUANT_SQG_EXTENSION_SHA256"
++_REQUIRE_PREBUILT_ENV = "KQUANT_SQG_REQUIRE_PREBUILT"
++_SHA256_RE = re.compile(r"[0-9a-f]{64}")
++
++
++def _sha256_file(path: Path) -> str:
++    digest = hashlib.sha256()
++    with path.open("rb") as handle:
++        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
++            digest.update(chunk)
++    return digest.hexdigest()
++
++
++def _import_prebuilt_extension(path: Path):
++    """Import one exact pybind extension without invoking the JIT builder."""
++
++    spec = importlib.util.spec_from_file_location(_EXTENSION_NAME, path)
++    if spec is None or spec.loader is None:
++        raise ImportError(f"cannot construct an import spec for SQG extension: {path}")
++    previous = sys.modules.get(_EXTENSION_NAME)
++    module = importlib.util.module_from_spec(spec)
++    sys.modules[_EXTENSION_NAME] = module
++    try:
++        spec.loader.exec_module(module)
++    except BaseException:
++        if previous is None:
++            sys.modules.pop(_EXTENSION_NAME, None)
++        else:
++            sys.modules[_EXTENSION_NAME] = previous
++        raise
++    return module
++
++
++def _load_prebuilt_extension(raw_path: str, expected_sha256: str):
++    candidate = Path(raw_path)
++    if not candidate.is_absolute():
++        raise ValueError(f"{_EXTENSION_PATH_ENV} must be an absolute path")
++    if candidate.is_symlink():
++        raise ValueError(f"{_EXTENSION_PATH_ENV} may not name a symlink")
++    try:
++        path = candidate.resolve(strict=True)
++    except FileNotFoundError as exc:
++        raise FileNotFoundError(
++            f"sealed SQG extension does not exist: {candidate}"
++        ) from exc
++    if not path.is_file():
++        raise ValueError(f"sealed SQG extension is not a regular file: {path}")
++    allowed_names = {
++        f"{_EXTENSION_NAME}{suffix}"
++        for suffix in importlib.machinery.EXTENSION_SUFFIXES
++    }
++    if path.name not in allowed_names:
++        raise ValueError(
++            "sealed SQG extension filename must identify "
++            f"{_EXTENSION_NAME}: {path.name}"
++        )
++    if _SHA256_RE.fullmatch(expected_sha256) is None:
++        raise ValueError(f"{_EXTENSION_SHA256_ENV} must be 64 lowercase hex digits")
++    observed_sha256 = _sha256_file(path)
++    if observed_sha256 != expected_sha256:
++        raise RuntimeError(
++            "sealed SQG extension SHA256 mismatch: "
++            f"{observed_sha256} != {expected_sha256}"
++        )
++    module = _import_prebuilt_extension(path)
++    module_path = Path(str(getattr(module, "__file__", ""))).resolve()
++    if module_path != path:
++        raise RuntimeError(
++            f"loaded SQG extension origin differs: {module_path} != {path}"
++        )
++    if not callable(getattr(module, "quantize_tiles_procedural", None)) or not callable(
++        getattr(module, "quantize_tiles_sqg", None)
++    ):
++        raise RuntimeError(
++            "sealed SQG extension lacks required quantization entry points"
++        )
++    if _sha256_file(path) != expected_sha256:
++        raise RuntimeError("sealed SQG extension changed while it was being imported")
++    return module
++
++
+ @lru_cache(maxsize=1)
+ def _extension():
++    require_value = os.environ.get(_REQUIRE_PREBUILT_ENV, "")
++    if require_value not in ("", "0", "1"):
++        raise ValueError(f"{_REQUIRE_PREBUILT_ENV} must be unset, 0, or 1")
++    raw_path = os.environ.get(_EXTENSION_PATH_ENV, "")
++    expected_sha256 = os.environ.get(_EXTENSION_SHA256_ENV, "")
++    if require_value == "1" or raw_path or expected_sha256:
++        if not raw_path or not expected_sha256:
++            raise RuntimeError(
++                f"{_EXTENSION_PATH_ENV} and {_EXTENSION_SHA256_ENV} are both required "
++                "for sealed prebuilt SQG loading"
++            )
++        return _load_prebuilt_extension(raw_path, expected_sha256)
++
+     project = Path(__file__).resolve().parents[1]
+     return load(
+-        name="kquant_sqg_quantize_ext_v22",
++        name=_EXTENSION_NAME,
+         sources=[
+             str(project / "kquant/csrc/sqg_quantize.cpp"),
+             str(project / "kquant/csrc/sqg_quantize.cu"),
+@@ -42,7 +144,7 @@ def _sqg_temp_buffers(device: torch.device, bits: int):
+ 
+     multiprocessors = torch.cuda.get_device_properties(device).multi_processor_count
+     edges = 65536 >> bits
+-    decisions = edges // (4 if bits == 2 else 2)
++    decisions = edges // 4 if bits == 2 else (edges // 2 if bits <= 4 else edges)
+     free_bytes, _ = torch.cuda.mem_get_info(device)
+     decision_bytes_per_tile = 256 * decisions
+     affordable = max(256, int(free_bytes * 0.5) // decision_bytes_per_tile)
+@@ -94,13 +196,17 @@ def install_sqg_quantizer(quantizer_module) -> None:
+                 )
+             if not isinstance(rate_codebooks, Mapping):
+                 raise TypeError("sqg_e4m3_luts_by_bits must be a mapping")
+-            if set(rate_codebooks) - {2, 3, 4}:
+-                raise ValueError("rate-specific SQG LUT keys must be K2, K3, or K4")
++            if set(rate_codebooks) - set(range(2, 7)):
++                raise ValueError("rate-specific SQG LUT keys must be K2 through K6")
+             try:
+                 codebook = rate_codebooks[bits]
+             except KeyError as exc:
+                 raise ValueError(f"missing rate-specific SQG K{bits} LUT") from exc
+             if codebook is None:
++                if bits not in (2, 3, 4):
++                    raise ValueError(
++                        "non-SQG procedural controls remain limited to K2 through K4"
++                    )
+                 output = torch.empty_like(tiles)
+                 indices = torch.empty_like(tiles, dtype=torch.int16)
+                 costs, edges = _sqg_temp_buffers(tiles.device, bits)
+diff --git a/tests/test_qsrt.py b/tests/test_qsrt.py
+index 0d5189d..6e5c469 100644
+--- a/tests/test_qsrt.py
++++ b/tests/test_qsrt.py
+@@ -218,6 +218,8 @@ def test_matrix_rate_axis_tracks_shared_intermediate_dimension() -> None:
+         (2, [6939] * 6),
+         (3, [30469, 1337, 1337, 14711, 14711, 30469]),
+         (4, [17767, 291, -12817, -30293, 17767, 291]),
++        (5, [12820, 68, 21686, -14526, -31643, 13775]),
++        (6, [-31984, 16, 8338, 20871, -11377, -29904]),
+     ],
+ )
+ def test_trellis_edge_pack_matches_native_word_order(
+@@ -231,7 +233,7 @@ def test_trellis_edge_pack_matches_native_word_order(
+     torch.testing.assert_close(unpack_trellis_edges(packed, bits), indices)
+ 
+ 
+-@pytest.mark.parametrize("bits", [2, 3, 4])
++@pytest.mark.parametrize("bits", range(2, 7))
+ def test_trellis_edge_pack_round_trips_only_stored_low_bits(bits: int) -> None:
+     generator = torch.Generator().manual_seed(100 + bits)
+     indices = torch.randint(-32768, 32768, (7, 256), generator=generator)
+@@ -241,7 +243,7 @@ def test_trellis_edge_pack_round_trips_only_stored_low_bits(bits: int) -> None:
+     torch.testing.assert_close(unpacked.int(), indices.int() & ((1 << bits) - 1))
+ 
+ 
+-@pytest.mark.parametrize("bits", [2, 3, 4])
++@pytest.mark.parametrize("bits", range(2, 7))
+ def test_cyclic_trellis_states_round_trip_exactly(bits: int) -> None:
+     generator = torch.Generator().manual_seed(150 + bits)
+     edges = torch.randint(0, 1 << bits, (5, 256), generator=generator)
+@@ -251,7 +253,7 @@ def test_cyclic_trellis_states_round_trip_exactly(bits: int) -> None:
+     torch.testing.assert_close(unpack_trellis_states(packed, bits), states)
+ 
+ 
+-@pytest.mark.parametrize("bits", [2, 3, 4])
++@pytest.mark.parametrize("bits", range(2, 7))
+ def test_cyclic_trellis_state_closed_form_matches_scalar_recurrence(bits: int) -> None:
+     generator = torch.Generator().manual_seed(175 + bits)
+     edges = torch.randint(0, 1 << bits, (3, 5, 256), generator=generator)
+diff --git a/tests/test_sqg_e4m3.py b/tests/test_sqg_e4m3.py
+index f630827..669fd6d 100644
+--- a/tests/test_sqg_e4m3.py
++++ b/tests/test_sqg_e4m3.py
+@@ -34,6 +34,8 @@ _SQG_XOR_CHEB_T12_SHA256 = {
+     2: "62027916386245a84c86156a0a08b6cf07e41548871af0e56fb40780558f6293",
+     3: "afe7b3633e7d243b00b379b18ec4dca573722b3727cafef47fcb6470d7e7e6c9",
+     4: "5a9620f0c4d8f0a60d0b6fbea921dcecbd193e6febf31043aaa6403c20389c2f",
++    5: "a1caf5572c67d4face423478f423e1ed6149a8c920be5c28d4702946c93cae2f",
++    6: "8e0edc92e20eb92f5287dbdadc8e870fc500fc03c5eed996cae0781d197ebcb0",
+ }
+ 
+ 
+@@ -48,7 +50,7 @@ def test_primary_sqg_xor_cheb_t12_matches_runtime_contract() -> None:
+     assert _sha256(table) == _SQG_XOR_CHEB_T12_SHA256["t12"]
+ 
+ 
+-@pytest.mark.parametrize("bits", (2, 3, 4))
++@pytest.mark.parametrize("bits", range(2, 7))
+ def test_primary_sqg_xor_cheb_t12_graph_and_labels_match_runtime_contract(
+     bits: int,
+ ) -> None:
+@@ -69,7 +71,7 @@ def test_primary_sqg_xor_cheb_t12_graph_and_labels_match_runtime_contract(
+     assert _sha256(labels) == _SQG_XOR_CHEB_T12_SHA256[bits]
+ 
+ 
+-@pytest.mark.parametrize("bits", (2, 3, 4))
++@pytest.mark.parametrize("bits", range(2, 7))
+ def test_sqg_preprojection_mapping_is_a_permutation(bits: int) -> None:
+     ranks = sqg_rank_permutation(bits)
+     assert ranks.dtype == torch.int64
+@@ -77,7 +79,7 @@ def test_sqg_preprojection_mapping_is_a_permutation(bits: int) -> None:
+     assert torch.equal(torch.sort(ranks).values, torch.arange(1 << 16))
+ 
+ 
+-@pytest.mark.parametrize("bits", (2, 3, 4))
++@pytest.mark.parametrize("bits", range(2, 7))
+ def test_sqg_codebook_round_trips_exact_e4m3(bits: int) -> None:
+     raw = sqg_e4m3_bytes(bits, "normal")
+     codebook = sqg_e4m3_codebook(bits, "normal")
+@@ -88,7 +90,7 @@ def test_sqg_codebook_round_trips_exact_e4m3(bits: int) -> None:
+     assert bool(torch.isfinite(codebook).all())
+ 
+ 
+-@pytest.mark.parametrize("bits", (2, 3, 4))
++@pytest.mark.parametrize("bits", range(2, 7))
+ def test_sqg_each_state_spans_all_coarse_strata(bits: int) -> None:
+     ranks = sqg_rank_permutation(bits)
+     width = 16 - bits
+@@ -97,7 +99,7 @@ def test_sqg_each_state_spans_all_coarse_strata(bits: int) -> None:
+     assert torch.equal(torch.sort(strata, dim=1).values, expected)
+ 
+ 
+-@pytest.mark.parametrize("bits", (2, 3, 4))
++@pytest.mark.parametrize("bits", range(2, 7))
+ def test_shared_rank_lut_preserves_one_graph_for_every_rate(bits: int) -> None:
+     rank_lut = torch.arange(1 << 16, dtype=torch.int64).remainder(126).to(torch.uint8)
+     raw = sqg_e4m3_bytes_from_rank_lut(bits, rank_lut)
+@@ -118,7 +120,7 @@ def test_shared_rank_lut_rejects_nonfinite_e4m3() -> None:
+         sqg_e4m3_bytes_from_rank_lut(2, rank_lut)
+ 
+ 
+-@pytest.mark.parametrize("bits", (2, 3, 4))
++@pytest.mark.parametrize("bits", range(2, 7))
+ def test_sqg_cheb_normal_named_codebook_uses_exact_shared_rank_law(bits: int) -> None:
+     expected = sqg_cheb_normal_rank_e4m3_bytes().index_select(
+         0, sqg_rank_permutation(bits)
+@@ -132,7 +134,7 @@ def test_sqg_cheb_normal_named_codebook_uses_exact_shared_rank_law(bits: int) ->
+     )
+ 
+ 
+-@pytest.mark.parametrize("bits", (2, 3, 4))
++@pytest.mark.parametrize("bits", range(2, 7))
+ @pytest.mark.parametrize("rate_axis", ("k", "n"))
+ def test_w2_k2_q8h4_profile_changes_only_k_axis_k2(
+     bits: int, rate_axis: str
+diff --git a/tests/test_sqg_quantizer.py b/tests/test_sqg_quantizer.py
+index 8d99e86..c470ee2 100644
+--- a/tests/test_sqg_quantizer.py
++++ b/tests/test_sqg_quantizer.py
+@@ -1,11 +1,135 @@
+ from __future__ import annotations
+ 
++import hashlib
++from pathlib import Path
+ from types import SimpleNamespace
+ 
++import pytest
+ import torch
+ 
+ import kquant.sqg_quantizer as sqg_quantizer
++from kquant.exl3_reference import reconstruct_trellis_states
+ from kquant.sqg_quantizer import install_sqg_quantizer
++from kquant.sqg_e4m3 import sqg_xor_cheb_t12_bytes
++
++
++def _sealed_extension(
++    tmp_path: Path,
++    monkeypatch: pytest.MonkeyPatch,
++    *,
++    filename: str = "kquant_sqg_quantize_ext_v22.so",
++) -> tuple[Path, SimpleNamespace]:
++    path = tmp_path / filename
++    payload = b"sealed-sqg-extension-test-double"
++    path.write_bytes(payload)
++    module = SimpleNamespace(
++        __file__=str(path),
++        quantize_tiles_procedural=lambda *_args: None,
++        quantize_tiles_sqg=lambda *_args: None,
++    )
++    monkeypatch.setenv("KQUANT_SQG_REQUIRE_PREBUILT", "1")
++    monkeypatch.setenv("KQUANT_SQG_EXTENSION_PATH", str(path))
++    monkeypatch.setenv(
++        "KQUANT_SQG_EXTENSION_SHA256", hashlib.sha256(payload).hexdigest()
++    )
++    return path, module
++
++
++def test_required_prebuilt_extension_uses_exact_path_and_never_jit_builds(
++    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
++) -> None:
++    path, module = _sealed_extension(tmp_path, monkeypatch)
++    imported: list[Path] = []
++
++    def fake_import(actual: Path):
++        imported.append(actual)
++        return module
++
++    monkeypatch.setattr(sqg_quantizer, "_import_prebuilt_extension", fake_import)
++    monkeypatch.setattr(
++        sqg_quantizer,
++        "load",
++        lambda **_kwargs: pytest.fail("JIT builder must not run in prebuilt mode"),
++    )
++    sqg_quantizer._extension.cache_clear()
++    try:
++        assert sqg_quantizer._extension() is module
++        assert imported == [path.resolve()]
++    finally:
++        sqg_quantizer._extension.cache_clear()
++
++
++def test_required_prebuilt_extension_rejects_wrong_filename_before_import_or_jit(
++    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
++) -> None:
++    _sealed_extension(tmp_path, monkeypatch, filename="foreign_extension.so")
++    monkeypatch.setattr(
++        sqg_quantizer,
++        "_import_prebuilt_extension",
++        lambda _path: pytest.fail("foreign extension must not be imported"),
++    )
++    monkeypatch.setattr(
++        sqg_quantizer,
++        "load",
++        lambda **_kwargs: pytest.fail("JIT builder must not run in prebuilt mode"),
++    )
++    sqg_quantizer._extension.cache_clear()
++    try:
++        with pytest.raises(ValueError, match="filename must identify"):
++            sqg_quantizer._extension()
++    finally:
++        sqg_quantizer._extension.cache_clear()
++
++
++def test_required_prebuilt_extension_rejects_hash_mismatch_before_import_or_jit(
++    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
++) -> None:
++    _sealed_extension(tmp_path, monkeypatch)
++    monkeypatch.setenv("KQUANT_SQG_EXTENSION_SHA256", "0" * 64)
++    monkeypatch.setattr(
++        sqg_quantizer,
++        "_import_prebuilt_extension",
++        lambda _path: pytest.fail("mismatched extension must not be imported"),
++    )
++    monkeypatch.setattr(
++        sqg_quantizer,
++        "load",
++        lambda **_kwargs: pytest.fail("JIT builder must not run in prebuilt mode"),
++    )
++    sqg_quantizer._extension.cache_clear()
++    try:
++        with pytest.raises(RuntimeError, match="SHA256 mismatch"):
++            sqg_quantizer._extension()
++    finally:
++        sqg_quantizer._extension.cache_clear()
++
++
++@pytest.mark.parametrize(
++    ("missing", "message"),
++    (
++        ("KQUANT_SQG_EXTENSION_PATH", "are both required"),
++        ("KQUANT_SQG_EXTENSION_SHA256", "are both required"),
++    ),
++)
++def test_required_prebuilt_extension_rejects_partial_binding_without_jit(
++    tmp_path: Path,
++    monkeypatch: pytest.MonkeyPatch,
++    missing: str,
++    message: str,
++) -> None:
++    _sealed_extension(tmp_path, monkeypatch)
++    monkeypatch.delenv(missing)
++    monkeypatch.setattr(
++        sqg_quantizer,
++        "load",
++        lambda **_kwargs: pytest.fail("JIT builder must not run in prebuilt mode"),
++    )
++    sqg_quantizer._extension.cache_clear()
++    try:
++        with pytest.raises(RuntimeError, match=message):
++            sqg_quantizer._extension()
++    finally:
++        sqg_quantizer._extension.cache_clear()
+ 
+ 
+ def test_rate_specific_none_dispatches_to_kquant_mcg(monkeypatch) -> None:
+@@ -88,3 +212,51 @@ def test_sqg_luts_are_transposed_by_predecessor_for_vector_loads(
+         )
+         assert torch.equal(actual, expected)
+         assert calls[-1][6:] == (bits, 128)
++
++
++@pytest.mark.parametrize("bits", [5, 6])
++def test_k5_k6_sqg_luts_remain_in_direct_state_order(
++    monkeypatch, bits: int
++) -> None:
++    calls: list[tuple] = []
++    monkeypatch.setattr(
++        sqg_quantizer,
++        "_sqg_temp_buffers",
++        lambda _device, _bits: (torch.empty(0), torch.empty(0)),
++    )
++    monkeypatch.setattr(
++        sqg_quantizer,
++        "_extension",
++        lambda: SimpleNamespace(quantize_tiles_sqg=lambda *args: calls.append(args)),
++    )
++    module = SimpleNamespace(quantize_tiles=lambda *_args: "upstream")
++    install_sqg_quantizer(module)
++    raw = torch.arange(65536, dtype=torch.int64).to(torch.uint8)
++    module.quantize_tiles(
++        torch.zeros((1, 256), dtype=torch.float32),
++        {"K": bits, "devices": ["cuda:0"], "sqg_e4m3_lut": raw},
++    )
++    assert torch.equal(calls[-1][5], raw)
++    assert calls[-1][6:] == (bits, 128)
++
++
++@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is unavailable")
++@pytest.mark.parametrize("bits", range(2, 7))
++def test_sqg_quantizer_keeps_saturated_cost_paths_closed(bits: int) -> None:
++    module = SimpleNamespace(quantize_tiles=lambda *_args: "upstream")
++    install_sqg_quantizer(module)
++    tiles = torch.full((2, 256), 1000.0, dtype=torch.float32, device="cuda")
++
++    _output, states = module.quantize_tiles(
++        tiles,
++        {
++            "K": bits,
++            "devices": ["cuda:0"],
++            "sqg_e4m3_lut": sqg_xor_cheb_t12_bytes(bits),
++            "tailbite_context": 128,
++        },
++    )
++
++    edges = states.to(torch.int64) & ((1 << bits) - 1)
++    expected = reconstruct_trellis_states(edges, bits)
++    assert torch.equal(states, expected)

reproduction/closure/patches/qsrt/changed-files/qsrt/pack/qsrt_candidates.py

@@ -0,0 +1,1637 @@
+"""Offline-streamed all-expert candidates for the QSRT codec."""
+
+from __future__ import annotations
+
+import math
+import os
+from dataclasses import asdict, dataclass
+from collections.abc import Mapping
+from typing import Literal, Protocol, Sequence
+
+import torch
+import torch.nn.functional as F
+
+from qsrt.capture import LayerSamples
+from qsrt.exl3_reference import (
+    CODEBOOK_SQG_XOR_CHEB_T12,
+    reconstruct_trellis_states,
+)
+from qsrt.qsrt_candidates import (
+    PERMUTATION_POLICIES,
+    PermutationPolicy,
+    RatePairSelection,
+    RequestPartition,
+    RoutedRows,
+    build_expert_hessians,
+    deterministic_expert_seed,
+    functional_sse_by_request,
+    permutation_policy_contexts,
+    select_expert_rows,
+    select_phase1_rate_pair,
+)
+from qsrt.qsrt import (
+    H308,
+    K2,
+    R0,
+    SCHEMA,
+    ExpertFormatSpec,
+    PHASE1_MODE_IDS,
+    RATE_TRANSFER_MODES,
+    RECORDS_PER_EXPERT,
+    resolve_mode,
+)
+from qsrt.qsrt_coupled import (
+    CoupledHadamardExecution,
+    CoupledHadamardSpec,
+    ExpertActivation,
+    coupled_execution,
+    encode_coupled_weights,
+    expert_activation,
+)
+from qsrt.pack.qsrt_encoder import (
+    Layout,
+    MATRICES,
+    QSRTExpertEncoding,
+    QSRTMatrixCandidate,
+    QSRTMatrixEncoding,
+    QSRTTransformSeeds,
+    SearchLayout,
+    finalize_qsrt_matrix_candidate,
+    quantize_qsrt_matrix_candidates_batched,
+    quantize_qsrt_matrix_expert_batch,
+)
+from qsrt.tp_simulator import situ
+
+
+CANDIDATE_POOL_KIND = "qsrt_kimi_k3_qsrt_candidate_pool"
+CANDIDATE_POOL_SCHEMA_VERSION = 7
+OFFICIAL_SOURCE_DAMAGE_METRIC = "official_source_excess_sse"
+HessianPolicy = Literal["captured_blend", "identity"]
+HESSIAN_POLICIES: tuple[HessianPolicy, ...] = ("captured_blend", "identity")
+FOLDED_SCALE_GRID: tuple[float, ...] = (0.85, 0.925, 1.0, 1.075, 1.15)
+
+
+def _debug_assert_candidate_state_closure(
+    candidate: QSRTMatrixCandidate,
+    *,
+    layer: int,
+    expert: int,
+    stage: str,
+) -> None:
+    encoded = candidate.encoded
+    plan = candidate.plan
+    tile_bits = tuple(int(value) for value in plan.encoder_tile_bits)
+    rate_dim = 0 if plan.rate_axis == "k" else 1
+    for bits in sorted(set(tile_bits)):
+        positions = torch.tensor(
+            [index for index, value in enumerate(tile_bits) if value == bits],
+            dtype=torch.long,
+            device=encoded.device,
+        )
+        selected = encoded.index_select(rate_dim, positions)
+        expected = reconstruct_trellis_states(selected, bits)
+        mismatch = selected != expected
+        if bool(torch.any(mismatch)):
+            count = int(torch.count_nonzero(mismatch))
+            first = tuple(
+                int(value)
+                for value in torch.nonzero(mismatch, as_tuple=False)[0]
+                .cpu()
+                .tolist()
+            )
+            raise RuntimeError(
+                f"layer {layer} expert {expert} {plan.matrix} R{plan.mode.mode_id} "
+                f"failed {stage} state closure at K{bits}: {count} states "
+                f"differ; first grouped coordinate {first}"
+            )
+
+
+def _debug_assert_candidate_tree_closure(
+    candidates,
+    experts: Sequence[int],
+    *,
+    layer: int,
+    stage: str,
+) -> None:
+    if os.environ.get("QSRT_DEBUG_CANDIDATE_CLOSURE") != "1":
+        return
+    for expert, tree in zip(experts, candidates):
+        stack = [tree]
+        while stack:
+            value = stack.pop()
+            if isinstance(value, QSRTMatrixCandidate):
+                _debug_assert_candidate_state_closure(
+                    value, layer=layer, expert=expert, stage=stage
+                )
+            elif isinstance(value, Mapping):
+                stack.extend(value.values())
+
+
+class MatrixStore(Protocol):
+    """Read-only source that materializes one official matrix per call."""
+
+    def load_matrix(
+        self,
+        layer: int,
+        expert: int,
+        matrix: str,
+        *,
+        device: torch.device | str | None = None,
+    ) -> torch.Tensor: ...
+
+
+@dataclass
+class QSRTCandidateEncoding:
+    """Selected payload plus the complete phase-1 selection evidence."""
+
+    expert: QSRTExpertEncoding
+    selection: RatePairSelection
+    block_contexts: torch.Tensor
+    block_scores: torch.Tensor
+    context_basis: str
+    covariance: dict[str, object]
+    evaluated_modes: tuple[int, ...]
+    evaluated_formats: tuple[tuple[int, int], ...]
+    fit_sse: dict[tuple[int, int], torch.Tensor]
+    confirmation_sse: dict[tuple[int, int], torch.Tensor]
+    fit_reference_energy: torch.Tensor
+    confirmation_reference_energy: torch.Tensor
+    fit_counts: torch.Tensor
+    confirmation_counts: torch.Tensor
+    mode_coding: dict[tuple[int, int], dict[str, object]]
+
+    def metadata(self) -> dict[str, object]:
+        return {
+            "selection": asdict(self.selection),
+            "context_basis": self.context_basis,
+            "covariance": self.covariance,
+            "evaluated_modes": list(self.evaluated_modes),
+            "mode_coding": {
+                (resolve_mode(r13).name if r13 == r2 else f"R{r13}/R{r2}"): value
+                for (r13, r2), value in self.mode_coding.items()
+            },
+        }
+
+
+def _source_middle(
+    gate_projection: torch.Tensor,
+    up_projection: torch.Tensor,
+    *,
+    activation: ExpertActivation = "situ",
+) -> torch.Tensor:
+    return expert_activation(
+        gate_projection,
+        up_projection,
+        activation=activation,
+    )
+
+
+def _activation_basis(activation: ExpertActivation) -> str:
+    if activation not in ("situ", "silu"):
+        raise ValueError(f"unsupported expert activation: {activation}")
+    return f"post_{activation}"
+
+
+def _folded_intermediate_conditioning(
+    block_scores: torch.Tensor,
+    block_contexts: torch.Tensor,
+    *,
+    power: float,
+) -> torch.Tensor | None:
+    """Derive a metadata-free per-record output conditioning profile.
+
+    The source rows are divided by this profile before w1/w3 quantization and
+    the exact inverse is folded into their existing expert-local ``svh``
+    vectors.  Factors are constant over each final 128-channel Hadamard block,
+    so the fold commutes with the transform and requires no new runtime tensor.
+    """
+
+    if not math.isfinite(power):
+        raise ValueError("folded scale power must be finite")
+    if power == 0.0:
+        return None
+    if block_scores.ndim != 1 or block_scores.numel() != 3072 // 4:
+        raise ValueError("folded conditioning requires 768 group scores")
+    if not bool(torch.all(torch.isfinite(block_scores))) or bool(
+        torch.any(block_scores < 0)
+    ):
+        raise ValueError("folded conditioning scores must be finite and nonnegative")
+    contexts = block_contexts.to(device=block_scores.device, dtype=torch.long)
+    group_energy = block_scores.float()
+    record_energy = torch.zeros(
+        RECORDS_PER_EXPERT, device=block_scores.device, dtype=torch.float32
+    ).scatter_add_(0, contexts, group_energy.float())
+    counts = torch.bincount(contexts, minlength=RECORDS_PER_EXPERT).float()
+    record_rms = (record_energy / counts).clamp_min(1e-30).sqrt()
+    geometric_mean = record_rms.log().mean().exp()
+    # A smaller folded scale magnifies important rows inside the trellis
+    # search and shrinks their absolute error again when folded back into svh.
+    # This is an output-metric conditioning heuristic, not a new runtime scale.
+    desired = (record_rms / geometric_mean).pow(-power)
+    grid = torch.tensor(FOLDED_SCALE_GRID, device=desired.device)
+    selected = grid[(desired[:, None] - grid[None, :]).abs().argmin(dim=1)]
+    return selected.index_select(0, contexts).repeat_interleave(4).contiguous()
+
+
+def _load_source_matrix(
+    store: MatrixStore,
+    layer: int,
+    expert: int,
+    matrix: str,
+    device: torch.device,
+) -> torch.Tensor:
+    source = store.load_matrix(layer, expert, matrix, device=device)
+    if source.device != device:
+        raise ValueError(
+            f"official source store returned {source.device}, expected {device}"
+        )
+    if source.dtype != torch.float32:
+        raise ValueError(
+            f"official source store returned {source.dtype}, expected torch.float32"
+        )
+    return source.contiguous()
+
+
+def _split_rows(rows: RoutedRows, requests: dict[int, str]) -> torch.Tensor:
+    allowed = torch.tensor(
+        sorted(requests), dtype=rows.request_steps.dtype, device=rows.request_steps.device
+    )
+    return torch.isin(rows.request_steps, allowed)
+
+
+def _empty_metrics(requests: dict[int, str]) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    count = len(requests)
+    return (
+        torch.zeros(count, dtype=torch.float64),
+        torch.zeros(count, dtype=torch.float64),
+        torch.zeros(count, dtype=torch.int64),
+    )
+
+
+@dataclass(frozen=True)
+class _DeferredFunctionalRows:
+    """Static routed-metric state shared by every candidate for one fold."""
+
+    device_mask: torch.Tensor
+    reference: torch.Tensor
+    route_weights: torch.Tensor
+    request_indices: torch.Tensor
+    reference_energy: torch.Tensor
+    counts: torch.Tensor
+    request_count: int
+
+
+def _prepare_deferred_functional_rows(
+    reference: torch.Tensor,
+    gates: torch.Tensor,
+    request_steps: torch.Tensor,
+    mask: torch.Tensor,
+    requests: dict[int, str],
+    *,
+    device: torch.device,
+) -> _DeferredFunctionalRows:
+    """Prepare invariant metric terms without changing their arithmetic."""
+
+    device_mask = mask.to(device=device)
+    selected_reference = reference[device_mask].double()
+    selected_gates = gates[device_mask]
+    selected_steps = request_steps[mask]
+    ordered_steps = sorted(map(int, requests))
+    positions = {step: index for index, step in enumerate(ordered_steps)}
+    try:
+        request_indices = torch.tensor(
+            [positions[int(step)] for step in selected_steps.cpu().tolist()],
+            dtype=torch.long,
+        )
+    except KeyError as exc:
+        raise ValueError(
+            f"metric row has unmapped request epoch {exc.args[0]}"
+        ) from exc
+    route_weights = selected_gates.double().square().cpu()
+    row_reference = (
+        selected_reference.square().sum(dim=1).cpu() * route_weights
+    )
+    reference_energy = torch.zeros(len(ordered_steps), dtype=torch.float64)
+    counts = torch.zeros(len(ordered_steps), dtype=torch.int64)
+    reference_energy.scatter_add_(0, request_indices, row_reference)
+    counts.scatter_add_(
+        0,
+        request_indices,
+        torch.ones_like(request_indices, dtype=torch.int64),
+    )
+    return _DeferredFunctionalRows(
+        device_mask=device_mask,
+        reference=selected_reference,
+        route_weights=route_weights,
+        request_indices=request_indices,
+        reference_energy=reference_energy,
+        counts=counts,
+        request_count=len(ordered_steps),
+    )
+
+
+def _defer_functional_row_sse(
+    plan: _DeferredFunctionalRows,
+    candidate: torch.Tensor,
+) -> torch.Tensor:
+    """Compute per-row SSE on CUDA and defer its host synchronization."""
+
+    delta = candidate[plan.device_mask].double() - plan.reference
+    return delta.square().sum(dim=1)
+
+
+def _finish_deferred_functional_sse(
+    plan: _DeferredFunctionalRows,
+    keys: Sequence[tuple[int, int]],
+    row_sse: Sequence[torch.Tensor],
+) -> dict[tuple[int, int], torch.Tensor]:
+    """Transfer all candidate rows once, then reproduce the CPU reductions."""
+
+    if len(keys) != len(row_sse):
+        raise ValueError("deferred functional metric keys and rows differ")
+    if not keys:
+        return {}
+    rows_cpu = torch.stack(tuple(row_sse)).cpu()
+    result: dict[tuple[int, int], torch.Tensor] = {}
+    for key, values in zip(keys, rows_cpu):
+        sse = torch.zeros(plan.request_count, dtype=torch.float64)
+        sse.scatter_add_(
+            0,
+            plan.request_indices,
+            values * plan.route_weights,
+        )
+        result[key] = sse
+    return result
+
+
+def _mode_evidence(encodings: dict[str, QSRTMatrixEncoding]) -> dict[str, object]:
+    trellis_bytes = sum(
+        value.tensors["trellis"].numel()
+        * value.tensors["trellis"].element_size()
+        for value in encodings.values()
+    )
+    scale_bytes = sum(
+        value.tensors[name].numel() * value.tensors[name].element_size()
+        for value in encodings.values()
+        for name in ("suh", "svh")
+    )
+    return {
+        "trellis_bytes": trellis_bytes,
+        "scale_bytes": scale_bytes,
+        "total_bytes_before_layer_deduplication": trellis_bytes + scale_bytes,
+        "matrices": {
+            matrix: {
+                "proxy": float(value.coding["proxy"]),
+                "mode": value.coding["mode"],
+                "trellis_descriptor": value.coding["trellis_descriptor"],
+            }
+            for matrix, value in encodings.items()
+        },
+    }
+
+
+def _candidate_mode_evidence(encodings: dict[str, object]) -> dict[str, object]:
+    trellis_bytes = sum(
+        value.reconstruction.numel()
+        * sum(value.plan.mode.context_bits)
+        // (len(value.plan.mode.context_bits) * 8)
+        for value in encodings.values()
+    )
+    scale_bytes = sum(
+        value.tensors[name].numel() * value.tensors[name].element_size()
+        for value in encodings.values()
+        for name in ("suh", "svh")
+    )
+    return {
+        "trellis_bytes": trellis_bytes,
+        "scale_bytes": scale_bytes,
+        "total_bytes_before_layer_deduplication": trellis_bytes + scale_bytes,
+        "matrices": {
+            matrix: {
+                "proxy": float(value.proxy),
+                "mode": value.plan.mode.name,
+                "mode_id": value.plan.mode.mode_id,
+                "rate_axis": value.plan.rate_axis,
+            }
+            for matrix, value in encodings.items()
+        },
+    }
+
+
+def _candidate_middle_by_r13(
+    inputs: torch.Tensor,
+    upstream: dict[str, dict[int, QSRTMatrixCandidate]],
+    evaluated_modes: Sequence[int],
+    *,
+    activation: ExpertActivation = "situ",
+) -> dict[int, torch.Tensor]:
+    """Replay exact candidate reconstructions through the model activation."""
+
+    return {
+        int(r13): _source_middle(
+            F.linear(inputs, upstream["w1"][int(r13)].reconstruction),
+            F.linear(inputs, upstream["w3"][int(r13)].reconstruction),
+            activation=activation,
+        )
+        for r13 in evaluated_modes
+    }
+
+
+def _conditional_h2_by_r13(
+    inputs: torch.Tensor,
+    gates: torch.Tensor,
+    fit_mask_device: torch.Tensor,
+    middle_by_r13: dict[int, torch.Tensor],
+    *,
+    global_h13: torch.Tensor,
+    global_h2: torch.Tensor,
+    device: torch.device,
+    hessian_policy: HessianPolicy,
+    have_local_support: bool,
+) -> tuple[dict[int, torch.Tensor], dict[str, object]]:
+    """Build one support-aware down-projection covariance per upstream mode."""
+
+    h2_by_r13: dict[int, torch.Tensor] = {}
+    evidence_by_r13: dict[str, object] = {}
+    identity_h2 = torch.eye(
+        global_h2.shape[0], dtype=torch.float32, device=device
+    )
+    for r13, middle in middle_by_r13.items():
+        if hessian_policy == "identity" or not have_local_support:
+            h2_by_r13[r13] = identity_h2
+            evidence_by_r13[f"R{r13}"] = {
+                "h2_shrinkage_policy": "identity_fallback",
+                "h2_local_alpha": 0.0,
+                "h2_effective_sample_size": 0.0,
+                "h2_identity_scale": 1.0,
+            }
+            continue
+        _, h2, evidence = build_expert_hessians(
+            inputs[fit_mask_device],
+            gates[fit_mask_device],
+            middle[fit_mask_device],
+            global_h13=global_h13,
+            global_h2=global_h2,
+            device=device,
+        )
+        h2_by_r13[r13] = h2
+        evidence_by_r13[f"R{r13}"] = evidence
+    return h2_by_r13, evidence_by_r13
+
+
+def _coupled_logical_block_contexts(
+    mode_id: int,
+    groups: int,
+) -> torch.Tensor:
+    """Label coupled coordinates without changing their LDLQ traversal order."""
+
+    context_count = resolve_mode(mode_id).context_count
+    if groups <= 0 or groups % context_count:
+        raise ValueError("coupled contexts do not divide neuron groups")
+    return torch.arange(context_count, dtype=torch.long).repeat_interleave(
+        groups // context_count
+    )
+
+
+def _quantize_conditional_down_grid(
+    source_w2_by_expert: Sequence[torch.Tensor],
+    block_contexts_by_expert: Sequence[torch.Tensor],
+    mode_specs_by_expert: Sequence[tuple[object, ...]],
+    h2_by_r13_by_expert: Sequence[dict[int, torch.Tensor]],
+    *,
+    layer: int,
+    device: torch.device,
+    shared_scale_scope: object,
+    quantizer_module: object | None,
+    codebook: str,
+    layout: SearchLayout,
+    ldlq_tf32: bool,
+    tailbite_context: int,
+    transform_seeds_by_expert: Sequence[
+        Mapping[str, QSRTTransformSeeds]
+    ] | None = None,
+) -> list[dict[int, dict[int, QSRTMatrixCandidate]]]:
+    """Encode the exhaustive ``w2[r13, r2]`` grid in flattened GPU batches."""
+
+    count = len(source_w2_by_expert)
+    if not (
+        count
+        == len(block_contexts_by_expert)
+        == len(mode_specs_by_expert)
+        == len(h2_by_r13_by_expert)
+    ):
+        raise ValueError("conditional down-grid inputs have inconsistent lengths")
+    if transform_seeds_by_expert is None:
+        transform_seeds_by_expert = [{} for _ in range(count)]
+    elif len(transform_seeds_by_expert) != count:
+        raise ValueError("conditional down-grid transform seeds have the wrong length")
+
+    flat_keys: list[tuple[int, int]] = []
+    flat_sources: list[dict[str, torch.Tensor]] = []
+    flat_contexts: list[torch.Tensor] = []
+    flat_modes: list[tuple[object, ...]] = []
+    flat_hessians: list[dict[str, torch.Tensor]] = []
+    flat_transform_seeds: list[Mapping[str, QSRTTransformSeeds]] = []
+    for expert_index, (source, contexts, modes, h2_by_r13) in enumerate(
+        zip(
+            source_w2_by_expert,
+            block_contexts_by_expert,
+            mode_specs_by_expert,
+            h2_by_r13_by_expert,
+        )
+    ):
+        mode_ids = tuple(int(mode.mode_id) for mode in modes)
+        h2_bases = tuple(h2_by_r13)
+        if h2_bases != mode_ids:
+            raise ValueError("conditional H2 modes differ from upstream modes")
+        for r13 in h2_bases:
+            flat_keys.append((expert_index, r13))
+            flat_sources.append({"w2": source})
+            flat_contexts.append(contexts)
+            flat_modes.append(modes)
+            flat_hessians.append({"w2": h2_by_r13[r13]})
+            flat_transform_seeds.append(transform_seeds_by_expert[expert_index])
+
+    guarded_reuse = layout == "qsrt_guarded_reuse"
+    if guarded_reuse:
+        flat_candidates = quantize_qsrt_matrix_expert_batch(
+            flat_sources,
+            flat_contexts,
+            modes_by_expert=[(modes[0],) for modes in flat_modes],
+            hessians_by_expert=flat_hessians,
+            layer=layer,
+            device=device,
+            shared_scale_scope=shared_scale_scope,
+            quantizer_module=quantizer_module,
+            codebook=codebook,
+            layout="importance_ordered",
+            ldlq_tf32=ldlq_tf32,
+            tailbite_context=tailbite_context,
+            transform_seeds_by_expert=flat_transform_seeds,
+        )
+        shifted_indices = [
+            index for index, modes in enumerate(flat_modes) if len(modes) > 1
+        ]
+        if shifted_indices:
+            shifted = quantize_qsrt_matrix_expert_batch(
+                [flat_sources[index] for index in shifted_indices],
+                [flat_contexts[index] for index in shifted_indices],
+                modes_by_expert=[flat_modes[index][1:] for index in shifted_indices],
+                hessians_by_expert=[flat_hessians[index] for index in shifted_indices],
+                layer=layer,
+                device=device,
+                shared_scale_scope=shared_scale_scope,
+                quantizer_module=quantizer_module,
+                codebook=codebook,
+                layout="qsrt_pair_prefix_reuse",
+                ldlq_tf32=ldlq_tf32,
+                tailbite_context=tailbite_context,
+                transform_seeds_by_expert=[
+                    flat_transform_seeds[index] for index in shifted_indices
+                ],
+            )
+            for index, extra in zip(shifted_indices, shifted):
+                flat_candidates[index]["w2"].update(extra["w2"])
+    else:
+        flat_candidates = quantize_qsrt_matrix_expert_batch(
+            flat_sources,
+            flat_contexts,
+            modes_by_expert=flat_modes,
+            hessians_by_expert=flat_hessians,
+            layer=layer,
+            device=device,
+            shared_scale_scope=shared_scale_scope,
+            quantizer_module=quantizer_module,
+            codebook=codebook,
+            layout=layout,
+            ldlq_tf32=ldlq_tf32,
+            tailbite_context=tailbite_context,
+            transform_seeds_by_expert=flat_transform_seeds,
+        )
+
+    result: list[dict[int, dict[int, QSRTMatrixCandidate]]] = [
+        {} for _ in range(count)
+    ]
+    for (expert_index, r13), candidates in zip(flat_keys, flat_candidates):
+        result[expert_index][r13] = candidates["w2"]
+    return result
+
+
+def encode_phase1_expert(
+    store: MatrixStore,
+    samples: LayerSamples,
+    *,
+    layer: int,
+    expert: int,
+    partition: RequestPartition,
+    global_h13: torch.Tensor,
+    global_h2: torch.Tensor,
+    fallback_block_contexts: torch.Tensor,
+    fallback_block_scores: torch.Tensor,
+    device: torch.device,
+    shared_scale_scope: object,
+    mode_ids: Sequence[int] = PHASE1_MODE_IDS,
+    quantizer_module: object | None = None,
+    min_fit_documents: int = 6,
+    min_confirmation_documents: int = 4,
+    bootstrap_replicates: int = 2_000,
+    minimum_improvement: float = 0.0,
+    logical_trellis_schema: str = SCHEMA,
+    codebook: str = CODEBOOK_SQG_XOR_CHEB_T12,
+    layout: SearchLayout = "importance_ordered",
+    ldlq_tf32: bool = False,
+    tailbite_context: int = 128,
+    hessian_policy: HessianPolicy = "captured_blend",
+    permutation_policy: PermutationPolicy = "h2_reverse",
+    folded_scale_power: float = 0.0,
+    transform_seeds: Mapping[str, QSRTTransformSeeds] | None = None,
+    activation: ExpertActivation = "situ",
+) -> QSRTCandidateEncoding:
+    """Encode and conservatively select one official-source expert.
+
+    The official model is never instantiated.  ``w1`` and ``w3`` are streamed
+    together so their same-shape R candidates can share one batched traversal;
+    they are released before ``w2`` is loaded.
+    """
+
+    if device.type != "cuda":
+        raise ValueError("QSRT candidate encoding requires CUDA")
+    if not 1 <= layer <= 92 or not 0 <= expert < 896:
+        raise ValueError("invalid K3 layer/expert assignment")
+    modes = tuple(int(mode) for mode in mode_ids)
+    fixed_k3 = modes == (R0.mode_id,)
+    fixed_high_rate = modes == (H308.mode_id,)
+    fixed_k2 = modes == (K2.mode_id,)
+    fixed_profile = fixed_k3 or fixed_high_rate or fixed_k2
+    if (
+        not modes
+        or len(set(modes)) != len(modes)
+        or (not fixed_profile and modes[0] != 0)
+    ):
+        raise ValueError(
+            "mode_ids must be H308/K2 alone or unique and begin with R0"
+        )
+    if any(mode not in (0, 1, 2, H308.mode_id, K2.mode_id) for mode in modes):
+        raise ValueError("mode_ids contain an unsupported schedule")
+    if hessian_policy not in HESSIAN_POLICIES:
+        raise ValueError(f"unsupported Hessian policy: {hessian_policy}")
+    if permutation_policy not in PERMUTATION_POLICIES:
+        raise ValueError(f"unsupported permutation policy: {permutation_policy}")
+
+    all_rows = select_expert_rows(samples, expert, partition.all)
+    fit_mask = _split_rows(all_rows, partition.fit)
+    confirmation_mask = _split_rows(all_rows, partition.confirmation)
+    fit_documents = int(torch.unique(all_rows.request_steps[fit_mask]).numel())
+    confirmation_documents = int(
+        torch.unique(all_rows.request_steps[confirmation_mask]).numel()
+    )
+    inputs = all_rows.inputs.float().to(device=device)
+    gates = all_rows.gates.float().to(device=device)
+
+    source_w1 = _load_source_matrix(store, layer, expert, "w1", device)
+    source_w3 = _load_source_matrix(store, layer, expert, "w3", device)
+    if all_rows.rows:
+        source_gate = F.linear(inputs, source_w1)
+        source_up = F.linear(inputs, source_w3)
+        source_middle = _source_middle(
+            source_gate,
+            source_up,
+            activation=activation,
+        )
+        del source_gate, source_up
+    else:
+        source_middle = torch.empty(
+            (0, global_h2.shape[0]), dtype=torch.float32, device=device
+        )
+
+    have_local_support = fit_documents >= min_fit_documents
+    fit_mask_device = fit_mask.to(device=device)
+    if have_local_support:
+        block_contexts, block_scores = permutation_policy_contexts(
+            source_middle[fit_mask_device],
+            gates[fit_mask_device],
+            policy=permutation_policy,
+        )
+        h13 = global_h13
+        covariance: dict[str, object] = {
+            "gate_square_sum": float(
+                gates[fit_mask_device].square().double().sum()
+            ),
+            "h13_local_alpha": 0.0,
+        }
+        context_basis = (
+            f"official_source_{_activation_basis(activation)}_fit_documents:"
+            f"{permutation_policy}"
+        )
+        covariance = {
+            **covariance,
+            "fit_rows": int(fit_mask.sum()),
+            "fit_documents": fit_documents,
+        }
+    else:
+        block_contexts = fallback_block_contexts.clone()
+        block_scores = fallback_block_scores.clone()
+        h13 = global_h13
+        context_basis = "fallback_context_order_no_local_support"
+        covariance = {
+            "fit_rows": int(fit_mask.sum()),
+            "fit_documents": fit_documents,
+            "gate_square_sum": float(gates[fit_mask_device].square().sum()),
+            "h13_local_alpha": 0.0,
+        }
+
+    if fixed_k2:
+        block_contexts = torch.zeros_like(block_contexts)
+        context_basis = "uniform_k2_identity_transformed_coordinate_order"
+    can_confirm = confirmation_documents >= min_confirmation_documents
+    evaluated_modes = (
+        modes if fixed_profile or (have_local_support and can_confirm) else (0,)
+    )
+    evaluated_formats = tuple(
+        (r13, r2) for r13 in evaluated_modes for r2 in evaluated_modes
+    )
+    block_contexts_device = block_contexts.to(device=device, dtype=torch.long)
+    intermediate_conditioning = _folded_intermediate_conditioning(
+        block_scores.to(device=device),
+        block_contexts_device,
+        power=folded_scale_power,
+    )
+    if intermediate_conditioning is not None:
+        covariance["folded_scale_power"] = float(folded_scale_power)
+        covariance["folded_scale_min"] = float(intermediate_conditioning.min())
+        covariance["folded_scale_max"] = float(intermediate_conditioning.max())
+        context_basis += "+folded_128_channel_conditioning"
+
+    mode_specs = tuple(resolve_mode(mode) for mode in evaluated_modes)
+    guarded_reuse = layout == "qsrt_guarded_reuse"
+    if guarded_reuse and any(mode > 2 for mode in evaluated_modes):
+        raise ValueError("qsrt_guarded_reuse supports only R0 through R2")
+    upstream_layout: Layout = "importance_ordered" if guarded_reuse else layout
+    upstream_candidates = quantize_qsrt_matrix_candidates_batched(
+        {"w1": source_w1, "w3": source_w3},
+        block_contexts_device,
+        modes=mode_specs,
+        hessians={"w1": h13, "w3": h13},
+        layer=layer,
+        device=device,
+        shared_scale_scope=shared_scale_scope,
+        quantizer_module=quantizer_module,
+        codebook=codebook,
+        layout=upstream_layout,
+        ldlq_tf32=ldlq_tf32,
+        tailbite_context=tailbite_context,
+        transform_seeds=transform_seeds,
+        intermediate_conditioning=intermediate_conditioning,
+    )
+    del source_w1, source_w3
+
+    middle_by_r13 = _candidate_middle_by_r13(
+        inputs,
+        upstream_candidates,
+        evaluated_modes,
+        activation=activation,
+    )
+    h2_by_r13, h2_evidence = _conditional_h2_by_r13(
+        inputs,
+        gates,
+        fit_mask_device,
+        middle_by_r13,
+        global_h13=global_h13,
+        global_h2=global_h2,
+        device=device,
+        hessian_policy=hessian_policy,
+        have_local_support=have_local_support,
+    )
+    covariance.update(
+        {
+            "basis": (
+                "identity"
+                if hessian_policy == "identity"
+                else (
+                    f"decoded_candidate_{_activation_basis(activation)}_"
+                    "adaptive_shrinkage"
+                    if have_local_support
+                    else "identity_support_fallback"
+                )
+            ),
+            "h2_by_r13": h2_evidence,
+        }
+    )
+    context_basis += f"+decoded_candidate_{_activation_basis(activation)}_h2"
+
+    source_w2 = _load_source_matrix(store, layer, expert, "w2", device)
+    if all_rows.rows:
+        reference_output = F.linear(source_middle, source_w2)
+    else:
+        reference_output = torch.empty(
+            (0, global_h13.shape[0]), dtype=torch.float32, device=device
+        )
+    down_candidates = _quantize_conditional_down_grid(
+        [source_w2],
+        [block_contexts_device],
+        [mode_specs],
+        [h2_by_r13],
+        layer=layer,
+        device=device,
+        shared_scale_scope=shared_scale_scope,
+        quantizer_module=quantizer_module,
+        codebook=codebook,
+        layout=layout,
+        ldlq_tf32=ldlq_tf32,
+        tailbite_context=tailbite_context,
+        transform_seeds_by_expert=[transform_seeds or {}],
+    )[0]
+    fit_sse: dict[tuple[int, int], torch.Tensor] = {}
+    confirmation_sse: dict[tuple[int, int], torch.Tensor] = {}
+    fit_reference = fit_counts = confirmation_reference = confirmation_counts = None
+    mode_coding: dict[tuple[int, int], dict[str, object]] = {}
+    for rate_pair in evaluated_formats:
+        r13, r2 = rate_pair
+        candidate_output = F.linear(
+            middle_by_r13[r13], down_candidates[r13][r2].reconstruction
+        )
+        fit_metric = (
+            functional_sse_by_request(
+                reference_output[fit_mask.to(device=device)],
+                candidate_output[fit_mask.to(device=device)],
+                gates[fit_mask.to(device=device)],
+                all_rows.request_steps[fit_mask],
+                partition.fit,
+            )
+            if int(fit_mask.sum())
+            else _empty_metrics(partition.fit)
+        )
+        confirmation_metric = (
+            functional_sse_by_request(
+                reference_output[confirmation_mask.to(device=device)],
+                candidate_output[confirmation_mask.to(device=device)],
+                gates[confirmation_mask.to(device=device)],
+                all_rows.request_steps[confirmation_mask],
+                partition.confirmation,
+            )
+            if int(confirmation_mask.sum())
+            else _empty_metrics(partition.confirmation)
+        )
+        fit_sse[rate_pair] = fit_metric[0]
+        confirmation_sse[rate_pair] = confirmation_metric[0]
+        if fit_reference is None:
+            fit_reference, fit_counts = fit_metric[1], fit_metric[2]
+            confirmation_reference, confirmation_counts = (
+                confirmation_metric[1],
+                confirmation_metric[2],
+            )
+        elif not (
+            torch.equal(fit_reference, fit_metric[1])
+            and torch.equal(fit_counts, fit_metric[2])
+            and torch.equal(confirmation_reference, confirmation_metric[1])
+            and torch.equal(confirmation_counts, confirmation_metric[2])
+        ):
+            raise ValueError("mode candidates changed metric support or reference energy")
+        mode_coding[rate_pair] = _candidate_mode_evidence(
+            {
+                "w1": upstream_candidates["w1"][r13],
+                "w3": upstream_candidates["w3"][r13],
+                "w2": down_candidates[r13][r2],
+            }
+        )
+        del candidate_output
+
+    del source_w2
+
+    assert fit_reference is not None and fit_counts is not None
+    assert confirmation_reference is not None and confirmation_counts is not None
+    if fixed_profile:
+        fixed_mode = H308 if fixed_high_rate else (K2 if fixed_k2 else R0)
+        selection = RatePairSelection(
+            proposed_r13=fixed_mode.mode_id,
+            proposed_r2=fixed_mode.mode_id,
+            selected_r13=fixed_mode.mode_id,
+            selected_r2=fixed_mode.mode_id,
+            accepted=True,
+            reason=(
+                "fixed_22k3_2k4_high_rate_profile"
+                if fixed_high_rate
+                else (
+                    "fixed_uniform_k2_profile"
+                    if fixed_k2
+                    else "fixed_uniform_k3_profile"
+                )
+            ),
+            fit_documents=int(torch.count_nonzero(fit_counts)),
+            confirmation_documents=int(torch.count_nonzero(confirmation_counts)),
+            confirmation_relative_improvement=None,
+            confirmation_ci95=(None, None),
+            bootstrap_replicates_valid=0,
+        )
+    else:
+        selection = select_phase1_rate_pair(
+            fit_sse,
+            confirmation_sse,
+            fit_counts=fit_counts,
+            confirmation_counts=confirmation_counts,
+            modes=evaluated_formats,
+            min_fit_documents=min_fit_documents,
+            min_confirmation_documents=min_confirmation_documents,
+            minimum_improvement=minimum_improvement,
+            bootstrap_replicates=bootstrap_replicates,
+            seed=deterministic_expert_seed(layer, expert),
+        )
+    selected_r13, selected_r2 = selection.selected
+    selected_candidates = {
+        "w1": upstream_candidates["w1"][selected_r13],
+        "w3": upstream_candidates["w3"][selected_r13],
+        "w2": down_candidates[selected_r13][selected_r2],
+    }
+    selected_encodings = {
+        matrix: finalize_qsrt_matrix_candidate(
+            value,
+            layer=layer,
+            logical_trellis_schema=logical_trellis_schema,
+            codebook=codebook,
+            tailbite_context=tailbite_context,
+        )
+        for matrix, value in selected_candidates.items()
+    }
+    physical_permutation = selected_encodings["w1"].plan.physical_permutation
+    for matrix in ("w3", "w2"):
+        if not torch.equal(
+            selected_encodings[matrix].plan.physical_permutation,
+            physical_permutation,
+        ):
+            raise ValueError("selected matrices do not share one physical permutation")
+    expert_encoding = QSRTExpertEncoding(
+        format=ExpertFormatSpec.compressed(selected_r13, selected_r2),
+        matrices=selected_encodings,
+        physical_permutation=physical_permutation,
+        coding={
+            "format": ExpertFormatSpec.compressed(
+                selected_r13, selected_r2
+            ).name,
+            "r13": selected_r13,
+            "r2": selected_r2,
+            "codebook": codebook,
+            "tailbite_context": int(tailbite_context),
+            **_mode_evidence(selected_encodings),
+        },
+    )
+    return QSRTCandidateEncoding(
+        expert=expert_encoding,
+        selection=selection,
+        block_contexts=block_contexts.detach().cpu().to(torch.uint8).contiguous(),
+        block_scores=block_scores.detach().cpu().float().contiguous(),
+        context_basis=context_basis,
+        covariance=covariance,
+        evaluated_modes=evaluated_modes,
+        evaluated_formats=evaluated_formats,
+        fit_sse=fit_sse,
+        confirmation_sse=confirmation_sse,
+        fit_reference_energy=fit_reference,
+        confirmation_reference_energy=confirmation_reference,
+        fit_counts=fit_counts,
+        confirmation_counts=confirmation_counts,
+        mode_coding=mode_coding,
+    )
+
+
+def encode_phase1_expert_batch(
+    store: MatrixStore,
+    samples: LayerSamples,
+    *,
+    layer: int,
+    experts: Sequence[int],
+    partition: RequestPartition,
+    global_h13: torch.Tensor,
+    global_h2: torch.Tensor,
+    fallback_block_contexts: torch.Tensor,
+    fallback_block_scores: torch.Tensor,
+    device: torch.device,
+    shared_scale_scope: object,
+    mode_ids: Sequence[int] = PHASE1_MODE_IDS,
+    quantizer_module: object | None = None,
+    min_fit_documents: int = 6,
+    min_confirmation_documents: int = 4,
+    bootstrap_replicates: int = 2_000,
+    minimum_improvement: float = 0.0,
+    logical_trellis_schema: str = SCHEMA,
+    codebook: str = CODEBOOK_SQG_XOR_CHEB_T12,
+    layout: SearchLayout = "importance_ordered",
+    ldlq_tf32: bool = False,
+    tailbite_context: int = 128,
+    hessian_policy: HessianPolicy = "captured_blend",
+    permutation_policy: PermutationPolicy = "h2_reverse",
+    folded_scale_power: float = 0.0,
+    transform_seeds_by_expert: Mapping[
+        int, Mapping[str, QSRTTransformSeeds]
+    ] | None = None,
+    coupled_intermediate_draws_by_expert: Mapping[int, int] | None = None,
+    activation: ExpertActivation = "situ",
+) -> list[QSRTCandidateEncoding]:
+    """Encode several experts while batching their independent trellis work.
+
+    This preserves the single-expert calibration, dense-H, candidate scoring,
+    and selection contracts.  Only the execution schedule changes: all
+    same-shape w1/w3 projections in the chunk are submitted together, followed
+    by all w2 projections.  The larger tile batches fill SM120 much more
+    effectively than one expert's R1/R2 replacement records.
+    """
+
+    expert_ids = tuple(int(expert) for expert in experts)
+    if not expert_ids or len(set(expert_ids)) != len(expert_ids):
+        raise ValueError("expert batch must be nonempty and unique")
+    if device.type != "cuda":
+        raise ValueError("QSRT candidate encoding requires CUDA")
+    if not 1 <= layer <= 92 or any(not 0 <= expert < 896 for expert in expert_ids):
+        raise ValueError("invalid K3 layer/expert assignment")
+    modes = tuple(int(mode) for mode in mode_ids)
+    fixed_k3 = modes == (R0.mode_id,)
+    fixed_high_rate = modes == (H308.mode_id,)
+    fixed_k2 = modes == (K2.mode_id,)
+    fixed_profile = fixed_k3 or fixed_high_rate or fixed_k2
+    coupled = coupled_intermediate_draws_by_expert is not None
+    _activation_basis(activation)
+    if coupled:
+        if modes not in (
+            (R0.mode_id,),
+            (H308.mode_id,),
+            (K2.mode_id,),
+        ):
+            raise ValueError(
+                "coupled activation-boundary transforms currently support "
+                "only a fixed K2, K3, or 22-K3/2-K4 schedule"
+            )
+        if folded_scale_power != 0.0:
+            raise ValueError(
+                "coupled activation-boundary transforms require "
+                "folded_scale_power zero"
+            )
+        missing = set(expert_ids) - set(coupled_intermediate_draws_by_expert)
+        extra = set(coupled_intermediate_draws_by_expert) - set(expert_ids)
+        if missing or extra:
+            raise ValueError("coupled draw map must cover exactly the expert batch")
+    if (
+        not modes
+        or len(set(modes)) != len(modes)
+        or (not fixed_profile and modes[0] != 0)
+    ):
+        raise ValueError(
+            "mode_ids must be H308/K2 alone or unique and begin with R0"
+        )
+    if any(mode not in (0, 1, 2, H308.mode_id, K2.mode_id) for mode in modes):
+        raise ValueError("mode_ids contain an unsupported schedule")
+    if hessian_policy not in HESSIAN_POLICIES:
+        raise ValueError(f"unsupported Hessian policy: {hessian_policy}")
+    if permutation_policy not in PERMUTATION_POLICIES:
+        raise ValueError(f"unsupported permutation policy: {permutation_policy}")
+
+    prepared: list[dict[str, object]] = []
+    upstream_sources: list[dict[str, torch.Tensor]] = []
+    upstream_contexts: list[torch.Tensor] = []
+    upstream_modes: list[tuple[object, ...]] = []
+    upstream_hessians: list[dict[str, torch.Tensor]] = []
+    upstream_conditioning: list[torch.Tensor | None] = []
+    seed_maps = [
+        (transform_seeds_by_expert or {}).get(expert, {}) for expert in expert_ids
+    ]
+
+    for expert in expert_ids:
+        all_rows = select_expert_rows(samples, expert, partition.all)
+        fit_mask = _split_rows(all_rows, partition.fit)
+        confirmation_mask = _split_rows(all_rows, partition.confirmation)
+        fit_documents = int(torch.unique(all_rows.request_steps[fit_mask]).numel())
+        confirmation_documents = int(
+            torch.unique(all_rows.request_steps[confirmation_mask]).numel()
+        )
+        inputs = all_rows.inputs.float().to(device=device)
+        gates = all_rows.gates.float().to(device=device)
+
+        source_w1 = _load_source_matrix(store, layer, expert, "w1", device)
+        source_w3 = _load_source_matrix(store, layer, expert, "w3", device)
+        if all_rows.rows:
+            source_gate = F.linear(inputs, source_w1)
+            source_up = F.linear(inputs, source_w3)
+            source_middle = _source_middle(
+                source_gate,
+                source_up,
+                activation=activation,
+            )
+            del source_gate, source_up
+        else:
+            source_middle = torch.empty(
+                (0, global_h2.shape[0]), dtype=torch.float32, device=device
+            )
+
+        have_local_support = fit_documents >= min_fit_documents
+        fit_mask_device = fit_mask.to(device=device)
+        if bool(fit_mask.any()):
+            block_contexts, block_scores = permutation_policy_contexts(
+                source_middle[fit_mask_device],
+                gates[fit_mask_device],
+                policy=permutation_policy,
+            )
+            context_basis = (
+                f"official_source_{_activation_basis(activation)}_fit_documents:"
+                f"{permutation_policy}"
+            )
+        else:
+            block_contexts = fallback_block_contexts.clone()
+            block_scores = fallback_block_scores.clone()
+            context_basis = "identity_logical_order_no_routed_fit_rows"
+
+        if have_local_support:
+            h13 = global_h13
+            covariance: dict[str, object] = {
+                "gate_square_sum": float(
+                    gates[fit_mask_device].square().double().sum()
+                ),
+                "h13_local_alpha": 0.0,
+            }
+            covariance = {
+                **covariance,
+                "fit_rows": int(fit_mask.sum()),
+                "fit_documents": fit_documents,
+            }
+        else:
+            h13 = global_h13
+            covariance = {
+                "fit_rows": int(fit_mask.sum()),
+                "fit_documents": fit_documents,
+                "gate_square_sum": float(
+                    gates[fit_mask.to(device=device)].square().sum()
+                ),
+                "h13_local_alpha": 0.0,
+            }
+
+        if coupled:
+            # The coupled preactivation transform is not coordinatewise, so an
+            # arbitrary neuron permutation cannot be pushed through it. Keep
+            # LDLQ in the transformed logical order. K2 has one context; R0 K3
+            # has 24 monotonically labelled equal-size contexts, whose stable
+            # importance order is the identity permutation.
+            context_count = resolve_mode(modes[0]).context_count
+            if any(
+                resolve_mode(mode).context_count != context_count
+                for mode in modes
+            ):
+                raise ValueError("coupled rate candidates change context cardinality")
+            block_contexts = _coupled_logical_block_contexts(
+                modes[0], block_contexts.numel()
+            )
+            context_basis = "coupled_transform_logical_coordinate_order"
+        elif fixed_k2:
+            block_contexts = torch.zeros_like(block_contexts)
+            context_basis = "uniform_k2_identity_coordinate_order"
+        can_confirm = confirmation_documents >= min_confirmation_documents
+        evaluated_modes = (
+            modes if fixed_profile or (have_local_support and can_confirm) else (0,)
+        )
+        mode_specs = tuple(resolve_mode(mode) for mode in evaluated_modes)
+        contexts_device = block_contexts.to(device=device, dtype=torch.long)
+        coupled_execution_basis: CoupledHadamardExecution | None = None
+        coupled_source_w2: torch.Tensor | None = None
+        coupled_reference_output: torch.Tensor | None = None
+        if coupled:
+            draw = int(coupled_intermediate_draws_by_expert[expert])
+            spec = CoupledHadamardSpec(
+                intermediate_draw=draw,
+                activation=activation,
+            )
+            source_w2 = _load_source_matrix(store, layer, expert, "w2", device)
+            transformed = encode_coupled_weights(
+                (source_w1, source_w3, source_w2), spec
+            )
+            coupled_execution_basis = coupled_execution(transformed, spec)
+            coupled_source_w2 = transformed[2]
+            coupled_reference_output = (
+                F.linear(source_middle, source_w2)
+                if all_rows.rows
+                else torch.empty(
+                    (0, global_h13.shape[0]),
+                    dtype=torch.float32,
+                    device=device,
+                )
+            )
+            source_w1, source_w3 = transformed[:2]
+            inputs = coupled_execution_basis.transform_inputs(inputs)
+            h13 = coupled_execution_basis.transform_h13(h13)
+            transformed_global_h2 = coupled_execution_basis.transform_h2(global_h2)
+            context_basis += f"+coupled_hadamard_intermediate_draw_{draw}"
+        else:
+            transformed_global_h2 = global_h2
+        intermediate_conditioning = _folded_intermediate_conditioning(
+            block_scores.to(device=device),
+            contexts_device,
+            power=folded_scale_power,
+        )
+        if intermediate_conditioning is not None:
+            covariance["folded_scale_power"] = float(folded_scale_power)
+            covariance["folded_scale_min"] = float(
+                intermediate_conditioning.min()
+            )
+            covariance["folded_scale_max"] = float(
+                intermediate_conditioning.max()
+            )
+            context_basis += "+folded_128_channel_conditioning"
+        prepared.append(
+            {
+                "expert": expert,
+                "all_rows": all_rows,
+                "fit_mask": fit_mask,
+                "confirmation_mask": confirmation_mask,
+                "inputs": inputs,
+                "gates": gates,
+                "source_middle": source_middle,
+                "block_contexts": block_contexts,
+                "block_scores": block_scores,
+                "context_basis": context_basis,
+                "covariance": covariance,
+                "evaluated_modes": evaluated_modes,
+                "evaluated_formats": tuple(
+                    (r13, r2)
+                    for r13 in evaluated_modes
+                    for r2 in evaluated_modes
+                ),
+                "h13": h13,
+                "global_h2_basis": transformed_global_h2,
+                "coupled_execution": coupled_execution_basis,
+                "coupled_source_w2": coupled_source_w2,
+                "coupled_reference_output": coupled_reference_output,
+            }
+        )
+        upstream_sources.append({"w1": source_w1, "w3": source_w3})
+        upstream_contexts.append(contexts_device)
+        upstream_modes.append(mode_specs)
+        upstream_hessians.append({"w1": h13, "w3": h13})
+        upstream_conditioning.append(intermediate_conditioning)
+
+    guarded_reuse = layout == "qsrt_guarded_reuse"
+    if guarded_reuse and any(mode > 2 for mode in modes):
+        raise ValueError("qsrt_guarded_reuse supports only R0 through R2")
+    upstream_layout: Layout = "importance_ordered" if guarded_reuse else layout
+    upstream_candidates = quantize_qsrt_matrix_expert_batch(
+        upstream_sources,
+        upstream_contexts,
+        modes_by_expert=upstream_modes,
+        hessians_by_expert=upstream_hessians,
+        layer=layer,
+        device=device,
+        shared_scale_scope=shared_scale_scope,
+        quantizer_module=quantizer_module,
+        codebook=codebook,
+        layout=upstream_layout,
+        ldlq_tf32=ldlq_tf32,
+        tailbite_context=tailbite_context,
+        transform_seeds_by_expert=seed_maps,
+        intermediate_conditioning_by_expert=upstream_conditioning,
+    )
+    _debug_assert_candidate_tree_closure(
+        upstream_candidates,
+        expert_ids,
+        layer=layer,
+        stage="post-upstream-encode",
+    )
+    del upstream_sources, upstream_hessians
+
+    down_sources: list[torch.Tensor] = []
+    conditional_h2s: list[dict[int, torch.Tensor]] = []
+    for item, upstream in zip(prepared, upstream_candidates):
+        expert = int(item["expert"])
+        source_w2 = item["coupled_source_w2"]
+        if source_w2 is None:
+            source_w2 = _load_source_matrix(store, layer, expert, "w2", device)
+        source_middle = item["source_middle"]
+        inputs = item["inputs"]
+        gates = item["gates"]
+        fit_mask = item["fit_mask"]
+        evaluated_modes = item["evaluated_modes"]
+        assert isinstance(source_middle, torch.Tensor)
+        assert isinstance(inputs, torch.Tensor)
+        assert isinstance(gates, torch.Tensor)
+        assert isinstance(fit_mask, torch.Tensor)
+        execution_basis = item["coupled_execution"]
+        if execution_basis is None:
+            middle_by_r13 = _candidate_middle_by_r13(
+                inputs,
+                upstream,
+                evaluated_modes,
+                activation=activation,
+            )
+        else:
+            middle_by_r13 = {
+                int(r13): execution_basis.decode_middle(
+                    inputs,
+                    upstream["w1"][int(r13)].reconstruction,
+                    upstream["w3"][int(r13)].reconstruction,
+                )
+                for r13 in evaluated_modes
+            }
+        h2_by_r13, h2_evidence = _conditional_h2_by_r13(
+            inputs,
+            gates,
+            fit_mask.to(device=device),
+            middle_by_r13,
+            global_h13=item["h13"],
+            global_h2=item["global_h2_basis"],
+            device=device,
+            hessian_policy=hessian_policy,
+            have_local_support=(
+                int(item["covariance"]["fit_documents"]) >= min_fit_documents
+            ),
+        )
+        item["middle_by_r13"] = middle_by_r13
+        item["covariance"].update(
+            {
+                "basis": (
+                    "identity"
+                    if hessian_policy == "identity"
+                    else (
+                        f"decoded_candidate_{_activation_basis(activation)}_"
+                        "adaptive_shrinkage"
+                        if int(item["covariance"]["fit_documents"])
+                        >= min_fit_documents
+                        else "identity_support_fallback"
+                    )
+                ),
+                "h2_by_r13": h2_evidence,
+            }
+        )
+        item["context_basis"] = (
+            str(item["context_basis"])
+            + f"+decoded_candidate_{_activation_basis(activation)}_h2"
+        )
+        all_rows = item["all_rows"]
+        if item["coupled_reference_output"] is not None:
+            reference_output = item["coupled_reference_output"]
+        elif all_rows.rows:
+            reference_output = F.linear(source_middle, source_w2)
+        else:
+            reference_output = torch.empty(
+                (0, global_h13.shape[0]), dtype=torch.float32, device=device
+            )
+        item["reference_output"] = reference_output
+        down_sources.append(source_w2)
+        conditional_h2s.append(h2_by_r13)
+
+    down_candidates = _quantize_conditional_down_grid(
+        down_sources,
+        upstream_contexts,
+        upstream_modes,
+        conditional_h2s,
+        layer=layer,
+        device=device,
+        shared_scale_scope=shared_scale_scope,
+        quantizer_module=quantizer_module,
+        codebook=codebook,
+        layout=layout,
+        ldlq_tf32=ldlq_tf32,
+        tailbite_context=tailbite_context,
+        transform_seeds_by_expert=seed_maps,
+    )
+    _debug_assert_candidate_tree_closure(
+        upstream_candidates,
+        expert_ids,
+        layer=layer,
+        stage="post-down-encode upstream recheck",
+    )
+    _debug_assert_candidate_tree_closure(
+        down_candidates,
+        expert_ids,
+        layer=layer,
+        stage="post-down-encode",
+    )
+    del down_sources, conditional_h2s
+
+    results: list[QSRTCandidateEncoding] = []
+    for item, upstream, down in zip(
+        prepared, upstream_candidates, down_candidates
+    ):
+        expert = int(item["expert"])
+        all_rows = item["all_rows"]
+        fit_mask = item["fit_mask"]
+        confirmation_mask = item["confirmation_mask"]
+        inputs = item["inputs"]
+        gates = item["gates"]
+        source_middle = item["source_middle"]
+        reference_output = item["reference_output"]
+        evaluated_modes = item["evaluated_modes"]
+        evaluated_formats = item["evaluated_formats"]
+        assert isinstance(inputs, torch.Tensor)
+        assert isinstance(gates, torch.Tensor)
+        assert isinstance(source_middle, torch.Tensor)
+        assert isinstance(reference_output, torch.Tensor)
+
+        fit_plan = _prepare_deferred_functional_rows(
+            reference_output,
+            gates,
+            all_rows.request_steps,
+            fit_mask,
+            partition.fit,
+            device=device,
+        )
+        confirmation_plan = _prepare_deferred_functional_rows(
+            reference_output,
+            gates,
+            all_rows.request_steps,
+            confirmation_mask,
+            partition.confirmation,
+            device=device,
+        )
+        deferred_fit_sse: list[torch.Tensor] = []
+        deferred_confirmation_sse: list[torch.Tensor] = []
+        format_order: list[tuple[int, int]] = []
+        mode_coding: dict[tuple[int, int], dict[str, object]] = {}
+        middle_by_r13 = item["middle_by_r13"]
+        for rate_pair in evaluated_formats:
+            r13, r2 = rate_pair
+            candidate_output = F.linear(
+                middle_by_r13[r13], down[r13][r2].reconstruction
+            )
+            execution_basis = item["coupled_execution"]
+            if execution_basis is not None:
+                candidate_output = execution_basis.decode_output(candidate_output)
+            deferred_fit_sse.append(
+                _defer_functional_row_sse(fit_plan, candidate_output)
+            )
+            deferred_confirmation_sse.append(
+                _defer_functional_row_sse(confirmation_plan, candidate_output)
+            )
+            format_order.append(rate_pair)
+            mode_coding[rate_pair] = _candidate_mode_evidence(
+                {
+                    "w1": upstream["w1"][r13],
+                    "w3": upstream["w3"][r13],
+                    "w2": down[r13][r2],
+                }
+            )
+            del candidate_output
+
+        fit_sse = _finish_deferred_functional_sse(
+            fit_plan, format_order, deferred_fit_sse
+        )
+        confirmation_sse = _finish_deferred_functional_sse(
+            confirmation_plan, format_order, deferred_confirmation_sse
+        )
+        fit_reference = fit_plan.reference_energy
+        fit_counts = fit_plan.counts
+        confirmation_reference = confirmation_plan.reference_energy
+        confirmation_counts = confirmation_plan.counts
+        if fixed_profile:
+            fixed_mode = H308 if fixed_high_rate else (K2 if fixed_k2 else R0)
+            selection = RatePairSelection(
+                proposed_r13=fixed_mode.mode_id,
+                proposed_r2=fixed_mode.mode_id,
+                selected_r13=fixed_mode.mode_id,
+                selected_r2=fixed_mode.mode_id,
+                accepted=True,
+                reason=(
+                    "fixed_22k3_2k4_high_rate_profile"
+                    if fixed_high_rate
+                    else (
+                        "fixed_uniform_k2_profile"
+                        if fixed_k2
+                        else "fixed_uniform_k3_profile"
+                    )
+                ),
+                fit_documents=int(torch.count_nonzero(fit_counts)),
+                confirmation_documents=int(torch.count_nonzero(confirmation_counts)),
+                confirmation_relative_improvement=None,
+                confirmation_ci95=(None, None),
+                bootstrap_replicates_valid=0,
+            )
+        else:
+            selection = select_phase1_rate_pair(
+                fit_sse,
+                confirmation_sse,
+                fit_counts=fit_counts,
+                confirmation_counts=confirmation_counts,
+                modes=evaluated_formats,
+                min_fit_documents=min_fit_documents,
+                min_confirmation_documents=min_confirmation_documents,
+                minimum_improvement=minimum_improvement,
+                bootstrap_replicates=bootstrap_replicates,
+                seed=deterministic_expert_seed(layer, expert),
+            )
+        selected_r13, selected_r2 = selection.selected
+        selected_candidates = {
+            "w1": upstream["w1"][selected_r13],
+            "w3": upstream["w3"][selected_r13],
+            "w2": down[selected_r13][selected_r2],
+        }
+        if os.environ.get("QSRT_DEBUG_CANDIDATE_CLOSURE") == "1":
+            for value in selected_candidates.values():
+                _debug_assert_candidate_state_closure(
+                    value,
+                    layer=layer,
+                    expert=expert,
+                    stage="post-scoring selected-candidate",
+                )
+        selected_encodings = {}
+        for matrix, value in selected_candidates.items():
+            try:
+                selected_encodings[matrix] = finalize_qsrt_matrix_candidate(
+                    value,
+                    layer=layer,
+                    logical_trellis_schema=logical_trellis_schema,
+                    codebook=codebook,
+                    tailbite_context=tailbite_context,
+                )
+            except Exception as exc:
+                raise RuntimeError(
+                    f"layer {layer} expert {expert} {matrix} finalization failed"
+                ) from exc
+        physical_permutation = selected_encodings["w1"].plan.physical_permutation
+        for matrix in ("w3", "w2"):
+            if not torch.equal(
+                selected_encodings[matrix].plan.physical_permutation,
+                physical_permutation,
+            ):
+                raise ValueError(
+                    "selected matrices do not share one physical permutation"
+                )
+        expert_encoding = QSRTExpertEncoding(
+            format=ExpertFormatSpec.compressed(selected_r13, selected_r2),
+            matrices=selected_encodings,
+            physical_permutation=physical_permutation,
+            coding={
+                "format": ExpertFormatSpec.compressed(
+                    selected_r13, selected_r2
+                ).name,
+                "r13": selected_r13,
+                "r2": selected_r2,
+                "codebook": codebook,
+                "tailbite_context": int(tailbite_context),
+                "coupled_hadamard": (
+                    None
+                    if item["coupled_execution"] is None
+                    else {
+                        "residual_block_size": item[
+                            "coupled_execution"
+                        ].spec.residual_block_size,
+                        "preactivation_block_size": item[
+                            "coupled_execution"
+                        ].spec.preactivation_block_size,
+                        "postactivation_block_size": item[
+                            "coupled_execution"
+                        ].spec.postactivation_block_size,
+                        "residual_draw": item[
+                            "coupled_execution"
+                        ].spec.residual_draw,
+                        "intermediate_draw": item[
+                            "coupled_execution"
+                        ].spec.intermediate_draw,
+                        "activation": item[
+                            "coupled_execution"
+                        ].spec.activation,
+                    }
+                ),
+                **_mode_evidence(selected_encodings),
+            },
+        )
+        results.append(
+            QSRTCandidateEncoding(
+                expert=expert_encoding,
+                selection=selection,
+                block_contexts=item["block_contexts"]
+                .detach()
+                .cpu()
+                .to(torch.uint8)
+                .contiguous(),
+                block_scores=item["block_scores"].detach().cpu().float().contiguous(),
+                context_basis=str(item["context_basis"]),
+                covariance=item["covariance"],
+                evaluated_modes=evaluated_modes,
+                evaluated_formats=evaluated_formats,
+                fit_sse=fit_sse,
+                confirmation_sse=confirmation_sse,
+                fit_reference_energy=fit_reference,
+                confirmation_reference_energy=confirmation_reference,
+                fit_counts=fit_counts,
+                confirmation_counts=confirmation_counts,
+                mode_coding=mode_coding,
+            )
+        )
+    return results
+
+
+def candidate_tensor_name(
+    layer: int, expert: int, matrix: str, part: str
+) -> str:
+    if matrix not in MATRICES or part not in ("trellis", "suh", "svh"):
+        raise ValueError("invalid QSRT candidate tensor component")
+    return (
+        f"language_model.model.layers.{layer}.block_sparse_moe.experts."
+        f"{expert}.{matrix}.qsrt_{part}"
+    )
+
+
+def selected_candidate_tensors(
+    layer: int, expert: int, candidate: QSRTCandidateEncoding
+) -> dict[str, torch.Tensor]:
+    """Copy one selected physical payload to CPU for layer-shard assembly."""
+
+    tensors: dict[str, torch.Tensor] = {}
+    for matrix, encoding in candidate.expert.matrices.items():
+        for part in ("trellis", "suh", "svh"):
+            tensors[candidate_tensor_name(layer, expert, matrix, part)] = (
+                encoding.tensors[part].detach().cpu().contiguous()
+            )
+    return tensors

reproduction/closure/patches/qsrt/changed-files/qsrt/qsrt_coupled.py

@@ -0,0 +1,381 @@
+"""Exact coupled activation-boundary coordinates for QSRT expert matrices.
+
+The transform interleaves gate and up rows before a normalized block-Hadamard
+rotation. The decoder joins the two stored projections, cancels that rotation,
+evaluates the coordinatewise SiTU activation, and rotates the resulting hidden
+coordinates into the matching down-projection basis. Residual-side transforms
+are layer-shared; the intermediate draw is expert-static.
+"""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+from typing import Literal
+
+import torch
+import torch.nn.functional as F
+
+from qsrt.tp_simulator import situ
+
+
+Tensor = torch.Tensor
+WeightTriplet = tuple[Tensor, Tensor, Tensor]
+ExpertActivation = Literal["situ", "silu"]
+
+
+def expert_activation(
+    gate: Tensor,
+    up: Tensor,
+    *,
+    activation: ExpertActivation,
+) -> Tensor:
+    """Evaluate the model's exact coordinatewise gated activation.
+
+    Coupled Hadamard construction is valid for either activation because the
+    preactivation transform is inverted before this operation and the
+    postactivation transform is applied afterward.  Keeping the choice in the
+    immutable transform specification prevents a Kimi SiTU candidate from
+    being replayed accidentally under GLM's standard SwiGLU contract.
+    """
+
+    if gate.shape != up.shape:
+        raise ValueError("gate and up projections must have matching shapes")
+    if activation == "situ":
+        return situ(gate, up)
+    if activation == "silu":
+        return F.silu(gate) * up
+    raise ValueError(f"unsupported expert activation: {activation}")
+
+
+@dataclass(frozen=True)
+class CoupledHadamardSpec:
+    """Immutable transform parameters carried by a coupled QSRT profile."""
+
+    residual_block_size: int = 512
+    preactivation_block_size: int = 128
+    postactivation_block_size: int = 128
+    residual_draw: int = 0
+    intermediate_draw: int = 0
+    activation: ExpertActivation = "situ"
+
+    def __post_init__(self) -> None:
+        for name, value in (
+            ("residual_block_size", self.residual_block_size),
+            ("preactivation_block_size", self.preactivation_block_size),
+            ("postactivation_block_size", self.postactivation_block_size),
+        ):
+            if value <= 0 or value & (value - 1):
+                raise ValueError(f"{name} must be a positive power of two")
+        if self.residual_draw != 0:
+            raise ValueError("the coupled profile fixes residual draw zero")
+        if not 0 <= self.intermediate_draw < 8:
+            raise ValueError("intermediate_draw must lie in 0..7")
+        if self.activation not in ("situ", "silu"):
+            raise ValueError("activation must be 'situ' or 'silu'")
+
+
+def block_hadamard(values: Tensor, *, block_size: int, dim: int = -1) -> Tensor:
+    """Apply a normalized self-inverse block Walsh-Hadamard transform."""
+
+    if values.ndim == 0 or not torch.is_floating_point(values):
+        raise TypeError("block Hadamard input must be floating point")
+    if block_size <= 0 or block_size & (block_size - 1):
+        raise ValueError("block_size must be a positive power of two")
+    axis = dim % values.ndim
+    if values.shape[axis] % block_size:
+        raise ValueError("Hadamard axis must be divisible by block_size")
+    # Naturally routed experts can have an empty fit or confirmation fold.
+    # The transform is linear, so its action on an empty batch is the empty
+    # float32 tensor with the same shape.  Returning before the butterfly
+    # reshape also avoids an ambiguous ``-1`` inference when numel is zero.
+    if values.numel() == 0:
+        return values.float().contiguous().clone()
+    output = values.float().movedim(axis, -1).contiguous().clone()
+    shape = output.shape
+    output = output.reshape(*shape[:-1], shape[-1] // block_size, block_size)
+    width = 1
+    while width < block_size:
+        paired = output.reshape(*output.shape[:-1], -1, 2, width)
+        left = paired[..., 0, :].clone()
+        right = paired[..., 1, :].clone()
+        paired[..., 0, :] = left + right
+        paired[..., 1, :] = left - right
+        output = paired.reshape(*output.shape)
+        width *= 2
+    return (
+        output.div_(math.sqrt(block_size))
+        .reshape(shape)
+        .movedim(-1, axis)
+        .contiguous()
+    )
+
+
+def rotation_signs(
+    length: int, *, draw: int, axis: int, device: torch.device
+) -> Tensor:
+    """Return one deterministic Rademacher vector from the frozen draw family."""
+
+    if length <= 0 or draw < 0:
+        raise ValueError("length must be positive and draw nonnegative")
+    if draw == 0:
+        return torch.ones(length, dtype=torch.float32, device=device)
+    generator = torch.Generator(device="cpu")
+    generator.manual_seed(
+        (0x6A09E667F3BCC909 * draw + 0xBB67AE8584CAA73B * axis)
+        & ((1 << 63) - 1)
+    )
+    return (
+        torch.randint(0, 2, (length,), generator=generator)
+        .mul_(2)
+        .sub_(1)
+        .float()
+        .to(device=device)
+    )
+
+
+def signed_block_hadamard(
+    values: Tensor,
+    *,
+    block_size: int,
+    signs: Tensor,
+    dim: int = -1,
+    inverse: bool = False,
+) -> Tensor:
+    """Apply ``D H`` to row vectors, or its exact inverse ``H D``."""
+
+    axis = dim % values.ndim
+    if signs.ndim != 1 or signs.numel() != values.shape[axis]:
+        raise ValueError("rotation signs do not align with the transformed axis")
+    shape = [1] * values.ndim
+    shape[axis] = signs.numel()
+    expanded = signs.to(device=values.device, dtype=torch.float32).reshape(shape)
+    if inverse:
+        return block_hadamard(values, block_size=block_size, dim=axis) * expanded
+    return block_hadamard(
+        values.float() * expanded, block_size=block_size, dim=axis
+    )
+
+
+def _validate_triplet(weights: WeightTriplet) -> tuple[int, int]:
+    w1, w3, w2 = weights
+    if w1.ndim != 2 or w3.ndim != 2 or w2.ndim != 2 or w1.shape != w3.shape:
+        raise ValueError("coupled weights require matching two-dimensional w1/w3")
+    intermediate, hidden = w1.shape
+    if tuple(w2.shape) != (hidden, intermediate):
+        raise ValueError("w2 must have shape [hidden, intermediate]")
+    if not all(torch.is_floating_point(value) for value in weights):
+        raise TypeError("coupled weights must be floating point")
+    if len({value.device for value in weights}) != 1:
+        raise ValueError("coupled weights must share one device")
+    return intermediate, hidden
+
+
+def encode_coupled_weights(
+    weights: WeightTriplet, spec: CoupledHadamardSpec
+) -> WeightTriplet:
+    """Reparameterize one full-precision expert into coupled coordinates."""
+
+    w1, w3, w2 = weights
+    intermediate, hidden = _validate_triplet(weights)
+    if hidden % spec.residual_block_size:
+        raise ValueError("hidden size does not close residual Hadamard blocks")
+    if 2 * intermediate % spec.preactivation_block_size:
+        raise ValueError("interleaved size does not close preactivation blocks")
+    if intermediate % spec.postactivation_block_size:
+        raise ValueError("intermediate size does not close postactivation blocks")
+    device = w1.device
+    input_signs = rotation_signs(hidden, draw=spec.residual_draw, axis=0, device=device)
+    pre_signs = rotation_signs(
+        2 * intermediate, draw=spec.intermediate_draw, axis=1, device=device
+    )
+    post_signs = rotation_signs(
+        intermediate, draw=spec.intermediate_draw, axis=2, device=device
+    )
+    output_signs = rotation_signs(
+        hidden, draw=spec.residual_draw, axis=3, device=device
+    )
+    interleaved = torch.stack((w1, w3), dim=1).reshape(2 * intermediate, hidden)
+    upstream = signed_block_hadamard(
+        signed_block_hadamard(
+            interleaved,
+            block_size=spec.residual_block_size,
+            signs=input_signs,
+            dim=1,
+        ),
+        block_size=spec.preactivation_block_size,
+        signs=pre_signs,
+        dim=0,
+    )
+    down = signed_block_hadamard(
+        signed_block_hadamard(
+            w2,
+            block_size=spec.postactivation_block_size,
+            signs=post_signs,
+            dim=1,
+        ),
+        block_size=spec.residual_block_size,
+        signs=output_signs,
+        dim=0,
+    )
+    return (
+        upstream[:intermediate].contiguous(),
+        upstream[intermediate:].contiguous(),
+        down.contiguous(),
+    )
+
+
+@dataclass(frozen=True)
+class CoupledHadamardExecution:
+    """Calibration and scoring operations for one transformed expert basis."""
+
+    hidden: int
+    intermediate: int
+    spec: CoupledHadamardSpec
+
+    def _signs(self, axis: int, *, device: torch.device) -> Tensor:
+        lengths = (self.hidden, 2 * self.intermediate, self.intermediate, self.hidden)
+        draw = self.spec.residual_draw if axis in (0, 3) else self.spec.intermediate_draw
+        return rotation_signs(lengths[axis], draw=draw, axis=axis, device=device)
+
+    def transform_inputs(self, rows: Tensor) -> Tensor:
+        if rows.ndim != 2 or rows.shape[1] != self.hidden:
+            raise ValueError("input rows do not match the hidden dimension")
+        return signed_block_hadamard(
+            rows,
+            block_size=self.spec.residual_block_size,
+            signs=self._signs(0, device=rows.device),
+            dim=1,
+        )
+
+    def transform_h13(self, hessian: Tensor) -> Tensor:
+        if tuple(hessian.shape) != (self.hidden, self.hidden):
+            raise ValueError("H13 does not match the hidden dimension")
+        signs = self._signs(0, device=hessian.device)
+        return signed_block_hadamard(
+            signed_block_hadamard(
+                hessian,
+                block_size=self.spec.residual_block_size,
+                signs=signs,
+                dim=0,
+            ),
+            block_size=self.spec.residual_block_size,
+            signs=signs,
+            dim=1,
+        )
+
+    def transform_h2(self, hessian: Tensor) -> Tensor:
+        if tuple(hessian.shape) != (self.intermediate, self.intermediate):
+            raise ValueError("H2 does not match the intermediate dimension")
+        signs = self._signs(2, device=hessian.device)
+        return signed_block_hadamard(
+            signed_block_hadamard(
+                hessian,
+                block_size=self.spec.postactivation_block_size,
+                signs=signs,
+                dim=0,
+            ),
+            block_size=self.spec.postactivation_block_size,
+            signs=signs,
+            dim=1,
+        )
+
+    def decode_middle(self, rows: Tensor, w1: Tensor, w3: Tensor) -> Tensor:
+        return self.decode_middle_outputs(F.linear(rows, w1), F.linear(rows, w3))
+
+    def decode_middle_outputs(
+        self, transformed_w1: Tensor, transformed_w3: Tensor
+    ) -> Tensor:
+        """Close the coupled boundary from already-computed FC1 outputs.
+
+        This is the runtime/calibration seam used when the two upstream GEMMs
+        have their own operand precision, reconstruction law, or fused kernel.
+        It applies exactly the same inverse preactivation rotation, model
+        activation, and postactivation rotation as :meth:`decode_middle`.
+        """
+
+        if (
+            transformed_w1.ndim != 2
+            or transformed_w3.ndim != 2
+            or transformed_w1.shape != transformed_w3.shape
+            or transformed_w1.shape[1] != self.intermediate
+            or transformed_w1.device != transformed_w3.device
+        ):
+            raise ValueError(
+                "coupled upstream outputs must be aligned [rows, intermediate]"
+            )
+        transformed = torch.cat((transformed_w1, transformed_w3), dim=1)
+        recovered = signed_block_hadamard(
+            transformed,
+            block_size=self.spec.preactivation_block_size,
+            signs=self._signs(1, device=transformed.device),
+            dim=1,
+            inverse=True,
+        )
+        middle = expert_activation(
+            recovered[:, 0::2],
+            recovered[:, 1::2],
+            activation=self.spec.activation,
+        )
+        return signed_block_hadamard(
+            middle,
+            block_size=self.spec.postactivation_block_size,
+            signs=self._signs(2, device=transformed.device),
+            dim=1,
+        )
+
+    def decode_output(self, transformed: Tensor) -> Tensor:
+        if transformed.ndim != 2 or transformed.shape[1] != self.hidden:
+            raise ValueError("output rows do not match the hidden dimension")
+        return signed_block_hadamard(
+            transformed,
+            block_size=self.spec.residual_block_size,
+            signs=self._signs(3, device=transformed.device),
+            dim=1,
+            inverse=True,
+        )
+
+    def transform_outputs(self, rows: Tensor) -> Tensor:
+        """Map caller-coordinate outputs into the coupled residual basis.
+
+        Calibration uses this forward counterpart of :meth:`decode_output`
+        when constructing a candidate-conditioned down cross term.  The
+        candidate operand and source-model teacher must share the transformed
+        output coordinates before forming ``B = Q.T @ Y``.
+        """
+
+        if rows.ndim != 2 or rows.shape[1] != self.hidden:
+            raise ValueError("output rows do not match the hidden dimension")
+        return signed_block_hadamard(
+            rows,
+            block_size=self.spec.residual_block_size,
+            signs=self._signs(3, device=rows.device),
+            dim=1,
+        )
+
+    def execute(self, rows: Tensor, weights: WeightTriplet) -> Tensor:
+        _validate_triplet(weights)
+        transformed_rows = self.transform_inputs(rows)
+        middle = self.decode_middle(transformed_rows, weights[0], weights[1])
+        return self.decode_output(F.linear(middle, weights[2]))
+
+
+def coupled_execution(
+    weights: WeightTriplet, spec: CoupledHadamardSpec
+) -> CoupledHadamardExecution:
+    intermediate, hidden = _validate_triplet(weights)
+    return CoupledHadamardExecution(hidden, intermediate, spec)
+
+
+__all__ = [
+    "CoupledHadamardExecution",
+    "CoupledHadamardSpec",
+    "ExpertActivation",
+    "WeightTriplet",
+    "block_hadamard",
+    "coupled_execution",
+    "encode_coupled_weights",
+    "expert_activation",
+    "rotation_signs",
+    "signed_block_hadamard",
+]

reproduction/closure/patches/qsrt/changed-files/tests/test_qsrt_coupled.py

@@ -0,0 +1,181 @@
+from __future__ import annotations
+
+import torch
+
+from qsrt.coupled_expert_study import (
+    CoupledTriplet,
+    encode_coupled_block_hadamard,
+    execute_coupled_block_hadamard,
+)
+from qsrt.qsrt_coupled import (
+    CoupledHadamardSpec,
+    block_hadamard,
+    coupled_execution,
+    encode_coupled_weights,
+    expert_activation,
+    rotation_signs,
+    signed_block_hadamard,
+)
+from qsrt.tp_simulator import situ
+
+
+def test_block_hadamard_preserves_empty_routed_batches() -> None:
+    values = torch.empty((0, 7, 16), dtype=torch.float16)
+    transformed = block_hadamard(values, block_size=8, dim=2)
+
+    assert transformed.shape == values.shape
+    assert transformed.dtype == torch.float32
+    assert transformed.numel() == 0
+
+
+def test_production_coupled_transform_closes_and_matches_research_oracle() -> None:
+    generator = torch.Generator().manual_seed(871)
+    rows, hidden, intermediate = 9, 16, 12
+    weights = (
+        torch.randn(intermediate, hidden, generator=generator),
+        torch.randn(intermediate, hidden, generator=generator),
+        torch.randn(hidden, intermediate, generator=generator),
+    )
+    inputs = torch.randn(rows, hidden, generator=generator)
+    spec = CoupledHadamardSpec(
+        residual_block_size=8,
+        preactivation_block_size=8,
+        postactivation_block_size=4,
+        intermediate_draw=3,
+    )
+
+    encoded = encode_coupled_weights(weights, spec)
+    execution = coupled_execution(encoded, spec)
+    output = execution.execute(inputs, encoded)
+    reference = torch.nn.functional.linear(
+        situ(
+            torch.nn.functional.linear(inputs, weights[0]),
+            torch.nn.functional.linear(inputs, weights[1]),
+        ),
+        weights[2],
+    )
+    research_encoded = encode_coupled_block_hadamard(
+        CoupledTriplet(*weights),
+        block_size=8,
+        preactivation_block_size=8,
+        postactivation_block_size=4,
+        intermediate_rotation_draw=3,
+    )
+    research_output = execute_coupled_block_hadamard(
+        inputs,
+        research_encoded,
+        block_size=8,
+        preactivation_block_size=8,
+        postactivation_block_size=4,
+        intermediate_rotation_draw=3,
+    )
+
+    for actual, expected in zip(encoded, research_encoded.tensors(), strict=True):
+        assert torch.equal(actual, expected)
+    assert torch.allclose(output, reference, rtol=2e-5, atol=2e-5)
+    assert torch.allclose(output, research_output, rtol=2e-5, atol=2e-5)
+
+
+def test_glm_silu_coupled_transform_closes_exact_expert_function() -> None:
+    generator = torch.Generator().manual_seed(5202)
+    rows, hidden, intermediate = 11, 24, 16
+    weights = (
+        torch.randn(intermediate, hidden, generator=generator),
+        torch.randn(intermediate, hidden, generator=generator),
+        torch.randn(hidden, intermediate, generator=generator),
+    )
+    inputs = torch.randn(rows, hidden, generator=generator)
+    spec = CoupledHadamardSpec(
+        residual_block_size=8,
+        preactivation_block_size=8,
+        postactivation_block_size=8,
+        intermediate_draw=6,
+        activation="silu",
+    )
+
+    encoded = encode_coupled_weights(weights, spec)
+    execution = coupled_execution(encoded, spec)
+    output = execution.execute(inputs, encoded)
+    gate = torch.nn.functional.linear(inputs, weights[0])
+    up = torch.nn.functional.linear(inputs, weights[1])
+    reference = torch.nn.functional.linear(
+        torch.nn.functional.silu(gate) * up,
+        weights[2],
+    )
+
+    assert torch.allclose(output, reference, rtol=2e-5, atol=2e-5)
+
+    transformed_input = execution.transform_inputs(inputs)
+    gate_output = torch.nn.functional.linear(transformed_input, encoded[0])
+    up_output = torch.nn.functional.linear(transformed_input, encoded[1])
+    middle_from_outputs = execution.decode_middle_outputs(
+        gate_output, up_output
+    )
+    middle_from_weights = execution.decode_middle(
+        transformed_input, encoded[0], encoded[1]
+    )
+    torch.testing.assert_close(
+        middle_from_outputs, middle_from_weights, rtol=0, atol=0
+    )
+
+    transformed_output = execution.transform_outputs(reference)
+    torch.testing.assert_close(
+        execution.decode_output(transformed_output), reference, rtol=2e-5, atol=2e-5
+    )
+
+
+def test_expert_activation_rejects_unknown_contract() -> None:
+    gate = torch.ones((2, 3))
+    up = torch.ones((2, 3))
+
+    try:
+        expert_activation(gate, up, activation="gelu")  # type: ignore[arg-type]
+    except ValueError as exc:
+        assert "unsupported expert activation" in str(exc)
+    else:
+        raise AssertionError("unknown expert activation was accepted")
+
+
+def test_coupled_hessian_transforms_preserve_quadratic_forms() -> None:
+    generator = torch.Generator().manual_seed(91)
+    hidden, intermediate = 16, 12
+    weights = (
+        torch.randn(intermediate, hidden, generator=generator),
+        torch.randn(intermediate, hidden, generator=generator),
+        torch.randn(hidden, intermediate, generator=generator),
+    )
+    spec = CoupledHadamardSpec(
+        residual_block_size=8,
+        preactivation_block_size=8,
+        postactivation_block_size=4,
+        intermediate_draw=5,
+    )
+    execution = coupled_execution(encode_coupled_weights(weights, spec), spec)
+    rows = torch.randn(11, hidden, generator=generator)
+    hidden_rows = torch.randn(11, intermediate, generator=generator)
+    h13 = rows.T @ rows
+    h2 = hidden_rows.T @ hidden_rows
+
+    transformed_rows = execution.transform_inputs(rows)
+    post_signs = rotation_signs(
+        intermediate, draw=spec.intermediate_draw, axis=2, device=rows.device
+    )
+    transformed_hidden_rows = signed_block_hadamard(
+        hidden_rows,
+        block_size=spec.postactivation_block_size,
+        signs=post_signs,
+        dim=1,
+    )
+    transformed_hidden = execution.transform_h2(h2)
+    assert torch.allclose(
+        execution.transform_h13(h13),
+        transformed_rows.T @ transformed_rows,
+        rtol=2e-5,
+        atol=2e-5,
+    )
+    assert torch.allclose(
+        transformed_hidden,
+        transformed_hidden_rows.T @ transformed_hidden_rows,
+        rtol=2e-5,
+        atol=2e-5,
+    )

reproduction/closure/patches/qsrt/changes.patch

@@ -0,0 +1,442 @@
+diff --git a/qsrt/pack/qsrt_candidates.py b/qsrt/pack/qsrt_candidates.py
+index d7464fc..864b5c2 100644
+--- a/qsrt/pack/qsrt_candidates.py
++++ b/qsrt/pack/qsrt_candidates.py
+@@ -43,8 +43,10 @@ from qsrt.qsrt import (
+ from qsrt.qsrt_coupled import (
+     CoupledHadamardExecution,
+     CoupledHadamardSpec,
++    ExpertActivation,
+     coupled_execution,
+     encode_coupled_weights,
++    expert_activation,
+ )
+ from qsrt.pack.qsrt_encoder import (
+     Layout,
+@@ -172,9 +174,22 @@ class QSRTCandidateEncoding:
+ 
+ 
+ def _source_middle(
+-    gate_projection: torch.Tensor, up_projection: torch.Tensor
++    gate_projection: torch.Tensor,
++    up_projection: torch.Tensor,
++    *,
++    activation: ExpertActivation = "situ",
+ ) -> torch.Tensor:
+-    return situ(gate_projection, up_projection)
++    return expert_activation(
++        gate_projection,
++        up_projection,
++        activation=activation,
++    )
++
++
++def _activation_basis(activation: ExpertActivation) -> str:
++    if activation not in ("situ", "silu"):
++        raise ValueError(f"unsupported expert activation: {activation}")
++    return f"post_{activation}"
+ 
+ 
+ def _folded_intermediate_conditioning(
+@@ -407,13 +422,16 @@ def _candidate_middle_by_r13(
+     inputs: torch.Tensor,
+     upstream: dict[str, dict[int, QSRTMatrixCandidate]],
+     evaluated_modes: Sequence[int],
++    *,
++    activation: ExpertActivation = "situ",
+ ) -> dict[int, torch.Tensor]:
+-    """Replay exact candidate reconstructions through Kimi's SiTU."""
++    """Replay exact candidate reconstructions through the model activation."""
+ 
+     return {
+         int(r13): _source_middle(
+             F.linear(inputs, upstream["w1"][int(r13)].reconstruction),
+             F.linear(inputs, upstream["w3"][int(r13)].reconstruction),
++            activation=activation,
+         )
+         for r13 in evaluated_modes
+     }
+@@ -627,6 +645,7 @@ def encode_phase1_expert(
+     permutation_policy: PermutationPolicy = "h2_reverse",
+     folded_scale_power: float = 0.0,
+     transform_seeds: Mapping[str, QSRTTransformSeeds] | None = None,
++    activation: ExpertActivation = "situ",
+ ) -> QSRTCandidateEncoding:
+     """Encode and conservatively select one official-source expert.
+ 
+@@ -674,7 +693,11 @@ def encode_phase1_expert(
+     if all_rows.rows:
+         source_gate = F.linear(inputs, source_w1)
+         source_up = F.linear(inputs, source_w3)
+-        source_middle = _source_middle(source_gate, source_up)
++        source_middle = _source_middle(
++            source_gate,
++            source_up,
++            activation=activation,
++        )
+         del source_gate, source_up
+     else:
+         source_middle = torch.empty(
+@@ -697,7 +720,8 @@ def encode_phase1_expert(
+             "h13_local_alpha": 0.0,
+         }
+         context_basis = (
+-            f"official_source_post_situ_fit_documents:{permutation_policy}"
++            f"official_source_{_activation_basis(activation)}_fit_documents:"
++            f"{permutation_policy}"
+         )
+         covariance = {
+             **covariance,
+@@ -762,7 +786,10 @@ def encode_phase1_expert(
+     del source_w1, source_w3
+ 
+     middle_by_r13 = _candidate_middle_by_r13(
+-        inputs, upstream_candidates, evaluated_modes
++        inputs,
++        upstream_candidates,
++        evaluated_modes,
++        activation=activation,
+     )
+     h2_by_r13, h2_evidence = _conditional_h2_by_r13(
+         inputs,
+@@ -781,7 +808,8 @@ def encode_phase1_expert(
+                 "identity"
+                 if hessian_policy == "identity"
+                 else (
+-                    "decoded_candidate_post_situ_adaptive_shrinkage"
++                    f"decoded_candidate_{_activation_basis(activation)}_"
++                    "adaptive_shrinkage"
+                     if have_local_support
+                     else "identity_support_fallback"
+                 )
+@@ -789,7 +817,7 @@ def encode_phase1_expert(
+             "h2_by_r13": h2_evidence,
+         }
+     )
+-    context_basis += "+decoded_candidate_post_situ_h2"
++    context_basis += f"+decoded_candidate_{_activation_basis(activation)}_h2"
+ 
+     source_w2 = _load_source_matrix(store, layer, expert, "w2", device)
+     if all_rows.rows:
+@@ -996,6 +1024,7 @@ def encode_phase1_expert_batch(
+         int, Mapping[str, QSRTTransformSeeds]
+     ] | None = None,
+     coupled_intermediate_draws_by_expert: Mapping[int, int] | None = None,
++    activation: ExpertActivation = "situ",
+ ) -> list[QSRTCandidateEncoding]:
+     """Encode several experts while batching their independent trellis work.
+ 
+@@ -1019,6 +1048,7 @@ def encode_phase1_expert_batch(
+     fixed_k2 = modes == (K2.mode_id,)
+     fixed_profile = fixed_k3 or fixed_high_rate or fixed_k2
+     coupled = coupled_intermediate_draws_by_expert is not None
++    _activation_basis(activation)
+     if coupled:
+         if modes not in (
+             (R0.mode_id,),
+@@ -1079,7 +1109,11 @@ def encode_phase1_expert_batch(
+         if all_rows.rows:
+             source_gate = F.linear(inputs, source_w1)
+             source_up = F.linear(inputs, source_w3)
+-            source_middle = _source_middle(source_gate, source_up)
++            source_middle = _source_middle(
++                source_gate,
++                source_up,
++                activation=activation,
++            )
+             del source_gate, source_up
+         else:
+             source_middle = torch.empty(
+@@ -1095,7 +1129,8 @@ def encode_phase1_expert_batch(
+                 policy=permutation_policy,
+             )
+             context_basis = (
+-                f"official_source_post_situ_fit_documents:{permutation_policy}"
++                f"official_source_{_activation_basis(activation)}_fit_documents:"
++                f"{permutation_policy}"
+             )
+         else:
+             block_contexts = fallback_block_contexts.clone()
+@@ -1156,7 +1191,10 @@ def encode_phase1_expert_batch(
+         coupled_reference_output: torch.Tensor | None = None
+         if coupled:
+             draw = int(coupled_intermediate_draws_by_expert[expert])
+-            spec = CoupledHadamardSpec(intermediate_draw=draw)
++            spec = CoupledHadamardSpec(
++                intermediate_draw=draw,
++                activation=activation,
++            )
+             source_w2 = _load_source_matrix(store, layer, expert, "w2", device)
+             transformed = encode_coupled_weights(
+                 (source_w1, source_w3, source_w2), spec
+@@ -1272,7 +1310,10 @@ def encode_phase1_expert_batch(
+         execution_basis = item["coupled_execution"]
+         if execution_basis is None:
+             middle_by_r13 = _candidate_middle_by_r13(
+-                inputs, upstream, evaluated_modes
++                inputs,
++                upstream,
++                evaluated_modes,
++                activation=activation,
+             )
+         else:
+             middle_by_r13 = {
+@@ -1303,7 +1344,8 @@ def encode_phase1_expert_batch(
+                     "identity"
+                     if hessian_policy == "identity"
+                     else (
+-                        "decoded_candidate_post_situ_adaptive_shrinkage"
++                        f"decoded_candidate_{_activation_basis(activation)}_"
++                        "adaptive_shrinkage"
+                         if int(item["covariance"]["fit_documents"])
+                         >= min_fit_documents
+                         else "identity_support_fallback"
+@@ -1313,7 +1355,8 @@ def encode_phase1_expert_batch(
+             }
+         )
+         item["context_basis"] = (
+-            str(item["context_basis"]) + "+decoded_candidate_post_situ_h2"
++            str(item["context_basis"])
++            + f"+decoded_candidate_{_activation_basis(activation)}_h2"
+         )
+         all_rows = item["all_rows"]
+         if item["coupled_reference_output"] is not None:
+@@ -1535,6 +1578,9 @@ def encode_phase1_expert_batch(
+                         "intermediate_draw": item[
+                             "coupled_execution"
+                         ].spec.intermediate_draw,
++                        "activation": item[
++                            "coupled_execution"
++                        ].spec.activation,
+                     }
+                 ),
+                 **_mode_evidence(selected_encodings),
+diff --git a/qsrt/qsrt_coupled.py b/qsrt/qsrt_coupled.py
+index 6159771..4febf9b 100644
+--- a/qsrt/qsrt_coupled.py
++++ b/qsrt/qsrt_coupled.py
+@@ -11,6 +11,7 @@ from __future__ import annotations
+ 
+ import math
+ from dataclasses import dataclass
++from typing import Literal
+ 
+ import torch
+ import torch.nn.functional as F
+@@ -20,6 +21,31 @@ from qsrt.tp_simulator import situ
+ 
+ Tensor = torch.Tensor
+ WeightTriplet = tuple[Tensor, Tensor, Tensor]
++ExpertActivation = Literal["situ", "silu"]
++
++
++def expert_activation(
++    gate: Tensor,
++    up: Tensor,
++    *,
++    activation: ExpertActivation,
++) -> Tensor:
++    """Evaluate the model's exact coordinatewise gated activation.
++
++    Coupled Hadamard construction is valid for either activation because the
++    preactivation transform is inverted before this operation and the
++    postactivation transform is applied afterward.  Keeping the choice in the
++    immutable transform specification prevents a Kimi SiTU candidate from
++    being replayed accidentally under GLM's standard SwiGLU contract.
++    """
++
++    if gate.shape != up.shape:
++        raise ValueError("gate and up projections must have matching shapes")
++    if activation == "situ":
++        return situ(gate, up)
++    if activation == "silu":
++        return F.silu(gate) * up
++    raise ValueError(f"unsupported expert activation: {activation}")
+ 
+ 
+ @dataclass(frozen=True)
+@@ -31,6 +57,7 @@ class CoupledHadamardSpec:
+     postactivation_block_size: int = 128
+     residual_draw: int = 0
+     intermediate_draw: int = 0
++    activation: ExpertActivation = "situ"
+ 
+     def __post_init__(self) -> None:
+         for name, value in (
+@@ -44,6 +71,8 @@ class CoupledHadamardSpec:
+             raise ValueError("the coupled profile fixes residual draw zero")
+         if not 0 <= self.intermediate_draw < 8:
+             raise ValueError("intermediate_draw must lie in 0..7")
++        if self.activation not in ("situ", "silu"):
++            raise ValueError("activation must be 'situ' or 'silu'")
+ 
+ 
+ def block_hadamard(values: Tensor, *, block_size: int, dim: int = -1) -> Tensor:
+@@ -252,19 +281,46 @@ class CoupledHadamardExecution:
+         )
+ 
+     def decode_middle(self, rows: Tensor, w1: Tensor, w3: Tensor) -> Tensor:
+-        transformed = torch.cat((F.linear(rows, w1), F.linear(rows, w3)), dim=1)
++        return self.decode_middle_outputs(F.linear(rows, w1), F.linear(rows, w3))
++
++    def decode_middle_outputs(
++        self, transformed_w1: Tensor, transformed_w3: Tensor
++    ) -> Tensor:
++        """Close the coupled boundary from already-computed FC1 outputs.
++
++        This is the runtime/calibration seam used when the two upstream GEMMs
++        have their own operand precision, reconstruction law, or fused kernel.
++        It applies exactly the same inverse preactivation rotation, model
++        activation, and postactivation rotation as :meth:`decode_middle`.
++        """
++
++        if (
++            transformed_w1.ndim != 2
++            or transformed_w3.ndim != 2
++            or transformed_w1.shape != transformed_w3.shape
++            or transformed_w1.shape[1] != self.intermediate
++            or transformed_w1.device != transformed_w3.device
++        ):
++            raise ValueError(
++                "coupled upstream outputs must be aligned [rows, intermediate]"
++            )
++        transformed = torch.cat((transformed_w1, transformed_w3), dim=1)
+         recovered = signed_block_hadamard(
+             transformed,
+             block_size=self.spec.preactivation_block_size,
+-            signs=self._signs(1, device=rows.device),
++            signs=self._signs(1, device=transformed.device),
+             dim=1,
+             inverse=True,
+         )
+-        middle = situ(recovered[:, 0::2], recovered[:, 1::2])
++        middle = expert_activation(
++            recovered[:, 0::2],
++            recovered[:, 1::2],
++            activation=self.spec.activation,
++        )
+         return signed_block_hadamard(
+             middle,
+             block_size=self.spec.postactivation_block_size,
+-            signs=self._signs(2, device=rows.device),
++            signs=self._signs(2, device=transformed.device),
+             dim=1,
+         )
+ 
+@@ -279,6 +335,24 @@ class CoupledHadamardExecution:
+             inverse=True,
+         )
+ 
++    def transform_outputs(self, rows: Tensor) -> Tensor:
++        """Map caller-coordinate outputs into the coupled residual basis.
++
++        Calibration uses this forward counterpart of :meth:`decode_output`
++        when constructing a candidate-conditioned down cross term.  The
++        candidate operand and source-model teacher must share the transformed
++        output coordinates before forming ``B = Q.T @ Y``.
++        """
++
++        if rows.ndim != 2 or rows.shape[1] != self.hidden:
++            raise ValueError("output rows do not match the hidden dimension")
++        return signed_block_hadamard(
++            rows,
++            block_size=self.spec.residual_block_size,
++            signs=self._signs(3, device=rows.device),
++            dim=1,
++        )
++
+     def execute(self, rows: Tensor, weights: WeightTriplet) -> Tensor:
+         _validate_triplet(weights)
+         transformed_rows = self.transform_inputs(rows)
+@@ -296,10 +370,12 @@ def coupled_execution(
+ __all__ = [
+     "CoupledHadamardExecution",
+     "CoupledHadamardSpec",
++    "ExpertActivation",
+     "WeightTriplet",
+     "block_hadamard",
+     "coupled_execution",
+     "encode_coupled_weights",
++    "expert_activation",
+     "rotation_signs",
+     "signed_block_hadamard",
+ ]
+diff --git a/tests/test_qsrt_coupled.py b/tests/test_qsrt_coupled.py
+index 8f3f6a3..2fef347 100644
+--- a/tests/test_qsrt_coupled.py
++++ b/tests/test_qsrt_coupled.py
+@@ -12,6 +12,7 @@ from qsrt.qsrt_coupled import (
+     block_hadamard,
+     coupled_execution,
+     encode_coupled_weights,
++    expert_activation,
+     rotation_signs,
+     signed_block_hadamard,
+ )
+@@ -75,6 +76,66 @@ def test_production_coupled_transform_closes_and_matches_research_oracle() -> No
+     assert torch.allclose(output, research_output, rtol=2e-5, atol=2e-5)
+ 
+ 
++def test_glm_silu_coupled_transform_closes_exact_expert_function() -> None:
++    generator = torch.Generator().manual_seed(5202)
++    rows, hidden, intermediate = 11, 24, 16
++    weights = (
++        torch.randn(intermediate, hidden, generator=generator),
++        torch.randn(intermediate, hidden, generator=generator),
++        torch.randn(hidden, intermediate, generator=generator),
++    )
++    inputs = torch.randn(rows, hidden, generator=generator)
++    spec = CoupledHadamardSpec(
++        residual_block_size=8,
++        preactivation_block_size=8,
++        postactivation_block_size=8,
++        intermediate_draw=6,
++        activation="silu",
++    )
++
++    encoded = encode_coupled_weights(weights, spec)
++    execution = coupled_execution(encoded, spec)
++    output = execution.execute(inputs, encoded)
++    gate = torch.nn.functional.linear(inputs, weights[0])
++    up = torch.nn.functional.linear(inputs, weights[1])
++    reference = torch.nn.functional.linear(
++        torch.nn.functional.silu(gate) * up,
++        weights[2],
++    )
++
++    assert torch.allclose(output, reference, rtol=2e-5, atol=2e-5)
++
++    transformed_input = execution.transform_inputs(inputs)
++    gate_output = torch.nn.functional.linear(transformed_input, encoded[0])
++    up_output = torch.nn.functional.linear(transformed_input, encoded[1])
++    middle_from_outputs = execution.decode_middle_outputs(
++        gate_output, up_output
++    )
++    middle_from_weights = execution.decode_middle(
++        transformed_input, encoded[0], encoded[1]
++    )
++    torch.testing.assert_close(
++        middle_from_outputs, middle_from_weights, rtol=0, atol=0
++    )
++
++    transformed_output = execution.transform_outputs(reference)
++    torch.testing.assert_close(
++        execution.decode_output(transformed_output), reference, rtol=2e-5, atol=2e-5
++    )
++
++
++def test_expert_activation_rejects_unknown_contract() -> None:
++    gate = torch.ones((2, 3))
++    up = torch.ones((2, 3))
++
++    try:
++        expert_activation(gate, up, activation="gelu")  # type: ignore[arg-type]
++    except ValueError as exc:
++        assert "unsupported expert activation" in str(exc)
++    else:
++        raise AssertionError("unknown expert activation was accepted")
++
++
+ def test_coupled_hessian_transforms_preserve_quadratic_forms() -> None:
+     generator = torch.Generator().manual_seed(91)
+     hidden, intermediate = 16, 12