Upload 5 files

__init__.py

@@ -0,0 +1,59 @@
+"""Transformers-compatible Multiscreen implementation.
+
+This package ports the core architecture from ``dieOD/multiscreen-pytorch`` to
+Hugging Face Transformers-style ``PreTrainedConfig`` / ``PreTrainedModel``
+classes.
+"""
+
+from .configuration_multiscreen import MultiscreenConfig
+from .compile_utils import find_msvc_cl, load_vcvars_env, setup_compile_env
+from .data import PackedTextDataset
+from .modeling_multiscreen import (
+    GatedScreeningBlock,
+    MultiscreenForCausalLM,
+    MultiscreenLayer,
+    MultiscreenModel,
+    MultiscreenPreTrainedModel,
+    ScreeningCache,
+    convert_original_state_dict_for_causal_lm,
+    convert_original_state_dict_for_model,
+)
+
+__version__ = "0.1.2"
+
+__all__ = [
+    "MultiscreenConfig",
+    "MultiscreenPreTrainedModel",
+    "MultiscreenModel",
+    "MultiscreenForCausalLM",
+    "MultiscreenLayer",
+    "GatedScreeningBlock",
+    "ScreeningCache",
+    "convert_original_state_dict_for_causal_lm",
+    "convert_original_state_dict_for_model",
+    "PackedTextDataset",
+    "find_msvc_cl",
+    "load_vcvars_env",
+    "setup_compile_env",
+    "register_multiscreen_auto_classes",
+]
+
+
+def register_multiscreen_auto_classes() -> None:
+    """Register Multiscreen with Transformers auto classes in this process.
+
+    Use this when loading local checkpoints without ``trust_remote_code`` and
+    without installing the model into a Transformers source tree::
+
+        from multiscreen_transformers import register_multiscreen_auto_classes
+        register_multiscreen_auto_classes()
+
+        from transformers import AutoModelForCausalLM
+        model = AutoModelForCausalLM.from_pretrained("./checkpoint")
+    """
+
+    from transformers import AutoConfig, AutoModel, AutoModelForCausalLM
+
+    AutoConfig.register(MultiscreenConfig.model_type, MultiscreenConfig)
+    AutoModel.register(MultiscreenConfig, MultiscreenModel)
+    AutoModelForCausalLM.register(MultiscreenConfig, MultiscreenForCausalLM)

compile_utils.py

@@ -0,0 +1,93 @@
+"""Helpers for torch.compile setup, especially on Windows/MSVC."""
+
+from __future__ import annotations
+
+import glob
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+
+def find_msvc_cl() -> str | None:
+    """Find MSVC ``cl.exe`` for Triton/torch.compile on Windows."""
+
+    if os.environ.get("CC"):
+        return os.environ["CC"]
+
+    bases = [
+        Path(r"C:\Program Files (x86)\Microsoft Visual Studio"),
+        Path(r"C:\Program Files\Microsoft Visual Studio"),
+    ]
+    for base in bases:
+        if not base.exists():
+            continue
+        for vs in sorted(base.iterdir(), reverse=True):
+            pattern = str(
+                vs / "BuildTools" / "VC" / "Tools" / "MSVC" / "*" / "bin" / "Hostx64" / "x64" / "cl.exe"
+            )
+            matches = sorted(glob.glob(pattern))
+            if matches:
+                return matches[-1]
+    return None
+
+
+def _find_vcvarsall() -> Path | None:
+    bases = [
+        Path(r"C:\Program Files (x86)\Microsoft Visual Studio"),
+        Path(r"C:\Program Files\Microsoft Visual Studio"),
+    ]
+    for base in bases:
+        if not base.exists():
+            continue
+        for vs in sorted(base.iterdir(), reverse=True):
+            candidate = vs / "BuildTools" / "VC" / "Auxiliary" / "Build" / "vcvarsall.bat"
+            if candidate.exists():
+                return candidate
+    return None
+
+
+def load_vcvars_env() -> bool:
+    """Load the full MSVC build environment into ``os.environ`` on Windows."""
+
+    if sys.platform != "win32":
+        return False
+    if os.environ.get("VSCMD_VER"):
+        return True
+
+    vcvarsall = _find_vcvarsall()
+    if vcvarsall is None:
+        return False
+
+    try:
+        result = subprocess.run(
+            f'"{vcvarsall}" x64 >nul && set',
+            shell=True,
+            capture_output=True,
+            text=True,
+            check=False,
+        )
+    except OSError:
+        return False
+    if result.returncode != 0:
+        return False
+
+    for line in result.stdout.splitlines():
+        if "=" in line:
+            key, value = line.split("=", 1)
+            os.environ[key] = value
+    return True
+
+
+def setup_compile_env() -> str | None:
+    """Auto-detect MSVC and set ``CC`` for ``torch.compile`` when needed."""
+
+    if sys.platform == "win32":
+        load_vcvars_env()
+        if os.environ.get("CC"):
+            return os.environ["CC"]
+        cl_path = find_msvc_cl()
+        if cl_path:
+            os.environ["CC"] = cl_path
+        return cl_path
+    return os.environ.get("CC")

configuration_multiscreen.py

@@ -0,0 +1,336 @@
+"""Configuration for the Transformers-compatible Multiscreen model."""
+
+from __future__ import annotations
+
+import math
+from typing import Any
+
+try:  # Transformers has historically used both spellings in examples.
+    from transformers import PreTrainedConfig
+except ImportError:  # pragma: no cover - compatibility fallback for old releases.
+    from transformers import PretrainedConfig as PreTrainedConfig  # type: ignore
+
+
+class MultiscreenConfig(PreTrainedConfig):
+    """Configuration for Multiscreen causal language models.
+
+    This mirrors the architecture knobs from ``dieOD/multiscreen-pytorch`` while
+    exposing the conventional Transformers names where possible.
+
+    Important aliases
+    -----------------
+    ``hidden_size`` <-> original ``hidden_dim``
+    ``num_hidden_layers`` <-> original ``num_layers``
+    ``num_attention_heads`` <-> original ``num_heads``
+    ``max_position_embeddings`` <-> original ``max_seq_len``
+
+    Reproducibility controls
+    ------------------------
+    ``mipe_compute_dtype`` and ``softmask_compute_dtype`` can be ``"fp32"``
+    for the numerically safer Transformers port behavior, or ``"reference"``
+    to use the incoming tensor dtype like the standalone PyTorch reference.
+    ``strict_position_ids`` rejects batch-specific or non-contiguous
+    ``position_ids`` because the reference cache API is based on a scalar
+    ``start_pos``. ``zero_pad_hidden_states`` can additionally zero padded
+    query states after each residual layer; it defaults to ``False`` to keep
+    original residual behavior.
+    """
+
+    model_type = "multiscreen"
+    keys_to_ignore_at_inference = ["past_key_values"]
+    _alias_to_primary = {
+        "hidden_dim": "hidden_size",
+        "num_layers": "num_hidden_layers",
+        "num_heads": "num_attention_heads",
+        "max_seq_len": "max_position_embeddings",
+    }
+
+    def __init__(
+        self,
+        vocab_size: int = 50_257,
+        hidden_size: int | None = None,
+        hidden_dim: int | None = None,
+        num_hidden_layers: int | None = None,
+        num_layers: int | None = None,
+        num_attention_heads: int | None = None,
+        num_heads: int | None = None,
+        key_dim: int = 16,
+        value_dim: int = 64,
+        max_position_embeddings: int | None = None,
+        max_seq_len: int | None = None,
+        mipe_threshold: float = 256.0,
+        gradient_checkpointing: bool = False,
+        use_cache: bool = True,
+        labels_are_shifted: bool = False,
+        mipe_compute_dtype: str = "fp32",
+        softmask_compute_dtype: str = "fp32",
+        strict_position_ids: bool = True,
+        zero_pad_hidden_states: bool = False,
+        initializer_range: float = 0.1,
+        bos_token_id: int | None = None,
+        eos_token_id: int | None = None,
+        pad_token_id: int | None = None,
+        tie_word_embeddings: bool = True,
+        **kwargs: Any,
+    ) -> None:
+        # Saved Transformers configs may contain these superclass fields in
+        # kwargs. Multiscreen is always a decoder-only model, and passing them
+        # through while also setting them below would duplicate keyword args.
+        is_decoder = kwargs.pop("is_decoder", True)
+        is_encoder_decoder = kwargs.pop("is_encoder_decoder", False)
+        kwargs.pop("model_type", None)
+        if is_decoder is not True:
+            raise ValueError("MultiscreenConfig requires is_decoder=True")
+        if is_encoder_decoder is not False:
+            raise ValueError("MultiscreenConfig requires is_encoder_decoder=False")
+
+        hidden_size = self._resolve_alias(
+            primary=hidden_size,
+            alias=hidden_dim,
+            default=256,
+            primary_name="hidden_size",
+            alias_name="hidden_dim",
+        )
+        num_hidden_layers = self._resolve_alias(
+            primary=num_hidden_layers,
+            alias=num_layers,
+            default=8,
+            primary_name="num_hidden_layers",
+            alias_name="num_layers",
+        )
+        num_attention_heads = self._resolve_alias(
+            primary=num_attention_heads,
+            alias=num_heads,
+            default=8,
+            primary_name="num_attention_heads",
+            alias_name="num_heads",
+        )
+        max_position_embeddings = self._resolve_alias(
+            primary=max_position_embeddings,
+            alias=max_seq_len,
+            default=256,
+            primary_name="max_position_embeddings",
+            alias_name="max_seq_len",
+        )
+
+        if not bool(tie_word_embeddings):
+            raise ValueError(
+                "Multiscreen uses normalized tied input/output embeddings; "
+                "tie_word_embeddings must be True."
+            )
+
+        self.vocab_size = int(vocab_size)
+        self.hidden_size = int(hidden_size)
+        self.hidden_dim = int(hidden_size)  # original repo alias
+        self.num_hidden_layers = int(num_hidden_layers)
+        self.num_layers = int(num_hidden_layers)  # original repo alias
+        self.num_attention_heads = int(num_attention_heads)
+        self.num_heads = int(num_attention_heads)  # original repo alias
+        self.key_dim = int(key_dim)
+        self.value_dim = int(value_dim)
+        self.max_position_embeddings = int(max_position_embeddings)
+        self.max_seq_len = int(max_position_embeddings)  # original repo alias
+        self.mipe_threshold = float(mipe_threshold)
+        self.gradient_checkpointing = bool(gradient_checkpointing)
+        self.use_cache = bool(use_cache)
+        self.labels_are_shifted = bool(labels_are_shifted)
+        self.mipe_compute_dtype = str(mipe_compute_dtype)
+        self.softmask_compute_dtype = str(softmask_compute_dtype)
+        self.strict_position_ids = bool(strict_position_ids)
+        self.zero_pad_hidden_states = bool(zero_pad_hidden_states)
+        self.initializer_range = float(initializer_range)
+
+        self._validate()
+
+        super().__init__(
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            pad_token_id=pad_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            is_decoder=True,
+            is_encoder_decoder=False,
+            use_cache=use_cache,
+            **kwargs,
+        )
+
+        # Useful when pushing a repo with these Python files to the Hub and
+        # loading with trust_remote_code=True.
+        if not getattr(self, "auto_map", None):
+            self.auto_map = {
+                "AutoConfig": "configuration_multiscreen.MultiscreenConfig",
+                "AutoModel": "modeling_multiscreen.MultiscreenModel",
+                "AutoModelForCausalLM": "modeling_multiscreen.MultiscreenForCausalLM",
+            }
+        if not getattr(self, "architectures", None):
+            self.architectures = ["MultiscreenForCausalLM"]
+
+    @staticmethod
+    def _resolve_alias(
+        *,
+        primary: int | None,
+        alias: int | None,
+        default: int,
+        primary_name: str,
+        alias_name: str,
+    ) -> int:
+        if primary is None and alias is None:
+            return default
+        if primary is None:
+            return int(alias)  # type: ignore[arg-type]
+        if alias is None:
+            return int(primary)
+        if int(primary) != int(alias):
+            raise ValueError(
+                f"Conflicting values for {primary_name}={primary} and "
+                f"{alias_name}={alias}. Use only one or make them equal."
+            )
+        return int(primary)
+
+    @classmethod
+    def from_psi(
+        cls,
+        psi: int,
+        vocab_size: int = 50_257,
+        max_seq_len: int = 256,
+        **overrides: Any,
+    ) -> "MultiscreenConfig":
+        """Build a paper-style config from the supraparameter Psi.
+
+        The scaling rule used in the reference repo is ``N_L = N_H = Psi`` and
+        ``d_E = Psi²``.
+        """
+
+        return cls(
+            vocab_size=vocab_size,
+            hidden_size=psi * psi,
+            num_hidden_layers=psi,
+            num_attention_heads=psi,
+            max_position_embeddings=max_seq_len,
+            **overrides,
+        )
+
+    def _validate(self) -> None:
+        if self.vocab_size <= 0:
+            raise ValueError("vocab_size must be positive")
+        if self.hidden_size <= 0:
+            raise ValueError("hidden_size/hidden_dim must be positive")
+        if self.num_hidden_layers <= 0:
+            raise ValueError("num_hidden_layers/num_layers must be positive")
+        if self.num_attention_heads <= 0:
+            raise ValueError("num_attention_heads/num_heads must be positive")
+        if self.key_dim < 2:
+            raise ValueError("key_dim must be at least 2 because MiPE rotates the first two coordinates")
+        if self.value_dim <= 0:
+            raise ValueError("value_dim must be positive")
+        if self.max_position_embeddings <= 0:
+            raise ValueError("max_position_embeddings/max_seq_len must be positive")
+        if self.mipe_threshold <= 0:
+            raise ValueError("mipe_threshold must be positive")
+        allowed_compute_dtypes = {"fp32", "reference"}
+        if self.mipe_compute_dtype not in allowed_compute_dtypes:
+            raise ValueError(
+                "mipe_compute_dtype must be either 'fp32' or 'reference', "
+                f"got {self.mipe_compute_dtype!r}"
+            )
+        if self.softmask_compute_dtype not in allowed_compute_dtypes:
+            raise ValueError(
+                "softmask_compute_dtype must be either 'fp32' or 'reference', "
+                f"got {self.softmask_compute_dtype!r}"
+            )
+        if self.initializer_range <= 0:
+            raise ValueError("initializer_range must be positive")
+
+    @classmethod
+    def _normalize_alias_updates(cls, updates: dict[str, Any]) -> dict[str, Any]:
+        """Map original-repository aliases to canonical Transformers field names."""
+
+        normalized = dict(updates)
+        for alias, primary in cls._alias_to_primary.items():
+            if alias not in normalized:
+                continue
+            alias_value = normalized.pop(alias)
+            if primary in normalized and int(normalized[primary]) != int(alias_value):
+                raise ValueError(
+                    f"Conflicting update values for {primary}={normalized[primary]} "
+                    f"and {alias}={alias_value}. Use only one or make them equal."
+                )
+            normalized[primary] = alias_value
+        return normalized
+
+    def clone(self, **updates: Any) -> "MultiscreenConfig":
+        """Return a config copy with updated fields.
+
+        ``PreTrainedConfig.to_dict()`` contains both Transformers field names and
+        original-repository aliases such as ``hidden_size``/``hidden_dim``.
+        Reusing that dictionary directly can create alias conflicts when callers
+        update only one spelling, so this method rebuilds from canonical fields
+        and canonicalizes alias-style updates.
+        """
+
+        normalized_updates = self._normalize_alias_updates(updates)
+
+        # ``is_decoder`` and ``is_encoder_decoder`` are forced by this class and
+        # are passed explicitly to ``PreTrainedConfig`` in ``__init__``.
+        # Silently accepting the matching values makes clone robust to dicts
+        # produced by Transformers; conflicting values should fail loudly.
+        if "is_decoder" in normalized_updates:
+            if normalized_updates.pop("is_decoder") is not True:
+                raise ValueError("MultiscreenConfig requires is_decoder=True")
+        if "is_encoder_decoder" in normalized_updates:
+            if normalized_updates.pop("is_encoder_decoder") is not False:
+                raise ValueError("MultiscreenConfig requires is_encoder_decoder=False")
+
+        data: dict[str, Any] = {
+            "vocab_size": self.vocab_size,
+            "hidden_size": self.hidden_size,
+            "num_hidden_layers": self.num_hidden_layers,
+            "num_attention_heads": self.num_attention_heads,
+            "key_dim": self.key_dim,
+            "value_dim": self.value_dim,
+            "max_position_embeddings": self.max_position_embeddings,
+            "mipe_threshold": self.mipe_threshold,
+            "gradient_checkpointing": self.gradient_checkpointing,
+            "use_cache": self.use_cache,
+            "labels_are_shifted": self.labels_are_shifted,
+            "mipe_compute_dtype": self.mipe_compute_dtype,
+            "softmask_compute_dtype": self.softmask_compute_dtype,
+            "strict_position_ids": self.strict_position_ids,
+            "zero_pad_hidden_states": self.zero_pad_hidden_states,
+            "initializer_range": self.initializer_range,
+            "bos_token_id": self.bos_token_id,
+            "eos_token_id": self.eos_token_id,
+            "pad_token_id": self.pad_token_id,
+            "tie_word_embeddings": True,
+        }
+
+        # Preserve useful PreTrainedConfig extras without reintroducing explicit
+        # constructor arguments, aliases, or forced superclass kwargs.
+        skip_keys = set(data) | set(self._alias_to_primary) | {
+            "model_type",
+            "is_decoder",
+            "is_encoder_decoder",
+            "transformers_version",
+        }
+        for key, value in self.to_dict().items():
+            if key not in skip_keys:
+                data[key] = value
+
+        data.update(normalized_updates)
+        return self.__class__(**data)
+
+    @property
+    def num_params_estimate(self) -> int:
+        """Approximate parameter count, following the reference implementation.
+
+        The estimate assumes tied input/output embeddings and ignores the small
+        learned scalar parameters.
+        """
+
+        embed = self.vocab_size * self.hidden_size
+        per_tile = self.hidden_size * (2 * self.key_dim + 3 * self.value_dim)
+        total_tiles = self.num_hidden_layers * self.num_attention_heads
+        return int(embed + total_tiles * per_tile)
+
+    @property
+    def sqrt_hidden_size(self) -> float:
+        return math.sqrt(self.hidden_size)

data.py

@@ -0,0 +1,136 @@
+"""Dataset utilities for Multiscreen causal LM training."""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+from typing import Optional
+
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+
+
+class PackedTextDataset(Dataset):
+    """In-memory packed dataset for autoregressive language-model training.
+
+    Texts are tokenized, separated by EOS, concatenated, and chunked into fixed
+    length sequences.
+
+    By default this dataset follows the original ``dieOD/multiscreen-pytorch``
+    trainer: each stored chunk has ``seq_len + 1`` tokens, ``input_ids`` are
+    ``chunk[:-1]``, and ``labels`` are ``chunk[1:]``.  The item also includes a
+    scalar ``labels_are_shifted=True`` flag so a standard Transformers data
+    collator/Trainer can forward it to ``MultiscreenForCausalLM`` and avoid a
+    second internal next-token shift.
+
+    Set ``legacy_shifted_labels=False`` for conventional Hugging Face causal-LM
+    batches where ``labels == input_ids`` and the model performs the standard
+    internal shift. In that mode the dataset emits ``labels_are_shifted=False``.
+    """
+
+    def __init__(
+        self,
+        texts: Iterable[str],
+        tokenizer,
+        seq_len: int = 256,
+        eos_token_id: Optional[int] = None,
+        max_tokens: Optional[int] = None,
+        legacy_shifted_labels: bool = True,
+        return_labels_are_shifted: bool = True,
+    ) -> None:
+        if seq_len <= 0:
+            raise ValueError("seq_len must be positive")
+        self.seq_len = int(seq_len)
+        self.legacy_shifted_labels = bool(legacy_shifted_labels)
+        self.return_labels_are_shifted = bool(return_labels_are_shifted)
+
+        if eos_token_id is None:
+            eos_token_id = getattr(tokenizer, "eos_token_id", None)
+        if eos_token_id is None:
+            eos_token_id = 0
+        self.eos_token_id = int(eos_token_id)
+
+        all_ids: list[int] = []
+        for text in texts:
+            if not text:
+                continue
+            ids = tokenizer.encode(text, add_special_tokens=False)
+            all_ids.extend(int(i) for i in ids)
+            all_ids.append(self.eos_token_id)
+            if max_tokens is not None and len(all_ids) >= max_tokens:
+                all_ids = all_ids[:max_tokens]
+                break
+
+        chunk_size = self.seq_len + 1 if self.legacy_shifted_labels else self.seq_len
+        usable = (len(all_ids) // chunk_size) * chunk_size
+        if usable == 0:
+            raise ValueError(f"Not enough tokens for one chunk (need {chunk_size}, got {len(all_ids)})")
+
+        self.tokens = np.array(all_ids[:usable], dtype=np.int64).reshape(-1, chunk_size)
+
+    def __len__(self) -> int:
+        return int(self.tokens.shape[0])
+
+    def __getitem__(self, idx: int) -> dict[str, torch.Tensor]:
+        chunk = self.tokens[idx]
+        if self.legacy_shifted_labels:
+            input_ids = torch.from_numpy(chunk[:-1].copy())
+            labels = torch.from_numpy(chunk[1:].copy())
+        else:
+            input_ids = torch.from_numpy(chunk.copy())
+            labels = input_ids.clone()
+
+        item = {
+            "input_ids": input_ids,
+            "labels": labels,
+            "attention_mask": torch.ones_like(input_ids, dtype=torch.long),
+        }
+        if self.return_labels_are_shifted:
+            item["labels_are_shifted"] = torch.tensor(self.legacy_shifted_labels, dtype=torch.bool)
+        return item
+
+    @classmethod
+    def from_hf_dataset(
+        cls,
+        dataset_name: str,
+        tokenizer,
+        seq_len: int = 256,
+        split: str = "train",
+        text_column: str = "text",
+        config_name: Optional[str] = None,
+        max_tokens: Optional[int] = None,
+        legacy_shifted_labels: bool = True,
+        return_labels_are_shifted: bool = True,
+        cache_dir: Optional[str] = None,
+        data_files: Optional[str | list[str] | dict[str, str | list[str]]] = None,
+        data_dir: Optional[str] = None,
+        revision: Optional[str] = None,
+    ) -> "PackedTextDataset":
+        """Load and pack a Hugging Face dataset.
+
+        ``cache_dir`` is forwarded to :func:`datasets.load_dataset`, which is
+        useful when training from TinyStories or other Hub datasets on machines
+        with a dedicated dataset cache volume.  ``data_files`` / ``data_dir`` /
+        ``revision`` are kept as narrow passthroughs for local or pinned data
+        sources while preserving the original in-memory packing behavior.
+        """
+
+        from datasets import load_dataset
+
+        load_kwargs = {
+            "split": split,
+            "cache_dir": cache_dir,
+            "data_files": data_files,
+            "data_dir": data_dir,
+            "revision": revision,
+        }
+        load_kwargs = {k: v for k, v in load_kwargs.items() if v is not None}
+        ds = load_dataset(dataset_name, config_name, **load_kwargs)
+        return cls(
+            texts=(row[text_column] for row in ds),
+            tokenizer=tokenizer,
+            seq_len=seq_len,
+            max_tokens=max_tokens,
+            legacy_shifted_labels=legacy_shifted_labels,
+            return_labels_are_shifted=return_labels_are_shifted,
+        )

modeling_multiscreen.py

@@ -0,0 +1,1084 @@
+"""Transformers-compatible Multiscreen model.
+
+The screening block is ported from ``dieOD/multiscreen-pytorch`` and wrapped in
+Hugging Face ``PreTrainedModel`` classes.
+"""
+
+from __future__ import annotations
+
+import math
+import weakref
+from collections.abc import Mapping, Sequence
+from typing import Any, Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.checkpoint import checkpoint as grad_checkpoint
+from transformers import PreTrainedModel
+try:  # Transformers >=4.50 separates generation helpers from PreTrainedModel.
+    from transformers.generation import GenerationMixin
+except ImportError:  # pragma: no cover - compatibility with older releases.
+    try:
+        from transformers.generation.utils import GenerationMixin
+    except ImportError:  # pragma: no cover
+        class GenerationMixin:  # type: ignore[no-redef]
+            pass
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+from transformers.utils import logging
+
+from .configuration_multiscreen import MultiscreenConfig
+
+logger = logging.get_logger(__name__)
+
+# Per-layer screening cache.
+# K: (batch, num_heads, cached_length, key_dim), post-MiPE and unit-normalized.
+# V: (batch, num_heads, cached_length, value_dim), unit-normalized.
+ScreeningCache = tuple[torch.Tensor, torch.Tensor]
+
+
+def convert_original_state_dict_for_causal_lm(
+    state_dict: Mapping[str, torch.Tensor],
+    *,
+    strip_module_prefix: bool = True,
+) -> dict[str, torch.Tensor]:
+    """Convert original ``dieOD/multiscreen-pytorch`` weights for HF CausalLM.
+
+    The original repository's language model stores parameters under bare keys
+    such as ``embed.weight`` and ``layers.0.block.q_proj.weight``.  This
+    Transformers port keeps those modules inside ``MultiscreenForCausalLM`` as
+    ``self.multiscreen``; state dict keys are prefixed with ``multiscreen.``.
+    Loading the original checkpoint into ``MultiscreenForCausalLM`` requires
+    that prefix. Already
+    prefixed keys are left unchanged, so the helper is safe to call twice.
+
+    Args:
+        state_dict: State dict from the original implementation, or an already
+            converted state dict.
+        strip_module_prefix: Strip a leading ``module.`` prefix often added by
+            DataParallel/DDP wrappers before conversion.
+    """
+
+    converted: dict[str, torch.Tensor] = {}
+    for key, value in state_dict.items():
+        converted_key = key
+        if strip_module_prefix and converted_key.startswith("module."):
+            converted_key = converted_key[len("module.") :]
+        if not converted_key.startswith("multiscreen."):
+            converted_key = f"multiscreen.{converted_key}"
+        converted[converted_key] = value
+    return converted
+
+
+def convert_original_state_dict_for_model(
+    state_dict: Mapping[str, torch.Tensor],
+    *,
+    strip_module_prefix: bool = True,
+) -> dict[str, torch.Tensor]:
+    """Convert original or CausalLM-prefixed weights for bare ``MultiscreenModel``.
+
+    Original ``dieOD/multiscreen-pytorch`` keys are already suitable for the bare
+    decoder. This helper mainly strips a leading ``module.`` or ``multiscreen.``
+    prefix when needed.
+    """
+
+    converted: dict[str, torch.Tensor] = {}
+    for key, value in state_dict.items():
+        converted_key = key
+        if strip_module_prefix and converted_key.startswith("module."):
+            converted_key = converted_key[len("module.") :]
+        if converted_key.startswith("multiscreen."):
+            converted_key = converted_key[len("multiscreen.") :]
+        converted[converted_key] = value
+    return converted
+
+
+class MultiscreenPreTrainedModel(PreTrainedModel):
+    """Base class for Multiscreen Transformers models."""
+
+    config_class = MultiscreenConfig
+    base_model_prefix = "multiscreen"
+    # Transformers 5 expects tied-weight metadata to be a mapping, while older
+    # model classes often used a list of regex keys. Multiscreen has no
+    # duplicated output-head Parameter to tie or drop from the state dict:
+    # logits are computed directly from the normalized input embedding.
+    _tied_weights_keys: dict[str, str] = {}
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["MultiscreenLayer"]
+    _skip_keys_device_placement = "past_key_values"
+    # Newer Transformers Trainer may pass ``num_items_in_batch`` to models whose
+    # forward signature has **kwargs. Multiscreen computes a standard mean CE
+    # loss internally and does not consume that normalization hint.
+    accepts_loss_kwargs = False
+
+    def get_expanded_tied_weights_keys(self, all_submodels: bool = False) -> dict[str, str]:
+        """Return no storage-level tied-parameter mapping for Multiscreen.
+
+        The input/output embedding relationship is implemented by construction
+        in ``_compute_logits`` / ``_NormalizedTiedLMHead`` instead of by
+        assigning a second registered output-head Parameter to the input
+        embedding Parameter. Returning an empty mapping keeps Transformers 5
+        tied-weight bookkeeping on the mapping code path and avoids legacy
+        list-vs-dict crashes.
+        """
+
+        return {}
+
+    def _init_weights(self, module: nn.Module) -> None:  # pragma: no cover - post_init hook.
+        """No-op because modules initialize with the original Multiscreen rules.
+
+        The reference implementation uses per-projection initializers rather than
+        a single global initializer. Those are applied in each module's ``__init__``.
+        """
+
+        return None
+
+    def _set_gradient_checkpointing(
+        self,
+        module: nn.Module | None = None,
+        value: bool = False,
+        enable: bool | None = None,
+        gradient_checkpointing_func: Any | None = None,
+    ) -> None:
+        # Accept both the older Transformers hook signature
+        #   _set_gradient_checkpointing(module, value=False)
+        # and the newer one
+        #   _set_gradient_checkpointing(enable=True, gradient_checkpointing_func=...).
+        flag = value if enable is None else enable
+        if module is None:
+            for child in self.modules():
+                if isinstance(child, MultiscreenModel):
+                    child.gradient_checkpointing = bool(flag)
+            return
+        if isinstance(module, MultiscreenModel):
+            module.gradient_checkpointing = bool(flag)
+
+
+class MultiscreenModel(MultiscreenPreTrainedModel):
+    """Bare Multiscreen decoder model.
+
+    This returns hidden states, not vocabulary logits. Use
+    :class:`MultiscreenForCausalLM` for the original language-model behavior.
+    """
+
+    def __init__(self, config: MultiscreenConfig) -> None:
+        super().__init__(config)
+        self.config = config
+        self.gradient_checkpointing = bool(config.gradient_checkpointing)
+        self.zero_pad_hidden_states = bool(config.zero_pad_hidden_states)
+
+        d_e = config.hidden_size
+        self.embed = nn.Embedding(config.vocab_size, d_e)
+        self.s_E = nn.Parameter(torch.tensor(0.0))
+        self.s_F = nn.Parameter(torch.tensor(math.log(math.sqrt(d_e))))
+        self.layers = nn.ModuleList(
+            [MultiscreenLayer(config, layer_idx=i) for i in range(config.num_hidden_layers)]
+        )
+
+        # Original embedding initialization: N(0, 0.1 / sqrt(d_E)).
+        nn.init.normal_(self.embed.weight, mean=0.0, std=config.initializer_range / math.sqrt(d_e))
+        self.post_init()
+
+    def get_input_embeddings(self) -> nn.Embedding:
+        return self.embed
+
+    def set_input_embeddings(self, value: nn.Embedding) -> None:
+        self.embed = value
+        self.config.vocab_size = value.num_embeddings
+
+    def get_output_embeddings(self) -> nn.Embedding:
+        # Output is tied by construction via normalized input embedding.
+        return self.embed
+
+    def set_output_embeddings(self, value: nn.Embedding) -> None:
+        self.set_input_embeddings(value)
+
+    def tie_weights(self, *args: Any, **kwargs: Any) -> None:
+        # We do not create a separate lm_head; logits use self.embed.weight.
+        # Recent Transformers releases call tie_weights with keyword arguments
+        # such as recompute_mapping=... or missing_keys=...; accept and ignore
+        # them because Multiscreen ties weights by construction.
+        return None
+
+    def count_parameters(self) -> int:
+        return sum(p.numel() for p in self.parameters() if p.requires_grad)
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Sequence[ScreeningCache] | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        use_cache: bool | None = None,
+        output_attentions: bool | None = None,
+        output_hidden_states: bool | None = None,
+        return_dict: bool | None = None,
+        start_pos: int | None = None,
+        **kwargs: Any,
+    ) -> BaseModelOutputWithPast | tuple[torch.Tensor, ...]:
+        """Run the Multiscreen decoder.
+
+        Args follow Transformers conventions. ``start_pos`` is kept as an
+        explicit compatibility escape hatch for the original cache API.
+        """
+
+        kv_caches = kwargs.pop("kv_caches", None)
+        if kv_caches is not None:
+            if past_key_values is not None:
+                raise ValueError("Pass only one of `past_key_values` or original-api `kv_caches`, not both.")
+            past_key_values = kv_caches
+
+        # Compatibility with Trainer/TRL/Transformers call paths.
+        use_return_dict = kwargs.pop("use_return_dict", None)
+        if return_dict is None and use_return_dict is not None:
+            return_dict = bool(use_return_dict)
+        kwargs.pop("num_items_in_batch", None)
+
+        if past_key_values is not None and len(past_key_values) == 0:
+            past_key_values = None
+        if past_key_values is not None and len(past_key_values) != len(self.layers):
+            raise ValueError(
+                f"past_key_values must contain {len(self.layers)} layer caches, got {len(past_key_values)}."
+            )
+
+        if kwargs:
+            # Keep forward permissive, but surface likely typo/debug information.
+            # ``warning_once`` caches calls, so every argument must be hashable.
+            unused_kwargs = ", ".join(sorted(str(key) for key in kwargs.keys()))
+            logger.warning_once("Unused MultiscreenModel.forward kwargs: %s", unused_kwargs)
+
+        if output_attentions:
+            logger.warning_once(
+                "Multiscreen has no softmax attention weights; `output_attentions=True` returns None."
+            )
+
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        requested_cache = self.config.use_cache if use_cache is None else bool(use_cache)
+
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError("Pass either input_ids or inputs_embeds, not both.")
+        if input_ids is None and inputs_embeds is None:
+            raise ValueError("You must pass input_ids or inputs_embeds.")
+        if inputs_embeds is not None:
+            raise ValueError(
+                "Multiscreen does not accept `inputs_embeds` through the public Transformers API. "
+                "The reference architecture normalizes token embedding weights before lookup, so raw "
+                "embeddings from `get_input_embeddings()` are not equivalent to `input_ids`. "
+                "Pass `input_ids` instead."
+            )
+
+        if input_ids is None:
+            raise ValueError("input_ids unexpectedly None")
+        input_shape = input_ids.shape
+        batch_size, seq_len = input_shape
+        W_norm = F.normalize(self.embed.weight, dim=-1)
+        hidden_states = F.embedding(input_ids, W_norm) * self.s_E.exp()
+
+        if past_key_values is not None and len(past_key_values) > 0:
+            past_length = int(past_key_values[0][0].shape[2])
+        else:
+            past_length = 0
+
+        if start_pos is None:
+            if position_ids is not None:
+                start_pos = self._start_pos_from_position_ids(
+                    position_ids=position_ids,
+                    seq_len=seq_len,
+                    strict=bool(self.config.strict_position_ids),
+                )
+            else:
+                start_pos = past_length
+        elif position_ids is not None:
+            logger.warning_once(
+                "Multiscreen consumes a scalar `start_pos`; `position_ids` are ignored when `start_pos` is provided."
+            )
+
+        use_cache = requested_cache and (not self.training)
+        if requested_cache and self.training:
+            logger.warning_once("Multiscreen disables cache materialization while model.training is True.")
+
+        if self.gradient_checkpointing and self.training and use_cache:
+            logger.warning_once("use_cache=True is incompatible with gradient checkpointing; disabling cache.")
+            use_cache = False
+
+        total_length = past_length + seq_len
+        key_attention_mask, query_attention_mask = self._prepare_attention_masks(
+            attention_mask=attention_mask,
+            batch_size=batch_size,
+            past_length=past_length,
+            seq_len=seq_len,
+            total_length=total_length,
+            device=hidden_states.device,
+        )
+        query_mask_3d = (
+            query_attention_mask.to(dtype=hidden_states.dtype).unsqueeze(-1)
+            if query_attention_mask is not None
+            else None
+        )
+        if self.zero_pad_hidden_states and query_mask_3d is not None:
+            hidden_states = hidden_states * query_mask_3d
+
+        all_hidden_states: tuple[torch.Tensor, ...] | None = () if output_hidden_states else None
+        new_key_values: list[ScreeningCache] = []
+
+        for layer_idx, layer in enumerate(self.layers):
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (hidden_states,)  # type: ignore[operator]
+
+            past_layer = past_key_values[layer_idx] if past_key_values is not None else None
+
+            if self.gradient_checkpointing and self.training:
+                def custom_forward(
+                    x: torch.Tensor,
+                    layer_ref: MultiscreenLayer = layer,
+                    start_pos_ref: int = start_pos,
+                    key_attention_mask_ref: torch.Tensor | None = key_attention_mask,
+                    query_attention_mask_ref: torch.Tensor | None = query_attention_mask,
+                ) -> torch.Tensor:
+                    y, _ = layer_ref(
+                        x,
+                        start_pos=start_pos_ref,
+                        past_kv=None,
+                        use_cache=False,
+                        key_attention_mask=key_attention_mask_ref,
+                        query_attention_mask=query_attention_mask_ref,
+                    )
+                    return y
+
+                hidden_states = grad_checkpoint(custom_forward, hidden_states, use_reentrant=False)
+                new_kv = None
+            else:
+                hidden_states, new_kv = layer(
+                    hidden_states,
+                    start_pos=start_pos,
+                    past_kv=past_layer,
+                    use_cache=use_cache,
+                    key_attention_mask=key_attention_mask,
+                    query_attention_mask=query_attention_mask,
+                )
+
+            if self.zero_pad_hidden_states and query_mask_3d is not None:
+                hidden_states = hidden_states * query_mask_3d
+
+            if use_cache:
+                if new_kv is None:
+                    raise RuntimeError("Layer did not return a cache while use_cache=True")
+                new_key_values.append(new_kv)
+
+        if output_hidden_states:
+            all_hidden_states = all_hidden_states + (hidden_states,)  # type: ignore[operator]
+
+        past = tuple(new_key_values) if use_cache else None
+
+        if not return_dict:
+            outputs: tuple[Any, ...] = (hidden_states,)
+            if past is not None:
+                outputs += (past,)
+            if all_hidden_states is not None:
+                outputs += (all_hidden_states,)
+            return outputs  # type: ignore[return-value]
+
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past,
+            hidden_states=all_hidden_states,
+            attentions=None,
+        )
+
+    @staticmethod
+    def _start_pos_from_position_ids(
+        *,
+        position_ids: torch.LongTensor,
+        seq_len: int,
+        strict: bool,
+    ) -> int:
+        """Extract the scalar reference-style ``start_pos`` from position IDs.
+
+        Multiscreen's reference implementation uses one scalar ``start_pos`` for
+        every batch item. Arbitrary per-token or per-batch ``position_ids`` would
+        misalign MiPE and the distance softmask, so strict mode fails loudly.
+        """
+
+        if position_ids.dim() != 2:
+            raise ValueError("position_ids must have shape (batch, sequence_length)")
+        if int(position_ids.shape[1]) != seq_len:
+            raise ValueError(
+                f"position_ids length {position_ids.shape[1]} does not match input sequence length {seq_len}"
+            )
+        if seq_len == 0:
+            return 0
+
+        start_pos = int(position_ids[0, 0].item())
+        expected = torch.arange(
+            start_pos,
+            start_pos + seq_len,
+            device=position_ids.device,
+            dtype=position_ids.dtype,
+        ).unsqueeze(0).expand(position_ids.shape[0], -1)
+
+        if not torch.equal(position_ids, expected):
+            message = (
+                "Multiscreen only supports batch-shared contiguous position_ids, "
+                "because the reference cache API is based on a scalar start_pos. "
+                "Pass start_pos explicitly for reference-style decoding, or disable "
+                "config.strict_position_ids only if you intentionally want to use "
+                "position_ids[0, 0] and ignore the rest."
+            )
+            if strict:
+                raise ValueError(message)
+            logger.warning_once(message)
+        return start_pos
+
+    @staticmethod
+    def _prepare_attention_masks(
+        *,
+        attention_mask: torch.Tensor | None,
+        batch_size: int,
+        past_length: int,
+        seq_len: int,
+        total_length: int,
+        device: torch.device,
+    ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
+        if attention_mask is None:
+            return None, None
+
+        if attention_mask.dim() != 2:
+            raise ValueError("attention_mask must have shape (batch, sequence_length)")
+        if attention_mask.shape[0] != batch_size:
+            raise ValueError(
+                f"attention_mask batch size {attention_mask.shape[0]} does not match input batch {batch_size}"
+            )
+
+        mask = attention_mask.to(device=device)
+        mask_len = int(mask.shape[1])
+
+        if past_length > 0 and mask_len != total_length:
+            logger.warning_once(
+                "Cached Multiscreen decoding received an attention_mask whose length (%s) "
+                "does not cover the full cache length (%s). Omitted past cache positions are "
+                "treated as valid. Pass a full-length attention_mask when cached prefixes "
+                "contain padding.",
+                mask_len,
+                total_length,
+            )
+
+        if mask_len == total_length:
+            key_mask = mask
+            query_mask = mask[:, -seq_len:]
+        elif mask_len == seq_len:
+            if past_length > 0:
+                prefix = torch.ones(batch_size, past_length, device=device, dtype=mask.dtype)
+                key_mask = torch.cat([prefix, mask], dim=1)
+            else:
+                key_mask = mask
+            query_mask = mask
+        elif mask_len > total_length:
+            key_mask = mask[:, -total_length:]
+            query_mask = key_mask[:, -seq_len:]
+        elif mask_len < total_length:
+            # If only a shorter mask is supplied during cached decoding, assume
+            # the missing older cache positions are valid.
+            prefix = torch.ones(batch_size, total_length - mask_len, device=device, dtype=mask.dtype)
+            key_mask = torch.cat([prefix, mask], dim=1)
+            query_mask = key_mask[:, -seq_len:]
+        else:  # pragma: no cover - unreachable, kept for clarity.
+            key_mask = mask
+            query_mask = mask[:, -seq_len:]
+
+        return key_mask, query_mask
+
+
+
+
+class _NormalizedTiedLMHead(nn.Module):
+    """Parameter-free lm_head proxy for trainers that expect ``model.lm_head``.
+
+    Multiscreen computes logits with the unit-normalized input embedding matrix
+    and the learned scalar ``s_F`` instead of a standalone Linear layer.  Some
+    Hugging Face/TRL training paths look for ``model.lm_head.weight`` to avoid
+    materializing full logits.  This proxy exposes the mathematically equivalent
+    dynamic weight while keeping the true parameters tied to ``embed.weight`` and
+    ``s_F``.
+    """
+
+    def __init__(self, owner: "MultiscreenForCausalLM") -> None:
+        super().__init__()
+        self._owner_ref = weakref.ref(owner)
+
+    @property
+    def weight(self) -> torch.Tensor:
+        owner = self._owner_ref()
+        if owner is None:  # pragma: no cover - defensive only.
+            raise RuntimeError("Multiscreen lm_head owner has been garbage-collected")
+        W_norm = F.normalize(owner.multiscreen.embed.weight, dim=-1)
+        return W_norm * owner.multiscreen.s_F.exp()
+
+    @property
+    def bias(self) -> None:
+        return None
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return F.linear(hidden_states, self.weight)
+
+
+class MultiscreenForCausalLM(MultiscreenPreTrainedModel, GenerationMixin):
+    """Multiscreen decoder with normalized tied LM head."""
+
+    # ``self.lm_head`` is a parameter-free proxy whose dynamic ``weight``
+    # property is computed from ``multiscreen.embed.weight`` and ``s_F``. There
+    # is no registered ``lm_head.weight`` Parameter, so report no explicit tied
+    # parameter pair to Transformers.
+    _tied_weights_keys: dict[str, str] = {}
+
+    def __init__(self, config: MultiscreenConfig) -> None:
+        if not bool(getattr(config, "tie_word_embeddings", True)):
+            raise ValueError(
+                "Multiscreen uses normalized tied input/output embeddings; "
+                "tie_word_embeddings must be True."
+            )
+        super().__init__(config)
+        self.multiscreen = MultiscreenModel(config)
+        # Compatibility shim for TRL/SFTTrainer paths that expect a ``lm_head``
+        # attribute. It has no parameters; it dynamically reuses the normalized
+        # tied input embeddings exactly like ``_compute_logits``.
+        self.lm_head = _NormalizedTiedLMHead(self)
+        self.vocab_size = config.vocab_size
+        self.post_init()
+
+    def get_input_embeddings(self) -> nn.Embedding:
+        return self.multiscreen.get_input_embeddings()
+
+    def set_input_embeddings(self, value: nn.Embedding) -> None:
+        self.multiscreen.set_input_embeddings(value)
+        self.config.vocab_size = value.num_embeddings
+        self.vocab_size = value.num_embeddings
+
+    def get_output_embeddings(self) -> nn.Embedding:
+        return self.multiscreen.get_output_embeddings()
+
+    def set_output_embeddings(self, value: nn.Embedding) -> None:
+        self.set_input_embeddings(value)
+
+    def tie_weights(self, *args: Any, **kwargs: Any) -> None:
+        # Output logits directly reuse the normalized input embedding matrix.
+        # Recent Transformers releases call tie_weights with keyword arguments
+        # such as recompute_mapping=... or missing_keys=...; accept and ignore
+        # them because Multiscreen ties weights by construction.
+        return None
+
+    @staticmethod
+    def convert_original_state_dict(state_dict: Mapping[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+        """Convert original ``multiscreen-pytorch`` checkpoint keys for this class."""
+
+        return convert_original_state_dict_for_causal_lm(state_dict)
+
+    def _compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.lm_head(hidden_states)
+
+    @staticmethod
+    def _coerce_optional_bool(value: Any, name: str) -> bool | None:
+        """Convert Python or collated tensor booleans to a scalar bool.
+
+        ``PackedTextDataset`` can emit a scalar ``labels_are_shifted`` flag so a
+        standard Transformers data collator/Trainer forwards it to the model.
+        After collation this arrives as a batch tensor; mixed True/False values
+        in one batch are rejected because a single loss path must be chosen.
+        """
+
+        if value is None:
+            return None
+        if isinstance(value, torch.Tensor):
+            if value.numel() == 0:
+                return None
+            bool_values = value.detach().to(dtype=torch.bool).flatten()
+            has_true = bool(bool_values.any().item())
+            has_false = bool((~bool_values).any().item())
+            if has_true and has_false:
+                raise ValueError(f"{name} must be the same for every item in a batch.")
+            return has_true
+        return bool(value)
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Sequence[ScreeningCache] | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        labels: torch.LongTensor | None = None,
+        use_cache: bool | None = None,
+        output_attentions: bool | None = None,
+        output_hidden_states: bool | None = None,
+        return_dict: bool | None = None,
+        start_pos: int | None = None,
+        labels_are_shifted: bool | None = None,
+        legacy_shifted_labels: bool | None = None,
+        logits_to_keep: int = 0,
+        **kwargs: Any,
+    ) -> CausalLMOutputWithPast | tuple[torch.Tensor, ...]:
+        # Compatibility with Trainer/TRL/Transformers call paths.
+        # ``use_return_dict`` is a deprecated alias that may still be forwarded,
+        # and ``num_items_in_batch`` can be injected by recent Trainer versions
+        # when a model forward has **kwargs. Multiscreen does not consume it.
+        use_return_dict = kwargs.pop("use_return_dict", None)
+        if return_dict is None and use_return_dict is not None:
+            return_dict = bool(use_return_dict)
+        kwargs.pop("num_items_in_batch", None)
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        labels_are_shifted_value = self._coerce_optional_bool(labels_are_shifted, "labels_are_shifted")
+        legacy_shifted_labels_value = self._coerce_optional_bool(
+            legacy_shifted_labels, "legacy_shifted_labels"
+        )
+        if labels_are_shifted_value is not None and legacy_shifted_labels_value is not None:
+            if labels_are_shifted_value != legacy_shifted_labels_value:
+                raise ValueError("labels_are_shifted and legacy_shifted_labels disagree.")
+        if labels_are_shifted_value is None:
+            labels_are_shifted = (
+                legacy_shifted_labels_value
+                if legacy_shifted_labels_value is not None
+                else bool(getattr(self.config, "labels_are_shifted", False))
+            )
+        else:
+            labels_are_shifted = labels_are_shifted_value
+
+        kv_caches = kwargs.pop("kv_caches", None)
+        if kv_caches is not None:
+            if past_key_values is not None:
+                raise ValueError("Pass only one of `past_key_values` or original-api `kv_caches`, not both.")
+            past_key_values = kv_caches
+
+        if kwargs:
+            # Keep forward permissive for HF/TRL extras, but do not pass them to
+            # the bare decoder where they would only generate duplicate warnings.
+            unused_kwargs = ", ".join(sorted(str(key) for key in kwargs.keys()))
+            logger.warning_once("Unused MultiscreenForCausalLM.forward kwargs: %s", unused_kwargs)
+
+        model_outputs = self.multiscreen(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=True,
+            start_pos=start_pos,
+        )
+        hidden_states = model_outputs.last_hidden_state
+
+        if labels is None and logits_to_keep and logits_to_keep > 0:
+            logits_hidden_states = hidden_states[:, -logits_to_keep:, :]
+        else:
+            logits_hidden_states = hidden_states
+
+        logits = self._compute_logits(logits_hidden_states)
+        loss = None
+
+        if labels is not None:
+            if logits.shape[1] != labels.shape[1]:
+                # This can only happen if a caller forced logits_to_keep with labels.
+                logits = self._compute_logits(hidden_states)
+
+            loss_labels = labels.to(device=logits.device).clone()
+            loss_attention_mask = None
+            if attention_mask is not None:
+                loss_attention_mask = self._slice_loss_attention_mask(
+                    attention_mask=attention_mask,
+                    target_length=loss_labels.shape[1],
+                    device=loss_labels.device,
+                )
+                loss_labels = loss_labels.masked_fill(loss_attention_mask == 0, -100)
+
+            loss_fct = nn.CrossEntropyLoss(ignore_index=-100)
+            if labels_are_shifted:
+                loss = loss_fct(
+                    logits.reshape(-1, self.config.vocab_size),
+                    loss_labels.reshape(-1),
+                )
+            else:
+                if logits.shape[1] < 2:
+                    loss = logits.new_zeros(())
+                else:
+                    shift_logits = logits[..., :-1, :].contiguous()
+                    shift_labels = loss_labels[..., 1:].contiguous()
+                    if loss_attention_mask is not None:
+                        # Ignore predictions where either the query token or the
+                        # target token is padding. This avoids a left-padding edge
+                        # case where the last pad token predicts the first real token.
+                        valid_shift = (loss_attention_mask[..., :-1] != 0) & (
+                            loss_attention_mask[..., 1:] != 0
+                        )
+                        shift_labels = shift_labels.masked_fill(~valid_shift, -100)
+                    loss = loss_fct(
+                        shift_logits.reshape(-1, self.config.vocab_size),
+                        shift_labels.reshape(-1),
+                    )
+
+        if not return_dict:
+            output: tuple[Any, ...] = (logits,)
+            if model_outputs.past_key_values is not None:
+                output += (model_outputs.past_key_values,)
+            if model_outputs.hidden_states is not None:
+                output += (model_outputs.hidden_states,)
+            return ((loss,) + output) if loss is not None else output  # type: ignore[return-value]
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=model_outputs.past_key_values,
+            hidden_states=model_outputs.hidden_states,
+            attentions=None,
+        )
+
+    @staticmethod
+    def _slice_loss_attention_mask(
+        *,
+        attention_mask: torch.Tensor,
+        target_length: int,
+        device: torch.device,
+    ) -> torch.Tensor:
+        if attention_mask.dim() != 2:
+            raise ValueError("attention_mask must have shape (batch, sequence_length)")
+        mask = attention_mask.to(device=device)
+        mask_length = int(mask.shape[1])
+        if mask_length == target_length:
+            return mask
+        if mask_length > target_length:
+            return mask[:, -target_length:]
+        prefix = torch.ones(
+            mask.shape[0],
+            target_length - mask_length,
+            device=device,
+            dtype=mask.dtype,
+        )
+        return torch.cat([prefix, mask], dim=1)
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids: torch.LongTensor,
+        past_key_values: Sequence[ScreeningCache] | None = None,
+        attention_mask: torch.Tensor | None = None,
+        cache_position: torch.LongTensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        start_pos: int | None = None,
+        use_cache: bool | None = True,
+        **kwargs: Any,
+    ) -> dict[str, Any]:
+        """Prepare inputs for ``GenerationMixin.generate``.
+
+        Multiscreen keeps a simple tuple cache. When a cache is present, the
+        method slices ``input_ids`` to the new suffix and sets a scalar
+        ``start_pos`` equal to the cached sequence length.
+        """
+
+        kv_caches = kwargs.pop("kv_caches", None)
+        if kv_caches is not None:
+            if past_key_values is not None:
+                raise ValueError("Pass only one of `past_key_values` or original-api `kv_caches`, not both.")
+            past_key_values = kv_caches
+
+        if past_key_values is not None and len(past_key_values) > 0:
+            past_length = int(past_key_values[0][0].shape[2])
+            if input_ids.shape[1] > past_length:
+                input_ids = input_ids[:, past_length:]
+            else:
+                input_ids = input_ids[:, -1:]
+            # Cache length is the source of truth during generation. A stale
+            # explicit start_pos or arbitrary position_ids would misalign
+            # MiPE/softmask positions.
+            start_pos = past_length
+        else:
+            if start_pos is None:
+                if cache_position is not None and cache_position.numel() > 0:
+                    start_pos = int(cache_position[0].item())
+                elif position_ids is not None:
+                    start_pos = MultiscreenModel._start_pos_from_position_ids(
+                        position_ids=position_ids,
+                        seq_len=int(input_ids.shape[1]),
+                        strict=bool(self.config.strict_position_ids),
+                    )
+                else:
+                    start_pos = 0
+
+        # The model consumes scalar `start_pos`; forwarding arbitrary
+        # `position_ids` would give the false impression that batch-specific
+        # offsets are honored.
+        position_ids = None
+
+        model_inputs = {
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "position_ids": position_ids,
+            "past_key_values": past_key_values,
+            "use_cache": use_cache,
+            "start_pos": start_pos,
+        }
+        return model_inputs
+
+    @staticmethod
+    def _reorder_cache(
+        past_key_values: Sequence[ScreeningCache], beam_idx: torch.LongTensor
+    ) -> tuple[ScreeningCache, ...]:
+        """Beam-search cache reordering."""
+
+        reordered: list[ScreeningCache] = []
+        for key_cache, value_cache in past_key_values:
+            beam_idx_device = beam_idx.to(key_cache.device)
+            reordered.append(
+                (
+                    key_cache.index_select(0, beam_idx_device),
+                    value_cache.index_select(0, beam_idx_device.to(value_cache.device)),
+                )
+            )
+        return tuple(reordered)
+
+
+class MultiscreenLayer(nn.Module):
+    """Single residual Multiscreen layer."""
+
+    def __init__(self, config: MultiscreenConfig, layer_idx: int) -> None:
+        super().__init__()
+        self.block = GatedScreeningBlock(config, layer_idx)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        start_pos: int = 0,
+        past_kv: ScreeningCache | None = None,
+        use_cache: bool = False,
+        key_attention_mask: torch.Tensor | None = None,
+        query_attention_mask: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, ScreeningCache | None]:
+        block_out, new_kv = self.block(
+            x,
+            start_pos=start_pos,
+            past_kv=past_kv,
+            use_cache=use_cache,
+            key_attention_mask=key_attention_mask,
+            query_attention_mask=query_attention_mask,
+        )
+        if query_attention_mask is not None:
+            block_out = block_out * query_attention_mask.to(dtype=block_out.dtype).unsqueeze(-1)
+        return x + block_out, new_kv
+
+
+class GatedScreeningBlock(nn.Module):
+    """Parallel gated screening tiles for one Multiscreen layer.
+
+    Each tile performs Q/K/V/G projection, unit normalization, MiPE, independent
+    screening, TanhNorm, bounded gating, per-head scaling, and output projection.
+    """
+
+    def __init__(self, config: MultiscreenConfig, layer_idx: int) -> None:
+        super().__init__()
+        d_e = config.hidden_size
+        d_k = config.key_dim
+        d_v = config.value_dim
+        num_heads = config.num_attention_heads
+        num_layers = config.num_hidden_layers
+
+        self.layer_idx = layer_idx
+        self.NH = num_heads
+        self.dK = d_k
+        self.dV = d_v
+        self.wth = float(config.mipe_threshold)
+        self.max_seq_len = int(config.max_position_embeddings)
+        self.mipe_compute_dtype = str(config.mipe_compute_dtype)
+        self.softmask_compute_dtype = str(config.softmask_compute_dtype)
+
+        self.q_proj = nn.Linear(d_e, num_heads * d_k, bias=False)
+        self.k_proj = nn.Linear(d_e, num_heads * d_k, bias=False)
+        self.v_proj = nn.Linear(d_e, num_heads * d_v, bias=False)
+        self.g_proj = nn.Linear(d_e, num_heads * d_v, bias=False)
+        self.o_proj = nn.Linear(num_heads * d_v, d_e, bias=False)
+
+        # Per-head learned parameters from the reference implementation.
+        self.sw = nn.Parameter(torch.linspace(0, math.log(self.wth), num_heads))
+        self.sr = nn.Parameter(torch.zeros(num_heads))
+        self.sO = nn.Parameter(torch.full((num_heads,), math.log(1.0 / math.sqrt(num_heads * num_layers))))
+
+        init = config.initializer_range
+        nn.init.normal_(self.q_proj.weight, mean=0.0, std=init / math.sqrt(d_k))
+        nn.init.normal_(self.k_proj.weight, mean=0.0, std=init / math.sqrt(d_k))
+        nn.init.normal_(self.v_proj.weight, mean=0.0, std=init / math.sqrt(d_v))
+        nn.init.normal_(self.g_proj.weight, mean=0.0, std=init)
+        nn.init.normal_(self.o_proj.weight, mean=0.0, std=init / math.sqrt(d_e))
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        start_pos: int = 0,
+        past_kv: ScreeningCache | None = None,
+        use_cache: bool = False,
+        key_attention_mask: torch.Tensor | None = None,
+        query_attention_mask: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, ScreeningCache | None]:
+        batch_size, seq_len, _ = x.shape
+
+        q = self.q_proj(x).view(batch_size, seq_len, self.NH, self.dK)
+        k_new = self.k_proj(x).view(batch_size, seq_len, self.NH, self.dK)
+        v_new = self.v_proj(x).view(batch_size, seq_len, self.NH, self.dV)
+        g = self.g_proj(x).view(batch_size, seq_len, self.NH, self.dV)
+
+        u, new_kv = self._screening(
+            q=q,
+            k_new=k_new,
+            v_new=v_new,
+            start_pos=start_pos,
+            past_kv=past_kv,
+            use_cache=use_cache,
+            key_attention_mask=key_attention_mask,
+        )
+
+        g_hat = torch.tanh(F.silu(g))
+        h = u * g_hat
+        if query_attention_mask is not None:
+            h = h * query_attention_mask.to(dtype=h.dtype).view(batch_size, seq_len, 1, 1)
+        h = h * self.sO.exp().view(1, 1, self.NH, 1)
+        h = h.reshape(batch_size, seq_len, self.NH * self.dV)
+        return self.o_proj(h), new_kv
+
+    def _screening(
+        self,
+        q: torch.Tensor,
+        k_new: torch.Tensor,
+        v_new: torch.Tensor,
+        start_pos: int = 0,
+        past_kv: ScreeningCache | None = None,
+        use_cache: bool = False,
+        key_attention_mask: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, ScreeningCache | None]:
+        """Screening unit with optional per-layer KV cache."""
+
+        q = F.normalize(q, dim=-1)
+        k_new = F.normalize(k_new, dim=-1)
+        v_new = F.normalize(v_new, dim=-1)
+
+        w = self.sw.exp() + 1.0
+        r = self.sr.exp() + 1.0
+
+        q, k_new = self._apply_mipe(q, k_new, w, start_pos=start_pos)
+
+        q = q.transpose(1, 2)
+        k_new = k_new.transpose(1, 2)
+        v_new = v_new.transpose(1, 2)
+
+        if past_kv is not None:
+            past_k, past_v = past_kv
+            full_k = torch.cat([past_k, k_new], dim=2)
+            full_v = torch.cat([past_v, v_new], dim=2)
+        else:
+            full_k = k_new
+            full_v = v_new
+
+        seq_len = q.shape[2]
+        total_length = full_k.shape[2]
+
+        sim = torch.matmul(q, full_k.transpose(-2, -1))
+        mask = self._softmask(
+            T_new=seq_len,
+            T_total=total_length,
+            start_pos=start_pos,
+            w=w,
+            device=sim.device,
+            dtype=sim.dtype,
+            key_attention_mask=key_attention_mask,
+        )
+
+        rho_d = torch.clamp(
+            1.0 - r.view(1, -1, 1, 1).to(dtype=sim.dtype) * (1.0 - sim),
+            min=0.0,
+        ).square_().mul_(mask)
+
+        h = torch.matmul(rho_d, full_v)
+        h_norm = h.norm(dim=-1, keepdim=True).clamp(min=1e-8)
+        u = (torch.tanh(h_norm) / h_norm) * h
+
+        new_kv = (full_k, full_v) if use_cache else None
+        return u.transpose(1, 2), new_kv
+
+    @staticmethod
+    def _select_compute_dtype(input_dtype: torch.dtype, mode: str) -> torch.dtype:
+        if mode == "fp32":
+            return torch.float32
+        if mode == "reference":
+            return input_dtype
+        raise ValueError(f"Unknown Multiscreen compute dtype mode: {mode!r}")
+
+    def _apply_mipe(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        w: torch.Tensor,
+        start_pos: int = 0,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Minimal positional encoding on the first two Q/K coordinates."""
+
+        seq_len = q.shape[1]
+        compute_dtype = self._select_compute_dtype(q.dtype, self.mipe_compute_dtype)
+        w_float = w.to(device=q.device, dtype=compute_dtype)
+        phi = torch.where(
+            w_float < self.wth,
+            0.5 * (torch.cos(math.pi * w_float / self.wth) + 1.0),
+            torch.zeros_like(w_float),
+        )
+
+        positions = torch.arange(start_pos, start_pos + seq_len, device=q.device, dtype=compute_dtype)
+        pos_2d = positions.unsqueeze(1)
+        w_2d = w_float.unsqueeze(0)
+        effective_pos = torch.where(pos_2d >= self.max_seq_len, pos_2d % w_2d, pos_2d)
+        angles = effective_pos * (math.pi * phi / w_float).unsqueeze(0)
+
+        cos_a = torch.cos(angles).to(dtype=q.dtype)
+        sin_a = torch.sin(angles).to(dtype=q.dtype)
+
+        q0, q1 = q[..., 0], q[..., 1]
+        k0, k1 = k[..., 0], k[..., 1]
+
+        q_rot = torch.empty_like(q)
+        q_rot[..., 0] = q0 * cos_a - q1 * sin_a
+        q_rot[..., 1] = q0 * sin_a + q1 * cos_a
+        q_rot[..., 2:] = q[..., 2:]
+
+        k_rot = torch.empty_like(k)
+        k_rot[..., 0] = k0 * cos_a - k1 * sin_a
+        k_rot[..., 1] = k0 * sin_a + k1 * cos_a
+        k_rot[..., 2:] = k[..., 2:]
+        return q_rot, k_rot
+
+    def _softmask(
+        self,
+        T_new: int,
+        T_total: int,
+        start_pos: int,
+        w: torch.Tensor,
+        device: torch.device,
+        dtype: torch.dtype,
+        key_attention_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """Causal distance-aware softmask for new queries over all keys."""
+
+        compute_dtype = self._select_compute_dtype(dtype, self.softmask_compute_dtype)
+        q_pos = torch.arange(start_pos, start_pos + T_new, device=device, dtype=compute_dtype)
+        k_pos = torch.arange(T_total, device=device, dtype=compute_dtype)
+        rel = k_pos.unsqueeze(0) - q_pos.unsqueeze(1)
+
+        w_exp = w.to(device=device, dtype=compute_dtype).view(-1, 1, 1)
+        valid = (rel <= 0) & (rel > -w_exp)
+        mask = (0.5 * (torch.cos(math.pi * rel / w_exp) + 1.0)) * valid
+        mask = mask.unsqueeze(0).to(dtype=dtype)
+
+        if key_attention_mask is not None:
+            if key_attention_mask.shape[1] != T_total:
+                raise ValueError(
+                    f"key_attention_mask length {key_attention_mask.shape[1]} must equal total key length {T_total}"
+                )
+            mask = mask * key_attention_mask.to(device=device, dtype=dtype).view(-1, 1, 1, T_total)
+
+        return mask