Update modeling_neollm.py

modeling_neollm.py

@@ -12,6 +12,12 @@ factorised attention log-priors (Litman & Guo, 2026), optional StackMemory
 (Zhang et al., NeurIPS 2025), and optional HOLA-derived episodic memory for
 the FlashAttention-2 IHA seq-expand path (Cui, 2026).
 
+When enabled, Leviathan-JTok/JTok-M uses reference PyTorch operations by
+default. The runtime-only ``config.jtok_kernel_backend = "triton"`` opt-in
+dispatches the same module parameters through the compatible CCE JTok/JTok-M
+adapter. The default ``"torch"`` path is unchanged and the backend choice
+does not add parameters or buffers to the checkpoint.
+
 The FlashAttention training paths deliberately use dense right-padded
 attention with a shared dtype and compilation contract across IHA and non-IHA.
 See ``NeoLLMModel.forward`` for its padding/packing constraints.
@@ -56,6 +62,8 @@ References:
         Towards Flexible and Stable Transformer Attention." arXiv:2602.23057.
     Leviathan Generator: Batley & Saha (2026). "A Separable Architecture for
         Continuous Token Representation in Language Models." arXiv:2601.22040.
+    JTok: Yang et al. (2026). "JTok: On Token Embedding as another Axis of
+        Scaling Law via Joint Token Self-modulation." arXiv:2602.00800.
     KHRONOS: Batley & Saha (2025). "KHRONOS: a Kernel-Based Neural Architecture
         for Rapid, Resource-Efficient Scientific Computation." arXiv:2505.13315.
     Spelling Bee Embeddings: Rabe, Clymo & Dong (2026). "Spelling Bee
@@ -76,6 +84,7 @@ References:
 
 import inspect
 import math
+from dataclasses import dataclass
 from typing import Optional, Union, Tuple
 
 import torch
@@ -127,6 +136,7 @@ except (ImportError, ModuleNotFoundError, AttributeError):
 
 _LEV_KERNEL_AVAILABLE = False
 _leviathan_embedding_compiler_safe = None
+_leviathan_embedding_with_seed_compiler_safe = None
 _leviathan_supports = None
 try:
     from cut_cross_entropy.leviathan import (  # type: ignore[import-not-found]
@@ -136,12 +146,40 @@ try:
     )
 
     _LEV_KERNEL_AVAILABLE = bool(_LEV_HAS_KERNEL_FORWARD)
+    try:
+        from cut_cross_entropy.leviathan import (  # type: ignore[import-not-found]
+            leviathan_embedding_with_seed_compiler_safe as _leviathan_embedding_with_seed_compiler_safe,
+        )
+    except (ImportError, ModuleNotFoundError, AttributeError):
+        # An older CCE package may provide the legacy Leviathan kernel without
+        # the opt-in JTok geometry bridge. Keep the legacy path available.
+        pass
 except (ImportError, ModuleNotFoundError, AttributeError):
     # The model must remain usable with a plain CCE installation.  In that
     # environment the original reference LeviathanGenerator.forward below is
     # the complete embedding path.
     _LEV_HAS_KERNEL_FORWARD = False
 
+_LEV_GEOMETRY_KERNEL_AVAILABLE = (
+    _LEV_KERNEL_AVAILABLE
+    and _leviathan_embedding_with_seed_compiler_safe is not None
+)
+
+# JTok/JTok-M is a separate, strict opt-in path. Keep this import separate
+# from the legacy Leviathan import above: an older CCE installation may expose
+# the original Leviathan kernel but not the newer JTok adapter, and that must
+# not silently disable the legacy path.
+_JTOK_KERNEL_AVAILABLE = False
+_apply_neollm_jtok = None
+try:
+    from cut_cross_entropy.leviathan import (  # type: ignore[import-not-found]
+        apply_neollm_jtok as _apply_neollm_jtok,
+    )
+
+    _JTOK_KERNEL_AVAILABLE = True
+except (ImportError, ModuleNotFoundError, AttributeError):
+    pass
+
 
 def _leviathan_config_supported(config) -> bool:
     """Check LEV shapes without depending on NeoLLM's optional dtype field."""
@@ -185,6 +223,23 @@ from .configuration_neollm import NeoLLMConfig
 
 from transformers import AutoConfig, AutoModel, AutoModelForCausalLM
 
+
+@dataclass
+class NeoLLMBaseModelOutputWithPast(BaseModelOutputWithPast):
+    """Backbone output with compact, differentiable JTok-M statistics.
+
+    When JTok-M is active during training, ``jtokm_aux_stats`` contains one
+    ``(P, n, T)`` tuple per decoder layer.  ``P`` is the differentiable sum of
+    the router probabilities, ``n`` is the detached hard Top-K count, and
+    ``T`` is the number of valid routed tokens.  Keeping only these reductions
+    in the model output avoids retaining per-token router tensors until the
+    language-model loss is assembled.
+    """
+
+    jtokm_aux_stats: Optional[
+        Tuple[Tuple[torch.Tensor, torch.Tensor, torch.Tensor], ...]
+    ] = None
+
 logger = logging.get_logger(__name__)
 _NEOLLM_FA_USE_TOP_LEFT_MASK = flash_attn_supports_top_left_mask()
 
@@ -443,6 +498,26 @@ class LeviathanGenerator(nn.Module):
             torch.empty(self.num_modes, self.krank, self.hidden_size)
         )
 
+        # Shared continuous coordinate for Leviathan-JTok/JTok-M.  The
+        # original JTok paper stores a separate V×D table per layer; NeoLLM
+        # replaces that table with a compact surface over one Leviathan seed.
+        # These bridge parameters are present only when JTok is requested, so
+        # the default Leviathan checkpoint/state-dict contract is unchanged.
+        if bool(getattr(config, "use_jtok", False)):
+            self.jtok_seed_proj_weight = nn.Parameter(
+                torch.empty(self.d_seed, self.d_seed)
+            )
+            self.jtok_seed_proj_bias = nn.Parameter(torch.zeros(self.d_seed))
+            self.jtok_seed_norm_weight = nn.Parameter(torch.ones(self.d_seed))
+            self.jtok_seed_norm_bias = nn.Parameter(torch.zeros(self.d_seed))
+            self.jtok_seed_norm_eps = 1e-5
+        else:
+            self.register_parameter("jtok_seed_proj_weight", None)
+            self.register_parameter("jtok_seed_proj_bias", None)
+            self.register_parameter("jtok_seed_norm_weight", None)
+            self.register_parameter("jtok_seed_norm_bias", None)
+            self.jtok_seed_norm_eps = 1e-5
+
     def _base_k_decompose(self, token_ids: torch.Tensor) -> torch.Tensor:
         """
         Deterministic base-b decomposition: i -> (i_0, ..., i_{k-1}).
@@ -463,6 +538,48 @@ class LeviathanGenerator(nn.Module):
             ids = ids // self.b
         return coords
 
+    def _codebook_seed(self, token_ids: torch.Tensor) -> torch.Tensor:
+        """Return the differentiable compositional Leviathan seed ``z``."""
+        orig_shape = token_ids.shape
+        N = token_ids.numel()
+        coords = self._base_k_decompose(token_ids).reshape(N, self.k)
+        z = torch.zeros(
+            N,
+            self.d_seed,
+            device=token_ids.device,
+            dtype=self.codebooks.dtype,
+        )
+        for r in range(self.k):
+            z = z + self.codebooks[r][coords[:, r]]
+        return z.reshape(*orig_shape, self.d_seed)
+
+    def _jtok_coordinate_from_seed(self, z: torch.Tensor) -> torch.Tensor:
+        r"""Map a raw Leviathan seed to the shared bounded JTok coordinate.
+
+        The bridge is a single learned ``Linear → LayerNorm → sigmoid`` path
+        shared by all decoder layers.  It preserves the smooth compositional
+        geometry of Leviathan while giving every JTok surface the same
+        coordinate system.
+        """
+        if self.jtok_seed_proj_weight is None:
+            raise RuntimeError(
+                "JTok geometry requested but `use_jtok` was not enabled at "
+                "model construction time."
+            )
+        orig_shape = z.shape[:-1]
+        flat = z.reshape(-1, self.d_seed)
+        y = F.linear(
+            flat.to(self.jtok_seed_proj_weight.dtype),
+            self.jtok_seed_proj_weight,
+            self.jtok_seed_proj_bias,
+        ).float()
+        mean = y.mean(dim=-1, keepdim=True)
+        var = y.var(dim=-1, keepdim=True, unbiased=False)
+        y = (y - mean) * torch.rsqrt(var + self.jtok_seed_norm_eps)
+        y = y * self.jtok_seed_norm_weight.float() + self.jtok_seed_norm_bias.float()
+        y = torch.sigmoid(0.5 * y).clamp(0.0, 1.0)
+        return y.to(z.dtype).reshape(*orig_shape, self.d_seed)
+
     @staticmethod
     def _normalize_bspline_basis(B: torch.Tensor) -> torch.Tensor:
         """
@@ -533,7 +650,8 @@ class LeviathanGenerator(nn.Module):
         *,
         meap_mask_embedding: Optional[torch.Tensor] = None,
         meap_mask_token_id: Optional[int] = None,
-    ) -> torch.Tensor:
+        return_jtok_geometry: bool = False,
+    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
         """
         Generate embeddings from discrete token indices.
 
@@ -542,7 +660,10 @@ class LeviathanGenerator(nn.Module):
         Args:
             token_ids: (batch, seq_len) or (seq_len,).
         Returns:
-            embeddings [*token_ids.shape, hidden_size].
+            embeddings [*token_ids.shape, hidden_size].  When
+            ``return_jtok_geometry=True``, returns ``(embeddings, z_tilde)``
+            where ``z_tilde`` is the shared differentiable Leviathan/JTok
+            coordinate with shape ``[*token_ids.shape, d_seed]``.
         """
         has_meap = (
             meap_mask_embedding is not None or meap_mask_token_id is not None
@@ -562,7 +683,7 @@ class LeviathanGenerator(nn.Module):
                     f"({self.hidden_size},)"
                 )
 
-        if (
+        leviathan_kernel_ready = (
             self.use_leviathan_triton
             and _LEV_KERNEL_AVAILABLE
             and _leviathan_embedding_compiler_safe is not None
@@ -580,6 +701,37 @@ class LeviathanGenerator(nn.Module):
                     self.head_out_weight,
                 )
             )
+        )
+        require_leviathan_triton = bool(
+            getattr(self.config, "_require_leviathan_triton", False)
+        )
+        if require_leviathan_triton and not leviathan_kernel_ready:
+            raise RuntimeError(
+                "This run requires the Triton Leviathan kernel, but the "
+                "installed CCE package or the current geometry is not "
+                "supported. Refusing a reference fallback."
+            )
+        if (
+            return_jtok_geometry
+            and str(getattr(self.config, "jtok_kernel_backend", "torch"))
+            .strip()
+            .lower()
+            == "triton"
+            and not (leviathan_kernel_ready and _LEV_GEOMETRY_KERNEL_AVAILABLE)
+        ):
+            raise RuntimeError(
+                "JTok Triton requires a supported CUDA Leviathan kernel and "
+                "the CCE geometry adapter; refusing a silent reference fallback."
+            )
+
+        # Ordinary Leviathan calls keep their existing embedding-only custom
+        # op.  JTok additionally needs the raw compositional seed.  The CCE
+        # adapter keeps the legacy kernel for the embedding and exposes only a
+        # small differentiable base-b/codebook bridge for the geometry; this
+        # avoids both a full reference-Leviathan detour and a silent cut in the
+        # JTok -> codebook gradient.
+        if (
+            leviathan_kernel_ready
         ):
             params = {
                 "codebooks": self.codebooks,
@@ -590,6 +742,23 @@ class LeviathanGenerator(nn.Module):
                 "head_out_weight": self.head_out_weight,
             }
             try:
+                if return_jtok_geometry:
+                    if not _LEV_GEOMETRY_KERNEL_AVAILABLE:
+                        raise RuntimeError(
+                            "JTok geometry requested with the Leviathan kernel, "
+                            "but the installed CCE package lacks the geometry "
+                            "adapter."
+                        )
+                    embedding, seed = _leviathan_embedding_with_seed_compiler_safe(
+                        token_ids,
+                        params,
+                        self.config,
+                        self.knot_grid,
+                        mask_embedding=meap_mask_embedding,
+                        mask_token_id=meap_mask_token_id,
+                    )
+                    embedding = embedding.reshape(*token_ids.shape, self.hidden_size)
+                    return embedding, self._jtok_coordinate_from_seed(seed)
                 return _leviathan_embedding_compiler_safe(
                     token_ids,
                     params,
@@ -603,6 +772,8 @@ class LeviathanGenerator(nn.Module):
                 # implementation. A CUDA launch error must propagate: doing
                 # reference work after a failed launch can poison the next
                 # CUDA-graph/FlashAttention partition.
+                if require_leviathan_triton:
+                    raise
                 pass
 
         target_dtype = self.codebooks.dtype
@@ -610,11 +781,7 @@ class LeviathanGenerator(nn.Module):
         N = token_ids.numel()
 
         # -- Stage 1: compositional codebook indexing -------------------------
-        coords = self._base_k_decompose(token_ids)
-        coords_flat = coords.reshape(N, self.k)
-        z = torch.zeros(N, self.d_seed, device=token_ids.device, dtype=target_dtype)
-        for r in range(self.k):
-            z = z + self.codebooks[r][coords_flat[:, r]]
+        z = self._codebook_seed(token_ids).reshape(N, self.d_seed)
 
         # -- Stages 2-4: per-head projection, spline expansion, sum ----------
         e = torch.zeros(
@@ -655,11 +822,393 @@ class LeviathanGenerator(nn.Module):
                 meap_mask_embedding.to(device=e.device, dtype=e.dtype),
                 e,
             )
+        if return_jtok_geometry:
+            z_jtok = self._jtok_coordinate_from_seed(
+                z.reshape(*orig_shape, self.d_seed)
+            )
+            return e, z_jtok
         return e
 
 
 
 
+# ==================== LEVIATHAN-JTOK / JTOK-M MODULATION ====================
+
+
+class LeviathanJTok(nn.Module):
+    r"""Pure-PyTorch Leviathan-coupled implementation of JTok and JTok-M.
+
+    ``use_jtok=True, use_jtokm=False`` activates JTok.  Enabling
+    ``use_jtokm=True`` upgrades the same module to JTok-M; the configuration
+    validates that JTok and the Leviathan generator are active first.
+
+    The original JTok paper stores a token table per decoder layer.  NeoLLM
+    keeps the JTok modulation/routing equations but replaces every table row
+    with a continuous CP-separable B-spline surface over the shared
+    Leviathan-derived coordinate ``z_tilde``:
+
+    .. math::
+        \hat{\Delta m}_x^l = \Delta m_x^l \odot
+        (1+s^l\odot\operatorname{Norm}(S_l(\tilde z_x)))
+
+    JTok-M evaluates the selected Top-K surfaces and injects their normalized
+    mixture as ``Delta_m + Delta_r`` with the paper's
+    ``1/sqrt(2L)`` depth factor.  The default implementation is reference
+    PyTorch. Setting the runtime-only ``config.jtok_kernel_backend`` to
+    ``"triton"`` dispatches through the optional CCE JTok/JTok-M adapter while
+    preserving this module's parameters, state-dict layout, and decoder return
+    contract. The existing optional Leviathan embedding kernel remains a
+    separate path for ordinary embedding calls. The Triton choice is strict:
+    if the adapter is unavailable, forward raises instead of silently
+    changing the training or benchmark path back to Torch.
+
+    The temporary ``surface_chunk_tokens`` value bounds inference/evaluation
+    work tiles only.  It is not a context-length, batch-size, or VRAM limit.
+    """
+
+    def __init__(self, config: NeoLLMConfig, layer_idx: int):
+        super().__init__()
+        self.layer_idx = int(layer_idx)
+        self.use_mixture = bool(getattr(config, "use_jtokm", False))
+        self.jtok_kernel_backend = str(
+            getattr(config, "jtok_kernel_backend", "torch")
+        ).strip().lower()
+        if self.jtok_kernel_backend not in {"torch", "triton"}:
+            raise ValueError(
+                "`jtok_kernel_backend` must be either `torch` or `triton`."
+            )
+        if self.jtok_kernel_backend == "triton" and not _JTOK_KERNEL_AVAILABLE:
+            raise RuntimeError(
+                "`jtok_kernel_backend='triton'` was requested, but the "
+                "installed CCE package does not expose `apply_neollm_jtok`. "
+                "Install the compatible JTok/JTok-M CCE implementation or "
+                "set `jtok_kernel_backend='torch'`."
+            )
+        self.hidden_size = int(config.hidden_size)
+        self.d_seed = int(config.generator_d_seed)
+        self.num_knots = int(config.generator_num_knots)
+        self.num_modes = int(getattr(config, "jtok_num_modes", 4))
+        self.num_experts = (
+            int(getattr(config, "jtok_num_experts", 5))
+            if self.use_mixture
+            else 1
+        )
+        self.top_k = (
+            int(getattr(config, "jtok_top_k", 2))
+            if self.use_mixture
+            else 1
+        )
+        self.norm_eps = float(getattr(config, "jtok_norm_eps", 1e-6))
+        # Work-tile size only.  The optional attribute is intentionally read
+        # with getattr so older config.json files remain loadable.
+        self.surface_chunk_tokens = max(
+            int(getattr(config, "jtok_surface_chunk_tokens", 2048)), 1
+        )
+        self.jtokm_residual_scale = 1.0 / math.sqrt(
+            2.0 * float(config.num_hidden_layers)
+        )
+
+        self.register_buffer(
+            "knot_grid",
+            torch.linspace(0.0, 1.0, self.num_knots),
+            persistent=True,
+        )
+
+        # Continuous replacement for the paper's per-token tables.  These
+        # remain ordinary Parameters so both their surface and the shared
+        # Leviathan codebooks receive gradients during pretraining.
+        self.spline_coeff = nn.Parameter(
+            torch.empty(
+                self.num_experts,
+                self.num_modes,
+                self.d_seed,
+                self.num_knots,
+            )
+        )
+        self.W_out = nn.Parameter(
+            torch.empty(self.num_experts, self.num_modes, self.hidden_size)
+        )
+        self.W_res = nn.Parameter(
+            torch.empty(self.num_experts, self.d_seed, self.hidden_size)
+        )
+        self.scaler = nn.Parameter(torch.ones(self.hidden_size))
+        self.router = (
+            nn.Linear(self.hidden_size, self.num_experts, bias=False)
+            if self.use_mixture
+            else None
+        )
+
+    def _bspline_basis(self, z_flat: torch.Tensor) -> torch.Tensor:
+        """Evaluate the normalized quadratic B-spline basis in FP32."""
+        scale = float(self.num_knots - 1)
+        z32 = z_flat.float().unsqueeze(-1)
+        grid = self.knot_grid.float().view(1, 1, -1)
+        d = (z32 - grid).abs() * scale
+        B = torch.where(
+            d < 0.5,
+            0.75 - d.square(),
+            torch.where(
+                d < 1.5,
+                0.5 * (1.5 - d).square(),
+                torch.zeros_like(d),
+            ),
+        )
+        return B / B.sum(dim=-1, keepdim=True).clamp_min(1e-12)
+
+    def _eval_surfaces(
+        self,
+        z_flat: torch.Tensor,
+        target_dtype: torch.dtype,
+    ) -> torch.Tensor:
+        """Evaluate all continuous surfaces, returning ``[N, E, D]``."""
+        B = self._bspline_basis(z_flat)
+        phi = torch.einsum("emrg,nrg->nemr", self.spline_coeff.float(), B)
+        log_mag = torch.log(phi.abs() + 1e-9).sum(dim=-1)
+        num_neg = (phi < 0).to(torch.int32).sum(dim=-1)
+        prod_sign = 1.0 - 2.0 * (num_neg % 2).float()
+        modes = (prod_sign * torch.exp(log_mag)).to(target_dtype)
+        out_modes = torch.einsum(
+            "nem,emd->ned", modes, self.W_out.to(target_dtype)
+        )
+        out_res = torch.einsum(
+            "nr,erd->ned",
+            z_flat.to(target_dtype),
+            self.W_res.to(target_dtype),
+        )
+        return out_modes + out_res
+
+    def _eval_selected_surfaces(
+        self,
+        z_flat: torch.Tensor,
+        expert_idx: torch.Tensor,
+        target_dtype: torch.dtype,
+    ) -> torch.Tensor:
+        """Evaluate only the selected expert surfaces for each token."""
+        if expert_idx.ndim != 2:
+            raise ValueError("expert_idx must have shape [tokens, selected_experts]")
+        N, selected_k = expert_idx.shape
+        basis_all = self._bspline_basis(z_flat)
+        selected_flat = torch.zeros(
+            N * selected_k,
+            self.hidden_size,
+            dtype=target_dtype,
+            device=z_flat.device,
+        )
+        # Top-K indices are unique per token, so each flattened destination is
+        # written once.  This avoids materializing [N,K,M,D] or [N,K,R,D].
+        for expert_idx_value in range(self.num_experts):
+            token_idx, slot_idx = torch.where(expert_idx == expert_idx_value)
+            flat_idx = token_idx * selected_k + slot_idx
+            z_selected = z_flat[token_idx]
+            B = basis_all[token_idx]
+            coeff = self.spline_coeff[expert_idx_value].float()
+            phi = torch.einsum("sdg,mdg->smd", B, coeff)
+            log_mag = torch.log(phi.abs() + 1e-9).sum(dim=-1)
+            num_neg = (phi < 0).to(torch.int32).sum(dim=-1)
+            prod_sign = 1.0 - 2.0 * (num_neg % 2).float()
+            modes = prod_sign * torch.exp(log_mag)
+            values = modes.to(target_dtype) @ self.W_out[
+                expert_idx_value
+            ].to(target_dtype)
+            values = values + z_selected.to(target_dtype) @ self.W_res[
+                expert_idx_value
+            ].to(target_dtype)
+            selected_flat = selected_flat.index_copy(0, flat_idx, values)
+        return selected_flat.reshape(N, selected_k, self.hidden_size)
+
+    def _eval_selected_surfaces_tiled(
+        self,
+        z_flat: torch.Tensor,
+        expert_idx: torch.Tensor,
+        target_dtype: torch.dtype,
+    ) -> torch.Tensor:
+        """Route-first surface evaluation with bounded temporary token tiles."""
+        if z_flat.shape[0] != expert_idx.shape[0]:
+            raise ValueError("z_flat and expert_idx must contain the same token count")
+        chunks = []
+        for start in range(0, z_flat.shape[0], self.surface_chunk_tokens):
+            stop = min(start + self.surface_chunk_tokens, z_flat.shape[0])
+            chunks.append(
+                self._eval_selected_surfaces(
+                    z_flat[start:stop], expert_idx[start:stop], target_dtype
+                )
+            )
+        if not chunks:
+            return z_flat.new_empty((0, expert_idx.shape[1], self.hidden_size))
+        return torch.cat(chunks, dim=0)
+
+    def _eval_selected_surfaces_dense_tiled(
+        self,
+        z_flat: torch.Tensor,
+        expert_idx: torch.Tensor,
+        target_dtype: torch.dtype,
+    ) -> torch.Tensor:
+        """Compile-friendly selected evaluation over bounded dense tiles."""
+        if z_flat.shape[0] != expert_idx.shape[0]:
+            raise ValueError("z_flat and expert_idx must contain the same token count")
+        chunks = []
+        for start in range(0, z_flat.shape[0], self.surface_chunk_tokens):
+            stop = min(start + self.surface_chunk_tokens, z_flat.shape[0])
+            surfaces = self._eval_surfaces(z_flat[start:stop], target_dtype)
+            gather_idx = expert_idx[start:stop].unsqueeze(-1).expand(
+                stop - start, expert_idx.shape[1], self.hidden_size
+            )
+            chunks.append(surfaces.gather(1, gather_idx))
+        if not chunks:
+            return z_flat.new_empty((0, expert_idx.shape[1], self.hidden_size))
+        return torch.cat(chunks, dim=0)
+
+    def _l2_normalize(self, x: torch.Tensor) -> torch.Tensor:
+        return x / (x.norm(dim=-1, keepdim=True) + self.norm_eps)
+
+    def _valid_mask(
+        self,
+        valid_mask: Optional[torch.Tensor],
+        N: int,
+        device: torch.device,
+    ) -> torch.Tensor:
+        if valid_mask is None:
+            return torch.ones(N, device=device, dtype=torch.bool)
+        return valid_mask.reshape(N).to(device=device, dtype=torch.bool)
+
+    def _jtokm_aux_stats(
+        self,
+        logits: torch.Tensor,
+        topk_idx: torch.Tensor,
+        mask: torch.Tensor,
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Return differentiable probability and hard Top-K count reductions."""
+        aux_sigmoid = torch.sigmoid(logits.float())
+        # Normalize only the auxiliary probability branch.  The forward
+        # mixture remains selected-only Sigmoid+TopK, as in the source design.
+        prob_all = aux_sigmoid / aux_sigmoid.sum(
+            dim=-1, keepdim=True
+        ).clamp_min(self.norm_eps)
+        valid32 = mask.to(prob_all.dtype).unsqueeze(-1)
+        p_sum = (prob_all * valid32).sum(dim=0)
+        onehot = torch.zeros_like(prob_all).scatter(1, topk_idx, 1.0)
+        f_sum = (onehot * valid32).sum(dim=0).detach()
+        T = mask.sum().to(dtype=torch.float32).detach()
+        return p_sum, f_sum, T
+
+    def forward(
+        self,
+        delta_m: torch.Tensor,
+        z_tilde: torch.Tensor,
+        *,
+        router_state: Optional[torch.Tensor] = None,
+        valid_mask: Optional[torch.Tensor] = None,
+        compute_aux: bool = False,
+    ) -> Tuple[
+        torch.Tensor,
+        Optional[Tuple[torch.Tensor, torch.Tensor, torch.Tensor]],
+    ]:
+        """Apply JTok/JTok-M and optionally return training-only router stats."""
+        if self.jtok_kernel_backend == "triton":
+            # Dispatch only: this branch does not change compiler mode,
+            # CUDA-graph policy, training state, optimizer behavior, or any
+            # model-global state. The adapter returns the same
+            # ``(updated_delta, compact_router_stats)`` contract as the Torch
+            # implementation, so the surrounding pretraining graph remains
+            # unchanged when the backend is switched.
+            if _apply_neollm_jtok is None:
+                raise RuntimeError(
+                    "JTok Triton adapter became unavailable after module "
+                    "construction; refusing an implicit Torch fallback."
+                )
+            return _apply_neollm_jtok(
+                self,
+                delta_m,
+                z_tilde,
+                router_state=router_state,
+                valid_mask=valid_mask,
+                compute_aux=bool(compute_aux and self.use_mixture),
+                backend="triton",
+            )
+
+        orig_shape = delta_m.shape
+        N = delta_m.numel() // self.hidden_size
+        dm = delta_m.reshape(N, self.hidden_size)
+        z = z_tilde.reshape(N, self.d_seed)
+        mask = self._valid_mask(valid_mask, N, dm.device)
+        mask_f = mask.to(dm.dtype).unsqueeze(-1)
+
+        if not self.use_mixture:
+            surfaces = self._eval_surfaces(z, dm.dtype)
+            token_direction = self._l2_normalize(surfaces[:, 0, :])
+            gate = 1.0 + self.scaler.to(dm.dtype) * token_direction
+            gate = torch.where(mask.unsqueeze(-1), gate, torch.ones_like(gate))
+            return (dm * gate).reshape(orig_shape), None
+
+        if router_state is None:
+            raise ValueError("JTok-M requires a pre-attention `router_state`.")
+        h = router_state.reshape(N, self.hidden_size)
+        h = h * torch.rsqrt(h.square().mean(dim=-1, keepdim=True) + self.norm_eps)
+        logits = self.router(h)
+        topk_vals, topk_idx = torch.topk(logits, self.top_k, dim=-1)
+        selected_prob = torch.sigmoid(topk_vals)
+        weights = selected_prob / selected_prob.sum(
+            dim=-1, keepdim=True
+        ).clamp_min(self.norm_eps)
+
+        # Match the Torch reference semantics: keep the small dense path for
+        # gradient-enabled training and use route-first tiles for large eager
+        # evaluation.  Both routes implement the same surface equations.
+        if N <= self.surface_chunk_tokens or torch.is_grad_enabled():
+            surfaces = self._eval_surfaces(z, dm.dtype)
+            gather_idx = topk_idx.unsqueeze(-1).expand(
+                N, self.top_k, self.hidden_size
+            )
+            selected = surfaces.gather(1, gather_idx)
+        elif getattr(torch.compiler, "is_compiling", lambda: False)():
+            selected = self._eval_selected_surfaces_dense_tiled(
+                z, topk_idx, dm.dtype
+            )
+        else:
+            selected = self._eval_selected_surfaces_tiled(
+                z, topk_idx, dm.dtype
+            )
+
+        mixed = (weights.unsqueeze(-1) * selected).sum(dim=1)
+        delta_r = (
+            self.jtokm_residual_scale
+            * self.scaler.to(dm.dtype)
+            * self._l2_normalize(mixed)
+        )
+        delta_r = delta_r * mask_f
+        shared_update = dm + delta_r
+        aux_stats = self._jtokm_aux_stats(logits, topk_idx, mask) if compute_aux else None
+        return shared_update.reshape(orig_shape), aux_stats
+
+
+def compute_jtokm_aux_loss(
+    aux_stats_list: Tuple[Tuple[torch.Tensor, torch.Tensor, torch.Tensor], ...],
+    *,
+    n_e: int,
+    top_k: int,
+    weight: float,
+) -> torch.Tensor:
+    r"""Compute the JTok-M load-balancing objective from compact reductions.
+
+    For each layer, ``p_i=P_i/T`` and ``f_i=n_i/(T*K)``; the returned value is
+    the layer-average of ``weight * n_e * sum_i(p_i*f_i)``.  ``P_i`` remains
+    attached to the router graph while hard Top-K counts are detached.
+    """
+    if not aux_stats_list:
+        raise ValueError("JTok-M auxiliary loss requires non-empty router stats.")
+    layer_losses = []
+    for p_sum, f_sum, T in aux_stats_list:
+        denom_t = T.to(device=p_sum.device, dtype=torch.float32).clamp_min(1.0)
+        p_i = p_sum.float() / denom_t
+        f_i = f_sum.to(device=p_sum.device, dtype=torch.float32) / (
+            denom_t * float(top_k)
+        )
+        layer_losses.append(float(n_e) * (p_i * f_i).sum())
+    mean_loss = torch.stack(layer_losses).mean()
+    lambda_t = torch.as_tensor(weight, device=mean_loss.device, dtype=torch.float32)
+    return lambda_t * mean_loss
+
+
 # ==================== ORIGINAL COMPONENTS ====================
 
 
@@ -4676,6 +5225,23 @@ class NeoLLMDecoderLayer(GradientCheckpointingLayer):
         )
         self.self_attn = NeoLLMAttention(config, layer_idx)
         self.mlp = NeoLLMMLP(config)
+        # JTok/JTok-M operates on the shared MLP update.  It is instantiated
+        # only when the corresponding config flag is active; the default
+        # training graph therefore keeps the original residual topology.
+        self.use_jtok = bool(getattr(config, "use_jtok", False))
+        self.use_jtokm = bool(getattr(config, "use_jtokm", False))
+        if self.use_jtokm and not self.use_jtok:
+            # ``NeoLLMConfig`` validates this relationship, but decoder layers
+            # are also constructed directly by a few checkpoint conversion and
+            # test paths.  Fail explicitly there instead of silently dropping
+            # JTok-M while leaving the rest of the model apparently valid.
+            raise ValueError(
+                "`use_jtokm=True` requires `use_jtok=True`; enable the base "
+                "JTok path before enabling JTok-M."
+            )
+        self.jtok = (
+            LeviathanJTok(config, layer_idx) if self.use_jtok else None
+        )
         if self.use_siamesenorm:
             self.input_layernorm = None
             self.post_attention_layernorm = None
@@ -4971,6 +5537,9 @@ class NeoLLMDecoderLayer(GradientCheckpointingLayer):
         output_attentions: Optional[bool] = False,
         repo_rope_args: Optional[Tuple[torch.Tensor, float]] = None,
         position_ids: Optional[torch.LongTensor] = None,
+        jtok_z_tilde: Optional[torch.Tensor] = None,
+        jtok_valid_mask: Optional[torch.Tensor] = None,
+        jtok_compute_aux: bool = False,
         **kwargs: Unpack[FlashAttentionKwargs],
     ) -> Tuple:
         # SiameseNorm keeps two coupled streams with shared Attention/MLP
@@ -4990,6 +5559,7 @@ class NeoLLMDecoderLayer(GradientCheckpointingLayer):
             h_attn, stack_memory, stack_memory_mask = self.apply_stack_memory(
                 h_attn, stack_memory, stack_memory_mask
             )
+        jtok_router_state = h_attn
 
         h_lns = self.lns_attn(h_attn) if self.use_lns else h_attn
         attn_out, attn_weights, self.current_layer_fan = self.self_attn(
@@ -5020,7 +5590,21 @@ class NeoLLMDecoderLayer(GradientCheckpointingLayer):
         h_lns2 = self.lns_mlp(h_mlp) if self.use_lns else h_mlp
         delta_m = self.mlp(h_lns2)
 
-        shared_update = delta_m
+        jtok_aux_stats = None
+        if self.jtok is not None:
+            if jtok_z_tilde is None:
+                raise ValueError(
+                    "Active JTok/JTok-M requires Leviathan `jtok_z_tilde`."
+                )
+            shared_update, jtok_aux_stats = self.jtok(
+                delta_m,
+                jtok_z_tilde,
+                router_state=jtok_router_state if self.use_jtokm else None,
+                valid_mask=jtok_valid_mask,
+                compute_aux=bool(jtok_compute_aux and self.use_jtokm),
+            )
+        else:
+            shared_update = delta_m
 
         x_next = x_after_attn + stream_scale * shared_update
         y_next = y_after_attn + shared_update
@@ -5031,6 +5615,8 @@ class NeoLLMDecoderLayer(GradientCheckpointingLayer):
         outputs = (x_next, y_next)
         if self.use_stack_memory:
             outputs += (stack_memory, stack_memory_mask)
+        if jtok_aux_stats is not None:
+            outputs += jtok_aux_stats
         if output_attentions:
             outputs += (attn_weights,)
         return outputs
@@ -5046,6 +5632,9 @@ class NeoLLMDecoderLayer(GradientCheckpointingLayer):
         output_attentions: Optional[bool] = False,
         repo_rope_args: Optional[Tuple[torch.Tensor, float]] = None,
         position_ids: Optional[torch.LongTensor] = None,
+        jtok_z_tilde: Optional[torch.Tensor] = None,
+        jtok_valid_mask: Optional[torch.Tensor] = None,
+        jtok_compute_aux: bool = False,
         **kwargs: Unpack[FlashAttentionKwargs],
     ) -> Tuple:
         # ── Snapshot input ────────────────────────────────────────────────
@@ -5083,6 +5672,9 @@ class NeoLLMDecoderLayer(GradientCheckpointingLayer):
             h_normed = h_attn
         else:
             h_normed = _apply_norm(self.input_layernorm, h_attn)
+        # JTok-M follows the source design and routes from the fused
+        # pre-attention contextual state, before the attention output is added.
+        jtok_router_state = h_normed
 
         h_lns = self.lns_attn(h_normed) if self.use_lns else h_normed
         hidden_states, attn_weights, self.current_layer_fan = self.self_attn(
@@ -5131,6 +5723,20 @@ class NeoLLMDecoderLayer(GradientCheckpointingLayer):
         h_lns2 = self.lns_mlp(h_normed2) if self.use_lns else h_normed2
         delta_m = self.mlp(h_lns2)
 
+        jtok_aux_stats = None
+        if self.jtok is not None:
+            if jtok_z_tilde is None:
+                raise ValueError(
+                    "Active JTok/JTok-M requires Leviathan `jtok_z_tilde`."
+                )
+            delta_m, jtok_aux_stats = self.jtok(
+                delta_m,
+                jtok_z_tilde,
+                router_state=jtok_router_state if self.use_jtokm else None,
+                valid_mask=jtok_valid_mask,
+                compute_aux=bool(jtok_compute_aux and self.use_jtokm),
+            )
+
         # ── Residual junction (MLP sublayer) ────��────────────────────────
         # LAuReL augments the base MLP residual (residual_mlp + delta_m).
         if self.use_laurel:
@@ -5148,6 +5754,8 @@ class NeoLLMDecoderLayer(GradientCheckpointingLayer):
         hidden_states = self.gpas_mlp(mlp_aug) if self.use_gpas else mlp_aug
 
         outputs = (hidden_states,)
+        if jtok_aux_stats is not None:
+            outputs += jtok_aux_stats
         if output_attentions:
             outputs += (attn_weights,)
         return outputs
@@ -5658,6 +6266,37 @@ class NeoLLMPreTrainedModel(PreTrainedModel):
             # outputs starts with the same variance as a single projection.
             out_std = self.config.initializer_range / math.sqrt(module.num_modes)
             nn.init.normal_(module.head_out_weight, mean=0.0, std=out_std)
+            if module.jtok_seed_proj_weight is not None:
+                nn.init.normal_(
+                    module.jtok_seed_proj_weight,
+                    mean=0.0,
+                    std=self.config.initializer_range,
+                )
+                nn.init.zeros_(module.jtok_seed_proj_bias)
+                module.jtok_seed_norm_weight.data.fill_(1.0)
+                module.jtok_seed_norm_bias.data.zero_()
+
+        elif isinstance(module, LeviathanJTok):
+            # Start continuous JTok surfaces close to a neutral residual
+            # modulation while leaving all parameters trainable immediately.
+            nn.init.normal_(module.spline_coeff, mean=1.0, std=0.1)
+            nn.init.normal_(
+                module.W_out,
+                mean=0.0,
+                std=self.config.initializer_range,
+            )
+            nn.init.normal_(
+                module.W_res,
+                mean=0.0,
+                std=self.config.initializer_range,
+            )
+            module.scaler.data.fill_(1.0)
+            if module.router is not None:
+                nn.init.normal_(
+                    module.router.weight,
+                    mean=0.0,
+                    std=self.config.initializer_range,
+                )
 
         elif isinstance(module, StackMemory):
             nn.init.normal_(
@@ -5892,6 +6531,15 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
         if (input_ids is None) ^ (inputs_embeds is not None):
             raise ValueError("Specify exactly one of input_ids or inputs_embeds")
 
+        use_jtok = bool(getattr(self.config, "use_jtok", False))
+        use_jtokm = bool(getattr(self.config, "use_jtokm", False))
+        if use_jtok and input_ids is None:
+            raise ValueError(
+                "Active JTok/JTok-M requires `input_ids`; precomputed "
+                "`inputs_embeds` do not carry the Leviathan token coordinate."
+            )
+        jtok_z_tilde = None
+
         # ── Embedding stage ────────────────────────────────────────────────
         meap_mask_embedding = None
         meap_mask_token_id = None
@@ -5914,13 +6562,33 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
                 # selection so byte features cannot alter masked positions.
                 fuse_meap_in_leviathan = meap_active and self.spelling_bee is None
                 if fuse_meap_in_leviathan:
-                    inputs_embeds = self.token_generator(
-                        input_ids,
-                        meap_mask_embedding=meap_mask_embedding,
-                        meap_mask_token_id=meap_mask_token_id,
-                    )
+                    if use_jtok:
+                        generator_result = self.token_generator(
+                            input_ids,
+                            meap_mask_embedding=meap_mask_embedding,
+                            meap_mask_token_id=meap_mask_token_id,
+                            return_jtok_geometry=True,
+                        )
+                    else:
+                        # Keep the pre-JTok call signature for custom or
+                        # checkpoint-restored token generators.
+                        generator_result = self.token_generator(
+                            input_ids,
+                            meap_mask_embedding=meap_mask_embedding,
+                            meap_mask_token_id=meap_mask_token_id,
+                        )
+                else:
+                    if use_jtok:
+                        generator_result = self.token_generator(
+                            input_ids,
+                            return_jtok_geometry=True,
+                        )
+                    else:
+                        generator_result = self.token_generator(input_ids)
+                if use_jtok:
+                    inputs_embeds, jtok_z_tilde = generator_result
                 else:
-                    inputs_embeds = self.token_generator(input_ids)
+                    inputs_embeds = generator_result
                 meap_consumed = fuse_meap_in_leviathan
             else:
                 inputs_embeds = self.embed_tokens(input_ids)
@@ -5954,6 +6622,26 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
                 inputs_embeds,
             )
 
+        # The valid mask serves two separate purposes: JTok/JTok-M must not
+        # modulate padding or MEAP-masked positions, and JTok-M must exclude
+        # those positions from its load-balancing counts.  The original
+        # attention mask remains unchanged for positions/loss computation.
+        if use_jtok:
+            jtok_valid_mask = (
+                attention_mask.to(
+                    device=input_ids.device,
+                    dtype=torch.bool,
+                )
+                if attention_mask is not None
+                else torch.ones_like(input_ids, dtype=torch.bool)
+            )
+            if meap_active:
+                jtok_valid_mask = jtok_valid_mask & input_ids.ne(
+                    int(meap_mask_token_id)
+                )
+        else:
+            jtok_valid_mask = None
+
         if position_ids is None:
             if attention_mask is not None:
                 position_ids = attention_mask.to(dtype=torch.long).cumsum(-1) - 1
@@ -6018,6 +6706,10 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
         all_hidden_states = () if output_hidden_states else None
         all_attentions = () if output_attentions else None
         tweo_activations = () if output_tweo_activations else None
+        collect_jtokm_aux = bool(
+            use_jtokm and self.training and torch.is_grad_enabled()
+        )
+        all_jtokm_aux_stats = () if collect_jtokm_aux else None
 
         # ── StackMemory state ─────────────────────────────────────────────
         # Same source-level idea as StackTrans:
@@ -6130,6 +6822,9 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
                     output_attentions=output_attentions,
                     repo_rope_args=repo_rope_args,
                     position_ids=position_ids,
+                    jtok_z_tilde=jtok_z_tilde,
+                    jtok_valid_mask=jtok_valid_mask,
+                    jtok_compute_aux=collect_jtokm_aux,
                     **kwargs,
                 )
                 hidden_states = layer_outputs[0]
@@ -6151,6 +6846,9 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
                     output_attentions=output_attentions,
                     repo_rope_args=repo_rope_args,
                     position_ids=position_ids,
+                    jtok_z_tilde=jtok_z_tilde,
+                    jtok_valid_mask=jtok_valid_mask,
+                    jtok_compute_aux=collect_jtokm_aux,
                     **kwargs,
                 )
                 hidden_states = layer_outputs[0]
@@ -6160,6 +6858,16 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
             if use_attn_res:
                 attn_res_partial = hidden_states
 
+            if collect_jtokm_aux:
+                all_jtokm_aux_stats = all_jtokm_aux_stats + (
+                    (
+                        layer_outputs[extras_start],
+                        layer_outputs[extras_start + 1],
+                        layer_outputs[extras_start + 2],
+                    ),
+                )
+                extras_start += 3
+
             if output_attentions:
                 all_attentions = all_attentions + (layer_outputs[extras_start],)
                 extras_start += 1
@@ -6191,15 +6899,18 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
                 for v in [hidden_states, None, all_hidden_states, all_attentions]
                 if v is not None
             )
+            if all_jtokm_aux_stats is not None:
+                base_outputs = base_outputs + (all_jtokm_aux_stats,)
             if output_tweo_activations:
                 return base_outputs + (tweo_activations,)
             return base_outputs
 
-        outputs = BaseModelOutputWithPast(
+        outputs = NeoLLMBaseModelOutputWithPast(
             last_hidden_state=hidden_states,
             past_key_values=None,
             hidden_states=all_hidden_states,
             attentions=all_attentions,
+            jtokm_aux_stats=all_jtokm_aux_stats,
         )
         if output_tweo_activations:
             return outputs, tweo_activations
@@ -8031,6 +8742,7 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
         self._last_nextlat_loss = None
         self._last_nextlat_mse_loss = None
         self._last_nextlat_kl_loss = None
+        self._last_jtokm_aux_loss = None
         self._last_total_loss = None
 
         if config.use_token_generator:
@@ -8211,17 +8923,35 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
         )
 
         # Unpack: model returns BaseModelOutputWithPast, or a tuple when
-        # return_dict=False (with tweo_activations appended when enabled)
+        # return_dict=False (with compact JTok-M stats and/or tweo activations
+        # appended when enabled).
         tweo_activations = None
         hidden_states_tuple = None
+        jtokm_aux_stats = None
         if isinstance(model_out, tuple):
             outputs = model_out[0]
             if tweo_enabled and len(model_out) >= 2:
                 tweo_activations = model_out[-1]
             if isinstance(outputs, torch.Tensor):
                 hidden_states = outputs
-                if len(model_out) > 1 and isinstance(model_out[1], tuple):
-                    hidden_states_tuple = model_out[1]
+                if model_output_hidden_states:
+                    tuple_candidates = [
+                        candidate
+                        for candidate in model_out[1:]
+                        if isinstance(candidate, tuple)
+                        and (not candidate or not isinstance(candidate[0], tuple))
+                    ]
+                    if tuple_candidates:
+                        hidden_states_tuple = tuple_candidates[0]
+                for candidate in model_out[1:]:
+                    if (
+                        isinstance(candidate, tuple)
+                        and candidate
+                        and isinstance(candidate[0], tuple)
+                        and len(candidate[0]) == 3
+                    ):
+                        jtokm_aux_stats = candidate
+                        break
             elif isinstance(outputs, tuple):
                 hidden_states = outputs[0]
                 if len(outputs) > 2:
@@ -8229,10 +8959,12 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
             else:
                 hidden_states = outputs.last_hidden_state
                 hidden_states_tuple = outputs.hidden_states
+                jtokm_aux_stats = getattr(outputs, "jtokm_aux_stats", None)
         else:
             outputs = model_out
             hidden_states = outputs.last_hidden_state
             hidden_states_tuple = outputs.hidden_states
+            jtokm_aux_stats = getattr(outputs, "jtokm_aux_stats", None)
 
         loss = None
         ntp_loss = None
@@ -8255,6 +8987,7 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
         nextlat_loss = None
         nextlat_mse_loss = None
         nextlat_kl_loss = None
+        jtokm_aux_loss = None
         self._last_ntp_loss = None
         self._last_ntp_ce_unweighted = None
         self._last_mile_reweighting_delta = None
@@ -8280,6 +9013,7 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
         self._last_nextlat_loss = None
         self._last_nextlat_mse_loss = None
         self._last_nextlat_kl_loss = None
+        self._last_jtokm_aux_loss = None
         self._last_total_loss = None
         if labels is not None:
             if self.ntp_loss_backend == "liger":
@@ -8513,6 +9247,15 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
                 )
                 loss = loss + nextlat_loss
 
+            if bool(getattr(self.config, "use_jtokm", False)) and jtokm_aux_stats:
+                jtokm_aux_loss = compute_jtokm_aux_loss(
+                    jtokm_aux_stats,
+                    n_e=int(getattr(self.config, "jtok_num_experts", 5)),
+                    top_k=int(getattr(self.config, "jtok_top_k", 2)),
+                    weight=float(getattr(self.config, "jtok_aux_loss_weight", 1e-4)),
+                )
+                loss = loss + jtokm_aux_loss
+
             self._last_ntp_loss = ntp_loss.detach() if ntp_loss is not None else None
             self._last_tweo_loss = tweo_loss.detach() if tweo_loss is not None else None
             self._last_nitp_loss = nitp_loss.detach() if nitp_loss is not None else None
@@ -8608,6 +9351,9 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
             self._last_nextlat_kl_loss = (
                 nextlat_kl_loss.detach() if nextlat_kl_loss is not None else None
             )
+            self._last_jtokm_aux_loss = (
+                jtokm_aux_loss.detach() if jtokm_aux_loss is not None else None
+            )
             self._last_total_loss = loss.detach() if loss is not None else None
             logits = None
         else:
@@ -8635,6 +9381,9 @@ __all__ = [
     "NeoLLMPreTrainedModel",
     "NeoLLMConfig",
     "LeviathanGenerator",
+    "LeviathanJTok",
+    "NeoLLMBaseModelOutputWithPast",
+    "compute_jtokm_aux_loss",
     "SpellingBeeEmbedding",
     "StackMemory",
     "FANLayer",