Add KV-BSS (Key-Value Binding Softmax Sharpening) attention hook

HADAMARD_QUANT_INTEGRATION.md

@@ -0,0 +1,34 @@
+# MiniCPM5-2B-Hadamard-GSQ: KV-BSS integration
+
+This pull request contains a self-contained implementation of the custom
+MiniCPM Hadamard model path and the KV-BSS attention hook.
+
+## Included source files
+
+1. `configuration_minicpm_hadamard.py` defines the model and quantization
+   parameters, including the KV-BSS controls.
+2. `modeling_minicpm_hadamard.py` implements the model, causal masking, RoPE,
+   grouped-query attention, and legacy tuple cache support.
+3. `kv_bss.py` implements focus scaling and haze-floor filtering with explicit
+   shape validation and finite handling for fully masked rows.
+4. `test_inference.py` exercises a finite forward pass, GQA validation,
+   2D/4D attention-mask behavior, and cached-versus-uncached logit parity.
+
+## Verification
+
+Run from the model-code directory:
+
+```bash
+python -m unittest discover -s . -p 'test_inference.py' -v
+```
+
+The test is intentionally small and CPU-only. It validates implementation
+behavior without downloading a checkpoint and does not claim benchmark
+accuracy or long-context quality.
+
+## Scope
+
+The PR contains code and tests only. The separately published quantized
+checkpoint and its calibration report are linked from the model card:
+
+https://huggingface.co/F-Labs/MiniCPM5-2B-Hadamard-GSQ

configuration_minicpm_hadamard.py

@@ -0,0 +1,75 @@
+"""MiniCPMHadamardConfig — explicit bifurcation_rank added."""
+from transformers.configuration_utils import PretrainedConfig
+
+
+class MiniCPMHadamardConfig(PretrainedConfig):
+    model_type = "minicpm_hadamard"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        vocab_size=130560,
+        hidden_size=2048,
+        intermediate_size=6144,
+        num_hidden_layers=42,
+        num_attention_heads=16,
+        num_key_value_heads=2,
+        head_dim=128,
+        hidden_act="silu",
+        max_position_embeddings=131072,
+        initializer_range=0.02,
+        rms_norm_eps=1e-6,
+        use_cache=True,
+        pad_token_id=1,
+        bos_token_id=0,
+        eos_token_id=[1, 130073],
+        tie_word_embeddings=False,
+        rope_theta=5000000.0,
+        bits=4,
+        group_size=64,
+        hadamard_block_size=128,
+        residual_rank=16,
+        bifurcation_rank=24,
+        k_proj_rank=32,
+        layer_rank_map=None,
+        dense_tensor_names=None,
+        int8_tensor_names=None,
+        rotation_mode="fixed_hadamard",
+        rotation_seed=1729,
+        tau_focus=1.10,
+        haze_floor_margin=12.0,
+        **kwargs,
+    ):
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+        self.vocab_size = vocab_size
+        self.max_position_embeddings = max_position_embeddings
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.head_dim = head_dim
+        self.hidden_act = hidden_act
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
+        self.use_cache = use_cache
+        self.rope_theta = rope_theta
+        self.bits = bits
+        self.group_size = group_size
+        self.hadamard_block_size = hadamard_block_size
+        self.residual_rank = residual_rank
+        self.bifurcation_rank = bifurcation_rank
+        self.k_proj_rank = k_proj_rank
+        self.layer_rank_map = dict(layer_rank_map or {})
+        self.dense_tensor_names = list(dense_tensor_names or [])
+        self.int8_tensor_names = list(int8_tensor_names or [])
+        self.rotation_mode = rotation_mode
+        self.rotation_seed = int(rotation_seed)
+        self.tau_focus = tau_focus
+        self.haze_floor_margin = haze_floor_margin

kv_bss.py

@@ -0,0 +1,128 @@
+"""
+KV-BSS: Key-Value Binding Softmax Sharpening.
+"""
+
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class KVBSSAttentionHook(nn.Module):
+    def __init__(self, tau_focus: float = 1.10, haze_floor_margin: float = 12.0):
+        super().__init__()
+        if not math.isfinite(float(tau_focus)) or float(tau_focus) <= 0:
+            raise ValueError("tau_focus must be finite and greater than zero")
+        if not math.isfinite(float(haze_floor_margin)) or float(haze_floor_margin) < 0:
+            raise ValueError("haze_floor_margin must be finite and non-negative")
+        self.tau_focus = float(tau_focus)
+        self.haze_floor_margin = float(haze_floor_margin)
+
+    @staticmethod
+    def _prepare_mask(attention_mask, scores: torch.Tensor):
+        """Normalize 2D/4D masks and return additive mask plus hard-valid mask."""
+        b, h, q_len, kv_len = scores.shape
+        min_value = torch.finfo(scores.dtype).min
+
+        if attention_mask is None:
+            shape = (1, 1, 1, kv_len)
+            additive = torch.zeros(shape, dtype=scores.dtype, device=scores.device)
+            valid = torch.ones(shape, dtype=torch.bool, device=scores.device)
+            return additive, valid
+
+        mask = attention_mask.to(device=scores.device)
+        if mask.dim() == 2:
+            if tuple(mask.shape) != (b, kv_len):
+                raise ValueError(
+                    "2D attention_mask must have shape "
+                    f"({b}, {kv_len}), got {tuple(mask.shape)}"
+                )
+            valid = mask if mask.dtype == torch.bool else mask > 0
+            valid = valid[:, None, None, :]
+            additive = torch.where(
+                valid,
+                torch.zeros((), dtype=scores.dtype, device=scores.device),
+                torch.full((), min_value, dtype=scores.dtype, device=scores.device),
+            )
+            return additive, valid
+
+        if mask.dim() != 4:
+            raise ValueError("attention_mask must be None, 2D, or 4D")
+        try:
+            mask = torch.broadcast_to(mask, (b, h, q_len, kv_len))
+        except RuntimeError as exc:
+            raise ValueError(
+                "4D attention_mask is not broadcastable to "
+                f"{tuple(scores.shape)}, got {tuple(mask.shape)}"
+            ) from exc
+
+        if mask.dtype == torch.bool:
+            valid = mask
+            additive = torch.where(
+                valid,
+                torch.zeros((), dtype=scores.dtype, device=scores.device),
+                torch.full((), min_value, dtype=scores.dtype, device=scores.device),
+            )
+        else:
+            additive = mask.to(dtype=scores.dtype)
+            if torch.isnan(additive).any() or torch.isposinf(additive).any():
+                raise ValueError("attention_mask contains NaN or positive infinity")
+            # Hugging Face additive causal masks use finfo.min or -inf for
+            # blocked positions. Finite negative biases remain valid scores.
+            valid = torch.isfinite(additive) & (additive > min_value / 2)
+        return additive, valid
+
+    def forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attention_mask=None, scaling=None):
+        if query.dim() != 4 or key.dim() != 4 or value.dim() != 4:
+            raise ValueError("query, key, and value must all be 4D tensors")
+        b, h, q_len, d = query.shape
+        kb, kv_h, kv_len, kd = key.shape
+        vb, value_kv_h, value_kv_len, vd = value.shape
+        if (kb, value_kv_h, value_kv_len, vd) != (b, kv_h, kv_len, kd):
+            raise ValueError("key and value shapes must match in batch, heads, length, and depth")
+        if h == 0 or kv_h == 0 or h % kv_h != 0:
+            raise ValueError(f"query heads ({h}) must be a positive multiple of KV heads ({kv_h})")
+
+        if h != kv_h:
+            num_repeat = h // kv_h
+            key = key.repeat_interleave(num_repeat, dim=1)
+            value = value.repeat_interleave(num_repeat, dim=1)
+
+        if scaling is None:
+            scaling = 1.0 / math.sqrt(d)
+        if not math.isfinite(float(scaling)) or float(scaling) <= 0:
+            raise ValueError("scaling must be finite and greater than zero")
+
+        # Attention score accumulation in BF16 can overflow even when most
+        # individual activations are finite. Compute the score path in FP32 and
+        # isolate invalid query/key positions instead of taking down generation.
+        query_valid = torch.isfinite(query).all(dim=-1, keepdim=True)
+        key_valid = torch.isfinite(key).all(dim=-1).unsqueeze(-2)
+        query_fp32 = torch.nan_to_num(query.float(), nan=0.0, posinf=0.0, neginf=0.0)
+        key_fp32 = torch.nan_to_num(key.float(), nan=0.0, posinf=0.0, neginf=0.0)
+        scores = torch.matmul(query_fp32, key_fp32.transpose(-1, -2)) * float(scaling)
+        scores = scores * self.tau_focus
+        score_valid = torch.isfinite(scores)
+        scores = torch.nan_to_num(scores, nan=0.0, posinf=0.0, neginf=0.0)
+
+        additive_mask, hard_valid = self._prepare_mask(attention_mask, scores)
+        scores = scores + additive_mask
+        hard_valid = hard_valid.expand_as(scores)
+        hard_valid = hard_valid & score_valid & query_valid & key_valid
+
+        row_has_valid = hard_valid.any(dim=-1, keepdim=True)
+        neg_inf = torch.tensor(float("-inf"), dtype=scores.dtype, device=scores.device)
+        max_scores = scores.masked_fill(~hard_valid, neg_inf).amax(dim=-1, keepdim=True)
+        keep = hard_valid & (scores >= (max_scores - self.haze_floor_margin))
+        # Use the dtype minimum instead of a finite magic number so masked
+        # positions cannot receive probability mass at any supported dtype.
+        scores = scores.masked_fill(~keep, torch.finfo(scores.dtype).min)
+
+        probs = F.softmax(scores, dim=-1, dtype=torch.float32)
+        # A fully masked row is invalid input for ordinary softmax. Returning
+        # a finite zero vector keeps the failure contained and avoids NaN
+        # propagation through a whole decoder stack.
+        probs = probs * row_has_valid.to(dtype=probs.dtype)
+        value_fp32 = torch.nan_to_num(value.float(), nan=0.0, posinf=0.0, neginf=0.0)
+        out = torch.matmul(probs, value_fp32)
+        return torch.nan_to_num(out, nan=0.0, posinf=0.0, neginf=0.0).to(query.dtype)

modeling_minicpm_hadamard.py

@@ -0,0 +1,590 @@
+"""
+PyTorch modeling implementation for MiniCPM5-2B-Hadamard-GSQ.
+Engineered at F-Labs.
+
+Features:
+- HadamardLinear4bit: Group-wise INT4 with Walsh-Hadamard spin and Low-Rank SVD (SRC)
+- Int8Linear / DenseBF16Linear: explicit mixed-precision paths for measured risk tiers
+- ZeroCompressionShield: Pure BF16 RMSNorms, Biases, and Embeddings
+- Adaptive L0-L3 rank allocation from calibration, spectral spikes, jumps, and curvature
+- Key-Projection Sensitivity Defense (r=k_proj_rank on k_proj)
+- KVBSSAttentionHook: Key-Value Binding Softmax Sharpening for 128K context
+- RoPE (GPT-NeoX style, theta=rope_theta, head_dim=128) + causal mask + KV-cache
+- Full compliance with Hugging Face PreTrainedModel standards.
+"""
+
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Optional, Tuple
+from transformers.modeling_utils import PreTrainedModel
+from transformers.generation import GenerationMixin
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+
+try:
+    from .configuration_minicpm_hadamard import MiniCPMHadamardConfig
+    from .kv_bss import KVBSSAttentionHook
+except (ImportError, ValueError):
+    from configuration_minicpm_hadamard import MiniCPMHadamardConfig
+    from kv_bss import KVBSSAttentionHook
+
+_H_CACHE = {}
+
+
+def rademacher_signs(size: int, seed: int = 1729, dtype=torch.float32, device=None):
+    """Reproduce the calibration-time +/-1 input spin without storing a vector."""
+    generator = torch.Generator(device="cpu")
+    generator.manual_seed((int(seed) + 1009 * int(size)) % (2**63 - 1))
+    bits = torch.randint(0, 2, (size,), generator=generator, dtype=torch.int8)
+    return bits.to(torch.float32).mul_(2.0).sub_(1.0).to(device=device, dtype=dtype)
+
+
+def layer_residual_rank(config, layer_idx: int, fallback: int = 16) -> int:
+    """Read a measured per-layer rank map, retaining the legacy fallback."""
+    rank_map = getattr(config, "layer_rank_map", None)
+    if isinstance(rank_map, dict):
+        value = rank_map.get(str(layer_idx), rank_map.get(layer_idx))
+        if value is not None:
+            return max(0, int(value))
+    if 14 <= layer_idx <= 27:
+        return int(getattr(config, "bifurcation_rank", 24))
+    return int(getattr(config, "residual_rank", fallback))
+
+
+def get_hadamard_matrix(n: int, dtype=torch.float32, device=None):
+    key = (n, dtype, str(device))
+    if key in _H_CACHE:
+        return _H_CACHE[key]
+    if n == 1:
+        h = torch.tensor([[1.0]], dtype=dtype, device=device)
+    else:
+        h_half = get_hadamard_matrix(n // 2, dtype=dtype, device=device)
+        top = torch.cat([h_half, h_half], dim=1)
+        bottom = torch.cat([h_half, -h_half], dim=1)
+        h = torch.cat([top, bottom], dim=0) / math.sqrt(2.0)
+    _H_CACHE[key] = h
+    return h
+
+
+def apply_hadamard_rot(x: torch.Tensor, block_size: int = 128) -> torch.Tensor:
+    orig_shape = x.shape
+    d = orig_shape[-1]
+    if d % block_size != 0:
+        return x
+    h = get_hadamard_matrix(block_size, dtype=x.dtype, device=x.device)
+    reshaped = x.view(-1, d // block_size, block_size)
+    rotated = torch.matmul(reshaped, h)
+    return rotated.view(orig_shape)
+
+
+def apply_runtime_input_rotation(
+    x: torch.Tensor,
+    block_size: int,
+    rotation_mode: str,
+    rotation_signs: torch.Tensor,
+) -> torch.Tensor:
+    if x.shape[-1] % block_size != 0:
+        return x
+    if rotation_mode == "rademacher_hadamard":
+        x = x * rotation_signs.to(dtype=x.dtype, device=x.device)
+    return apply_hadamard_rot(x, block_size=block_size)
+
+
+def resolve_rotation_signs(
+    signs: torch.Tensor,
+    size: int,
+    seed: int,
+    x: torch.Tensor,
+) -> torch.Tensor:
+    """Recover non-persistent signs after low_cpu_mem_usage meta init."""
+    if signs.device.type == "meta" or signs.numel() != size or not bool(signs.any()):
+        return rademacher_signs(size, seed, dtype=x.dtype, device=x.device)
+    return signs.to(dtype=x.dtype, device=x.device)
+
+
+class DenseBF16Linear(nn.Module):
+    """Mixed-precision escape hatch for measured high-sensitivity projections."""
+
+    def __init__(self, in_features: int, out_features: int, block_size: int = 128,
+                 rotation_mode: str = "fixed_hadamard", rotation_seed: int = 1729):
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.block_size = block_size
+        self.rotation_mode = rotation_mode
+        self.rotation_seed = int(rotation_seed)
+        if rotation_mode == "rademacher_hadamard":
+            signs = rademacher_signs(in_features, self.rotation_seed, dtype=torch.float32)
+        else:
+            signs = torch.ones(in_features, dtype=torch.float32)
+        self.register_buffer("rotation_signs", signs, persistent=False)
+        self.register_buffer(
+            "weight_bf16",
+            torch.zeros((out_features, in_features), dtype=torch.bfloat16),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        signs = resolve_rotation_signs(
+            self.rotation_signs, self.in_features, self.rotation_seed, x
+        )
+        x_rot = apply_runtime_input_rotation(
+            x, self.block_size, self.rotation_mode, signs
+        )
+        return F.linear(x_rot, self.weight_bf16.to(dtype=x.dtype, device=x.device))
+
+
+class Int8Linear(nn.Module):
+    """Groupwise INT8 path for medium-risk projections."""
+
+    def __init__(self, in_features: int, out_features: int, group_size: int = 64,
+                 rank: int = 16, block_size: int = 128,
+                 rotation_mode: str = "fixed_hadamard", rotation_seed: int = 1729):
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.group_size = group_size
+        self.rank = rank
+        self.block_size = block_size
+        self.rotation_mode = rotation_mode
+        self.rotation_seed = int(rotation_seed)
+        if rotation_mode == "rademacher_hadamard":
+            signs = rademacher_signs(in_features, self.rotation_seed, dtype=torch.float32)
+        else:
+            signs = torch.ones(in_features, dtype=torch.float32)
+        self.register_buffer("rotation_signs", signs, persistent=False)
+        self.register_buffer("qweight_int8", torch.zeros((out_features, in_features), dtype=torch.int8))
+        self.register_buffer("scales_int8", torch.zeros((out_features, in_features // group_size), dtype=torch.bfloat16))
+        self.register_buffer("svd_a", torch.zeros((out_features, rank), dtype=torch.bfloat16))
+        self.register_buffer("svd_b", torch.zeros((rank, in_features), dtype=torch.bfloat16))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        signs = resolve_rotation_signs(
+            self.rotation_signs, self.in_features, self.rotation_seed, x
+        )
+        x_rot = apply_runtime_input_rotation(
+            x, self.block_size, self.rotation_mode, signs
+        )
+        m, n = self.qweight_int8.shape
+        w_float = (
+            self.qweight_int8.float().view(m, n // self.group_size, self.group_size)
+            * self.scales_int8.to(dtype=torch.float32).unsqueeze(-1)
+        ).view(m, n).to(dtype=x.dtype)
+        out_base = F.linear(x_rot, w_float)
+        x_svd = F.linear(x_rot, self.svd_b.to(dtype=x.dtype, device=x.device))
+        out_res = F.linear(x_svd, self.svd_a.to(dtype=x.dtype, device=x.device))
+        return out_base + out_res
+
+
+class HadamardLinear4bit(nn.Module):
+    """
+    4-bit Group-Scale Quantized Linear layer with:
+    1. Walsh-Hadamard Input Coordinate Spin (Outlier suppression)
+    2. Group-Scale INT4 quantization (group size 64)
+    3. Low-Rank Residual SVD Compensation: Y = X' W_quant^T + (X' B^T) A^T
+    """
+
+    def __init__(self, in_features: int, out_features: int, group_size: int = 64,
+                 rank: int = 16, block_size: int = 128, rotation_mode: str = "fixed_hadamard",
+                 rotation_seed: int = 1729):
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.group_size = group_size
+        self.rank = rank
+        self.block_size = block_size
+        self.rotation_mode = rotation_mode
+        self.rotation_seed = int(rotation_seed)
+
+        if self.rotation_mode == "rademacher_hadamard":
+            signs = rademacher_signs(
+                in_features, self.rotation_seed, dtype=torch.float32
+            )
+        else:
+            signs = torch.ones(in_features, dtype=torch.float32)
+        self.register_buffer("rotation_signs", signs, persistent=False)
+
+        self.register_buffer("qweight_packed", torch.zeros((out_features, in_features // 2), dtype=torch.uint8))
+        self.register_buffer("scales", torch.zeros((out_features, in_features // group_size), dtype=torch.bfloat16))
+        self.register_buffer("svd_a", torch.zeros((out_features, rank), dtype=torch.bfloat16))
+        self.register_buffer("svd_b", torch.zeros((rank, in_features), dtype=torch.bfloat16))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        signs = resolve_rotation_signs(
+            self.rotation_signs, self.in_features, self.rotation_seed, x
+        )
+        x_rot = apply_runtime_input_rotation(
+            x, self.block_size, self.rotation_mode, signs
+        )
+
+        low = (self.qweight_packed & 0x0F).to(torch.int8)
+        low = torch.where(low >= 8, low - 16, low)
+        high = ((self.qweight_packed >> 4) & 0x0F).to(torch.int8)
+        high = torch.where(high >= 8, high - 16, high)
+
+        m, half_n = self.qweight_packed.shape
+        w_int8 = torch.empty((m, half_n * 2), dtype=torch.int8, device=x.device)
+        w_int8[:, 0::2] = low
+        w_int8[:, 1::2] = high
+
+        w_float = (w_int8.float().view(m, self.in_features // self.group_size, self.group_size) * self.scales.unsqueeze(-1)).view(m, self.in_features).to(x.dtype)
+        out_base = F.linear(x_rot, w_float)
+
+        x_svd = F.linear(x_rot, self.svd_b.to(x.dtype))
+        out_res = F.linear(x_svd, self.svd_a.to(x.dtype))
+
+        return out_base + out_res
+
+
+def make_projection(config, full_name: str, in_features: int, out_features: int,
+                    rank: int, linear_kwargs):
+    dense_names = getattr(config, "dense_tensor_names", ())
+    int8_names = getattr(config, "int8_tensor_names", ())
+    if full_name in dense_names:
+        return DenseBF16Linear(in_features, out_features, **{
+            key: linear_kwargs[key]
+            for key in ("block_size", "rotation_mode", "rotation_seed")
+        })
+    if full_name in int8_names:
+        return Int8Linear(in_features, out_features, rank=rank, **linear_kwargs)
+    return HadamardLinear4bit(
+        in_features, out_features, rank=rank, **linear_kwargs
+    )
+
+
+class MiniCPMRMSNorm(nn.Module):
+    # ZeroCompressionShield: stays BF16, never quantized.
+    def __init__(self, hidden_size, eps=1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size, dtype=torch.bfloat16))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states):
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return (self.weight.to(torch.float32) * hidden_states).to(input_dtype)
+
+
+class MiniCPMHadamardRotaryEmbedding(nn.Module):
+    # GPT-NeoX style RoPE. inv_freq built from rope_theta; cached cos/sin up to max_pos.
+    def __init__(self, head_dim: int = 128, rope_theta: float = 5000000.0, max_position_embeddings: int = 131072):
+        super().__init__()
+        self.head_dim = head_dim
+        inv_freq = 1.0 / (rope_theta ** (torch.arange(0, head_dim, 2).float() / head_dim))
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self.max_position_embeddings = max_position_embeddings
+        self._cached_cos = None
+        self._cached_sin = None
+        self._cached_len = 0
+
+    def _build_cache(self, seq_len: int, device, dtype):
+        t = torch.arange(seq_len, device=device, dtype=self.inv_freq.dtype)
+        freqs = torch.outer(t, self.inv_freq)  # (seq_len, head_dim/2)
+        emb = torch.cat([freqs, freqs], dim=-1)  # (seq_len, head_dim)
+        self._cached_cos = emb.cos().to(dtype)
+        self._cached_sin = emb.sin().to(dtype)
+        self._cached_len = seq_len
+
+    def forward(self, x: torch.Tensor, position_ids: torch.Tensor):
+        # x: (b, heads, seq, head_dim) — only used for device/dtype.
+        seq_len = int(position_ids.max().item()) + 1
+        if self._cached_cos is None or seq_len > self._cached_len or self._cached_cos.device != x.device:
+            self._build_cache(seq_len, x.device, torch.float32)
+        cos = self._cached_cos[position_ids].to(x.dtype)  # (b, seq, head_dim)
+        sin = self._cached_sin[position_ids].to(x.dtype)
+        return cos.unsqueeze(1), sin.unsqueeze(1)  # (b, 1, seq, head_dim) — broadcast over heads
+
+
+def rotate_half(x: torch.Tensor) -> torch.Tensor:
+    d = x.shape[-1]
+    x1 = x[..., : d // 2]
+    x2 = x[..., d // 2 :]
+    return torch.cat([-x2, x1], dim=-1)
+
+
+def apply_rotary_pos_emb(q, k, cos, sin):
+    # cos/sin: (b, 1, seq, head_dim), broadcast over heads.
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+
+def _repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
+    if n_rep == 1:
+        return x
+    return x.repeat_interleave(n_rep, dim=1)
+
+
+def _make_causal_mask_4d(attention_mask_2d, q_len: int, kv_len: int, device, dtype) -> Optional[torch.Tensor]:
+    # Returns additive (b, 1, q_len, kv_len) mask: 0.0 keep / -inf block.
+    past_len = kv_len - q_len
+    causal = torch.full((q_len, kv_len), torch.finfo(dtype).min, device=device)
+    # allow j <= past_len + i
+    causal = torch.where(
+        torch.ones(q_len, kv_len, device=device, dtype=torch.bool).tril(diagonal=past_len),
+        torch.zeros((), device=device, dtype=dtype),
+        causal.to(dtype),
+    )  # (q, kv)
+    mask_4d = causal.unsqueeze(0).unsqueeze(0)  # (1, 1, q, kv)
+    if attention_mask_2d is not None:
+        am = attention_mask_2d.to(dtype)  # (b, kv_len)
+        additive = (1.0 - am) * torch.finfo(dtype).min  # pad -> -inf
+        mask_4d = mask_4d + additive[:, None, None, :]
+    return mask_4d
+
+
+class MiniCPMAttention(nn.Module):
+    def __init__(self, config: MiniCPMHadamardConfig, layer_idx: int = 0):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.head_dim = config.head_dim
+        self.num_key_value_heads = config.num_key_value_heads
+
+        res_rank = layer_residual_rank(config, layer_idx)
+        k_rank = getattr(config, "k_proj_rank", 32)
+
+        linear_kwargs = dict(
+            group_size=config.group_size,
+            block_size=getattr(config, "hadamard_block_size", 128),
+            rotation_mode=getattr(config, "rotation_mode", "fixed_hadamard"),
+            rotation_seed=getattr(config, "rotation_seed", 1729),
+        )
+        prefix = f"model.layers.{layer_idx}.self_attn"
+        self.k_proj = make_projection(config, f"{prefix}.k_proj", self.hidden_size, self.num_key_value_heads * self.head_dim, k_rank, linear_kwargs)
+        self.q_proj = make_projection(config, f"{prefix}.q_proj", self.hidden_size, self.num_heads * self.head_dim, res_rank, linear_kwargs)
+        self.v_proj = make_projection(config, f"{prefix}.v_proj", self.hidden_size, self.num_key_value_heads * self.head_dim, res_rank, linear_kwargs)
+        self.o_proj = make_projection(config, f"{prefix}.o_proj", self.num_heads * self.head_dim, self.hidden_size, res_rank, linear_kwargs)
+
+        self.kv_bss = KVBSSAttentionHook(tau_focus=config.tau_focus, haze_floor_margin=config.haze_floor_margin)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,  # 4D additive, built by model
+        position_ids: Optional[torch.Tensor] = None,
+        past_key_values=None,  # DynamicCache or tuple[(k,v)]
+        use_cache: bool = False,
+        cache_position: Optional[torch.Tensor] = None,
+        rotary_emb: Optional[MiniCPMHadamardRotaryEmbedding] = None,
+    ) -> Tuple[torch.Tensor, object]:
+        b, q_len, _ = hidden_states.shape
+        q = self.q_proj(hidden_states).view(b, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(hidden_states).view(b, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(hidden_states).view(b, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+
+        # RoPE on fresh q/k only (past already rotated).
+        if rotary_emb is not None and position_ids is not None:
+            cos, sin = rotary_emb(q, position_ids)
+            q, k = apply_rotary_pos_emb(q, k, cos, sin)
+
+        # Append KV cache.
+        if past_key_values is not None:
+            if hasattr(past_key_values, "update"):  # transformers DynamicCache / Cache
+                k, v = past_key_values.update(k, v, self.layer_idx)
+            else:  # legacy tuple / list
+                pk, pv = past_key_values[self.layer_idx] if past_key_values[self.layer_idx] is not None else (None, None)
+                if pk is not None:
+                    k = torch.cat([pk, k], dim=2)
+                    v = torch.cat([pv, v], dim=2)
+                    past_key_values[self.layer_idx] = (k, v)
+                elif use_cache:
+                    past_key_values[self.layer_idx] = (k, v)
+        attn_out = self.kv_bss(q, k, v, attention_mask=attention_mask)
+        attn_out = attn_out.transpose(1, 2).contiguous().view(b, q_len, -1)
+        out = self.o_proj(attn_out)
+        return out, past_key_values
+
+
+class MiniCPMMLP(nn.Module):
+    def __init__(self, config: MiniCPMHadamardConfig, layer_idx: int = 0):
+        super().__init__()
+        res_rank = layer_residual_rank(config, layer_idx)
+
+        linear_kwargs = dict(
+            group_size=config.group_size,
+            block_size=getattr(config, "hadamard_block_size", 128),
+            rotation_mode=getattr(config, "rotation_mode", "fixed_hadamard"),
+            rotation_seed=getattr(config, "rotation_seed", 1729),
+        )
+        prefix = f"model.layers.{layer_idx}.mlp"
+        self.gate_proj = make_projection(config, f"{prefix}.gate_proj", config.hidden_size, config.intermediate_size, res_rank, linear_kwargs)
+        self.up_proj = make_projection(config, f"{prefix}.up_proj", config.hidden_size, config.intermediate_size, res_rank, linear_kwargs)
+        self.down_proj = make_projection(config, f"{prefix}.down_proj", config.intermediate_size, config.hidden_size, res_rank, linear_kwargs)
+
+    def forward(self, x):
+        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class MiniCPMDecoderLayer(nn.Module):
+    def __init__(self, config: MiniCPMHadamardConfig, layer_idx: int):
+        super().__init__()
+        self.input_layernorm = MiniCPMRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.self_attn = MiniCPMAttention(config, layer_idx=layer_idx)
+        self.post_attention_layernorm = MiniCPMRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.mlp = MiniCPMMLP(config, layer_idx=layer_idx)
+
+    def forward(self, hidden_states, attention_mask=None, position_ids=None,
+                past_key_values=None, use_cache=False, cache_position=None, rotary_emb=None):
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        hidden_states, past_key_values = self.self_attn(
+            hidden_states, attention_mask=attention_mask, position_ids=position_ids,
+            past_key_values=past_key_values, use_cache=use_cache,
+            cache_position=cache_position, rotary_emb=rotary_emb,
+        )
+        hidden_states = residual + hidden_states
+
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+        return hidden_states, past_key_values
+
+
+class MiniCPMHadamardPreTrainedModel(PreTrainedModel):
+    config_class = MiniCPMHadamardConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = False
+    _no_split_modules = ["MiniCPMDecoderLayer"]
+
+    def _init_weights(self, module):
+        """Initialize scratch models while leaving quantized buffers untouched."""
+        if isinstance(module, (nn.Linear, nn.Embedding)):
+            std = float(getattr(self.config, "initializer_range", 0.02))
+            with torch.no_grad():
+                module.weight.normal_(mean=0.0, std=std)
+                if getattr(module, "bias", None) is not None:
+                    module.bias.zero_()
+        elif isinstance(module, MiniCPMRMSNorm):
+            with torch.no_grad():
+                module.weight.fill_(1.0)
+
+
+class MiniCPMHadamardModel(MiniCPMHadamardPreTrainedModel):
+    def __init__(self, config: MiniCPMHadamardConfig):
+        super().__init__(config)
+        self.padding_idx = getattr(config, "pad_token_id", 1)
+        self.vocab_size = config.vocab_size
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+        self.layers = nn.ModuleList([
+            MiniCPMDecoderLayer(config, idx) for idx in range(config.num_hidden_layers)
+        ])
+        self.norm = MiniCPMRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.rotary_emb = MiniCPMHadamardRotaryEmbedding(
+            head_dim=config.head_dim, rope_theta=config.rope_theta,
+            max_position_embeddings=config.max_position_embeddings,
+        )
+        self.post_init()
+
+    def forward(self, input_ids=None, attention_mask=None, position_ids=None,
+                past_key_values=None, use_cache=None, cache_position=None,
+                inputs_embeds=None, **kwargs):
+        use_cache = self.config.use_cache if use_cache is None else use_cache
+        if inputs_embeds is not None:
+            x = inputs_embeds
+            b, q_len = x.shape[:2]
+            device = x.device
+        else:
+            b, q_len = input_ids.shape
+            device = input_ids.device
+            x = self.embed_tokens(input_ids)
+
+        # Past length for position ids / mask.
+        if past_key_values is not None and hasattr(past_key_values, "get_seq_length"):
+            past_len = past_key_values.get_seq_length()
+        elif past_key_values is not None and isinstance(past_key_values, (list, tuple)) and len(past_key_values) > 0 and past_key_values[0] is not None:
+            past_len = past_key_values[0][0].shape[2]
+        else:
+            past_len = 0
+
+        # NOTE (transformers>=5 compat): generate() passes trimmed input_ids with
+        # FULL-length position_ids (and sometimes a stale cache_position). Trusting
+        # them blindly lets RoPE broadcast-expand the query length, which breaks
+        # KV-BSS masking (scores kv != mask kv). Rebuild whenever shapes disagree.
+        if cache_position is None or cache_position.shape[0] != q_len:
+            cache_position = torch.arange(past_len, past_len + q_len, device=device)
+        if position_ids is None or position_ids.shape[-1] != q_len:
+            position_ids = cache_position.unsqueeze(0).expand(b, -1)
+
+        kv_len = past_len + q_len
+        mask_4d = None
+        if q_len > 1 or attention_mask is not None:
+            # Full 2D mask over kv window if user passed one, else causal only.
+            am_2d = None
+            if attention_mask is not None:
+                if attention_mask.dim() == 4:
+                    mask_4d = attention_mask
+                elif attention_mask.dim() == 2:
+                    am_2d = attention_mask
+            if mask_4d is None:
+                mask_4d = _make_causal_mask_4d(am_2d, q_len, kv_len, device, x.dtype)
+
+        # Legacy tuple cache init on first use_cache call.
+        if use_cache and past_key_values is None:
+            past_key_values = [None] * self.config.num_hidden_layers
+
+        for layer in self.layers:
+            x, past_key_values = layer(
+                x, attention_mask=mask_4d, position_ids=position_ids,
+                past_key_values=past_key_values, use_cache=use_cache,
+                cache_position=cache_position, rotary_emb=self.rotary_emb,
+            )
+        x = self.norm(x)
+        return BaseModelOutputWithPast(last_hidden_state=x, past_key_values=past_key_values)
+
+
+class MiniCPMHadamardForCausalLM(MiniCPMHadamardPreTrainedModel, GenerationMixin):
+    def __init__(self, config: MiniCPMHadamardConfig):
+        super().__init__(config)
+        self.model = MiniCPMHadamardModel(config)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def forward(self, input_ids=None, attention_mask=None, position_ids=None,
+                past_key_values=None, use_cache=None, cache_position=None,
+                inputs_embeds=None, labels=None, **kwargs):
+        hidden = self.model(
+            input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids,
+            past_key_values=past_key_values, use_cache=use_cache,
+            cache_position=cache_position, inputs_embeds=inputs_embeds, **kwargs,
+        )
+        logits = self.lm_head(hidden.last_hidden_state)
+        loss = None
+        if labels is not None:
+            loss = F.cross_entropy(logits.view(-1, logits.shape[-1]), labels.view(-1), ignore_index=-100)
+        return CausalLMOutputWithPast(
+            loss=loss, logits=logits, past_key_values=hidden.past_key_values,
+        )
+
+    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, attention_mask=None, **kwargs):
+        # Trim to last token when cache is active.
+        if past_key_values is not None:
+            if isinstance(past_key_values, (list, tuple)):
+                past_len = past_key_values[0][0].shape[2] if past_key_values[0] is not None else 0
+            elif hasattr(past_key_values, "get_seq_length"):
+                past_len = past_key_values.get_seq_length()
+            else:
+                past_len = 0
+            if input_ids.shape[1] > past_len:
+                input_ids = input_ids[:, past_len:]
+        cache_position = kwargs.get("cache_position", None)
+        return {"input_ids": input_ids, "past_key_values": past_key_values,
+                "attention_mask": attention_mask, "cache_position": cache_position,
+                "position_ids": kwargs.get("position_ids", None), "use_cache": True}

test_inference.py

@@ -0,0 +1,93 @@
+"""Self-contained numerical checks for the MiniCPM5-2B-Hadamard-GSQ code."""
+
+import unittest
+import torch
+from configuration_minicpm_hadamard import MiniCPMHadamardConfig
+from modeling_minicpm_hadamard import MiniCPMHadamardForCausalLM
+from kv_bss import KVBSSAttentionHook
+
+
+def tiny_config(**overrides):
+    values = dict(
+        hidden_size=32,
+        intermediate_size=64,
+        num_hidden_layers=2,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        head_dim=8,
+        vocab_size=97,
+        group_size=8,
+        hadamard_block_size=8,
+        max_position_embeddings=64,
+        use_cache=True,
+    )
+    values.update(overrides)
+    return MiniCPMHadamardConfig(**values)
+
+
+class MiniCPMNumericalTests(unittest.TestCase):
+    def test_architecture_forward_is_finite(self):
+        torch.manual_seed(0)
+        model = MiniCPMHadamardForCausalLM(tiny_config()).eval()
+        input_ids = torch.randint(0, 97, (1, 8))
+        with torch.no_grad():
+            out = model(input_ids, use_cache=False)
+        self.assertEqual(tuple(out.logits.shape), (1, 8, 97))
+        self.assertTrue(torch.isfinite(out.logits).all())
+
+    def test_kv_bss_mask_and_gqa_validation(self):
+        hook = KVBSSAttentionHook(tau_focus=1.0, haze_floor_margin=100.0)
+        query = torch.zeros(1, 2, 2, 4)
+        key = torch.zeros(1, 1, 3, 4)
+        value = torch.arange(12, dtype=torch.float32).view(1, 1, 3, 4)
+        # A 2D keep/pad mask must never route the padded value to the output.
+        out = hook(query, key, value, attention_mask=torch.tensor([[1, 1, 0]]))
+        expected = value[:, :, :2].mean(dim=2).unsqueeze(2).expand_as(out)
+        self.assertTrue(torch.allclose(out, expected))
+        # Fully masked rows remain finite and resolve to a zero vector.
+        blocked = torch.zeros(1, 3)
+        out_blocked = hook(query, key, value, attention_mask=blocked)
+        self.assertTrue(torch.isfinite(out_blocked).all())
+        self.assertEqual(float(out_blocked.abs().sum()), 0.0)
+        with self.assertRaises(ValueError):
+            bad_key = torch.zeros(1, 3, 3, 4)
+            bad_value = torch.zeros(1, 3, 3, 4)
+            hook(query, bad_key, bad_value)
+
+    def test_cached_and_uncached_logits_match(self):
+        torch.manual_seed(1)
+        model = MiniCPMHadamardForCausalLM(tiny_config()).eval()
+        input_ids = torch.tensor([[3, 7, 11, 19, 23]])
+        with torch.no_grad():
+            full = model(input_ids, use_cache=False).logits
+            prefix = model(input_ids[:, :3], use_cache=True)
+            suffix = model(
+                input_ids[:, 3:],
+                past_key_values=prefix.past_key_values,
+                use_cache=True,
+            )
+        self.assertTrue(torch.isfinite(suffix.logits).all())
+        torch.testing.assert_close(
+            suffix.logits, full[:, 3:], rtol=2e-4, atol=2e-5
+        )
+
+    def test_invalid_mask_is_rejected(self):
+        hook = KVBSSAttentionHook()
+        q = torch.zeros(1, 1, 1, 4)
+        k = torch.zeros(1, 1, 2, 4)
+        v = torch.zeros(1, 1, 2, 4)
+        with self.assertRaises(ValueError):
+            hook(q, k, v, attention_mask=torch.ones(1, 3))
+
+    def test_nonfinite_scores_are_contained(self):
+        hook = KVBSSAttentionHook(tau_focus=1.0, haze_floor_margin=12.0)
+        q = torch.zeros(1, 2, 2, 4, dtype=torch.bfloat16)
+        k = torch.zeros(1, 1, 2, 4, dtype=torch.bfloat16)
+        v = torch.ones(1, 1, 2, 4, dtype=torch.bfloat16)
+        q[0, 0, 0, 0] = float("nan")
+        k[0, 0, 1, 0] = float("inf")
+        out = hook(q, k, v)
+        self.assertTrue(torch.isfinite(out).all())
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)