dburner/Qwen3.5-4B-Q8_0-FlashNgram-NativeBridgeV3

Qwen3.5-4B-Q8_0-FlashNgram-NativeBridgeV3-budget042-step2500-fixed.gguf

key[t,l] = RMSNorm_key38(W_key38(p[t]))[l]

q_shared[t] = W_query35_to_PLE(h35[t])
query[t,l] = RMSNorm_query38(q_shared[t] * lane_scale[l] + lane_bias[l])

gate[t,l] = sigmoid((signed_sqrt(dot(key[t,l], query[t,l]) / sqrt(2560))
                     + gate_bias) / temperature)

z[t,l] = gate[t,l] * RMSNorm_conv38(W_value38(p[t]))
feature[t,l] = z[t,l] + SiLU(DepthwiseCausalConv38(z)[t,l])

feature35[t] = W_outputPLE_to_35(sum_l(softmax(lane_logits)[l] * feature[t,l]))
delta[t] = alpha * residual_budget * RMS(h35[t]) * feature35[t]
h35'[t] = h35[t] + delta[t]

flowchart TB
    P["N-gram lookup<br/>pₜ ∈ R²⁵⁶⁰"]

    subgraph Q38["Original Qwen 3.8 Flash-Next PLE"]
        direction TB

        H38["Native HyperConnection residual<br/>h₃₈,t ∈ R¹⁰²⁴⁰"]

        H38 --> Q38L["Reshape into four native lanes"]
        Q38L --> Q380["query lane 0 · R²⁵⁶⁰"]
        Q38L --> Q381["query lane 1 · R²⁵⁶⁰"]
        Q38L --> Q382["query lane 2 · R²⁵⁶⁰"]
        Q38L --> Q383["query lane 3 · R²⁵⁶⁰"]

        P --> WK38["Frozen packed W_key38<br/>R²⁵⁶⁰ → R¹⁰²⁴⁰"]
        WK38 --> K38L["Four key lanes"]
        K38L --> K380["key lane 0 · R²⁵⁶⁰"]
        K38L --> K381["key lane 1 · R²⁵⁶⁰"]
        K38L --> K382["key lane 2 · R²⁵⁶⁰"]
        K38L --> K383["key lane 3 · R²⁵⁶⁰"]

        Q380 --> G38["four lane-wise<br/>key/query gates"]
        Q381 --> G38
        Q382 --> G38
        Q383 --> G38
        K380 --> G38
        K381 --> G38
        K382 --> G38
        K383 --> G38

        G38 --> F38["Native four-lane PLE feature"]
    end

    subgraph V3["Our Qwen 3.5 → Qwen 3.8 NativeBridge V3"]
        direction TB

        H35["Frozen Qwen 3.5 residual<br/>h₃₅,t ∈ R²⁵⁶⁰"]
        H35 --> WQ["Trained shared query projection<br/>W_query35→PLE · R²⁵⁶⁰ → R²⁵⁶⁰"]
        WQ --> QS["shared query · R²⁵⁶⁰"]

        QS --> A0["lane 0: q × scale₀ + bias₀"]
        QS --> A1["lane 1: q × scale₁ + bias₁"]
        QS --> A2["lane 2: q × scale₂ + bias₂"]
        QS --> A3["lane 3: q × scale₃ + bias₃"]

        A0 --> NQ["Frozen Qwen 3.8 query RMSNorm"]
        A1 --> NQ
        A2 --> NQ
        A3 --> NQ

        P --> WK["Frozen packed W_key38<br/>R²⁵⁶⁰ → R¹⁰²⁴⁰"]
        WK --> KL["four frozen key lanes"]

        NQ --> G["four gate values<br/>one per token × lane"]
        KL --> G

        P --> WV["Frozen W_value38<br/>R²⁵⁶⁰ → R²⁵⁶⁰"]
        WV --> RV["repeat value into four lanes"]

        G --> Z["gate × normalized value<br/>four lanes"]
        RV --> Z
        Z --> C["Frozen 4-lane causal<br/>depthwise convolution"]
        C --> F["z + SiLU(conv(z))<br/>four PLE feature lanes"]

        F --> MIX["Trained softmax lane mix<br/>R¹⁰²⁴⁰ → R²⁵⁶⁰"]
        MIX --> WO["Trained output bridge<br/>W_outputPLE→35"]
        WO --> D["deltaₜ = alpha × budget × RMS(h₃₅,t) × feature"]
        H35 --> ADD["Add before Qwen 3.5 block 2"]
        D --> ADD
    end

    style Q38 fill:#274860,color:#fff
    style V3 fill:#4b365b,color:#fff
    style WQ fill:#7c4d2f,color:#fff
    style MIX fill:#7c4d2f,color:#fff
    style WO fill:#7c4d2f,color:#fff
    style WK38 fill:#31594a,color:#fff
    style WK fill:#31594a,color:#fff
    style WV fill:#31594a,color:#fff
    style C fill:#31594a,color:#fff

gate mean ≈ 0.33–0.35
gate standard deviation ≈ 0.32–0.33
gate ≤ 0.1: ≈ 31–35% of token × lane values
gate ≥ 0.5: ≈ 30–34% of token × lane values
gate ≥ 0.9: ≈ 6–8% of token × lane values

delta / hidden = RMS(delta_t) / RMS(h35_t)

https://github.com/dburner/llama.cpp/tree/feature/qwen35-native-ple-bridge-v3

.\llama-server.exe `
  -m .\Qwen3.5-4B-Q8_0-FlashNgram-NativeBridgeV3-budget042-step2500-fixed.gguf `
  --n-gpu-layers 99 `
  --ctx-size 48192 `
  --flash-attn on `
  --jinja `
  --reasoning on `
  --ngram-gate 1 `
  --spec-type none `
  --host 127.0.0.1 `
  --port 8080

"""Frozen Qwen 3.8 PLE components and a trainable Qwen 3.5 bridge.

Only the compact PLE projections/norms/convolution are loaded here.  The huge
n-gram table remains in :mod:`ple_store` and is supplied as 2560-wide PLE
embeddings by the existing lookup wrapper.
"""
from __future__ import annotations

import math
from pathlib import Path

import torch
from safetensors.torch import load_file
from torch import nn
from torch.nn import functional as F


HIDDEN_SIZE = 2560
HC_COUNT = 4
HC_HIDDEN_SIZE = HIDDEN_SIZE * HC_COUNT


def _rms(x: torch.Tensor) -> torch.Tensor:
    return x.float().square().mean(dim=-1, keepdim=True).sqrt()


def _rms_norm(x: torch.Tensor, weight: torch.Tensor, eps: float = 1e-6) -> torch.Tensor:
    inv_rms = torch.rsqrt(x.float().square().mean(dim=-1, keepdim=True) + eps)
    return (x.float() * inv_rms) * weight.float()


class FrozenQwen38Ple(nn.Module):
    """Non-persistent frozen source PLE tensors extracted from Flash-Next."""

    REQUIRED = {
        "key_weight", "value_weight", "key_norm_weight", "query_norm_weight",
        "conv_norm_weight", "conv_weight",
    }

    def __init__(self, tensors: dict[str, torch.Tensor]) -> None:
        super().__init__()
        missing = self.REQUIRED - set(tensors)
        unexpected = set(tensors) - self.REQUIRED
        if missing or unexpected:
            raise ValueError(f"native PLE tensors mismatch; missing={sorted(missing)}, unexpected={sorted(unexpected)}")
        expected = {
            "key_weight": (HC_HIDDEN_SIZE, HIDDEN_SIZE),
            "value_weight": (HIDDEN_SIZE, HIDDEN_SIZE),
            "key_norm_weight": (HC_HIDDEN_SIZE,),
            "query_norm_weight": (HC_HIDDEN_SIZE,),
            "conv_norm_weight": (HC_HIDDEN_SIZE,),
            "conv_weight": (HC_HIDDEN_SIZE, 1, 4),
        }
        for name, shape in expected.items():
            value = tensors[name]
            if tuple(value.shape) != shape:
                raise ValueError(f"{name} shape {tuple(value.shape)} != {shape}")
            if not value.is_floating_point() or not torch.isfinite(value).all():
                raise ValueError(f"{name} must be finite floating point")
            # These buffers move with .to(), but are deliberately excluded
            # from checkpoints: each run names its immutable source file.
            self.register_buffer(name, value.contiguous(), persistent=False)

    @classmethod
    def from_safetensors(cls, path: Path) -> "FrozenQwen38Ple":
        if not path.is_file():
            raise FileNotFoundError(f"native Qwen 3.8 PLE source does not exist: {path}")
        return cls(load_file(path, device="cpu"))

    def key(self, ple: torch.Tensor) -> torch.Tensor:
        return _rms_norm(F.linear(ple.float(), self.key_weight.float()), self.key_norm_weight)

    def query_norm(self, query: torch.Tensor) -> torch.Tensor:
        return _rms_norm(query, self.query_norm_weight)

    def value(self, ple: torch.Tensor) -> torch.Tensor:
        return F.linear(ple.float(), self.value_weight.float())

    def conv_norm(self, value_hc: torch.Tensor) -> torch.Tensor:
        return _rms_norm(value_hc, self.conv_norm_weight)

    def causal_conv(self, x: torch.Tensor) -> torch.Tensor:
        """Apply the source depthwise kernel with explicit left-only padding."""
        if x.ndim != 3 or x.shape[-1] != HC_HIDDEN_SIZE:
            raise ValueError(f"native conv expects [batch, tokens, {HC_HIDDEN_SIZE}]")
        channels = x.transpose(1, 2)
        padded = F.pad(channels, (9, 0))  # dilation 3 * (kernel 4 - 1)
        return F.conv1d(padded, self.conv_weight.float(), dilation=3,
                        groups=HC_HIDDEN_SIZE).transpose(1, 2)


class Qwen35NativePleBridge(nn.Module):
    """Trainable Qwen 3.5 <-> frozen Qwen 3.8 four-lane PLE bridge (v3).

    ``gate_bias`` is deliberately a single trainable scalar, rather than a
    modification of a Qwen 3.8 tensor.  It calibrates how often the frozen
    source PLE is used for this new Qwen 3.5 residual stream.  Temperature is
    fixed per run so that the bridge cannot evade a gate-usage objective merely
    by making its logits arbitrarily sharp.
    """

    def __init__(self, native_ple: FrozenQwen38Ple, *, hidden_size: int = HIDDEN_SIZE,
                 ple_embed_dim: int = HIDDEN_SIZE, residual_budget: float = 0.05,
                 gate_bias_init: float = 0.0, gate_temperature: float = 1.0) -> None:
        super().__init__()
        if hidden_size != HIDDEN_SIZE or ple_embed_dim != HIDDEN_SIZE:
            raise ValueError("native bridge requires 2560-wide Qwen 3.5 residuals and PLE embeddings")
        if not math.isfinite(residual_budget) or residual_budget <= 0:
            raise ValueError("residual_budget must be finite and positive")
        if not math.isfinite(gate_bias_init):
            raise ValueError("gate_bias_init must be finite")
        if not math.isfinite(gate_temperature) or gate_temperature <= 0:
            raise ValueError("gate_temperature must be finite and positive")
        self.hidden_size = hidden_size
        self.ple_embed_dim = ple_embed_dim
        self.hc_count = HC_COUNT
        # Unlike a Python float, the budget is checkpointed.  That is required
        # for a continuation phase that ramps it: evaluation must reconstruct
        # the exact budget of the selected checkpoint, not just the run's
        # starting value.
        self.register_buffer("residual_budget", torch.tensor(residual_budget, dtype=torch.float32), persistent=True)
        self.native_ple = native_ple
        self.query_proj = nn.Linear(hidden_size, hidden_size, bias=False)
        self.query_lane_scale = nn.Parameter(torch.ones(HC_COUNT, hidden_size, dtype=torch.float32))
        self.query_lane_bias = nn.Parameter(torch.zeros(HC_COUNT, hidden_size, dtype=torch.float32))
        self.output_lane_logits = nn.Parameter(torch.zeros(HC_COUNT, hidden_size, dtype=torch.float32))
        self.output_proj = nn.Linear(hidden_size, hidden_size, bias=False)
        self.gate_bias = nn.Parameter(torch.tensor(gate_bias_init, dtype=torch.float32))
        # This is persistent so a checkpoint carries its exact calibration,
        # but is intentionally not trainable.
        self.register_buffer("gate_temperature", torch.tensor(gate_temperature, dtype=torch.float32), persistent=True)
        self.alpha = nn.Parameter(torch.zeros((), dtype=torch.float32))
        self.register_buffer("architecture_version", torch.tensor(3, dtype=torch.int32), persistent=True)
        self.last_activation_metrics: dict[str, torch.Tensor] = {}
        self.gate_mean_for_regularization: torch.Tensor | None = None

    def load_state_dict(self, state_dict: dict[str, torch.Tensor], strict: bool = True):
        """Load pre-calibration v3 checkpoints as an explicitly neutral gate.

        Early v3 runs did not contain calibration tensors.  Their behaviour is
        exactly represented by bias=0 and temperature=1, which also lets them
        remain valid alpha-ablation baselines after this extension.
        """
        compatible_state = dict(state_dict)
        compatible_state.setdefault("gate_bias", torch.zeros_like(self.gate_bias))
        compatible_state.setdefault("gate_temperature", torch.ones_like(self.gate_temperature))
        compatible_state.setdefault("residual_budget", self.residual_budget.detach().clone())
        return super().load_state_dict(compatible_state, strict=strict)

    def set_residual_budget(self, value: float) -> None:
        """Set the frozen global residual multiplier for the current step."""
        if not math.isfinite(value) or value <= 0:
            raise ValueError("residual budget must be finite and positive")
        with torch.no_grad():
            self.residual_budget.fill_(value)

    def forward(self, hidden_states: torch.Tensor, ple_embeddings: torch.Tensor) -> torch.Tensor:
        if hidden_states.ndim != 3 or ple_embeddings.ndim != 3:
            raise ValueError("hidden_states and ple_embeddings must have shape [batch, tokens, width]")
        if hidden_states.shape[:2] != ple_embeddings.shape[:2]:
            raise ValueError("hidden_states and ple_embeddings must have matching batch and token dimensions")
        if hidden_states.shape[-1] != self.hidden_size or ple_embeddings.shape[-1] != self.ple_embed_dim:
            raise ValueError("native bridge width mismatch")
        batch, tokens, _ = hidden_states.shape

        key = self.native_ple.key(ple_embeddings).reshape(batch, tokens, HC_COUNT, HIDDEN_SIZE)
        query_shared = self.query_proj(hidden_states.float())
        query_hc = query_shared.unsqueeze(-2) * self.query_lane_scale + self.query_lane_bias
        query = self.native_ple.query_norm(query_hc.reshape(batch, tokens, HC_HIDDEN_SIZE))
        query = query.reshape(batch, tokens, HC_COUNT, HIDDEN_SIZE)
        score = (key.float() * query.float()).sum(dim=-1, keepdim=True) / math.sqrt(HIDDEN_SIZE)
        signed_sqrt = score.sign() * score.abs().clamp_min(torch.finfo(torch.float32).tiny).sqrt()
        raw_gate = torch.sigmoid(signed_sqrt)
        gate_logits = (signed_sqrt + self.gate_bias) / self.gate_temperature
        gate = torch.sigmoid(gate_logits)
        # Retain the differentiable batch mean only until the trainer has
        # formed its optional usage regularizer for this forward pass.
        self.gate_mean_for_regularization = gate.mean()

        value = self.native_ple.value(ple_embeddings)
        value_hc = value.unsqueeze(-2).expand(-1, -1, HC_COUNT, -1)
        value_norm = self.native_ple.conv_norm(value_hc.reshape(batch, tokens, HC_HIDDEN_SIZE))
        z = gate.to(dtype=value_norm.dtype) * value_norm.reshape(batch, tokens, HC_COUNT, HIDDEN_SIZE)
        conv = F.silu(self.native_ple.causal_conv(z.reshape(batch, tokens, HC_HIDDEN_SIZE)))
        conv_hc = conv.reshape(batch, tokens, HC_COUNT, HIDDEN_SIZE)
        feature_hc = z + conv_hc
        lane_weights = torch.softmax(self.output_lane_logits, dim=0)
        mixed = (feature_hc.float() * lane_weights.unsqueeze(0).unsqueeze(0)).sum(dim=-2)
        feature35 = self.output_proj(mixed)
        # Gate is already inside feature35.  Do not divide by its RMS here.
        hidden_scale = _rms(hidden_states).detach()
        injected = (self.alpha.to(dtype=feature35.dtype) * self.residual_budget.to(dtype=feature35.dtype)
                    * hidden_scale.to(dtype=feature35.dtype) * feature35)

        with torch.no_grad():
            hidden_global_rms = hidden_states.detach().float().square().mean().sqrt()
            eps = torch.finfo(torch.float32).tiny
            def relative_rms(x: torch.Tensor) -> torch.Tensor:
                return x.detach().float().square().mean().sqrt() / hidden_global_rms.clamp_min(eps)
            gate_float = gate.detach().float()
            raw_gate_float = raw_gate.detach().float()
            injected_float = injected.detach().float()
            injected_rms = injected_float.square().mean().sqrt()
            self.last_activation_metrics = {
                "gate_mean": gate_float.mean(),
                "gate_std": gate_float.std(unbiased=False),
                "raw_gate_mean": raw_gate_float.mean(),
                "gate_bias": self.gate_bias.detach().float(),
                "gate_temperature": self.gate_temperature.detach().float(),
                "residual_budget": self.residual_budget.detach().float(),
                "gate_fraction_le_0_1": (gate_float <= 0.1).float().mean(),
                "gate_fraction_ge_0_5": (gate_float >= 0.5).float().mean(),
                "gate_fraction_ge_0_9": (gate_float >= 0.9).float().mean(),
                "value_to_hidden_rms": relative_rms(value),
                "z_to_hidden_rms": relative_rms(z),
                "conv_to_hidden_rms": relative_rms(conv_hc),
                "feature_hc_to_hidden_rms": relative_rms(feature_hc),
                "bridge_feature_rms": feature35.detach().float().square().mean().sqrt(),
                "injected_rms": injected_rms,
                "injected_to_hidden_rms": injected_rms / hidden_global_rms.clamp_min(eps),
            }
        return injected.to(dtype=hidden_states.dtype)