Publish Modilify Mk2 Preview

README.md

@@ -14,49 +14,91 @@ tags:
 
 ![LOGO](assets/01-LOGO.jpg)
 
-# Modilify Mk2 Preview
+# Modilify Mk2 Preview · 26B-A5B
 
-A 26B-A4B multimodal block-diffusion model with dual-timescale latent deliberation.
+**Refine a whole canvas. Think in latent space. Carry memory forward.**
 
-Mk2 does not dump chain-of-thought into extra visible tokens. Each heavy denoise runs a latent Transformer over a packed trajectory history and a denoise-time tape, then writes persistent memory only when the canvas actually commits. Working state is recomputed every step. Persistent slots survive the rolling window. The exclusive excess-entropy commit formula still decides how many tokens lock in; temperature and failure budget are first-class inference knobs.
+Modilify Mk2 is a **26B-A5B multimodal block-diffusion model** that brings parallel token refinement, latent deliberation, and persistent trajectory memory into one generation loop. A rolling 256-token canvas gives the model room to revise upcoming text together. A dedicated latent Transformer turns the history of those revisions into context for the next step. As output advances, a learned memory writer carries information across canvas boundaries.
 
-This repository is the first public Mk2 preview checkpoint: merged BF16 weights, remote code, processor, and tokenizer. The text trunk and dual-timescale latent stack come from schema23 training step 900 (~12.4 million adaptation tokens). The Gemma 4 vision tower is restored from DiffusionGemma so text, image, and video share one decoder. This is not a full benchmark release.
+The central idea is simple: **give each answer a workspace, a memory, and a variable compute budget.** Mk2 can refine its internal state over multiple passes before committing text, and release multiple tokens together when its confidence-and-entropy policy permits. Deliberation happens in continuous hidden states, without requiring every internal update to become a visible reasoning token.
 
-## Architecture
+Built on DiffusionGemma, Mk2 combines sparse expert routing with a dual-timescale latent architecture. Text, images, and sampled video frames feed the same generation path.
 
-The heavy trunk is DiffusionGemma 26B-A4B. Inside every denoise, a 4-layer latent Transformer reads the noisy 256-token canvas plus:
+## What makes Mk2 different
 
-1. **Packed history** (T=16). Four views of each canvas position's recent deliberation, projected into rank-1024 space. New canvas positions start empty; they do not inherit the previous token's thought.
-2. **Denoise tape**. Row-level probes written on a time ring. They do not shift when tokens commit.
-3. **Persistent slots** (256 × 2816). Updated only at commit by a Transformer writer. Full-attention decoder layers read working and persistent buses.
+| Architecture choice | What it enables |
+| --- | --- |
+| **Parallel block diffusion** | Revise a 256-token canvas jointly and commit a variable-length prefix, allowing multiple output tokens per denoising pass. |
+| **Trajectory-aware latent deliberation** | Condition the next revision on how hidden states have evolved, including their changes, acceleration, and residuals relative to the latest state. |
+| **Memory at two timescales** | Rebuild working state every pass while retaining persistent slots across rolling-window shifts within a generation. |
+| **Memory inside the decoder** | Feed working and persistent memory directly into full-attention decoder layers so both can influence token refinement. |
+| **Adaptive commitment** | Use confidence and entropy to decide how much text to release, with explicit budgets for continued refinement and forced progress. |
+| **Sparse multimodal foundation** | Select 8 of 128 experts per token and bring text, image, and video context into a shared decoder. |
 
-Visible tokens are the product of that loop, not the workspace. Easy prompts commit a long prefix. Hard prompts keep pondering.
+## Inside the generation loop
 
-The encoder is the official Gemma 4 multimodal encoder. Image and video tokens condition prefix KV the same way text does; the rolling canvas and latent stack stay text-side.
+### A canvas built for revision
+
+Mk2 maintains a rolling canvas of 256 candidate tokens. Each denoising pass updates the candidates using the encoded prompt, the current canvas, and latent context. Attention lets positions within the canvas inform one another before the output prefix is finalized. After a commit, the canvas shifts forward and opens space for new candidates.
+
+This creates two useful degrees of freedom: **how many times to refine** and **how many tokens to release**. A pass can commit a longer prefix when the policy permits, or spend additional computation refining an uncertain frontier.
+
+### Deliberation that reads its own trajectory
+
+A 4-layer latent Transformer, 2,816 dimensions wide, builds working context before each decoder denoising pass. It reads the current canvas alongside three complementary sources of state:
+
+- **Per-token trajectory history:** 16 recent frames represented through four views—hidden state, first difference, second difference, and residual from the latest state—projected to rank 1,024. These give the processor access to the direction and stability of recent revisions. History follows surviving canvas tokens; new positions start empty.
+- **Denoising tape:** 16 pooled probes per frame summarize canvas activity in a time-indexed ring. The tape preserves recent step-level context as token positions move through the window.
+- **Persistent memory:** 256 slots of 2,816 dimensions carry learned summaries across commits within the generation.
+
+The resulting working context enters the decoder through its self-conditioning bridge and working-memory bus. Full-attention layers also read a separate persistent-memory bus. **The refinement history becomes an input to the next refinement.**
+
+### Fast working state, lasting commit memory
+
+Working state is recomputed on every denoising pass. Persistent slots update only when tokens commit, using a Transformer writer with a separate gate for each slot. That writer draws on the committed region's trajectory, working state, and final decoder representation.
+
+This separates rapid revision from memory consolidation: the canvas can keep changing while persistent memory stays stable between commits. When the window advances, the memory slots remain available to later tokens.
+
+### Compute that follows the commit frontier
+
+Mk2 combines proposal confidence with an excess-entropy penalty and selects the longest prefix whose cumulative failure score stays below the configured budget. A tighter budget requires stronger evidence before normal commitment; a looser budget admits longer prefixes for the same scores. Stagnation handling and a pondering watchdog bound continued refinement.
+
+Temperature and commit budget are exposed directly at inference time, making the generation policy adjustable per request. These scores govern commitment; they are not calibrated guarantees of factual correctness.
+
+### One generation path for text, images, and video
+
+The Gemma 4 vision tower supplies visual features through the DiffusionGemma multimodal encoder. Text, image, and sampled video-frame inputs are encoded into the prefix KV cache that conditions the decoder. The rolling text canvas then uses the same latent deliberation and memory loop across all three input modalities.
 
 ## Model Summary
 
 | | |
 | --- | ---: |
 | Architecture | Mixture-of-Experts block diffusion + dual-timescale latent Transformer |
-| Total Parameters | 26.139B text trunk + 569.550M vision encoder + latent stack |
-| Text Heavy-Denoise Activated Parameters | 4.159B |
-| Vision Encoder | Gemma 4 Vision, 569.550M |
+| Model Size | 26B-A5B |
+| Vision Encoder | Gemma 4 Vision |
 | Layers | 30 |
 | Number of Experts | 128 |
 | Selected Experts per Token | 8 |
 | Vocabulary Size | 262,144 |
-| Context Length | 262,144 tokens |
+| Configured Context Length | 262,144 tokens |
 | Sliding Window | 1024 |
 | Canvas Length | 256 |
 | Latent Width | 2,816 |
-| Latent Memory | 256 slots × 2,816-d, 4 layers |
+| Latent Transformer | 4 layers, 16 attention heads |
+| Persistent Memory | 256 slots × 2,816 dimensions |
+| Memory Writer | 2-layer commit-sequence Transformer + per-slot gated writer |
 | Trajectory History | 16 frames, 4 views, rank 1,024 |
-| Denoise Tape | 16 probes |
+| Denoise Tape | 16 probes per frame |
 | Modality | Text, Image, Video |
-| Preview checkpoint | schema23 step 900 |
+| Preview checkpoint | Training step 900 |
 | Adaptation tokens | ~12.4 million |
 
+## Preview release
+
+This first public preview includes **merged BF16 weights, inference code, processor, and tokenizer**. The text backbone and latent stack are exported from training step 900, after approximately 12.4 million adaptation tokens. The Gemma 4 vision tower is restored from DiffusionGemma.
+
+The release makes the architecture available for hands-on exploration and evaluation. Comprehensive benchmark results are not included; measured speed, reasoning quality, and multimodal reliability remain to be established for specific workloads.
+
 ## Getting Started
 
 Transformers 5.14.1 is the minimum supported version.
@@ -157,12 +199,12 @@ output = model.generate(**inputs, max_new_tokens=256)
 
 ## Thinking mode
 
-The chat template controls the prompt, not the model's first generated tokens.
+Latent deliberation runs inside the generation loop regardless of the chat template's thinking flag. The flag controls the prompt's request for a textual thought channel:
 
 - `enable_thinking=True` inserts a system turn that contains `<|think|>` and still ends the prompt at `<|turn>model`.
 - `enable_thinking=False` does **not** inject an empty thought channel. The prompt ends at `<|turn>model`.
 
-The model may still open `<|channel>thought` on its own. That is generation, not a template artifact. Applications should not assume hidden reasoning is complete, correct, or appropriate to expose to end users.
+The model may still open `<|channel>thought` on its own when thinking is disabled. The template flag does not guarantee suppression of generated thought-channel text. Applications should handle that channel explicitly before displaying an answer.
 
 ## Configurable inference
 
@@ -208,13 +250,13 @@ model = AutoModelForMultimodalLM.from_pretrained(
 )
 ```
 
-Lower temperature and a tighter budget make the model more cautious and usually slower. Higher temperature and a looser budget commit more tokens per denoise. These are compute-control decisions, not guarantees of correctness.
+A tighter commit budget allows fewer tokens for the same confidence-and-entropy scores; a looser budget allows more. Temperature changes the sampling distribution and also affects those scores, so its effect on throughput depends on the prompt and generation trajectory. Measure latency and output quality together when tuning these controls.
 
 Generation supports left-padded batches with independent stopping. Batch prompts of similar lengths together for the best throughput. Streaming and caller-supplied KV caches remain limited to batch size 1.
 
 ## Evaluation status, limitations, and risks
 
-This is a preview. It does not include a complete accuracy, robustness, calibration, fairness, or safety evaluation. Structural export checks and a small graduate-level qualitative probe, if present, do not establish fitness for use.
+This preview does not include a complete accuracy, robustness, calibration, fairness, or safety evaluation. Architectural features describe how Mk2 generates; they do not establish benchmark superiority or fitness for a particular deployment. The configured context limit is not a validated long-context quality result.
 
 The model can hallucinate facts, citations, visual details, or temporal relationships; reproduce bias, unsafe content, personal information, or copyrighted material; and consume substantial time and memory during long iterative generation. Confidence-based commits control compute. They do not certify that a prefix is true. Visual performance can degrade with poor resolution, motion, occlusion, unusual aspect ratios, or domain shift.
 
@@ -228,7 +270,7 @@ Released under the [Modilify Open Model License 1.0](LICENSE), subject to its re
 
 ```bibtex
 @software{modilify_mk2_preview_2026,
-  title = {Modilify Mk2 Preview},
+  title = {Modilify Mk2 Preview: 26B-A5B},
   author = {Modilify},
   year = {2026},
   note = {A multimodal dual-timescale latent-deliberation derivative of DiffusionGemma}

model-00008-of-00011.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0b52a8c906686648eb04b3e28bc5c597e9b687de0a473ec093ad339dbfd8da77
+size 4884578046

model-00009-of-00011.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d425f9f2a800dd424e997db11933f6e4602bab648e2731a3d76fbe778554b2c4
+size 4913414718

model-00010-of-00011.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:802add3f8c81ff7c44dbf96632682b90a1a0a08f35e80be867d4f0294654c03b
+size 4884577974

model-00011-of-00011.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1e3a84ff63b8e779f287a5b227aa8252383554d7b3527fc895aeb623b0c998f1
+size 3959052840

modeling_modilify_mk2.py

@@ -0,0 +1,735 @@
+# Copyright 2026 Modilify
+# SPDX-License-Identifier: LicenseRef-Modilify-Open-Model-1.0
+"""Multimodal Modilify Mk2 model with recurrent latent deliberation."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import dataclass, replace
+from typing import Any
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+from transformers.cache_utils import Cache
+from transformers.modeling_outputs import BaseModelOutputWithPast
+from transformers.utils import ModelOutput
+
+from transformers.masking_utils import (
+    ALL_MASK_ATTENTION_FUNCTIONS,
+    bidirectional_mask_function,
+)
+from transformers.models.diffusion_gemma import (
+    DiffusionGemmaDecoderModel,
+    DiffusionGemmaEncoderModel,
+    DiffusionGemmaPreTrainedModel,
+)
+from transformers.models.diffusion_gemma.modeling_diffusion_gemma import (
+    DiffusionGemmaRMSNorm,
+    DiffusionGemmaTextRouter,
+)
+
+from .mps_ops import mps_segmented_experts_forward as _mps_segmented_experts_forward  # noqa: F401
+from .configuration_modilify_mk2 import ModilifyMk2Config
+from .generation_modilify_mk2 import ModilifyMk2GenerationConfig, ModilifyMk2GenerationMixin
+from .latent_deliberation import (
+    LatentDeliberationState,
+    LatentDeliberationTransformer,
+    LatentProcessorOutput,
+    TrajectoryHistory,
+    TrajectoryTape,
+    empty_trajectory_tape,
+)
+from .vocab_ops import chunked_vocab_statistics
+
+
+@dataclass
+class ModilifyMk2DecoderOutput(BaseModelOutputWithPast):
+    token_embeddings: torch.FloatTensor | None = None
+    latent_residual_diagnostics: dict[str, torch.Tensor] | None = None
+
+
+@dataclass
+class ModilifyMk2ModelOutput(BaseModelOutputWithPast):
+    token_embeddings: torch.FloatTensor | None = None
+    encoder_last_hidden_state: torch.FloatTensor | None = None
+    latent_residual_diagnostics: dict[str, torch.Tensor] | None = None
+
+
+@dataclass
+class ModilifyMk2BlockDiffusionOutput(ModelOutput):
+    """Inference output used by the rolling diffusion generator."""
+
+    logits: torch.FloatTensor | None = None
+    heavy_hidden_state: torch.FloatTensor | None = None
+    next_latent_state: LatentDeliberationState | None = None
+    past_key_values: Cache | None = None
+    encoder_last_hidden_state: torch.FloatTensor | None = None
+    temporal_context: torch.FloatTensor | None = None
+    history_projected: torch.FloatTensor | None = None
+    working_state: torch.FloatTensor | None = None
+    latent_residual_diagnostics: dict[str, torch.Tensor] | None = None
+    proposal: torch.LongTensor | None = None
+    proposal_confidence: torch.FloatTensor | None = None
+    token_entropy: torch.FloatTensor | None = None
+    greedy_proposal: torch.LongTensor | None = None
+    greedy_confidence: torch.FloatTensor | None = None
+
+
+class ModilifyMk2RMSNorm(DiffusionGemmaRMSNorm):
+    """Official RMSNorm parameters, pre-fusion same-dtype forward."""
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        normed_output = self._norm(hidden_states)
+        if self.with_scale:
+            normed_output = normed_output * self.weight.to(dtype=normed_output.dtype)
+        return normed_output.type_as(hidden_states)
+
+
+class ModilifyMk2TextRouter(DiffusionGemmaTextRouter):
+    """Official router parameters with fp32 softmax so top-k weights stay finite."""
+
+    def __init__(self, config: Any) -> None:
+        super().__init__(config)
+        self.norm = ModilifyMk2RMSNorm(self.hidden_size, eps=self.eps, with_scale=False)
+
+    def forward(
+        self, hidden_states: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        hidden_states = self.norm(hidden_states)
+        hidden_states = hidden_states * self.scale * self.scalar_root_size
+        expert_scores = self.proj(hidden_states)
+        # Match the official DiffusionGemma router: bf16 softmax underflows to
+        # all-zero top-k rows, then 0/0 NaNs the MoE residual and every loss.
+        router_probabilities = F.softmax(expert_scores, dim=-1, dtype=torch.float32)
+        top_k_weights, top_k_index = torch.topk(
+            router_probabilities,
+            k=self.config.top_k_experts,
+            dim=-1,
+        )
+        top_k_weights = top_k_weights / top_k_weights.sum(dim=-1, keepdim=True).clamp_min(
+            torch.finfo(torch.float32).tiny
+        )
+        top_k_weights = top_k_weights * self.per_expert_scale[top_k_index]
+        return router_probabilities, top_k_weights, top_k_index
+
+
+def install_modilify_mk2_trunk_semantics(module: nn.Module) -> None:
+    """Swap official leaf modules on this instance. Never patch Transformers classes."""
+
+    for name, child in list(module.named_children()):
+        if type(child) is DiffusionGemmaRMSNorm:
+            dim = int(child.weight.shape[0]) if child.with_scale else 1
+            replacement = ModilifyMk2RMSNorm(
+                dim, eps=child.eps, with_scale=child.with_scale
+            )
+            replacement.load_state_dict(child.state_dict())
+            setattr(module, name, replacement)
+        elif type(child) is DiffusionGemmaTextRouter:
+            replacement = ModilifyMk2TextRouter(child.config)
+            replacement.load_state_dict(child.state_dict())
+            setattr(module, name, replacement)
+        else:
+            install_modilify_mk2_trunk_semantics(child)
+
+
+class ModilifyMk2EncoderModel(DiffusionGemmaEncoderModel):
+    """Unmodified Transformers DiffusionGemma multimodal encoder."""
+
+    config_class = ModilifyMk2Config
+
+
+class ModilifyMk2DecoderModel(DiffusionGemmaDecoderModel):
+    """Diffusion decoder accepting compact self-conditioning embeddings."""
+
+    config_class = ModilifyMk2Config
+    latent_residual_rms_ratio_cap = 0.5
+
+    def __init__(self, config: ModilifyMk2Config):
+        super().__init__(config)
+        install_modilify_mk2_trunk_semantics(self)
+
+    @staticmethod
+    def create_diffusion_decoder_attention_mask(
+        config: Any,
+        inputs_embeds: torch.Tensor,
+        past_key_values: Cache,
+        decoder_attention_mask: torch.Tensor | dict | None = None,
+    ) -> dict[str, torch.Tensor | None]:
+        """Official mask builder without the all-True sliding-window skip."""
+
+        if past_key_values is None:
+            raise ValueError(
+                "The diffusion mask requires `past_key_values` to construct the next attention mask correctly"
+            )
+        if (
+            decoder_attention_mask is None
+            or config._attn_implementation
+            not in ALL_MASK_ATTENTION_FUNCTIONS._global_mapping
+        ):
+            return {"full_attention": None, "sliding_attention": None}
+        if isinstance(decoder_attention_mask, dict) and all(
+            mask.ndim == 4 for mask in decoder_attention_mask.values()
+        ):
+            return decoder_attention_mask
+
+        text_config = config.get_text_config() if hasattr(config, "get_text_config") else config
+        q_length = inputs_embeds.shape[1]
+        q_offset = past_key_values.get_seq_length()
+        if isinstance(q_offset, torch.Tensor):
+            q_offset = q_offset.to(inputs_embeds.device)
+        additional_kv_length = (
+            getattr(config, "canvas_length", 0) if past_key_values.is_compileable else 0
+        )
+        mask_mapping: dict[str, torch.Tensor | None] = {}
+        for layer_pattern in set(text_config.layer_types):
+            layer_idx = past_key_values.is_sliding.index(
+                layer_pattern == "sliding_attention"
+            )
+            kv_length, kv_offset = past_key_values.get_mask_sizes(q_length, layer_idx)
+            kv_length += additional_kv_length
+            if layer_pattern == "sliding_attention" and past_key_values.is_compileable:
+                sliding_layer = past_key_values.layers[layer_idx]
+                max_length = sliding_layer.get_max_length() + additional_kv_length
+                if kv_length >= max_length:
+                    kv_length = max_length
+            mask_mapping[layer_pattern] = ALL_MASK_ATTENTION_FUNCTIONS[
+                config._attn_implementation
+            ](
+                batch_size=inputs_embeds.shape[0],
+                q_length=q_length,
+                kv_length=kv_length,
+                q_offset=q_offset,
+                kv_offset=kv_offset,
+                mask_function=bidirectional_mask_function,
+                attention_mask=decoder_attention_mask,
+                allow_is_causal_skip=False,
+                allow_is_bidirectional_skip=True,
+                local_size=getattr(text_config, "sliding_window", None),
+                dtype=inputs_embeds.dtype,
+                config=text_config,
+                use_vmap=False,
+                device=inputs_embeds.device,
+            )
+        return mask_mapping
+
+    def merge_latent_context(
+        self,
+        token_embeddings: torch.Tensor,
+        latent_context: torch.Tensor | None,
+        *,
+        collect_diagnostics: bool = True,
+    ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
+        """Map latent context through the frozen native self-conditioning bridge."""
+
+        if latent_context is None:
+            context = torch.zeros_like(token_embeddings)
+        else:
+            if latent_context.shape != token_embeddings.shape:
+                raise ValueError("Temporal context must match canvas hidden-state shape.")
+            context = latent_context.to(token_embeddings)
+        mapper = self.self_conditioning
+        normalized_context = mapper.pre_norm(context)
+        mapped_context = mapper.down_proj(
+            mapper.act_fn(mapper.gate_proj(normalized_context))
+            * mapper.up_proj(normalized_context)
+        )
+        mapped_fp32 = mapped_context.float()
+        # Use the energy directly in the cap denominator.  Computing
+        # sqrt(E[x²]) and immediately squaring it again has an undefined
+        # backward at the identity-init point x=0 (0/0 in d(sqrt)/dx), which
+        # poisoned the detached temporal-context VJP on every denoise.
+        mapped_token_energy = mapped_fp32.square().mean(dim=-1, keepdim=True)
+        token_rms_per_token = (
+            token_embeddings.float().square().mean(dim=-1, keepdim=True).sqrt()
+        )
+        residual_cap = self.latent_residual_rms_ratio_cap * token_rms_per_token
+        soft_cap_scale = residual_cap / torch.sqrt(
+            mapped_token_energy + residual_cap.square() + 1.0e-12
+        )
+        mapped_context = (mapped_fp32 * soft_cap_scale).to(dtype=mapped_context.dtype)
+        combined = token_embeddings + mapped_context
+        diagnostics: dict[str, torch.Tensor] = {}
+        if collect_diagnostics:
+            token_rms = token_embeddings.detach().float().square().mean().sqrt()
+            mapped_rms = mapped_context.detach().float().square().mean().sqrt()
+            diagnostics = {
+                "token_embedding_rms": token_rms,
+                "latent_context_rms": context.detach().float().square().mean().sqrt(),
+                "mapped_sc_rms": mapped_rms,
+                "actual_residual_rms": mapped_rms,
+                "latent_to_embedding_rms_ratio": (
+                    mapped_rms / token_rms.clamp_min(1.0e-12)
+                ),
+                "merged_input_rms": combined.detach().float().square().mean().sqrt(),
+            }
+        return mapper.post_norm(combined), diagnostics
+
+    def _run_stack(
+        self,
+        inputs_embeds: torch.Tensor,
+        *,
+        past_key_values: Cache | None,
+        decoder_attention_mask: torch.Tensor | dict | None,
+        decoder_position_ids: torch.LongTensor | None,
+        **kwargs: Any,
+    ) -> torch.Tensor:
+        if decoder_position_ids is None:
+            prefix = past_key_values.get_seq_length(layer_idx=0) if past_key_values is not None else 0
+            decoder_position_ids = torch.arange(
+                prefix, prefix + inputs_embeds.shape[1], device=inputs_embeds.device
+            ).unsqueeze(0)
+        if not isinstance(mask_mapping := decoder_attention_mask, dict):
+            mask_mapping = self.create_diffusion_decoder_attention_mask(
+                config=self.text_config,
+                inputs_embeds=inputs_embeds,
+                past_key_values=past_key_values,
+                decoder_attention_mask=decoder_attention_mask,
+            )
+        working_bus = kwargs.pop("working_bus", None)
+        working_state = kwargs.pop("working_state", None)
+        persistent_bus = kwargs.pop("persistent_bus", None)
+        memory_bus = kwargs.pop("memory_bus", None)
+        if memory_bus is None:
+            memory_bus = persistent_bus
+        memory_slots = kwargs.pop("memory_slots", None)
+        slot_identity = kwargs.pop("slot_identity", None)
+        cache = past_key_values
+        hidden = inputs_embeds
+        positions = {
+            layer_type: self.rotary_emb(hidden, decoder_position_ids, layer_type)
+            for layer_type in self.unique_layer_types
+        }
+        working_kv = None
+        if working_bus is not None and working_state is not None:
+            working_kv = working_bus.prepare_kv(working_state)
+        memory_kv = None
+        if memory_bus is not None and memory_slots is not None:
+            memory_kv = memory_bus.prepare_kv(memory_slots, slot_identity)
+        working_reader = 0
+        reader_index = 0
+        for index in range(self.text_config.num_hidden_layers):
+            layer = self.layers[index]
+            layer_type = self.text_config.layer_types[index]
+            hidden = layer(
+                hidden,
+                position_embeddings=positions[layer_type],
+                attention_mask=mask_mapping[layer_type],
+                position_ids=decoder_position_ids,
+                past_key_values=cache,
+                **kwargs,
+            )
+            if layer_type == "full_attention":
+                if (
+                    working_kv is not None
+                    and working_reader < working_bus.num_readers
+                ):
+                    hidden = working_bus.read(
+                        hidden, working_reader, working_kv[0], working_kv[1]
+                    )
+                    working_reader += 1
+                if (
+                    memory_kv is not None
+                    and reader_index < memory_bus.num_readers
+                ):
+                    hidden = memory_bus.read(
+                        hidden, reader_index, memory_kv[0], memory_kv[1]
+                    )
+                    reader_index += 1
+        return self.norm(hidden)
+
+    def forward(
+        self,
+        decoder_input_ids: torch.LongTensor,
+        past_key_values: Cache | None = None,
+        decoder_token_embeddings: torch.FloatTensor | None = None,
+        temporal_context_embeddings: torch.FloatTensor | None = None,
+        decoder_attention_mask: torch.Tensor | dict | None = None,
+        decoder_position_ids: torch.LongTensor | None = None,
+        collect_latent_diagnostics: bool = True,
+        memory_bus: Any | None = None,
+        memory_slots: torch.Tensor | None = None,
+        working_bus: Any | None = None,
+        working_state: torch.Tensor | None = None,
+        persistent_bus: Any | None = None,
+        slot_identity: torch.Tensor | None = None,
+        **kwargs: Any,
+    ) -> ModilifyMk2DecoderOutput:
+        if "use_cache" in kwargs:
+            raise ValueError("The diffusion decoder always reads the supplied cache.")
+        if decoder_token_embeddings is None:
+            token_embeddings = self.embed_tokens(decoder_input_ids)
+        else:
+            token_embeddings = decoder_token_embeddings
+            expected = (*decoder_input_ids.shape, self.text_config.hidden_size)
+            if token_embeddings.shape != expected:
+                raise ValueError("Precomputed decoder embeddings have the wrong shape.")
+        context_embeddings = (
+            torch.zeros_like(token_embeddings)
+            if temporal_context_embeddings is None
+            else temporal_context_embeddings.to(token_embeddings)
+        )
+        inputs_embeds, diagnostics = self.merge_latent_context(
+            token_embeddings,
+            context_embeddings if temporal_context_embeddings is not None else None,
+            collect_diagnostics=collect_latent_diagnostics,
+        )
+        hidden = self._run_stack(
+            inputs_embeds,
+            past_key_values=past_key_values,
+            decoder_attention_mask=decoder_attention_mask,
+            decoder_position_ids=decoder_position_ids,
+            memory_bus=memory_bus,
+            memory_slots=memory_slots,
+            working_bus=working_bus,
+            working_state=working_state,
+            persistent_bus=persistent_bus,
+            slot_identity=slot_identity,
+            **kwargs,
+        )
+        return ModilifyMk2DecoderOutput(
+            last_hidden_state=hidden,
+            past_key_values=past_key_values,
+            token_embeddings=token_embeddings,
+            latent_residual_diagnostics=diagnostics,
+        )
+
+
+class ModilifyMk2Model(DiffusionGemmaPreTrainedModel):
+    config_class = ModilifyMk2Config
+    _tied_weights_keys = {
+        "encoder.language_model.norm.weight": "decoder.norm.weight",
+        r"encoder.language_model.layers\.(?:[^.]+\.)*weight": r"decoder.layers\.(?:[^.]+\.)*weight",
+        r"encoder.language_model.layers\.(?:[^.]+\.)*scale": r"decoder.layers\.(?:[^.]+\.)*scale",
+        r"encoder.language_model.layers\.(?:[^.]+\.)*per_expert_scale": r"decoder.layers\.(?:[^.]+\.)*per_expert_scale",
+        r"encoder.language_model.layers\.(?:[^.]+\.)*gate_up_proj": r"decoder.layers\.(?:[^.]+\.)*gate_up_proj",
+        r"encoder.language_model.layers\.(?:[^.]+\.)*down_proj": r"decoder.layers\.(?:[^.]+\.)*down_proj",
+        "encoder.language_model.embed_tokens.weight": "decoder.embed_tokens.weight",
+    }
+
+    def __init__(self, config: ModilifyMk2Config):
+        super().__init__(config)
+        self.encoder = ModilifyMk2EncoderModel(config)
+        self.decoder = ModilifyMk2DecoderModel(config)
+        install_modilify_mk2_trunk_semantics(self)
+        self.post_init()
+
+    def get_encoder(self):
+        return self.encoder
+
+    def get_decoder(self):
+        return self.decoder
+
+    def get_input_embeddings(self):
+        return self.encoder.get_input_embeddings()
+
+    def set_input_embeddings(self, value):
+        self.encoder.set_input_embeddings(value)
+        self.decoder.embed_tokens = value
+
+    def forward(
+        self,
+        *,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | dict | None = None,
+        past_key_values: Cache | None = None,
+        position_ids: torch.LongTensor | None = None,
+        decoder_input_ids: torch.LongTensor,
+        decoder_token_embeddings: torch.FloatTensor | None = None,
+        temporal_context_embeddings: torch.FloatTensor | None = None,
+        decoder_attention_mask: torch.Tensor | dict | None = None,
+        decoder_position_ids: torch.LongTensor | None = None,
+        return_encoder_outputs: bool = True,
+        collect_latent_diagnostics: bool = True,
+        memory_bus: Any | None = None,
+        memory_slots: torch.Tensor | None = None,
+        working_bus: Any | None = None,
+        working_state: torch.Tensor | None = None,
+        persistent_bus: Any | None = None,
+        slot_identity: torch.Tensor | None = None,
+        **kwargs: Any,
+    ) -> ModilifyMk2ModelOutput:
+        encoder_hidden = None
+        encoder_keys = ("pixel_values", "mm_token_type_ids", "image_position_ids", "inputs_embeds")
+        encoder_kwargs = {key: kwargs.pop(key) for key in encoder_keys if key in kwargs}
+        if input_ids is not None:
+            encoded = self.encoder(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                past_key_values=past_key_values,
+                position_ids=position_ids,
+                **encoder_kwargs,
+            )
+            past_key_values = encoded.past_key_values
+            if return_encoder_outputs:
+                encoder_hidden = encoded.last_hidden_state
+        elif past_key_values is None:
+            raise ValueError("Either `input_ids` or `past_key_values` is required.")
+        decoded = self.decoder(
+            decoder_input_ids=decoder_input_ids,
+            decoder_token_embeddings=decoder_token_embeddings,
+            past_key_values=past_key_values,
+            temporal_context_embeddings=temporal_context_embeddings,
+            decoder_attention_mask=decoder_attention_mask,
+            decoder_position_ids=decoder_position_ids,
+            collect_latent_diagnostics=collect_latent_diagnostics,
+            memory_bus=memory_bus,
+            memory_slots=memory_slots,
+            working_bus=working_bus,
+            working_state=working_state,
+            persistent_bus=persistent_bus,
+            slot_identity=slot_identity,
+            **kwargs,
+        )
+        return ModilifyMk2ModelOutput(
+            last_hidden_state=decoded.last_hidden_state,
+            past_key_values=past_key_values,
+            token_embeddings=decoded.token_embeddings,
+            encoder_last_hidden_state=encoder_hidden,
+            latent_residual_diagnostics=decoded.latent_residual_diagnostics,
+        )
+
+
+class ModilifyMk2ForBlockDiffusion(DiffusionGemmaPreTrainedModel, ModilifyMk2GenerationMixin):
+    config_class = ModilifyMk2Config
+    _tied_weights_keys = {"lm_head.weight": "model.decoder.embed_tokens.weight"}
+    generation_config_class = ModilifyMk2GenerationConfig
+
+    @torch.no_grad()
+    def _init_weights(self, module: nn.Module) -> None:
+        super()._init_weights(module)
+        if isinstance(module, LatentDeliberationTransformer):
+            module.reset_identity_parameters()
+            module.working_memory_bus.freeze()
+            module.persistent_memory_bus.freeze()
+
+    def __init__(self, config: ModilifyMk2Config):
+        super().__init__(config)
+        self.model = ModilifyMk2Model(config)
+        layer_types = tuple(getattr(config.text_config, "layer_types", None) or ())
+        self.latent_deliberation = LatentDeliberationTransformer(
+            hidden_size=config.text_config.hidden_size,
+            vocab_size=config.text_config.vocab_size,
+            latent_dim=config.latent_dim,
+            ffn_dim=config.latent_ffn_dim,
+            memory_slots=config.latent_memory_slots,
+            num_layers=config.latent_num_layers,
+            num_heads=config.latent_num_heads,
+            local_attention_window=config.latent_local_attention_window,
+            dropout=config.latent_dropout,
+            history_length=config.latent_history_length,
+            tape_probes=config.latent_tape_probes,
+            history_kv_rank=config.latent_history_kv_rank,
+            num_memory_readers=sum(layer_type == "full_attention" for layer_type in layer_types),
+            num_working_readers=(
+                sum(layer_type == "full_attention" for layer_type in layer_types)
+                if config.working_memory_bus else 0
+            ),
+            num_persistent_readers=(
+                sum(layer_type == "full_attention" for layer_type in layer_types)
+                if config.persistent_memory_bus else 0
+            ),
+            working_last_block_global=config.latent_working_last_block_global,
+            experience_roles=config.experience_roles,
+            commit_sequence_layers=config.commit_sequence_layers,
+            commit_sequence_dim=config.commit_sequence_dim,
+            max_canvas_length=config.canvas_length,
+        )
+        self.lm_head = nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False)
+        self.final_logit_softcapping = config.text_config.final_logit_softcapping
+        self.post_init()
+        install_modilify_mk2_trunk_semantics(self)
+
+    def _finalize_logits(self, hidden: torch.Tensor) -> torch.Tensor:
+        logits = self.lm_head(hidden)
+        return torch.tanh(logits / self.final_logit_softcapping) * self.final_logit_softcapping
+
+    def _prepare_latent_context(
+        self,
+        decoder_input_ids: torch.LongTensor,
+        *,
+        history: TrajectoryHistory | None,
+        tape: TrajectoryTape | None,
+        confidence: torch.Tensor | None,
+        entropy: torch.Tensor | None,
+        age: torch.Tensor | None,
+        latent_state: LatentDeliberationState | None,
+    ) -> tuple[torch.Tensor, LatentDeliberationState, torch.Tensor, torch.Tensor]:
+        batch, canvas = decoder_input_ids.shape
+        dtype = self.model.decoder.embed_tokens.weight.dtype
+        if latent_state is None:
+            latent_state = LatentDeliberationState.empty(
+                batch_size=batch, canvas_length=canvas,
+                latent_dim=self.config.latent_dim, memory_slots=self.config.latent_memory_slots,
+                device=decoder_input_ids.device, dtype=dtype,
+            )
+        if history is None:
+            history = TrajectoryHistory.empty(
+                batch_size=batch,
+                canvas_length=canvas,
+                hidden_size=self.config.text_config.hidden_size,
+                history_length=self.config.latent_history_length,
+                device=decoder_input_ids.device,
+                dtype=dtype,
+            )
+        if tape is None:
+            tape = empty_trajectory_tape(
+                batch_size=batch,
+                config=self.config,
+                device=decoder_input_ids.device,
+                dtype=dtype,
+            )
+        confidence = latent_state.confidence if confidence is None else confidence.squeeze(-1).float()
+        entropy = latent_state.entropy if entropy is None else entropy.squeeze(-1).float()
+        if age is not None:
+            latent_state = replace(
+                latent_state,
+                age=age.to(device=decoder_input_ids.device, dtype=torch.int32),
+            )
+        token_embeddings = self.model.decoder.embed_tokens(decoder_input_ids)
+        processed: LatentProcessorOutput = self.latent_deliberation(
+            token_embeddings=token_embeddings,
+            confidence=confidence,
+            entropy=entropy,
+            state=latent_state,
+            history=history,
+            tape=tape,
+        )
+        return (
+            processed.context,
+            processed.state,
+            token_embeddings,
+            processed.history_projected,
+        )
+
+    def forward(
+        self,
+        *,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | dict | None = None,
+        past_key_values: Cache | None = None,
+        position_ids: torch.LongTensor | None = None,
+        decoder_input_ids: torch.LongTensor,
+        previous_confidence: torch.FloatTensor | None = None,
+        previous_entropy: torch.FloatTensor | None = None,
+        token_age: torch.Tensor | None = None,
+        latent_state: LatentDeliberationState | None = None,
+        history: TrajectoryHistory | None = None,
+        tape: TrajectoryTape | None = None,
+        history_hidden_state: TrajectoryHistory | torch.FloatTensor | None = None,
+        decoder_attention_mask: torch.Tensor | dict | None = None,
+        decoder_position_ids: torch.LongTensor | None = None,
+        return_encoder_outputs: bool = True,
+        compact_vocab: bool = False,
+        denoise_temperature: float | None = None,
+        repetition_token_mask: torch.BoolTensor | None = None,
+        repetition_penalty: float = 1.0,
+        sampling_generators: Sequence[torch.Generator] | None = None,
+        collect_latent_diagnostics: bool = True,
+        **kwargs: Any,
+    ) -> ModilifyMk2BlockDiffusionOutput:
+        if history is None and isinstance(history_hidden_state, TrajectoryHistory):
+            history = history_hidden_state
+        (
+            latent_context,
+            next_state,
+            decoder_token_embeddings,
+            history_projected,
+        ) = self._prepare_latent_context(
+            decoder_input_ids,
+            history=history,
+            tape=tape,
+            confidence=previous_confidence,
+            entropy=previous_entropy,
+            age=token_age,
+            latent_state=latent_state,
+        )
+        working_bus = self.latent_deliberation.working_memory_bus
+        persistent_bus = self.latent_deliberation.persistent_memory_bus
+        working_readers_active = working_bus.num_readers > 0
+        memory_readers_active = persistent_bus.num_readers > 0
+        outputs = self.model(
+            input_ids=input_ids, attention_mask=attention_mask,
+            past_key_values=past_key_values, position_ids=position_ids,
+            decoder_input_ids=decoder_input_ids,
+            decoder_token_embeddings=decoder_token_embeddings,
+            temporal_context_embeddings=latent_context,
+            decoder_attention_mask=decoder_attention_mask,
+            decoder_position_ids=decoder_position_ids,
+            return_encoder_outputs=return_encoder_outputs,
+            collect_latent_diagnostics=collect_latent_diagnostics,
+            working_bus=working_bus if working_readers_active else None,
+            working_state=latent_context if working_readers_active else None,
+            persistent_bus=persistent_bus if memory_readers_active else None,
+            memory_slots=next_state.memory_slots if memory_readers_active else None,
+            slot_identity=(
+                self.latent_deliberation.scaled_memory_slot_identity(
+                    batch_size=decoder_input_ids.shape[0],
+                    device=decoder_input_ids.device,
+                    dtype=latent_context.dtype,
+                )
+                if memory_readers_active else None
+            ),
+            **kwargs,
+        )
+        temperature = (
+            self.config.denoise_temperature
+            if denoise_temperature is None
+            else float(denoise_temperature)
+        )
+        proposal = proposal_confidence = token_entropy = None
+        greedy_proposal = greedy_confidence = None
+        logits = None
+        if compact_vocab:
+            (
+                proposal,
+                proposal_confidence,
+                token_entropy,
+                greedy_proposal,
+                greedy_confidence,
+            ) = chunked_vocab_statistics(
+                outputs.last_hidden_state.detach(),
+                self.lm_head.weight.detach(),
+                softcap=self.final_logit_softcapping,
+                temperature=temperature,
+                chunk_size=self.config.vocab_chunk_size,
+                repetition_token_mask=repetition_token_mask,
+                repetition_penalty=repetition_penalty,
+                sampling_generators=sampling_generators,
+            )
+        else:
+            logits = self._finalize_logits(outputs.last_hidden_state)
+        return ModilifyMk2BlockDiffusionOutput(
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            encoder_last_hidden_state=outputs.encoder_last_hidden_state,
+            heavy_hidden_state=outputs.last_hidden_state,
+            next_latent_state=next_state,
+            temporal_context=latent_context,
+            history_projected=history_projected,
+            working_state=latent_context,
+            latent_residual_diagnostics=outputs.latent_residual_diagnostics,
+            proposal=proposal,
+            proposal_confidence=proposal_confidence,
+            token_entropy=token_entropy,
+            greedy_proposal=greedy_proposal,
+            greedy_confidence=greedy_confidence,
+        )
+
+
+ModilifyMk2Model.register_for_auto_class("AutoModel")
+ModilifyMk2ForBlockDiffusion.register_for_auto_class("AutoModelForCausalLM")
+ModilifyMk2ForBlockDiffusion.register_for_auto_class("AutoModelForMultimodalLM")
+
+
+__all__ = [
+    "ModilifyMk2BlockDiffusionOutput",
+    "ModilifyMk2Config",
+    "ModilifyMk2DecoderModel",
+    "ModilifyMk2EncoderModel",
+    "ModilifyMk2ForBlockDiffusion",
+    "ModilifyMk2Model",
+    "ModilifyMk2RMSNorm",
+    "ModilifyMk2TextRouter",
+    "install_modilify_mk2_trunk_semantics",
+]

mps_ops.py

@@ -0,0 +1,618 @@
+# Copyright 2026 Modilify
+# SPDX-License-Identifier: LicenseRef-Modilify-Open-Model-1.0
+"""MPS-specific kernels that preserve the model's mathematical operations."""
+
+from __future__ import annotations
+
+import math
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+from transformers.integrations.moe import ALL_EXPERTS_FUNCTIONS, _grouped_linear
+
+
+def lora_aware_eager_experts_forward(
+    self: nn.Module,
+    hidden_states: torch.Tensor,
+    top_k_index: torch.Tensor,
+    top_k_weights: torch.Tensor,
+) -> torch.Tensor:
+    final_hidden_states = torch.zeros_like(hidden_states)
+    with torch.no_grad():
+        expert_mask = torch.nn.functional.one_hot(top_k_index, num_classes=self.num_experts)
+        expert_mask = expert_mask.permute(2, 1, 0)
+        expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero()
+
+    lora_dropout = getattr(self, "lora_dropout", nn.Identity())
+    lora_scaling = getattr(self, "lora_scaling", 1.0)
+
+    for expert_idx in expert_hit:
+        expert_idx = expert_idx[0]
+        if expert_idx == self.num_experts:
+            continue
+        top_k_pos, token_idx = torch.where(expert_mask[expert_idx])
+        current_state = hidden_states[token_idx]
+        gate_up = nn.functional.linear(current_state, self.gate_up_proj[expert_idx])
+        if hasattr(self, "lora_gate_up_a"):
+            update = nn.functional.linear(
+                nn.functional.linear(lora_dropout(current_state), self.lora_gate_up_a[expert_idx]),
+                self.lora_gate_up_b[expert_idx],
+            )
+            gate_up = gate_up + lora_scaling * update
+        gate, up = gate_up.chunk(2, dim=-1)
+        current_hidden_states = self.act_fn(gate) * up
+        expert_output = nn.functional.linear(current_hidden_states, self.down_proj[expert_idx])
+        if hasattr(self, "lora_down_a"):
+            update = nn.functional.linear(
+                nn.functional.linear(lora_dropout(current_hidden_states), self.lora_down_a[expert_idx]),
+                self.lora_down_b[expert_idx],
+            )
+            expert_output = expert_output + lora_scaling * update
+        current_hidden_states = expert_output * top_k_weights[token_idx, top_k_pos, None]
+        final_hidden_states.index_add_(0, token_idx, current_hidden_states.to(final_hidden_states.dtype))
+
+    return final_hidden_states
+
+
+def grouped_expert_offsets(
+    sorted_expert_ids: torch.Tensor,
+    num_experts: int,
+) -> torch.Tensor:
+    """Build grouped-mm offsets without MPS ``histc`` host synchronization."""
+    counts = torch.bincount(sorted_expert_ids, minlength=num_experts)[:num_experts]
+    return torch.cumsum(counts, dim=0, dtype=torch.int32)
+
+
+def mps_grouped_mm_experts_forward(
+    self: torch.nn.Module,
+    hidden_states: torch.Tensor,
+    top_k_index: torch.Tensor,
+    top_k_weights: torch.Tensor,
+) -> torch.Tensor:
+    """Transformers grouped-mm MoE with histogram-free MPS routing offsets."""
+    num_top_k = top_k_index.size(-1)
+    num_tokens = hidden_states.size(0)
+    hidden_dim = hidden_states.size(-1)
+    sample_weights = top_k_weights.reshape(-1)
+    expert_ids = top_k_index.reshape(-1)
+
+    expert_ids_g, perm = torch.sort(expert_ids)
+    selected_hidden_states_g = hidden_states[perm // num_top_k]
+    sample_weights_g = sample_weights[perm]
+    offsets = grouped_expert_offsets(expert_ids_g, self.num_experts)
+
+    sentinel_mask = (expert_ids_g >= self.num_experts).unsqueeze(-1)
+    expert_ids_g.clamp_(max=self.num_experts - 1)
+    selected_hidden_states_g.masked_fill_(sentinel_mask, 0.0)
+
+    gate_up_weights = self.gate_up_proj if self.has_gate else self.up_proj
+    gate_up_biases = (
+        self.gate_up_proj_bias[expert_ids_g]
+        if self.has_gate and self.has_bias
+        else self.up_proj_bias[expert_ids_g]
+        if self.has_bias
+        else None
+    )
+    projected = _grouped_linear(
+        selected_hidden_states_g,
+        gate_up_weights,
+        offsets,
+        bias=gate_up_biases,
+        is_transposed=self.is_transposed,
+    )
+    projected = self._apply_gate(projected) if self.has_gate else self.act_fn(projected)
+
+    down_biases = self.down_proj_bias[expert_ids_g] if self.has_bias else None
+    projected = _grouped_linear(
+        projected,
+        self.down_proj,
+        offsets,
+        bias=down_biases,
+        is_transposed=self.is_transposed,
+    )
+    weighted = projected * sample_weights_g.unsqueeze(-1)
+    weighted.masked_fill_(sentinel_mask, 0.0)
+
+    # Scatter directly back to the original top-k order.  Constructing the
+    # inverse permutation and then gathering performs the same permutation
+    # with an extra index tensor and an extra read pass over `weighted`.
+    reordered = torch.empty_like(weighted)
+    reordered.index_copy_(0, perm, weighted)
+    weighted = reordered
+    return weighted.view(num_tokens, num_top_k, hidden_dim).sum(dim=1).to(hidden_states.dtype)
+
+
+ALL_EXPERTS_FUNCTIONS.register("mps_grouped_mm", mps_grouped_mm_experts_forward)
+
+
+class _IndependentExpertLoRA(torch.autograd.Function):
+    """Capacity-padded reference path for independent expert LoRA."""
+
+    @staticmethod
+    def forward(
+        ctx,
+        inputs: torch.Tensor,
+        weight_a: torch.Tensor,
+        weight_b: torch.Tensor,
+        expert_ids: torch.LongTensor,
+        counts: torch.LongTensor,
+        capacity: int,
+    ) -> torch.Tensor:
+        offsets = counts.cumsum(0)
+        starts = offsets - counts
+        slots = torch.arange(inputs.shape[0], device=inputs.device) - starts[expert_ids]
+        padded = inputs.new_zeros((weight_a.shape[0], capacity, inputs.shape[-1]))
+        padded[expert_ids, slots] = inputs
+        low_rank = torch.bmm(padded, weight_a.transpose(1, 2))
+        updates = torch.bmm(low_rank, weight_b.transpose(1, 2))
+        ctx.save_for_backward(
+            padded,
+            low_rank,
+            weight_a,
+            weight_b,
+            expert_ids,
+            slots,
+        )
+        return updates[expert_ids, slots]
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor):
+        padded, low_rank, weight_a, weight_b, expert_ids, slots = ctx.saved_tensors
+        grad_padded = grad_output.new_zeros(
+            (weight_a.shape[0], padded.shape[1], grad_output.shape[-1])
+        )
+        grad_padded[expert_ids, slots] = grad_output
+        grad_weight_b = torch.bmm(grad_padded.transpose(1, 2), low_rank)
+        grad_low_rank = torch.bmm(grad_padded, weight_b)
+        grad_weight_a = torch.bmm(grad_low_rank.transpose(1, 2), padded)
+        grad_inputs_padded = torch.bmm(grad_low_rank, weight_a)
+        grad_inputs = grad_inputs_padded[expert_ids, slots]
+        return grad_inputs, grad_weight_a, grad_weight_b, None, None, None
+
+
+_INDEPENDENT_EXPERT_LORA_METAL_SOURCE = r"""
+#include <metal_stdlib>
+using namespace metal;
+
+constant uint kRank = 8;
+constant uint kThreads = 256;
+
+kernel void expert_lora_forward_a(
+    device const bfloat* inputs [[buffer(0)]],
+    device const bfloat* weight_a [[buffer(1)]],
+    device const long* expert_ids [[buffer(2)]],
+    device bfloat* low_rank [[buffer(3)]],
+    constant uint& row_count [[buffer(4)]],
+    constant uint& input_dim [[buffer(5)]],
+    uint row [[threadgroup_position_in_grid]],
+    uint lane [[thread_index_in_threadgroup]]) {
+  if (row >= row_count) return;
+  const uint expert = uint(expert_ids[row]);
+  threadgroup float partial[kThreads * kRank];
+  for (uint rank = 0; rank < kRank; ++rank) {
+    float value = 0.0f;
+    const uint a_base = (expert * kRank + rank) * input_dim;
+    const uint x_base = row * input_dim;
+    for (uint column = lane; column < input_dim; column += kThreads) {
+      value += float(inputs[x_base + column]) * float(weight_a[a_base + column]);
+    }
+    partial[lane * kRank + rank] = value;
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  for (uint stride = kThreads / 2; stride > 0; stride >>= 1) {
+    if (lane < stride) {
+      for (uint rank = 0; rank < kRank; ++rank) {
+        partial[lane * kRank + rank] +=
+            partial[(lane + stride) * kRank + rank];
+      }
+    }
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+  }
+  if (lane == 0) {
+    for (uint rank = 0; rank < kRank; ++rank) {
+      low_rank[row * kRank + rank] = bfloat(partial[rank]);
+    }
+  }
+}
+
+kernel void expert_lora_forward_b(
+    device const bfloat* low_rank [[buffer(0)]],
+    device const bfloat* weight_b [[buffer(1)]],
+    device const long* expert_ids [[buffer(2)]],
+    device bfloat* output [[buffer(3)]],
+    constant uint& row_count [[buffer(4)]],
+    constant uint& output_dim [[buffer(5)]],
+    uint index [[thread_position_in_grid]]) {
+  const uint total = row_count * output_dim;
+  if (index >= total) return;
+  const uint row = index / output_dim;
+  const uint column = index - row * output_dim;
+  const uint expert = uint(expert_ids[row]);
+  const uint b_base = (expert * output_dim + column) * kRank;
+  float value = 0.0f;
+  for (uint rank = 0; rank < kRank; ++rank) {
+    value += float(low_rank[row * kRank + rank])
+        * float(weight_b[b_base + rank]);
+  }
+  output[index] = bfloat(value);
+}
+
+kernel void expert_lora_backward_low_rank(
+    device const bfloat* grad_output [[buffer(0)]],
+    device const bfloat* weight_b [[buffer(1)]],
+    device const long* expert_ids [[buffer(2)]],
+    device bfloat* grad_low_rank [[buffer(3)]],
+    constant uint& row_count [[buffer(4)]],
+    constant uint& output_dim [[buffer(5)]],
+    uint row [[threadgroup_position_in_grid]],
+    uint lane [[thread_index_in_threadgroup]]) {
+  if (row >= row_count) return;
+  const uint expert = uint(expert_ids[row]);
+  threadgroup float partial[kThreads * kRank];
+  for (uint rank = 0; rank < kRank; ++rank) {
+    float value = 0.0f;
+    const uint grad_base = row * output_dim;
+    for (uint column = lane; column < output_dim; column += kThreads) {
+      const uint b_index =
+          (expert * output_dim + column) * kRank + rank;
+      value += float(grad_output[grad_base + column])
+          * float(weight_b[b_index]);
+    }
+    partial[lane * kRank + rank] = value;
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  for (uint stride = kThreads / 2; stride > 0; stride >>= 1) {
+    if (lane < stride) {
+      for (uint rank = 0; rank < kRank; ++rank) {
+        partial[lane * kRank + rank] +=
+            partial[(lane + stride) * kRank + rank];
+      }
+    }
+    threadgroup_barrier(mem_flags::mem_threadgroup);
+  }
+  if (lane == 0) {
+    for (uint rank = 0; rank < kRank; ++rank) {
+      grad_low_rank[row * kRank + rank] = bfloat(partial[rank]);
+    }
+  }
+}
+
+kernel void expert_lora_backward_inputs(
+    device const bfloat* grad_low_rank [[buffer(0)]],
+    device const bfloat* weight_a [[buffer(1)]],
+    device const long* expert_ids [[buffer(2)]],
+    device bfloat* grad_inputs [[buffer(3)]],
+    constant uint& row_count [[buffer(4)]],
+    constant uint& input_dim [[buffer(5)]],
+    uint index [[thread_position_in_grid]]) {
+  const uint total = row_count * input_dim;
+  if (index >= total) return;
+  const uint row = index / input_dim;
+  const uint column = index - row * input_dim;
+  const uint expert = uint(expert_ids[row]);
+  float value = 0.0f;
+  for (uint rank = 0; rank < kRank; ++rank) {
+    const uint a_index =
+        (expert * kRank + rank) * input_dim + column;
+    value += float(grad_low_rank[row * kRank + rank])
+        * float(weight_a[a_index]);
+  }
+  grad_inputs[index] = bfloat(value);
+}
+
+kernel void expert_lora_backward_a(
+    device const bfloat* grad_low_rank [[buffer(0)]],
+    device const bfloat* inputs [[buffer(1)]],
+    device const long* offsets [[buffer(2)]],
+    device bfloat* grad_weight_a [[buffer(3)]],
+    constant uint& expert_count [[buffer(4)]],
+    constant uint& input_dim [[buffer(5)]],
+    uint index [[thread_position_in_grid]]) {
+  const uint total = expert_count * kRank * input_dim;
+  if (index >= total) return;
+  const uint column = index % input_dim;
+  const uint rank_expert = index / input_dim;
+  const uint rank = rank_expert % kRank;
+  const uint expert = rank_expert / kRank;
+  const uint start = uint(offsets[expert]);
+  const uint stop = uint(offsets[expert + 1]);
+  float value = 0.0f;
+  for (uint row = start; row < stop; ++row) {
+    value += float(grad_low_rank[row * kRank + rank])
+        * float(inputs[row * input_dim + column]);
+  }
+  grad_weight_a[index] = bfloat(value);
+}
+
+kernel void expert_lora_backward_b(
+    device const bfloat* grad_output [[buffer(0)]],
+    device const bfloat* low_rank [[buffer(1)]],
+    device const long* offsets [[buffer(2)]],
+    device bfloat* grad_weight_b [[buffer(3)]],
+    constant uint& expert_count [[buffer(4)]],
+    constant uint& output_dim [[buffer(5)]],
+    uint index [[thread_position_in_grid]]) {
+  const uint total = expert_count * output_dim * kRank;
+  if (index >= total) return;
+  const uint rank = index % kRank;
+  const uint output_expert = index / kRank;
+  const uint column = output_expert % output_dim;
+  const uint expert = output_expert / output_dim;
+  const uint start = uint(offsets[expert]);
+  const uint stop = uint(offsets[expert + 1]);
+  float value = 0.0f;
+  for (uint row = start; row < stop; ++row) {
+    value += float(grad_output[row * output_dim + column])
+        * float(low_rank[row * kRank + rank]);
+  }
+  grad_weight_b[index] = bfloat(value);
+}
+"""
+
+
+_independent_expert_lora_metal_library = None
+
+
+def _get_independent_expert_lora_metal_library():
+    global _independent_expert_lora_metal_library
+    if _independent_expert_lora_metal_library is None:
+        _independent_expert_lora_metal_library = torch.mps.compile_shader(
+            _INDEPENDENT_EXPERT_LORA_METAL_SOURCE
+        )
+    return _independent_expert_lora_metal_library
+
+
+class _MetalIndependentExpertLoRA(torch.autograd.Function):
+    """Independent rank-8 expert A/B with direct BF16 Metal forward/backward."""
+
+    @staticmethod
+    def forward(
+        ctx,
+        inputs: torch.Tensor,
+        weight_a: torch.Tensor,
+        weight_b: torch.Tensor,
+        expert_ids: torch.LongTensor,
+        counts: torch.LongTensor,
+    ) -> torch.Tensor:
+        inputs = inputs.contiguous()
+        weight_a = weight_a.contiguous()
+        weight_b = weight_b.contiguous()
+        expert_ids = expert_ids.contiguous()
+        expert_count, rank, input_dim = weight_a.shape
+        row_count = inputs.shape[0]
+        output_dim = weight_b.shape[1]
+        if rank != 8:
+            raise ValueError("The direct Metal expert LoRA kernel requires rank 8.")
+        offsets = torch.cat(
+            (
+                counts.new_zeros(1, dtype=torch.int64),
+                counts.cumsum(0, dtype=torch.int64),
+            )
+        ).contiguous()
+        low_rank = inputs.new_empty((row_count, rank))
+        output = inputs.new_empty((row_count, output_dim))
+        library = _get_independent_expert_lora_metal_library()
+        library.expert_lora_forward_a(
+            inputs,
+            weight_a,
+            expert_ids,
+            low_rank,
+            row_count,
+            input_dim,
+            threads=(row_count * 256, 1, 1),
+            group_size=(256, 1, 1),
+        )
+        library.expert_lora_forward_b(
+            low_rank,
+            weight_b,
+            expert_ids,
+            output,
+            row_count,
+            output_dim,
+            threads=(row_count * output_dim, 1, 1),
+            group_size=(256, 1, 1),
+        )
+        ctx.save_for_backward(
+            inputs,
+            low_rank,
+            weight_a,
+            weight_b,
+            expert_ids,
+            offsets,
+        )
+        ctx.dimensions = (expert_count, row_count, input_dim, output_dim)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor):
+        inputs, low_rank, weight_a, weight_b, expert_ids, offsets = ctx.saved_tensors
+        expert_count, row_count, input_dim, output_dim = ctx.dimensions
+        grad_output = grad_output.contiguous()
+        grad_low_rank = low_rank.new_empty(low_rank.shape)
+        grad_inputs = inputs.new_empty(inputs.shape)
+        grad_weight_a = weight_a.new_empty(weight_a.shape)
+        grad_weight_b = weight_b.new_empty(weight_b.shape)
+        library = _get_independent_expert_lora_metal_library()
+        library.expert_lora_backward_low_rank(
+            grad_output,
+            weight_b,
+            expert_ids,
+            grad_low_rank,
+            row_count,
+            output_dim,
+            threads=(row_count * 256, 1, 1),
+            group_size=(256, 1, 1),
+        )
+        library.expert_lora_backward_inputs(
+            grad_low_rank,
+            weight_a,
+            expert_ids,
+            grad_inputs,
+            row_count,
+            input_dim,
+            threads=(row_count * input_dim, 1, 1),
+            group_size=(256, 1, 1),
+        )
+        library.expert_lora_backward_a(
+            grad_low_rank,
+            inputs,
+            offsets,
+            grad_weight_a,
+            expert_count,
+            input_dim,
+            threads=(expert_count * 8 * input_dim, 1, 1),
+            group_size=(256, 1, 1),
+        )
+        library.expert_lora_backward_b(
+            grad_output,
+            low_rank,
+            offsets,
+            grad_weight_b,
+            expert_count,
+            output_dim,
+            threads=(expert_count * output_dim * 8, 1, 1),
+            group_size=(256, 1, 1),
+        )
+        return grad_inputs, grad_weight_a, grad_weight_b, None, None
+
+
+def _independent_expert_lora(
+    inputs: torch.Tensor,
+    weight_a: torch.Tensor,
+    weight_b: torch.Tensor,
+    expert_ids: torch.LongTensor,
+    counts: torch.LongTensor,
+) -> torch.Tensor:
+    """Dispatch the compact independent-expert operator without sharing A/B."""
+
+    if inputs.shape[0] == 0:
+        return inputs.new_empty((0, weight_b.shape[1]))
+    if (
+        inputs.device.type == "mps"
+        and inputs.dtype == torch.bfloat16
+        and weight_a.dtype == torch.bfloat16
+        and weight_b.dtype == torch.bfloat16
+        and weight_a.shape[1] == 8
+        and hasattr(torch.mps, "compile_shader")
+    ):
+        return _MetalIndependentExpertLoRA.apply(
+            inputs,
+            weight_a,
+            weight_b,
+            expert_ids,
+            counts,
+        )
+    return _IndependentExpertLoRA.apply(
+        inputs,
+        weight_a,
+        weight_b,
+        expert_ids,
+        counts,
+        int(counts.max().detach().cpu()),
+    )
+
+
+def mps_segmented_experts_forward(
+    self: torch.nn.Module,
+    hidden_states: torch.Tensor,
+    top_k_index: torch.Tensor,
+    top_k_weights: torch.Tensor,
+) -> torch.Tensor:
+    """Run GPU-grouped base experts plus independent Metal LoRA A/B."""
+    if hidden_states.device.type != "mps":
+        return lora_aware_eager_experts_forward(
+            self,
+            hidden_states,
+            top_k_index,
+            top_k_weights,
+        )
+
+    num_top_k = top_k_index.size(-1)
+    num_tokens, hidden_dim = hidden_states.shape
+    sample_weights = top_k_weights.reshape(-1)
+    expert_ids = top_k_index.reshape(-1)
+    expert_ids_g, permutation = torch.sort(expert_ids)
+    selected_hidden_g = hidden_states[permutation // num_top_k]
+    sample_weights_g = sample_weights[permutation]
+
+    counts = torch.bincount(
+        expert_ids_g,
+        minlength=self.num_experts,
+    )[: self.num_experts]
+    offsets = torch.cumsum(counts, dim=0, dtype=torch.int32)
+    gate_up_weights = self.gate_up_proj if self.has_gate else self.up_proj
+    gate_up_bias = None
+    if self.has_bias:
+        packed_bias = self.gate_up_proj_bias if self.has_gate else self.up_proj_bias
+        gate_up_bias = packed_bias[expert_ids_g]
+    gate_up_g = _grouped_linear(
+        selected_hidden_g,
+        gate_up_weights,
+        offsets,
+        bias=gate_up_bias,
+        is_transposed=self.is_transposed,
+    )
+    if hasattr(self, "lora_gate_up_a") and self.training:
+        gate_up_update = _independent_expert_lora(
+            self.lora_dropout(selected_hidden_g),
+            self.lora_gate_up_a,
+            self.lora_gate_up_b,
+            expert_ids_g,
+            counts,
+        )
+        gate_up_g.add_(gate_up_update, alpha=self.lora_scaling)
+    elif hasattr(self, "lora_gate_up_a"):
+        gate_up_update = _independent_expert_lora(
+            selected_hidden_g,
+            self.lora_gate_up_a,
+            self.lora_gate_up_b,
+            expert_ids_g,
+            counts,
+        )
+        gate_up_g.add_(gate_up_update, alpha=self.lora_scaling)
+    activated_g = (
+        self._apply_gate(gate_up_g)
+        if self.has_gate
+        else self.act_fn(gate_up_g)
+    )
+    down_bias = self.down_proj_bias[expert_ids_g] if self.has_bias else None
+    projected_g = _grouped_linear(
+        activated_g,
+        self.down_proj,
+        offsets,
+        bias=down_bias,
+        is_transposed=self.is_transposed,
+    )
+    if hasattr(self, "lora_down_a") and self.training:
+        down_update = _independent_expert_lora(
+            self.lora_dropout(activated_g),
+            self.lora_down_a,
+            self.lora_down_b,
+            expert_ids_g,
+            counts,
+        )
+        projected_g.add_(down_update, alpha=self.lora_scaling)
+    elif hasattr(self, "lora_down_a"):
+        down_update = _independent_expert_lora(
+            activated_g,
+            self.lora_down_a,
+            self.lora_down_b,
+            expert_ids_g,
+            counts,
+        )
+        projected_g.add_(down_update, alpha=self.lora_scaling)
+    weighted_g = projected_g * sample_weights_g[:, None]
+    weighted = torch.empty_like(weighted_g)
+    weighted.index_copy_(0, permutation, weighted_g)
+    return weighted.view(num_tokens, num_top_k, hidden_dim).sum(dim=1).to(hidden_states.dtype)
+
+
+ALL_EXPERTS_FUNCTIONS.register("mps_segmented", mps_segmented_experts_forward)
+
+
+__all__ = [
+    "grouped_expert_offsets",
+    "mps_grouped_mm_experts_forward",
+    "mps_segmented_experts_forward",
+]

processor_config.json

@@ -0,0 +1,75 @@
+{
+  "audio_ms_per_token": 40,
+  "audio_seq_length": 750,
+  "feature_extractor": {
+    "dither": 0.0,
+    "feature_extractor_type": "Gemma4AudioFeatureExtractor",
+    "feature_size": 128,
+    "fft_length": 512,
+    "fft_overdrive": false,
+    "frame_length": 320,
+    "hop_length": 160,
+    "input_scale_factor": 1.0,
+    "max_frequency": 8000.0,
+    "mel_floor": 0.001,
+    "min_frequency": 0.0,
+    "padding_side": "right",
+    "padding_value": 0.0,
+    "per_bin_mean": null,
+    "per_bin_stddev": null,
+    "preemphasis": 0.0,
+    "preemphasis_htk_flavor": true,
+    "return_attention_mask": true,
+    "sampling_rate": 16000
+  },
+  "image_processor": {
+    "do_convert_rgb": true,
+    "do_normalize": false,
+    "do_rescale": true,
+    "do_resize": true,
+    "image_mean": [
+      0.0,
+      0.0,
+      0.0
+    ],
+    "image_processor_type": "Gemma4ImageProcessor",
+    "image_seq_length": 280,
+    "image_std": [
+      1.0,
+      1.0,
+      1.0
+    ],
+    "max_soft_tokens": 280,
+    "patch_size": 16,
+    "pooling_kernel_size": 3,
+    "resample": 3,
+    "rescale_factor": 0.00392156862745098
+  },
+  "image_seq_length": 280,
+  "processor_class": "Gemma4Processor",
+  "video_processor": {
+    "do_convert_rgb": true,
+    "do_normalize": true,
+    "do_rescale": true,
+    "do_resize": true,
+    "do_sample_frames": true,
+    "image_mean": [
+      0.0,
+      0.0,
+      0.0
+    ],
+    "image_std": [
+      1.0,
+      1.0,
+      1.0
+    ],
+    "max_soft_tokens": 70,
+    "num_frames": 32,
+    "patch_size": 16,
+    "pooling_kernel_size": 3,
+    "resample": 3,
+    "rescale_factor": 0.00392156862745098,
+    "return_metadata": false,
+    "video_processor_type": "Gemma4VideoProcessor"
+  }
+}

tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cc8d3a0ce36466ccc1278bf987df5f71db1719b9ca6b4118264f45cb627bfe0f
+size 32169626

tokenizer_config.json

@@ -0,0 +1,96 @@
+{
+  "audio_token": "<|audio|>",
+  "backend": "tokenizers",
+  "boa_token": "<|audio>",
+  "boi_token": "<|image>",
+  "bos_token": "<bos>",
+  "eoa_token": "<audio|>",
+  "eoc_token": "<channel|>",
+  "eoi_token": "<image|>",
+  "eos_token": "<eos>",
+  "eot_token": "<turn|>",
+  "escape_token": "<|\"|>",
+  "etc_token": "<tool_call|>",
+  "etd_token": "<tool|>",
+  "etr_token": "<tool_response|>",
+  "extra_special_tokens": [
+    "<|video|>"
+  ],
+  "image_token": "<|image|>",
+  "is_local": true,
+  "local_files_only": false,
+  "mask_token": "<mask>",
+  "model_max_length": 1000000000000000019884624838656,
+  "model_specific_special_tokens": {
+    "audio_token": "<|audio|>",
+    "boa_token": "<|audio>",
+    "boi_token": "<|image>",
+    "eoa_token": "<audio|>",
+    "eoc_token": "<channel|>",
+    "eoi_token": "<image|>",
+    "eot_token": "<turn|>",
+    "escape_token": "<|\"|>",
+    "etc_token": "<tool_call|>",
+    "etd_token": "<tool|>",
+    "etr_token": "<tool_response|>",
+    "image_token": "<|image|>",
+    "soc_token": "<|channel>",
+    "sot_token": "<|turn>",
+    "stc_token": "<|tool_call>",
+    "std_token": "<|tool>",
+    "str_token": "<|tool_response>",
+    "think_token": "<|think|>"
+  },
+  "pad_token": "<pad>",
+  "padding_side": "left",
+  "processor_class": "Gemma4Processor",
+  "response_schema": {
+    "properties": {
+      "content": {
+        "type": "string"
+      },
+      "role": {
+        "const": "assistant"
+      },
+      "thinking": {
+        "type": "string"
+      },
+      "tool_calls": {
+        "items": {
+          "properties": {
+            "function": {
+              "properties": {
+                "arguments": {
+                  "additionalProperties": {},
+                  "type": "object",
+                  "x-parser": "gemma4-tool-call"
+                },
+                "name": {
+                  "type": "string"
+                }
+              },
+              "type": "object",
+              "x-regex": "call\\:(?P<name>\\w+)(?P<arguments>\\{.*\\})"
+            },
+            "type": {
+              "const": "function"
+            }
+          },
+          "type": "object"
+        },
+        "type": "array",
+        "x-regex-iterator": "<\\|tool_call>(.*?)<tool_call\\|>"
+      }
+    },
+    "type": "object",
+    "x-regex": "(\\<\\|channel\\>thought\\n(?P<thinking>.*?)\\<channel\\|\\>)?(?P<tool_calls>\\<\\|tool_call\\>.*\\<tool_call\\|\\>)?(?P<content>(?:(?!\\<turn\\|\\>)(?!\\<\\|tool_response\\>).)+)?(?:\\<turn\\|\\>|\\<\\|tool_response\\>)?"
+  },
+  "soc_token": "<|channel>",
+  "sot_token": "<|turn>",
+  "stc_token": "<|tool_call>",
+  "std_token": "<|tool>",
+  "str_token": "<|tool_response>",
+  "think_token": "<|think|>",
+  "tokenizer_class": "GemmaTokenizer",
+  "unk_token": "<unk>"
+}

vocab_ops.py

@@ -0,0 +1,525 @@
+# Copyright 2026 Modilify
+# SPDX-License-Identifier: LicenseRef-Modilify-Open-Model-1.0
+"""Exact memory-bounded vocabulary projection and sampling."""
+
+from __future__ import annotations
+
+import math
+from collections.abc import Sequence
+
+import torch
+from torch.nn import functional as F
+
+
+def _stable_max(values: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """Return ``max`` indices that stay in-range on MPS NaN/-inf rows."""
+
+    best, index = values.max(dim=-1)
+    width = values.shape[-1]
+    if width <= 0:
+        return best, index
+    return best, index.clamp(0, width - 1)
+
+
+class _ChunkedVocabStatistics(torch.autograd.Function):
+    """One projection pass for CE and exact Gumbel-max categorical sampling."""
+
+    @staticmethod
+    def forward(
+        ctx,
+        hidden: torch.Tensor,
+        weight: torch.Tensor,
+        labels: torch.LongTensor,
+        softcap: float,
+        temperature: float,
+        chunk_size: int,
+    ) -> tuple[
+        torch.Tensor,
+        torch.Tensor,
+        torch.LongTensor,
+        torch.Tensor,
+        torch.Tensor,
+        torch.LongTensor,
+        torch.Tensor,
+    ]:
+        if hidden.ndim != 2 or weight.ndim != 2:
+            raise ValueError("Chunked vocabulary tensors must be matrices.")
+        if hidden.shape[1] != weight.shape[1]:
+            raise ValueError("Hidden and vocabulary projection dimensions differ.")
+        if temperature <= 0 or chunk_size <= 0:
+            raise ValueError("Temperature and vocabulary chunk size must be positive.")
+        labels = labels.to(device=hidden.device, dtype=torch.long)
+        if labels.shape != (hidden.shape[0],):
+            raise ValueError("Labels must contain one value per hidden row.")
+
+        rows = hidden.shape[0]
+        row_indices = torch.arange(rows, device=hidden.device)
+        valid = labels.ge(0)
+        raw_log_z = torch.full(
+            (rows,), -torch.inf, device=hidden.device, dtype=torch.float32
+        )
+        sample_log_z = torch.full_like(raw_log_z, -torch.inf)
+        gold_score = torch.zeros_like(raw_log_z)
+        best_gumbel = torch.full_like(raw_log_z, -torch.inf)
+        selected_score = torch.zeros_like(raw_log_z)
+        selected = torch.zeros(rows, device=hidden.device, dtype=torch.long)
+        greedy_score = torch.full_like(raw_log_z, -torch.inf)
+        greedy = torch.zeros(rows, device=hidden.device, dtype=torch.long)
+        moment_max = torch.full_like(raw_log_z, -torch.inf)
+        moment_sum = torch.zeros_like(raw_log_z)
+        moment_weighted = torch.zeros_like(raw_log_z)
+        vocab_size = weight.shape[0]
+
+        with torch.no_grad():
+            for start in range(0, vocab_size, int(chunk_size)):
+                stop = min(start + int(chunk_size), vocab_size)
+                raw = F.linear(hidden, weight[start:stop])
+                scores = (
+                    torch.tanh(raw.float() / float(softcap))
+                    * float(softcap)
+                )
+                raw_log_z = torch.logaddexp(
+                    raw_log_z,
+                    torch.logsumexp(scores, dim=-1),
+                )
+                sample_scores = scores / float(temperature)
+                sample_log_z = torch.logaddexp(
+                    sample_log_z,
+                    torch.logsumexp(sample_scores, dim=-1),
+                )
+
+                in_chunk = valid & labels.ge(start) & labels.lt(stop)
+                local_gold = (labels - start).clamp(0, stop - start - 1)
+                gold_score = torch.where(
+                    in_chunk,
+                    scores[row_indices, local_gold],
+                    gold_score,
+                )
+
+                uniform = torch.rand(
+                    sample_scores.shape,
+                    device=sample_scores.device,
+                    dtype=torch.float32,
+                ).clamp_(
+                    min=torch.finfo(torch.float32).tiny,
+                    max=1.0 - torch.finfo(torch.float32).eps,
+                )
+                gumbel_scores = sample_scores - torch.log(-torch.log(uniform))
+                chunk_best, chunk_index = _stable_max(gumbel_scores)
+                replace_best = chunk_best.gt(best_gumbel)
+                candidate_score = sample_scores.gather(
+                    1, chunk_index[:, None]
+                ).squeeze(-1)
+                best_gumbel = torch.maximum(best_gumbel, chunk_best)
+                selected = torch.where(
+                    replace_best,
+                    chunk_index + start,
+                    selected,
+                )
+                selected_score = torch.where(
+                    replace_best,
+                    candidate_score,
+                    selected_score,
+                )
+
+                chunk_max, chunk_argmax = _stable_max(sample_scores)
+                replace_greedy = chunk_max.gt(greedy_score)
+                greedy_score = torch.maximum(greedy_score, chunk_max)
+                greedy = torch.where(
+                    replace_greedy,
+                    chunk_argmax + start,
+                    greedy,
+                )
+                shifted = torch.exp(sample_scores - chunk_max[:, None])
+                chunk_sum = shifted.sum(dim=-1)
+                chunk_weighted = (shifted * sample_scores).sum(dim=-1)
+                merged_max = torch.maximum(moment_max, chunk_max)
+                previous_scale = torch.exp(moment_max - merged_max)
+                chunk_scale = torch.exp(chunk_max - merged_max)
+                moment_sum = (
+                    moment_sum * previous_scale + chunk_sum * chunk_scale
+                )
+                moment_weighted = (
+                    moment_weighted * previous_scale
+                    + chunk_weighted * chunk_scale
+                )
+                moment_max = merged_max
+
+        raw_nll = torch.where(
+            valid,
+            raw_log_z - gold_score,
+            torch.zeros_like(raw_log_z),
+        )
+        temperature_nll = torch.where(
+            valid,
+            sample_log_z - gold_score / float(temperature),
+            torch.zeros_like(sample_log_z),
+        )
+        confidence = torch.exp(selected_score - sample_log_z).clamp_(0.0, 1.0)
+        greedy_confidence = torch.exp(greedy_score - sample_log_z).clamp_(0.0, 1.0)
+        entropy = sample_log_z - moment_weighted / moment_sum.clamp_min(
+            torch.finfo(torch.float32).tiny
+        )
+
+        ctx.save_for_backward(
+            hidden,
+            weight,
+            labels,
+            selected,
+            raw_log_z,
+            sample_log_z,
+            confidence,
+            entropy,
+        )
+        ctx.softcap = float(softcap)
+        ctx.temperature = float(temperature)
+        ctx.chunk_size = int(chunk_size)
+        ctx.mark_non_differentiable(selected, greedy, greedy_confidence)
+        return (
+            raw_nll,
+            temperature_nll,
+            selected,
+            confidence,
+            entropy,
+            greedy,
+            greedy_confidence,
+        )
+
+    @staticmethod
+    def backward(
+        ctx,
+        grad_raw_nll: torch.Tensor | None,
+        grad_temperature_nll: torch.Tensor | None,
+        grad_selected: torch.Tensor | None,
+        grad_confidence: torch.Tensor | None,
+        grad_entropy: torch.Tensor | None,
+        grad_greedy: torch.Tensor | None,
+        grad_greedy_confidence: torch.Tensor | None,
+    ):
+        del grad_selected, grad_greedy, grad_greedy_confidence
+        (
+            hidden,
+            weight,
+            labels,
+            selected,
+            raw_log_z,
+            sample_log_z,
+            confidence,
+            entropy,
+        ) = ctx.saved_tensors
+        grad_raw_nll = (
+            torch.zeros_like(raw_log_z)
+            if grad_raw_nll is None
+            else grad_raw_nll.float()
+        )
+        grad_temperature_nll = (
+            torch.zeros_like(sample_log_z)
+            if grad_temperature_nll is None
+            else grad_temperature_nll.float()
+        )
+        grad_confidence = (
+            torch.zeros_like(confidence)
+            if grad_confidence is None
+            else grad_confidence.float()
+        )
+        valid = labels.ge(0)
+        row_indices = torch.arange(hidden.shape[0], device=hidden.device)
+        grad_hidden = torch.zeros_like(hidden, dtype=torch.float32)
+        confidence_scale = (
+            grad_confidence * confidence / ctx.temperature
+        )
+        for start in range(0, weight.shape[0], ctx.chunk_size):
+            stop = min(start + ctx.chunk_size, weight.shape[0])
+            raw = F.linear(hidden, weight[start:stop])
+            scores = torch.tanh(raw.float() / ctx.softcap) * ctx.softcap
+            raw_probability = torch.exp(scores - raw_log_z[:, None])
+            sample_probability = torch.exp(
+                scores / ctx.temperature - sample_log_z[:, None]
+            )
+            score_gradient = (
+                raw_probability * (grad_raw_nll * valid)[:, None]
+                + sample_probability
+                * (grad_temperature_nll * valid)[:, None]
+                / ctx.temperature
+                - sample_probability * confidence_scale[:, None]
+            )
+            if grad_entropy is not None:
+                grad_entropy_f = grad_entropy.float()
+                sample_scores = scores / ctx.temperature
+                entropy_grad = -(
+                    grad_entropy_f / ctx.temperature
+                )[:, None] * sample_probability * (
+                    sample_scores - sample_log_z[:, None] + entropy[:, None]
+                )
+                score_gradient = score_gradient + entropy_grad
+
+            gold_in_chunk = valid & labels.ge(start) & labels.lt(stop)
+            gold_local = (labels - start).clamp(0, stop - start - 1)
+            score_gradient[row_indices, gold_local] -= (
+                grad_raw_nll * gold_in_chunk
+                + grad_temperature_nll
+                * gold_in_chunk
+                / ctx.temperature
+            )
+            selected_in_chunk = selected.ge(start) & selected.lt(stop)
+            selected_local = (selected - start).clamp(0, stop - start - 1)
+            score_gradient[row_indices, selected_local] += (
+                confidence_scale * selected_in_chunk
+            )
+            score_gradient *= 1.0 - (scores / ctx.softcap).square()
+            # The B16 shape feeds 4,096 rows through a 32,768-wide vocabulary
+            # chunk.  MPS BF16 GEMM can return NaNs for this backward-only
+            # projection even when every score gradient is finite.  Accumulate
+            # the exact VJP in FP32; casting the finished hidden gradient back
+            # to the model dtype happens only once below.
+            chunk_gradient = F.linear(
+                score_gradient,
+                weight[start:stop].float().transpose(0, 1),
+            )
+            grad_hidden.add_(chunk_gradient)
+
+        return (
+            grad_hidden.to(hidden.dtype),
+            None,
+            None,
+            None,
+            None,
+            None,
+        )
+
+
+def chunked_vocab_training_statistics(
+    hidden: torch.Tensor,
+    weight: torch.Tensor,
+    labels: torch.LongTensor,
+    *,
+    softcap: float,
+    temperature: float,
+    chunk_size: int,
+) -> tuple[
+    torch.Tensor,
+    torch.Tensor,
+    torch.LongTensor,
+    torch.Tensor,
+    torch.Tensor,
+    torch.LongTensor,
+    torch.Tensor,
+]:
+    """Return exact losses plus sampled and greedy proposal statistics."""
+
+    shape = labels.shape
+    outputs = _ChunkedVocabStatistics.apply(
+        hidden.reshape(-1, hidden.shape[-1]),
+        weight,
+        labels.reshape(-1),
+        float(softcap),
+        float(temperature),
+        int(chunk_size),
+    )
+    (
+        raw_nll,
+        temperature_nll,
+        proposal,
+        confidence,
+        entropy,
+        greedy,
+        greedy_confidence,
+    ) = outputs
+    return (
+        raw_nll.view(shape),
+        temperature_nll.view(shape),
+        proposal.view(shape),
+        confidence.view(shape),
+        entropy.view(shape),
+        greedy.view(shape),
+        greedy_confidence.view(shape),
+    )
+
+
+@torch.no_grad()
+def chunked_vocab_statistics(
+    hidden: torch.Tensor,
+    weight: torch.Tensor,
+    *,
+    softcap: float,
+    temperature: float,
+    chunk_size: int,
+    repetition_token_mask: torch.BoolTensor | None = None,
+    repetition_penalty: float = 1.0,
+    sampling_generators: Sequence[torch.Generator] | None = None,
+) -> tuple[
+    torch.LongTensor,
+    torch.Tensor,
+    torch.Tensor,
+    torch.LongTensor,
+    torch.Tensor,
+]:
+    """Return exact inference statistics, optionally penalizing seen tokens."""
+
+    if not math.isfinite(repetition_penalty) or repetition_penalty <= 0:
+        raise ValueError("`repetition_penalty` must be a finite positive number.")
+
+    # Preserve the original custom-autograd path exactly when disabled. In
+    # particular, this keeps the same chunk-local RNG draws and default output.
+    if repetition_penalty != 1.0 or sampling_generators is not None:
+        if hidden.ndim < 2 or weight.ndim != 2:
+            raise ValueError("Chunked vocabulary tensors must have matrix features.")
+        if hidden.shape[-1] != weight.shape[1]:
+            raise ValueError("Hidden and vocabulary projection dimensions differ.")
+        if temperature <= 0 or chunk_size <= 0:
+            raise ValueError("Temperature and vocabulary chunk size must be positive.")
+        expected_mask_shape = (hidden.shape[0], weight.shape[0])
+        if repetition_penalty != 1.0:
+            if repetition_token_mask is None or repetition_token_mask.shape != expected_mask_shape:
+                raise ValueError(
+                    "`repetition_token_mask` must have shape [batch, vocabulary]."
+                )
+            if repetition_token_mask.dtype != torch.bool:
+                raise ValueError("`repetition_token_mask` must be a boolean tensor.")
+        if sampling_generators is not None and len(sampling_generators) != hidden.shape[0]:
+            raise ValueError(
+                "`sampling_generators` must contain one generator per batch row."
+            )
+
+        output_shape = hidden.shape[:-1]
+        flat_hidden = hidden.reshape(-1, hidden.shape[-1])
+        rows_per_batch = math.prod(hidden.shape[1:-1])
+        batch_indices = torch.arange(
+            flat_hidden.shape[0], device=hidden.device
+        ).div(rows_per_batch, rounding_mode="floor")
+        sample_log_z = torch.full(
+            (flat_hidden.shape[0],),
+            -torch.inf,
+            device=hidden.device,
+            dtype=torch.float32,
+        )
+        best_gumbel = torch.full_like(sample_log_z, -torch.inf)
+        selected_score = torch.zeros_like(sample_log_z)
+        selected = torch.zeros(
+            flat_hidden.shape[0], device=hidden.device, dtype=torch.long
+        )
+        greedy_score = torch.full_like(sample_log_z, -torch.inf)
+        greedy = torch.zeros_like(selected)
+        moment_max = torch.full_like(sample_log_z, -torch.inf)
+        moment_sum = torch.zeros_like(sample_log_z)
+        moment_weighted = torch.zeros_like(sample_log_z)
+
+        for start in range(0, weight.shape[0], int(chunk_size)):
+            stop = min(start + int(chunk_size), weight.shape[0])
+            raw = F.linear(flat_hidden, weight[start:stop])
+            scores = torch.tanh(raw.float() / float(softcap)) * float(softcap)
+            if repetition_penalty != 1.0:
+                assert repetition_token_mask is not None
+                seen = repetition_token_mask[:, start:stop].index_select(
+                    0, batch_indices
+                )
+                penalized = torch.where(
+                    scores < 0,
+                    scores * float(repetition_penalty),
+                    scores / float(repetition_penalty),
+                )
+                scores = torch.where(seen, penalized, scores)
+            sample_scores = scores / float(temperature)
+            sample_log_z = torch.logaddexp(
+                sample_log_z, torch.logsumexp(sample_scores, dim=-1)
+            )
+
+            if sampling_generators is None:
+                uniform = torch.rand(
+                    sample_scores.shape,
+                    device=sample_scores.device,
+                    dtype=torch.float32,
+                )
+            else:
+                # Drawing each request from its own generator makes sampling
+                # invariant to dynamic admission, slot changes, and batch order.
+                per_request_shape = (
+                    rows_per_batch,
+                    sample_scores.shape[-1],
+                )
+                uniform = torch.cat(
+                    [
+                        torch.rand(
+                            per_request_shape,
+                            device=sample_scores.device,
+                            dtype=torch.float32,
+                            generator=generator,
+                        )
+                        for generator in sampling_generators
+                    ],
+                    dim=0,
+                )
+            uniform = uniform.clamp_(
+                min=torch.finfo(torch.float32).tiny,
+                max=1.0 - torch.finfo(torch.float32).eps,
+            )
+            gumbel_scores = sample_scores - torch.log(-torch.log(uniform))
+            chunk_best, chunk_index = _stable_max(gumbel_scores)
+            replace_best = chunk_best.gt(best_gumbel)
+            candidate_score = sample_scores.gather(
+                1, chunk_index[:, None]
+            ).squeeze(-1)
+            best_gumbel = torch.maximum(best_gumbel, chunk_best)
+            selected = torch.where(replace_best, chunk_index + start, selected)
+            selected_score = torch.where(
+                replace_best, candidate_score, selected_score
+            )
+
+            chunk_max, chunk_argmax = _stable_max(sample_scores)
+            replace_greedy = chunk_max.gt(greedy_score)
+            greedy_score = torch.maximum(greedy_score, chunk_max)
+            greedy = torch.where(replace_greedy, chunk_argmax + start, greedy)
+            shifted = torch.exp(sample_scores - chunk_max[:, None])
+            chunk_sum = shifted.sum(dim=-1)
+            chunk_weighted = (shifted * sample_scores).sum(dim=-1)
+            merged_max = torch.maximum(moment_max, chunk_max)
+            previous_scale = torch.exp(moment_max - merged_max)
+            chunk_scale = torch.exp(chunk_max - merged_max)
+            moment_sum = moment_sum * previous_scale + chunk_sum * chunk_scale
+            moment_weighted = (
+                moment_weighted * previous_scale + chunk_weighted * chunk_scale
+            )
+            moment_max = merged_max
+
+        confidence = torch.exp(selected_score - sample_log_z).clamp_(0.0, 1.0)
+        greedy_confidence = torch.exp(greedy_score - sample_log_z).clamp_(0.0, 1.0)
+        entropy = sample_log_z - moment_weighted / moment_sum.clamp_min(
+            torch.finfo(torch.float32).tiny
+        )
+        return (
+            selected.view(output_shape),
+            confidence.view(output_shape),
+            entropy.view(output_shape),
+            greedy.view(output_shape),
+            greedy_confidence.view(output_shape),
+        )
+
+    labels = torch.full(
+        hidden.shape[:-1],
+        -100,
+        device=hidden.device,
+        dtype=torch.long,
+    )
+    (
+        _,
+        _,
+        proposal,
+        confidence,
+        entropy,
+        greedy,
+        greedy_confidence,
+    ) = chunked_vocab_training_statistics(
+        hidden,
+        weight,
+        labels,
+        softcap=softcap,
+        temperature=temperature,
+        chunk_size=chunk_size,
+    )
+    return proposal, confidence, entropy, greedy, greedy_confidence
+
+
+__all__ = [
+    "chunked_vocab_statistics",
+    "chunked_vocab_training_statistics",
+]