Update architecture and tokenizer

README.md

@@ -1,103 +1,104 @@
----
-language:
-- en
-license: mit
-library_name: transformers
-pipeline_tag: text-generation
-tags:
-- pytorch
-- research
-- sparse-attention
-- mixture-of-experts
----
-
-# SHRAM — Sparse Hybrid Token Routed Attention Mixture
-
-A research baseline implementing the SHRAM architecture from "An Examination of Sparse
-Attention for Long Context Purposes." No pretrained weights — pull the architecture from
-the Hub and instantiate a freshly initialised model from config. Every parameter is
-overridable at instantiation time via kwargs.
-
-> **Important:** `trust_remote_code=True` is required. It downloads the architecture
-> source files from the Hub and imports them into your Python process. Review the
-> source at [smithblack-0/SHRAM](https://huggingface.co/smithblack-0/SHRAM) before use.
-
-## Architecture
-
-SHRAM replaces every standard attention layer with a hybrid layer `H(x) = h_l(x) + h_s(x)`:
-
-- **h_l** — local sliding-window causal attention path.
-- **h_s** — MoSRAH sparse routed path. Each token selects K of L available expert heads
-  via token-choice routing. Bottlenecked Ensemble Attention (BEA) is applied per head.
-
-All other components follow the Llama 3 baseline (RMSNorm, SwiGLU FFN, RoPE).
-
-## Usage
-
-This repository contains no pretrained weights. The intended workflow is: pull the
-architecture config from the Hub, instantiate a model with fresh random weights, then
-train it yourself.
-
-```python
-from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
-
-# Step 1: pull the architecture config from the Hub.
-# AutoConfig.from_pretrained downloads config.json only — no weights are loaded.
-# Override any parameter via kwargs.
-config = AutoConfig.from_pretrained(
-    "smithblack-0/SHRAM",
-    trust_remote_code=True,
-    num_hidden_layers=16,       # example override
-    num_mosrah_heads=32,        # example override
-)
-
-# Step 2: instantiate with fresh random weights.
-# from_config never loads a checkpoint — it always produces a randomly initialised model.
-model = AutoModelForCausalLM.from_config(config, trust_remote_code=True)
-
-# Step 3: load the tokenizer.
-tokenizer = AutoTokenizer.from_pretrained("smithblack-0/SHRAM")
-```
-
-After training your own checkpoint, save and reload it in the standard way:
-
-```python
-model.save_pretrained("./my-checkpoint")
-model = AutoModelForCausalLM.from_pretrained("./my-checkpoint", trust_remote_code=True)
-```
-
-## Constructor Defaults
-
-The values below are the defaults you get if you call `AutoConfig.from_pretrained` with
-no overrides. They are not the parameters of a pretrained model — this repository
-contains no weights. All values are overridable via kwargs.
-
-| Parameter | Default |
-|-----------|---------|
-| `alpha` | 1.0 |
-| `attention_dropout` | 0.0 |
-| `beta` | 32.0 |
-| `dtype` | None |
-| `head_dim` | 16 |
-| `hidden_size` | 512 |
-| `inference_sequence_length` | 1024 |
-| `intermediate_size` | 1366 |
-| `local_rope_theta` | 10000.0 |
-| `mosrah_rope_theta` | 10000.0 |
-| `num_hidden_layers` | 12 |
-| `num_mosrah_heads` | 16 |
-| `num_selected_heads` | 16 |
-| `num_sliding_window_heads` | 16 |
-| `output_hidden_states` | False |
-| `rms_norm_eps` | 1e-05 |
-| `rope_mode` | main_sequence |
-| `tie_word_embeddings` | False |
-| `training_sequence_length` | 1024 |
-| `use_cache` | True |
-| `vocab_size` | 50277 |
-| `window_size` | 128 |
-
-## License
-
-MIT. Clean-room synthesis informed by the reference paper. Tokenizer is GPT-NeoX
-(`EleutherAI/gpt-neox-20b`, Apache 2.0).
+---
+language:
+- en
+license: mit
+library_name: transformers
+pipeline_tag: text-generation
+tags:
+- pytorch
+- research
+- sparse-attention
+- mixture-of-experts
+---
+
+# SHRAM — Sparse Hybrid Token Routed Attention Mixture
+
+A research baseline implementing the SHRAM architecture from "An Examination of Sparse
+Attention for Long Context Purposes." No pretrained weights — pull the architecture from
+the Hub and instantiate a freshly initialised model from config. Every parameter is
+overridable at instantiation time via kwargs.
+
+> **Important:** `trust_remote_code=True` is required. It downloads the architecture
+> source files from the Hub and imports them into your Python process. Review the
+> source at [smithblack-0/SHRAM](https://huggingface.co/smithblack-0/SHRAM) before use. Those interested can also
+> clone the git repository at https://github.com/smithblack-0/advanced-transformers-lib
+
+## Architecture
+
+SHRAM replaces every standard attention layer with a hybrid layer `H(x) = h_l(x) + h_s(x)`:
+
+- **h_l** — local sliding-window causal attention path.
+- **h_s** — MoSRAH sparse routed path. Each token selects K of L available expert heads
+  via token-choice routing. Bottlenecked Ensemble Attention (BEA) is applied per head.
+
+All other components follow the Llama 3 baseline (RMSNorm, SwiGLU FFN, RoPE).
+
+## Usage
+
+This repository contains no pretrained weights. The intended workflow is: pull the
+architecture config from the Hub, instantiate a model with fresh random weights, then
+train it yourself.
+
+```python
+from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
+
+# Step 1: pull the architecture config from the Hub.
+# AutoConfig.from_pretrained downloads config.json only — no weights are loaded.
+# Override any parameter via kwargs.
+config = AutoConfig.from_pretrained(
+    "smithblack-0/SHRAM",
+    trust_remote_code=True,
+    num_hidden_layers=16,       # example override
+    num_mosrah_heads=32,        # example override
+)
+
+# Step 2: instantiate with fresh random weights.
+# from_config never loads a checkpoint — it always produces a randomly initialised model.
+model = AutoModelForCausalLM.from_config(config, trust_remote_code=True)
+
+# Step 3: load the tokenizer.
+tokenizer = AutoTokenizer.from_pretrained("smithblack-0/SHRAM")
+```
+
+After training your own checkpoint, save and reload it in the standard way:
+
+```python
+model.save_pretrained("./my-checkpoint")
+model = AutoModelForCausalLM.from_pretrained("./my-checkpoint", trust_remote_code=True)
+```
+
+## Constructor Defaults
+
+The values below are the defaults you get if you call `AutoConfig.from_pretrained` with
+no overrides. They are not the parameters of a pretrained model — this repository
+contains no weights. All values are overridable via kwargs.
+
+| Parameter | Default |
+|-----------|---------|
+| `alpha` | 1.0 |
+| `attention_dropout` | 0.0 |
+| `beta` | 32.0 |
+| `dtype` | None |
+| `head_dim` | 16 |
+| `hidden_size` | 512 |
+| `inference_sequence_length` | 1024 |
+| `intermediate_size` | 1366 |
+| `local_rope_theta` | 10000.0 |
+| `mosrah_rope_theta` | 10000.0 |
+| `num_hidden_layers` | 12 |
+| `num_mosrah_heads` | 16 |
+| `num_selected_heads` | 16 |
+| `num_sliding_window_heads` | 16 |
+| `output_hidden_states` | False |
+| `rms_norm_eps` | 1e-05 |
+| `rope_mode` | main_sequence |
+| `tie_word_embeddings` | False |
+| `training_sequence_length` | 1024 |
+| `use_cache` | True |
+| `vocab_size` | 50277 |
+| `window_size` | 128 |
+
+## License
+
+MIT. Clean-room synthesis informed by the reference paper. Tokenizer is GPT-NeoX
+(`EleutherAI/gpt-neox-20b`, Apache 2.0).

__attention__bottlenecked_ensemble_attention.py

@@ -1,252 +1,252 @@
-"""Bottlenecked Ensemble Attention (BEA) for the MoSRAH sparse path.
-
-BEA is the packed expert-choice attention operator over the MoSRAH sparse path.
-It consumes packed expert-choice tensors, a supplied position tensor, an active
-token mask, and an optional layer-local MoSRAH cache. It returns outputs in the
-same packed expert-choice space expected by later unpacking.
-
-BEA does not compute positions and does not choose packed-position semantics.
-Those are supplied by the caller. If caching is used, BEA stores post-RoPE keys
-(K̃) and raw values (V) into the sparse cache and attends against the
-accumulated cached state returned by that cache.
-"""
-
-import math
-
-import torch
-import torch.nn as nn
-from torch.nn.attention.flex_attention import create_block_mask, flex_attention
-
-from .configuration import ShramConfig
-from .__cache__mosrah_cache import MoSRAHCache
-from .rope import RotaryEmbedding
-
-
-class BottleneckedEnsembleAttention(nn.Module):
-    """
-    Packed expert-choice attention operator for the MoSRAH sparse path.
-    Operates per-head independently on an ensemble of tokens.
-    FlexAttention saves flops on dead tokens.
-
-    Architectural properties:
-      - consumes packed expert-choice tensors of shape (B, L, T, d)
-      - uses independent per-head Q/K/V/O projection parameters
-      - applies YaRN-capable RoPE using supplied position_ids
-      - stores post-RoPE K̃ and raw V in MoSRAHCache when caching is enabled
-      - uses a fast fused attention path
-      - returns outputs in the same packed expert-choice space (B, L, T, d)
-
-    Args:
-        config: SHRAM config. Must expose `hidden_size`, `num_mosrah_heads`,
-            `head_dim`, `mosrah_rope_theta`, `training_sequence_length`,
-            `inference_sequence_length`, `alpha`, and `beta`.
-    """
-
-    def __init__(self, config: ShramConfig) -> None:
-        super().__init__()
-
-        self.hidden_size = config.hidden_size
-        self.num_heads = config.num_mosrah_heads
-        self.head_dim = config.head_dim
-
-        # Independent per-head projections. No cross-head parameter sharing.
-        self.q_proj = nn.Parameter(
-            torch.empty(self.num_heads, self.hidden_size, self.head_dim)
-        )
-        self.k_proj = nn.Parameter(
-            torch.empty(self.num_heads, self.hidden_size, self.head_dim)
-        )
-        self.v_proj = nn.Parameter(
-            torch.empty(self.num_heads, self.hidden_size, self.head_dim)
-        )
-        self.o_proj = nn.Parameter(
-            torch.empty(self.num_heads, self.head_dim, self.hidden_size)
-        )
-
-        self._reset_parameters()
-
-        # BEA uses the YaRN-capable RoPE path. The caller supplies the position tensor;
-        # this unit only consumes it. In training modes, dilation will be 1.0 and so
-        # no yarn dilation occurs.
-        self.rope = RotaryEmbedding(
-            mode="yarn",
-            head_dim=self.head_dim,
-            theta=config.mosrah_rope_theta,
-            initial_seq_length=config.training_sequence_length,
-            dilation=config.scale,
-            alpha=config.alpha,
-            beta=config.beta,
-        )
-
-    def forward(
-        self,
-        packed_embeddings: torch.Tensor,
-        position_ids: torch.Tensor,
-        active_mask: torch.Tensor,
-        cache: MoSRAHCache | None = None,
-    ) -> torch.Tensor:
-        """Apply BEA to packed expert-choice tensors.
-
-        Args:
-            packed_embeddings: Packed expert-choice hidden states of shape (B, L, T, d).
-            position_ids: Supplied packed positions of shape (B, L, T).
-            active_mask: Boolean active-token mask of shape (B, L, T).
-            cache: Optional layer-local MoSRAH cache.
-
-        Returns:
-            Packed expert-choice output tensor of shape (B, L, T, d).
-        """
-        batch_size, _, query_length, _ = packed_embeddings.shape
-        self._validate_tensor_shape(packed_embeddings)
-        self._validate_position_shape(packed_embeddings, position_ids)
-        self._validate_active_mask_shape(packed_embeddings, active_mask)
-
-        # Independent per-head projections:
-        # (B, L, T, d) x (L, d, u) -> (B, L, T, u)
-        query_states = torch.einsum("bltd,ldu->bltu", packed_embeddings, self.q_proj)
-        key_states = torch.einsum("bltd,ldu->bltu", packed_embeddings, self.k_proj)
-        value_states = torch.einsum("bltd,ldu->bltu", packed_embeddings, self.v_proj)
-
-        rotated_query_states, rotated_key_states, attention_scaling = self.rope(
-            query_states,
-            key_states,
-            position_ids,
-        )
-
-        if cache is not None:
-            # In cached execution, the current query tensor uses local tensor rows
-            # 0..Q-1, but the key tensor returned by the cache is the full accumulated
-            # packed sequence for each (batch, head) slot. The only additional data
-            # needed to align those two views is the pre-update cached prefix length.
-            # which will indicate how many queries were processed before now.
-            num_tokens_processed = cache.get_heads_lengths().clone()
-            key_states, value_states, key_active_mask = cache.update(
-                rotated_key_states,
-                value_states,
-                active_mask,
-            )
-        else:
-            num_tokens_processed = torch.zeros(
-                batch_size,
-                self.num_heads,
-                dtype=torch.long,
-                device=packed_embeddings.device,
-            )
-            key_states = rotated_key_states
-            key_active_mask = active_mask
-
-        block_mask = self._make_block_mask(
-             query_active_mask=active_mask,
-             key_active_mask=key_active_mask,
-             num_tokens_processed=num_tokens_processed,
-             query_length=query_length,
-             key_length=key_states.shape[2],
-             device=packed_embeddings.device,
-         )
-        attended_states = flex_attention(
-             rotated_query_states,
-             key_states,
-             value_states,
-             block_mask=block_mask,
-             scale=attention_scaling / math.sqrt(self.head_dim),
-        )
-
-        # Project back to model width:
-        # (B, L, T, u) x (L, u, d) -> (B, L, T, d)
-        return torch.einsum("bltu,lud->bltd", attended_states, self.o_proj)
-
-    def _reset_parameters(self) -> None:
-        """Initialize per-head projection weights."""
-        for weight in (self.q_proj, self.k_proj, self.v_proj, self.o_proj):
-            nn.init.xavier_uniform_(weight)
-
-    def _validate_tensor_shape(self, packed_embeddings: torch.Tensor) -> None:
-        """Validate the local packed-embedding shape contract required by BEA."""
-        if packed_embeddings.shape[1] != self.num_heads:
-            raise ValueError(
-                f"Expected packed_embeddings.shape[1] == num_mosrah_heads={self.num_heads}, "
-                f"got {packed_embeddings.shape[1]}."
-            )
-
-        if packed_embeddings.shape[-1] != self.hidden_size:
-            raise ValueError(
-                f"Expected packed_embeddings last dim == hidden_size={self.hidden_size}, "
-                f"got {packed_embeddings.shape[-1]}."
-            )
-
-    def _validate_position_shape(
-        self,
-        packed_embeddings: torch.Tensor,
-        position_ids: torch.Tensor,
-    ) -> None:
-        """Validate the supplied packed-position tensor shape."""
-        if position_ids.shape != packed_embeddings.shape[:3]:
-            raise ValueError(
-                f"position_ids must have shape {tuple(packed_embeddings.shape[:3])}, "
-                f"got {tuple(position_ids.shape)}."
-            )
-
-    def _validate_active_mask_shape(
-        self,
-        packed_embeddings: torch.Tensor,
-        active_mask: torch.Tensor,
-    ) -> None:
-        """Validate the supplied active-token mask shape."""
-        if active_mask.shape != packed_embeddings.shape[:3]:
-            raise ValueError(
-                f"active_mask must have shape {tuple(packed_embeddings.shape[:3])}, "
-                f"got {tuple(active_mask.shape)}."
-            )
-
-    def _make_block_mask(
-        self,
-        query_active_mask: torch.Tensor,
-        key_active_mask: torch.Tensor,
-        num_tokens_processed: torch.Tensor,
-        query_length: int,
-        key_length: int,
-        device: torch.device,
-    ):
-        """Create the packed-sequence causal mask for FlexAttention.
-
-        At the root, causality is still triangular. The only nuance is cached
-        execution: query rows are indexed locally as 0..Q-1 inside the current
-        query tensor, but the key tensor may already contain a cached prefix for
-        that (batch, head) slot. The causal horizon for query tensor row q is
-        therefore:
-
-            cached_prefix_lengths[b, h] + q
-
-        Query and key activity masks are then composed with that triangular rule
-        so FlexAttention can skip padded query rows and ignore inactive key slots.
-        """
-        batch_size, num_heads, _ = query_active_mask.shape
-
-        # Build the per-(batch, head, query_row) triangular horizon from a simple
-        # arange over query rows plus the cached prefix lengths for each slot.
-        relative_query_positions = torch.arange(
-            query_length,
-            device=device,
-            dtype=torch.long,
-        ).view(1, 1, query_length)
-        causal_query_positions = num_tokens_processed.unsqueeze(-1) + relative_query_positions
-
-        def packed_causal_mask(
-            batch_idx: torch.Tensor,
-            head_idx: torch.Tensor,
-            query_idx: torch.Tensor,
-            key_idx: torch.Tensor,
-        ) -> torch.Tensor:
-            query_is_active = query_active_mask[batch_idx, head_idx, query_idx]
-            key_is_active = key_active_mask[batch_idx, head_idx, key_idx]
-            is_causal = key_idx <= causal_query_positions[batch_idx, head_idx, query_idx]
-            return query_is_active & key_is_active & is_causal
-
-        return create_block_mask(
-            packed_causal_mask,
-            B=batch_size,
-            H=num_heads,
-            Q_LEN=query_length,
-            KV_LEN=key_length,
-            device=device,
+"""Bottlenecked Ensemble Attention (BEA) for the MoSRAH sparse path.
+
+BEA is the packed expert-choice attention operator over the MoSRAH sparse path.
+It consumes packed expert-choice tensors, a supplied position tensor, an active
+token mask, and an optional layer-local MoSRAH cache. It returns outputs in the
+same packed expert-choice space expected by later unpacking.
+
+BEA does not compute positions and does not choose packed-position semantics.
+Those are supplied by the caller. If caching is used, BEA stores post-RoPE keys
+(K̃) and raw values (V) into the sparse cache and attends against the
+accumulated cached state returned by that cache.
+"""
+
+import math
+
+import torch
+import torch.nn as nn
+from torch.nn.attention.flex_attention import create_block_mask, flex_attention
+
+from .configuration import ShramConfig
+from .__cache__mosrah_cache import MoSRAHCache
+from .rope import RotaryEmbedding
+
+
+class BottleneckedEnsembleAttention(nn.Module):
+    """
+    Packed expert-choice attention operator for the MoSRAH sparse path.
+    Operates per-head independently on an ensemble of tokens.
+    FlexAttention saves flops on dead tokens.
+
+    Architectural properties:
+      - consumes packed expert-choice tensors of shape (B, L, T, d)
+      - uses independent per-head Q/K/V/O projection parameters
+      - applies YaRN-capable RoPE using supplied position_ids
+      - stores post-RoPE K̃ and raw V in MoSRAHCache when caching is enabled
+      - uses a fast fused attention path
+      - returns outputs in the same packed expert-choice space (B, L, T, d)
+
+    Args:
+        config: SHRAM config. Must expose `hidden_size`, `num_mosrah_heads`,
+            `head_dim`, `mosrah_rope_theta`, `training_sequence_length`,
+            `inference_sequence_length`, `alpha`, and `beta`.
+    """
+
+    def __init__(self, config: ShramConfig) -> None:
+        super().__init__()
+
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_mosrah_heads
+        self.head_dim = config.head_dim
+
+        # Independent per-head projections. No cross-head parameter sharing.
+        self.q_proj = nn.Parameter(
+            torch.empty(self.num_heads, self.hidden_size, self.head_dim)
+        )
+        self.k_proj = nn.Parameter(
+            torch.empty(self.num_heads, self.hidden_size, self.head_dim)
+        )
+        self.v_proj = nn.Parameter(
+            torch.empty(self.num_heads, self.hidden_size, self.head_dim)
+        )
+        self.o_proj = nn.Parameter(
+            torch.empty(self.num_heads, self.head_dim, self.hidden_size)
+        )
+
+        self._reset_parameters()
+
+        # BEA uses the YaRN-capable RoPE path. The caller supplies the position tensor;
+        # this unit only consumes it. In training modes, dilation will be 1.0 and so
+        # no yarn dilation occurs.
+        self.rope = RotaryEmbedding(
+            mode="yarn",
+            head_dim=self.head_dim,
+            theta=config.mosrah_rope_theta,
+            initial_seq_length=config.training_sequence_length,
+            dilation=config.scale,
+            alpha=config.alpha,
+            beta=config.beta,
+        )
+
+    def forward(
+        self,
+        packed_embeddings: torch.Tensor,
+        position_ids: torch.Tensor,
+        active_mask: torch.Tensor,
+        cache: MoSRAHCache | None = None,
+    ) -> torch.Tensor:
+        """Apply BEA to packed expert-choice tensors.
+
+        Args:
+            packed_embeddings: Packed expert-choice hidden states of shape (B, L, T, d).
+            position_ids: Supplied packed positions of shape (B, L, T).
+            active_mask: Boolean active-token mask of shape (B, L, T).
+            cache: Optional layer-local MoSRAH cache.
+
+        Returns:
+            Packed expert-choice output tensor of shape (B, L, T, d).
+        """
+        batch_size, _, query_length, _ = packed_embeddings.shape
+        self._validate_tensor_shape(packed_embeddings)
+        self._validate_position_shape(packed_embeddings, position_ids)
+        self._validate_active_mask_shape(packed_embeddings, active_mask)
+
+        # Independent per-head projections:
+        # (B, L, T, d) x (L, d, u) -> (B, L, T, u)
+        query_states = torch.einsum("bltd,ldu->bltu", packed_embeddings, self.q_proj)
+        key_states = torch.einsum("bltd,ldu->bltu", packed_embeddings, self.k_proj)
+        value_states = torch.einsum("bltd,ldu->bltu", packed_embeddings, self.v_proj)
+
+        rotated_query_states, rotated_key_states, attention_scaling = self.rope(
+            query_states,
+            key_states,
+            position_ids,
+        )
+
+        if cache is not None:
+            # In cached execution, the current query tensor uses local tensor rows
+            # 0..Q-1, but the key tensor returned by the cache is the full accumulated
+            # packed sequence for each (batch, head) slot. The only additional data
+            # needed to align those two views is the pre-update cached prefix length.
+            # which will indicate how many queries were processed before now.
+            num_tokens_processed = cache.get_heads_lengths().clone()
+            key_states, value_states, key_active_mask = cache.update(
+                rotated_key_states,
+                value_states,
+                active_mask,
+            )
+        else:
+            num_tokens_processed = torch.zeros(
+                batch_size,
+                self.num_heads,
+                dtype=torch.long,
+                device=packed_embeddings.device,
+            )
+            key_states = rotated_key_states
+            key_active_mask = active_mask
+
+        block_mask = self._make_block_mask(
+             query_active_mask=active_mask,
+             key_active_mask=key_active_mask,
+             num_tokens_processed=num_tokens_processed,
+             query_length=query_length,
+             key_length=key_states.shape[2],
+             device=packed_embeddings.device,
+         )
+        attended_states = flex_attention(
+             rotated_query_states,
+             key_states,
+             value_states,
+             block_mask=block_mask,
+             scale=attention_scaling / math.sqrt(self.head_dim),
+        )
+
+        # Project back to model width:
+        # (B, L, T, u) x (L, u, d) -> (B, L, T, d)
+        return torch.einsum("bltu,lud->bltd", attended_states, self.o_proj)
+
+    def _reset_parameters(self) -> None:
+        """Initialize per-head projection weights."""
+        for weight in (self.q_proj, self.k_proj, self.v_proj, self.o_proj):
+            nn.init.xavier_uniform_(weight)
+
+    def _validate_tensor_shape(self, packed_embeddings: torch.Tensor) -> None:
+        """Validate the local packed-embedding shape contract required by BEA."""
+        if packed_embeddings.shape[1] != self.num_heads:
+            raise ValueError(
+                f"Expected packed_embeddings.shape[1] == num_mosrah_heads={self.num_heads}, "
+                f"got {packed_embeddings.shape[1]}."
+            )
+
+        if packed_embeddings.shape[-1] != self.hidden_size:
+            raise ValueError(
+                f"Expected packed_embeddings last dim == hidden_size={self.hidden_size}, "
+                f"got {packed_embeddings.shape[-1]}."
+            )
+
+    def _validate_position_shape(
+        self,
+        packed_embeddings: torch.Tensor,
+        position_ids: torch.Tensor,
+    ) -> None:
+        """Validate the supplied packed-position tensor shape."""
+        if position_ids.shape != packed_embeddings.shape[:3]:
+            raise ValueError(
+                f"position_ids must have shape {tuple(packed_embeddings.shape[:3])}, "
+                f"got {tuple(position_ids.shape)}."
+            )
+
+    def _validate_active_mask_shape(
+        self,
+        packed_embeddings: torch.Tensor,
+        active_mask: torch.Tensor,
+    ) -> None:
+        """Validate the supplied active-token mask shape."""
+        if active_mask.shape != packed_embeddings.shape[:3]:
+            raise ValueError(
+                f"active_mask must have shape {tuple(packed_embeddings.shape[:3])}, "
+                f"got {tuple(active_mask.shape)}."
+            )
+
+    def _make_block_mask(
+        self,
+        query_active_mask: torch.Tensor,
+        key_active_mask: torch.Tensor,
+        num_tokens_processed: torch.Tensor,
+        query_length: int,
+        key_length: int,
+        device: torch.device,
+    ):
+        """Create the packed-sequence causal mask for FlexAttention.
+
+        At the root, causality is still triangular. The only nuance is cached
+        execution: query rows are indexed locally as 0..Q-1 inside the current
+        query tensor, but the key tensor may already contain a cached prefix for
+        that (batch, head) slot. The causal horizon for query tensor row q is
+        therefore:
+
+            cached_prefix_lengths[b, h] + q
+
+        Query and key activity masks are then composed with that triangular rule
+        so FlexAttention can skip padded query rows and ignore inactive key slots.
+        """
+        batch_size, num_heads, _ = query_active_mask.shape
+
+        # Build the per-(batch, head, query_row) triangular horizon from a simple
+        # arange over query rows plus the cached prefix lengths for each slot.
+        relative_query_positions = torch.arange(
+            query_length,
+            device=device,
+            dtype=torch.long,
+        ).view(1, 1, query_length)
+        causal_query_positions = num_tokens_processed.unsqueeze(-1) + relative_query_positions
+
+        def packed_causal_mask(
+            batch_idx: torch.Tensor,
+            head_idx: torch.Tensor,
+            query_idx: torch.Tensor,
+            key_idx: torch.Tensor,
+        ) -> torch.Tensor:
+            query_is_active = query_active_mask[batch_idx, head_idx, query_idx]
+            key_is_active = key_active_mask[batch_idx, head_idx, key_idx]
+            is_causal = key_idx <= causal_query_positions[batch_idx, head_idx, query_idx]
+            return query_is_active & key_is_active & is_causal
+
+        return create_block_mask(
+            packed_causal_mask,
+            B=batch_size,
+            H=num_heads,
+            Q_LEN=query_length,
+            KV_LEN=key_length,
+            device=device,
         )

__attention__expert_packing.py

@@ -1,335 +1,335 @@
-"""Expert packing and unpacking for the MoSRAH path.
-
-This module implements the low-level token-choice -> expert-choice -> token-choice
-conversion boundary specified in the paper. The externally visible behavior is fixed:
-
-- setup_packing() prepares the auxiliary ordering data.
-- pack_experts() converts routed token-choice state into packed expert-choice state.
-- unpack_experts() restores token-choice ordering afterward.
-
-Stable sort is a correctness requirement. It preserves causal ordering inside each
-expert bucket, which is the foundation on which BEA's later triangular causal mask
-is correct.
-
-pack_experts() returns two distinct masks that serve different roles and must not
-be interchanged:
-
-- unpacking_mask: marks every packed slot that contains a routed token copy,
-  live or dead. Always has exactly B*N*K True entries. Required by unpack_experts
-  so its reshape invariant holds regardless of outer token liveness.
-- active_mask: marks only the packed slots whose source token was semantically
-  live. This is what BEA consumes for attention gating. Dead outer tokens must
-  not influence sparse attention outputs.
-"""
-
-import torch
-
-
-# ---------------------------------------------------------------------------
-# Setup
-# ---------------------------------------------------------------------------
-
-def setup_packing(
-    selected_heads: torch.Tensor,
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    """Prepare the auxiliary ordering data used by pack/unpack.
-
-    Routing produces token-choice state I of shape (B, N, K): for each token, which
-    K experts were selected. Packing needs the same routed token copies reordered into
-    expert-major order so each expert bucket becomes contiguous.
-
-    The paper's setup step does this by flattening (N, K) into one axis to produce
-    H in token-major order, then computing a stable argsort permutation Pi over the
-    expert indices stored in H. Applying Pi reorders the flattened routed copies into
-    expert-major order while preserving their original token order *within* each expert
-    bucket. That preservation is why stable sort is required for causality.
-
-    Args:
-        selected_heads: Routed token-choice head selections I of shape (B, N, K).
-
-    Returns:
-        Tuple of:
-          - flattened_selected_heads: H of shape (B, N*K)
-          - permutation: stable expert-major permutation Pi of shape (B, N*K)
-          - inverse_permutation: inverse permutation Pi^{-1} of shape (B, N*K)
-    """
-    batch_size, sequence_length, num_selected_heads = selected_heads.shape
-    flattened_selected_heads = selected_heads.reshape(
-        batch_size,
-        sequence_length * num_selected_heads,
-    )
-
-    permutation = torch.argsort(flattened_selected_heads, dim=-1, stable=True)
-    inverse_permutation = torch.argsort(permutation, dim=-1)
-
-    return flattened_selected_heads, permutation, inverse_permutation
-
-
-# ---------------------------------------------------------------------------
-# Packing
-# ---------------------------------------------------------------------------
-
-def pack_experts(
-    hidden_states: torch.Tensor,
-    position_ids: torch.Tensor,
-    selected_heads: torch.Tensor,
-    num_experts: int,
-    flattened_selected_heads: torch.Tensor,
-    permutation: torch.Tensor,
-    outer_active_mask: torch.Tensor,
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-    """Pack token-choice hidden states into expert-choice padded form.
-
-    The paper's packing path has two jobs:
-
-    1. Convert routed token-choice copies into expert-major order.
-    2. Materialize that expert-major order into a padded tensor layout BEA can consume.
-
-    The routed hidden-state copies are not stored explicitly in token-choice form.
-    Instead, the same token hidden state is conceptually copied once per selected expert.
-    The packing step reconstructs those copies by expanding local source-token indices,
-    reordering those indices with Pi, then gathering hidden states, positions, and outer
-    liveness in that packed order. All three are carried through the same expert-major
-    rearrangement so they remain aligned in the packed frame.
-
-    Packed positions are sourced from the authoritative upstream position_ids tensor
-    rather than synthesized locally from arange(N). This preserves advanced positions
-    correctly during cached inference while leaving training/full-sequence behavior
-    unchanged when position_ids is the ordinary sequential token positions.
-
-    Args:
-        hidden_states: Token-choice hidden states x of shape (B, N, d).
-        position_ids: Authoritative upstream token positions J of shape (B, N).
-        selected_heads: Routed head selections I of shape (B, N, K).
-        num_experts: Total number of experts L.
-        flattened_selected_heads: H from setup_packing(), shape (B, N*K).
-        permutation: Pi from setup_packing(), shape (B, N*K).
-        outer_active_mask: Current-chunk active mask of shape (B, N), where True
-            means the token is semantically live. Dead tokens do not become
-            semantically active in the packed sparse representation.
-
-    Returns:
-        Tuple of:
-          - packed_hidden_states: x' of shape (B, L, T, d)
-          - packed_positions: J' of shape (B, L, T)
-          - unpacking_mask: of shape (B, L, T). True where a slot contains any
-            routed token copy, live or dead. Always has exactly B*N*K True entries.
-            Pass this to unpack_experts — not active_mask.
-          - active_mask: of shape (B, L, T). True only where a slot contains a
-            copy of a live outer token. Pass this to BEA for attention gating.
-    """
-    batch_size, sequence_length, hidden_dim = hidden_states.shape
-    _, _, num_selected_heads = selected_heads.shape
-
-    # -----------------------------------------------------------------------
-    # Reconstruct routed local source-token indices in token-choice order.
-    #
-    # The internal arange(N) is no longer the packed position tensor. It is only
-    # the local source-row index object used to gather from the current chunk
-    # tensor x. Flattening this object gives a (B, N*K) tensor aligned with H's
-    # token-major routed-copy order.
-    # -----------------------------------------------------------------------
-    source_token_indices = torch.arange(
-        sequence_length,
-        device=hidden_states.device,
-        dtype=torch.long,
-    ).view(1, sequence_length, 1).expand(
-        batch_size,
-        sequence_length,
-        num_selected_heads,
-    )
-    flattened_source_indices = source_token_indices.reshape(
-        batch_size,
-        sequence_length * num_selected_heads,
-    )
-
-    # -----------------------------------------------------------------------
-    # Reorder source-token indices into expert-major order.
-    #
-    # Applying Pi yields the local source-token rows in the packed expert-major
-    # order required by the paper. Those same reordered source indices are then
-    # used to gather hidden states, authoritative upstream positions, and outer
-    # liveness so all three remain aligned under the exact same packing
-    # transformation.
-    # -----------------------------------------------------------------------
-    sorted_source_indices = flattened_source_indices.gather(
-        dim=1,
-        index=permutation,
-    )
-    sorted_hidden_states = hidden_states.gather(
-        dim=1,
-        index=sorted_source_indices.unsqueeze(-1).expand(-1, -1, hidden_dim),
-    )
-    sorted_positions = position_ids.gather(
-        dim=1,
-        index=sorted_source_indices,
-    )
-    sorted_active_mask = outer_active_mask.gather(
-        dim=1,
-        index=sorted_source_indices,
-    )
-
-    # -----------------------------------------------------------------------
-    # Count how many routed copies land in each expert bucket.
-    #
-    # S[b, l] is the number of routed token copies assigned to expert l in batch b.
-    # T is the maximum such count across all batches and experts; it determines the
-    # padded expert-length dimension of the packed representation.
-    # -----------------------------------------------------------------------
-    tokens_per_expert = _bincount_rows(
-        values=flattened_selected_heads,
-        num_bins=num_experts,
-    )
-    max_tokens_per_expert = int(tokens_per_expert.max().item())
-
-    # -----------------------------------------------------------------------
-    # Construct the active-token mask M.
-    #
-    # Each expert bucket is left-justified: if S[b, l] = s, then slots
-    # t = 0, ..., s-1 are active and all later slots are padding. The resulting
-    # mask therefore both identifies real packed tokens and enforces left-justified
-    # packing. This is the unpacking_mask — it marks slot occupancy regardless of
-    # outer token liveness, and always has exactly B*N*K True entries.
-    # -----------------------------------------------------------------------
-    time_axis = torch.arange(
-        max_tokens_per_expert,
-        device=hidden_states.device,
-        dtype=torch.long,
-    ).view(1, 1, max_tokens_per_expert)
-    unpacking_mask = time_axis < tokens_per_expert.unsqueeze(-1)
-
-    # -----------------------------------------------------------------------
-    # Materialize the padded packed tensors.
-    #
-    # The packed hidden states x', packed original-token positions J', and packed
-    # active-token mask are allocated as zero-filled tensors. Active entries are
-    # then written into those buffers in the expert-major order established above.
-    # Padding remains zero / inactive.
-    # -----------------------------------------------------------------------
-    packed_hidden_states = hidden_states.new_zeros(
-        batch_size,
-        num_experts,
-        max_tokens_per_expert,
-        hidden_dim,
-    )
-    packed_positions = position_ids.new_zeros(
-        batch_size,
-        num_experts,
-        max_tokens_per_expert,
-    )
-    active_mask = torch.zeros(
-        batch_size,
-        num_experts,
-        max_tokens_per_expert,
-        dtype=torch.bool,
-        device=hidden_states.device,
-    )
-
-    packed_hidden_states[unpacking_mask] = sorted_hidden_states.reshape(-1, hidden_dim)
-    packed_positions[unpacking_mask] = sorted_positions.reshape(-1)
-    active_mask[unpacking_mask] = sorted_active_mask.reshape(-1)
-
-    return packed_hidden_states, packed_positions, unpacking_mask, active_mask
-
-
-# ---------------------------------------------------------------------------
-# Unpacking
-# ---------------------------------------------------------------------------
-
-def unpack_experts(
-    expert_outputs: torch.Tensor,
-    selected_heads: torch.Tensor,
-    unpacking_mask: torch.Tensor,
-    inverse_permutation: torch.Tensor,
-) -> torch.Tensor:
-    """Restore token-choice ordering from BEA expert-choice output.
-
-    Unpacking inverts the packing path only on occupied entries. Padding does not
-    participate: the output tensor is first filtered by unpacking_mask to recover
-    only the real routed-token copies in expert-major order, then Pi^{-1} restores
-    the original token-choice ordering, and finally the tensor is reshaped back to
-    (B, N, K, d).
-
-    The unpacking_mask — not active_mask — must be used here. Even copies of dead
-    outer tokens occupy slots and must be un-scattered correctly for the inverse
-    permutation to hold. The total True entry count in unpacking_mask is always
-    B*N*K, which is exactly what the reshape to (B, N*K, d) requires.
-
-    Args:
-        expert_outputs: Expert-choice BEA output y of shape (B, L, T, d).
-        selected_heads: Routed head selections I of shape (B, N, K).
-        unpacking_mask: From pack_experts(), shape (B, L, T). Identifies all
-            occupied packed slots regardless of outer token liveness.
-        inverse_permutation: Pi^{-1} from setup_packing(), shape (B, N*K).
-
-    Returns:
-        Restored token-choice tensor y_tilde of shape (B, N, K, d).
-    """
-    batch_size, sequence_length, num_selected_heads = selected_heads.shape
-    hidden_dim = expert_outputs.shape[-1]
-
-    active_outputs = expert_outputs[unpacking_mask]
-    sorted_token_choice_outputs = active_outputs.reshape(
-        batch_size,
-        sequence_length * num_selected_heads,
-        hidden_dim,
-    )
-    restored_outputs = sorted_token_choice_outputs.gather(
-        dim=1,
-        index=inverse_permutation.unsqueeze(-1).expand(-1, -1, hidden_dim),
-    )
-
-    return restored_outputs.reshape(
-        batch_size,
-        sequence_length,
-        num_selected_heads,
-        hidden_dim,
-    )
-
-
-# ---------------------------------------------------------------------------
-# Helpers
-# ---------------------------------------------------------------------------
-
-def _bincount_rows(
-    values: torch.Tensor,
-    num_bins: int,
-) -> torch.Tensor:
-    """Count per-row integer occurrences for a 2D tensor.
-
-    torch.bincount operates on a flat 1D vector, but the packing algorithm needs
-    one bincount per batch row. The trick used here is to shift each row into its
-    own disjoint bin range before flattening:
-
-      row 0 uses bins [0, ..., num_bins - 1]
-      row 1 uses bins [num_bins, ..., 2*num_bins - 1]
-      row 2 uses bins [2*num_bins, ..., 3*num_bins - 1]
-      ...
-
-    After that shift, one global torch.bincount produces all row-local counts at
-    once. Reshaping the result back to (B, num_bins) recovers the per-row bincount.
-
-    This is a vectorized implementation detail only; externally visible behavior
-    remains exactly the paper's S tensor of per-batch per-expert token counts.
-
-    Args:
-        values: Integer tensor of shape (B, M) with entries in [0, num_bins).
-        num_bins: Number of bins.
-
-    Returns:
-        Counts tensor of shape (B, num_bins).
-    """
-    batch_size = values.shape[0]
-
-    row_offsets = torch.arange(
-        batch_size,
-        device=values.device,
-        dtype=values.dtype,
-    ).unsqueeze(1) * num_bins
-    shifted_values = values + row_offsets
-
-    counts = torch.bincount(
-        shifted_values.reshape(-1),
-        minlength=batch_size * num_bins,
-    )
-    return counts.reshape(batch_size, num_bins)
+"""Expert packing and unpacking for the MoSRAH path.
+
+This module implements the low-level token-choice -> expert-choice -> token-choice
+conversion boundary specified in the paper. The externally visible behavior is fixed:
+
+- setup_packing() prepares the auxiliary ordering data.
+- pack_experts() converts routed token-choice state into packed expert-choice state.
+- unpack_experts() restores token-choice ordering afterward.
+
+Stable sort is a correctness requirement. It preserves causal ordering inside each
+expert bucket, which is the foundation on which BEA's later triangular causal mask
+is correct.
+
+pack_experts() returns two distinct masks that serve different roles and must not
+be interchanged:
+
+- unpacking_mask: marks every packed slot that contains a routed token copy,
+  live or dead. Always has exactly B*N*K True entries. Required by unpack_experts
+  so its reshape invariant holds regardless of outer token liveness.
+- active_mask: marks only the packed slots whose source token was semantically
+  live. This is what BEA consumes for attention gating. Dead outer tokens must
+  not influence sparse attention outputs.
+"""
+
+import torch
+
+
+# ---------------------------------------------------------------------------
+# Setup
+# ---------------------------------------------------------------------------
+
+def setup_packing(
+    selected_heads: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Prepare the auxiliary ordering data used by pack/unpack.
+
+    Routing produces token-choice state I of shape (B, N, K): for each token, which
+    K experts were selected. Packing needs the same routed token copies reordered into
+    expert-major order so each expert bucket becomes contiguous.
+
+    The paper's setup step does this by flattening (N, K) into one axis to produce
+    H in token-major order, then computing a stable argsort permutation Pi over the
+    expert indices stored in H. Applying Pi reorders the flattened routed copies into
+    expert-major order while preserving their original token order *within* each expert
+    bucket. That preservation is why stable sort is required for causality.
+
+    Args:
+        selected_heads: Routed token-choice head selections I of shape (B, N, K).
+
+    Returns:
+        Tuple of:
+          - flattened_selected_heads: H of shape (B, N*K)
+          - permutation: stable expert-major permutation Pi of shape (B, N*K)
+          - inverse_permutation: inverse permutation Pi^{-1} of shape (B, N*K)
+    """
+    batch_size, sequence_length, num_selected_heads = selected_heads.shape
+    flattened_selected_heads = selected_heads.reshape(
+        batch_size,
+        sequence_length * num_selected_heads,
+    )
+
+    permutation = torch.argsort(flattened_selected_heads, dim=-1, stable=True)
+    inverse_permutation = torch.argsort(permutation, dim=-1)
+
+    return flattened_selected_heads, permutation, inverse_permutation
+
+
+# ---------------------------------------------------------------------------
+# Packing
+# ---------------------------------------------------------------------------
+
+def pack_experts(
+    hidden_states: torch.Tensor,
+    position_ids: torch.Tensor,
+    selected_heads: torch.Tensor,
+    num_experts: int,
+    flattened_selected_heads: torch.Tensor,
+    permutation: torch.Tensor,
+    outer_active_mask: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Pack token-choice hidden states into expert-choice padded form.
+
+    The paper's packing path has two jobs:
+
+    1. Convert routed token-choice copies into expert-major order.
+    2. Materialize that expert-major order into a padded tensor layout BEA can consume.
+
+    The routed hidden-state copies are not stored explicitly in token-choice form.
+    Instead, the same token hidden state is conceptually copied once per selected expert.
+    The packing step reconstructs those copies by expanding local source-token indices,
+    reordering those indices with Pi, then gathering hidden states, positions, and outer
+    liveness in that packed order. All three are carried through the same expert-major
+    rearrangement so they remain aligned in the packed frame.
+
+    Packed positions are sourced from the authoritative upstream position_ids tensor
+    rather than synthesized locally from arange(N). This preserves advanced positions
+    correctly during cached inference while leaving training/full-sequence behavior
+    unchanged when position_ids is the ordinary sequential token positions.
+
+    Args:
+        hidden_states: Token-choice hidden states x of shape (B, N, d).
+        position_ids: Authoritative upstream token positions J of shape (B, N).
+        selected_heads: Routed head selections I of shape (B, N, K).
+        num_experts: Total number of experts L.
+        flattened_selected_heads: H from setup_packing(), shape (B, N*K).
+        permutation: Pi from setup_packing(), shape (B, N*K).
+        outer_active_mask: Current-chunk active mask of shape (B, N), where True
+            means the token is semantically live. Dead tokens do not become
+            semantically active in the packed sparse representation.
+
+    Returns:
+        Tuple of:
+          - packed_hidden_states: x' of shape (B, L, T, d)
+          - packed_positions: J' of shape (B, L, T)
+          - unpacking_mask: of shape (B, L, T). True where a slot contains any
+            routed token copy, live or dead. Always has exactly B*N*K True entries.
+            Pass this to unpack_experts — not active_mask.
+          - active_mask: of shape (B, L, T). True only where a slot contains a
+            copy of a live outer token. Pass this to BEA for attention gating.
+    """
+    batch_size, sequence_length, hidden_dim = hidden_states.shape
+    _, _, num_selected_heads = selected_heads.shape
+
+    # -----------------------------------------------------------------------
+    # Reconstruct routed local source-token indices in token-choice order.
+    #
+    # The internal arange(N) is no longer the packed position tensor. It is only
+    # the local source-row index object used to gather from the current chunk
+    # tensor x. Flattening this object gives a (B, N*K) tensor aligned with H's
+    # token-major routed-copy order.
+    # -----------------------------------------------------------------------
+    source_token_indices = torch.arange(
+        sequence_length,
+        device=hidden_states.device,
+        dtype=torch.long,
+    ).view(1, sequence_length, 1).expand(
+        batch_size,
+        sequence_length,
+        num_selected_heads,
+    )
+    flattened_source_indices = source_token_indices.reshape(
+        batch_size,
+        sequence_length * num_selected_heads,
+    )
+
+    # -----------------------------------------------------------------------
+    # Reorder source-token indices into expert-major order.
+    #
+    # Applying Pi yields the local source-token rows in the packed expert-major
+    # order required by the paper. Those same reordered source indices are then
+    # used to gather hidden states, authoritative upstream positions, and outer
+    # liveness so all three remain aligned under the exact same packing
+    # transformation.
+    # -----------------------------------------------------------------------
+    sorted_source_indices = flattened_source_indices.gather(
+        dim=1,
+        index=permutation,
+    )
+    sorted_hidden_states = hidden_states.gather(
+        dim=1,
+        index=sorted_source_indices.unsqueeze(-1).expand(-1, -1, hidden_dim),
+    )
+    sorted_positions = position_ids.gather(
+        dim=1,
+        index=sorted_source_indices,
+    )
+    sorted_active_mask = outer_active_mask.gather(
+        dim=1,
+        index=sorted_source_indices,
+    )
+
+    # -----------------------------------------------------------------------
+    # Count how many routed copies land in each expert bucket.
+    #
+    # S[b, l] is the number of routed token copies assigned to expert l in batch b.
+    # T is the maximum such count across all batches and experts; it determines the
+    # padded expert-length dimension of the packed representation.
+    # -----------------------------------------------------------------------
+    tokens_per_expert = _bincount_rows(
+        values=flattened_selected_heads,
+        num_bins=num_experts,
+    )
+    max_tokens_per_expert = int(tokens_per_expert.max().item())
+
+    # -----------------------------------------------------------------------
+    # Construct the active-token mask M.
+    #
+    # Each expert bucket is left-justified: if S[b, l] = s, then slots
+    # t = 0, ..., s-1 are active and all later slots are padding. The resulting
+    # mask therefore both identifies real packed tokens and enforces left-justified
+    # packing. This is the unpacking_mask — it marks slot occupancy regardless of
+    # outer token liveness, and always has exactly B*N*K True entries.
+    # -----------------------------------------------------------------------
+    time_axis = torch.arange(
+        max_tokens_per_expert,
+        device=hidden_states.device,
+        dtype=torch.long,
+    ).view(1, 1, max_tokens_per_expert)
+    unpacking_mask = time_axis < tokens_per_expert.unsqueeze(-1)
+
+    # -----------------------------------------------------------------------
+    # Materialize the padded packed tensors.
+    #
+    # The packed hidden states x', packed original-token positions J', and packed
+    # active-token mask are allocated as zero-filled tensors. Active entries are
+    # then written into those buffers in the expert-major order established above.
+    # Padding remains zero / inactive.
+    # -----------------------------------------------------------------------
+    packed_hidden_states = hidden_states.new_zeros(
+        batch_size,
+        num_experts,
+        max_tokens_per_expert,
+        hidden_dim,
+    )
+    packed_positions = position_ids.new_zeros(
+        batch_size,
+        num_experts,
+        max_tokens_per_expert,
+    )
+    active_mask = torch.zeros(
+        batch_size,
+        num_experts,
+        max_tokens_per_expert,
+        dtype=torch.bool,
+        device=hidden_states.device,
+    )
+
+    packed_hidden_states[unpacking_mask] = sorted_hidden_states.reshape(-1, hidden_dim)
+    packed_positions[unpacking_mask] = sorted_positions.reshape(-1)
+    active_mask[unpacking_mask] = sorted_active_mask.reshape(-1)
+
+    return packed_hidden_states, packed_positions, unpacking_mask, active_mask
+
+
+# ---------------------------------------------------------------------------
+# Unpacking
+# ---------------------------------------------------------------------------
+
+def unpack_experts(
+    expert_outputs: torch.Tensor,
+    selected_heads: torch.Tensor,
+    unpacking_mask: torch.Tensor,
+    inverse_permutation: torch.Tensor,
+) -> torch.Tensor:
+    """Restore token-choice ordering from BEA expert-choice output.
+
+    Unpacking inverts the packing path only on occupied entries. Padding does not
+    participate: the output tensor is first filtered by unpacking_mask to recover
+    only the real routed-token copies in expert-major order, then Pi^{-1} restores
+    the original token-choice ordering, and finally the tensor is reshaped back to
+    (B, N, K, d).
+
+    The unpacking_mask — not active_mask — must be used here. Even copies of dead
+    outer tokens occupy slots and must be un-scattered correctly for the inverse
+    permutation to hold. The total True entry count in unpacking_mask is always
+    B*N*K, which is exactly what the reshape to (B, N*K, d) requires.
+
+    Args:
+        expert_outputs: Expert-choice BEA output y of shape (B, L, T, d).
+        selected_heads: Routed head selections I of shape (B, N, K).
+        unpacking_mask: From pack_experts(), shape (B, L, T). Identifies all
+            occupied packed slots regardless of outer token liveness.
+        inverse_permutation: Pi^{-1} from setup_packing(), shape (B, N*K).
+
+    Returns:
+        Restored token-choice tensor y_tilde of shape (B, N, K, d).
+    """
+    batch_size, sequence_length, num_selected_heads = selected_heads.shape
+    hidden_dim = expert_outputs.shape[-1]
+
+    active_outputs = expert_outputs[unpacking_mask]
+    sorted_token_choice_outputs = active_outputs.reshape(
+        batch_size,
+        sequence_length * num_selected_heads,
+        hidden_dim,
+    )
+    restored_outputs = sorted_token_choice_outputs.gather(
+        dim=1,
+        index=inverse_permutation.unsqueeze(-1).expand(-1, -1, hidden_dim),
+    )
+
+    return restored_outputs.reshape(
+        batch_size,
+        sequence_length,
+        num_selected_heads,
+        hidden_dim,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+def _bincount_rows(
+    values: torch.Tensor,
+    num_bins: int,
+) -> torch.Tensor:
+    """Count per-row integer occurrences for a 2D tensor.
+
+    torch.bincount operates on a flat 1D vector, but the packing algorithm needs
+    one bincount per batch row. The trick used here is to shift each row into its
+    own disjoint bin range before flattening:
+
+      row 0 uses bins [0, ..., num_bins - 1]
+      row 1 uses bins [num_bins, ..., 2*num_bins - 1]
+      row 2 uses bins [2*num_bins, ..., 3*num_bins - 1]
+      ...
+
+    After that shift, one global torch.bincount produces all row-local counts at
+    once. Reshaping the result back to (B, num_bins) recovers the per-row bincount.
+
+    This is a vectorized implementation detail only; externally visible behavior
+    remains exactly the paper's S tensor of per-batch per-expert token counts.
+
+    Args:
+        values: Integer tensor of shape (B, M) with entries in [0, num_bins).
+        num_bins: Number of bins.
+
+    Returns:
+        Counts tensor of shape (B, num_bins).
+    """
+    batch_size = values.shape[0]
+
+    row_offsets = torch.arange(
+        batch_size,
+        device=values.device,
+        dtype=values.dtype,
+    ).unsqueeze(1) * num_bins
+    shifted_values = values + row_offsets
+
+    counts = torch.bincount(
+        shifted_values.reshape(-1),
+        minlength=batch_size * num_bins,
+    )
+    return counts.reshape(batch_size, num_bins)

__attention__load_balance_loss.py

@@ -1,88 +1,88 @@
-"""Auxiliary-loss-free load balancing operator for MoSRAH routing.
-
-This module implements the custom autograd Function H(b, f) described in the paper's
-Implementation Concerns section. The operator bridges two requirements that are in
-tension: it must behave like a standard auxiliary loss (scalar output, scalable via
-multiplication) so that existing training loops remain compatible, while simultaneously
-implementing DeepSeek-style bias correction rather than the usual auxiliary-loss gradient
-path through the router weights.
-
-The resolution is a custom backward pass. The forward emits the load balance imbalance
-as a scalar loss. The backward, instead of differentiating that scalar with respect to
-its inputs, writes a bias-correction gradient directly to expert_bias. This gradient is
-then consumed by the main AdamW optimizer in the normal way, achieving DeepSeek-style
-correction without a standalone SGD update step.
-
-Paper ref: Appendix A.Implementation Concerns.
-"""
-
-import torch
-
-
-class LoadBalanceLoss(torch.autograd.Function):
-    """Custom autograd operator for DeepSeek-style auxiliary-loss-free load balancing.
-
-    Forward computes the load balance imbalance:
-
-        L_load_balance = H(b, f) = sum_l | f_l - 1/L |
-
-    Backward emits a bias-correction gradient to expert_bias:
-
-        grad_b = L_grad * sign(f_l - 1/L)
-
-    expert_bias (b) is included as a forward input so PyTorch registers it as a node
-    in the computation graph and routes gradients through it. routing_freqs (f) receives
-    no gradient — its origin is the discrete TopK operation which has no gradient, so
-    defining a gradient for f here would be mathematically incorrect.
-
-    Paper ref: Appendix A.Implementation Concerns.
-    """
-
-    @staticmethod
-    def forward(
-        ctx: torch.autograd.function.FunctionCtx,
-        expert_bias: torch.Tensor,
-        routing_freqs: torch.Tensor,
-    ) -> torch.Tensor:
-        """Compute the load balance loss.
-
-        Args:
-            ctx: Autograd context for saving state needed in backward.
-            expert_bias: Learned per-head bias b, shape (L,). Included as an input so
-                PyTorch tracks it as a computation graph node needing a gradient.
-            routing_freqs: Realized routing frequency f_l per head, shape (L,). Computed
-                from the discrete TopK selection — not differentiable.
-
-        Returns:
-            Scalar loss equal to sum_l |f_l - 1/L|.
-        """
-        L = expert_bias.shape[0]
-        # imbalance = f_l - 1/L for each head: positive means overloaded, negative means
-        # underloaded. Saved for backward where sign(imbalance) determines the direction
-        # of the bias-correction update.
-        imbalance = routing_freqs - 1.0 / L
-        ctx.save_for_backward(imbalance)
-        return imbalance.abs().sum()
-
-    @staticmethod
-    def backward(
-        ctx: torch.autograd.function.FunctionCtx,
-        grad_output: torch.Tensor,
-    ) -> tuple[torch.Tensor, None]:
-        """Emit the DeepSeek-style bias-correction gradient.
-
-        Args:
-            ctx: Autograd context carrying imbalance saved in forward.
-            grad_output: Incoming gradient L_grad (scalar). Any rescaling of the loss
-                by the training loop arrives here and is propagated to grad_b, so the
-                correction magnitude is proportional to the loss weight chosen by the
-                consumer.
-
-        Returns:
-            Gradient for expert_bias: L_grad * sign(f_l - 1/L), shape (L,).
-            None for routing_freqs: no gradient is defined for the discrete routing
-            frequency.
-        """
-        (imbalance,) = ctx.saved_tensors
-        grad_expert_bias = grad_output * imbalance.sign()
-        return grad_expert_bias, None
+"""Auxiliary-loss-free load balancing operator for MoSRAH routing.
+
+This module implements the custom autograd Function H(b, f) described in the paper's
+Implementation Concerns section. The operator bridges two requirements that are in
+tension: it must behave like a standard auxiliary loss (scalar output, scalable via
+multiplication) so that existing training loops remain compatible, while simultaneously
+implementing DeepSeek-style bias correction rather than the usual auxiliary-loss gradient
+path through the router weights.
+
+The resolution is a custom backward pass. The forward emits the load balance imbalance
+as a scalar loss. The backward, instead of differentiating that scalar with respect to
+its inputs, writes a bias-correction gradient directly to expert_bias. This gradient is
+then consumed by the main AdamW optimizer in the normal way, achieving DeepSeek-style
+correction without a standalone SGD update step.
+
+Paper ref: Appendix A.Implementation Concerns.
+"""
+
+import torch
+
+
+class LoadBalanceLoss(torch.autograd.Function):
+    """Custom autograd operator for DeepSeek-style auxiliary-loss-free load balancing.
+
+    Forward computes the load balance imbalance:
+
+        L_load_balance = H(b, f) = sum_l | f_l - 1/L |
+
+    Backward emits a bias-correction gradient to expert_bias:
+
+        grad_b = L_grad * sign(f_l - 1/L)
+
+    expert_bias (b) is included as a forward input so PyTorch registers it as a node
+    in the computation graph and routes gradients through it. routing_freqs (f) receives
+    no gradient — its origin is the discrete TopK operation which has no gradient, so
+    defining a gradient for f here would be mathematically incorrect.
+
+    Paper ref: Appendix A.Implementation Concerns.
+    """
+
+    @staticmethod
+    def forward(
+        ctx: torch.autograd.function.FunctionCtx,
+        expert_bias: torch.Tensor,
+        routing_freqs: torch.Tensor,
+    ) -> torch.Tensor:
+        """Compute the load balance loss.
+
+        Args:
+            ctx: Autograd context for saving state needed in backward.
+            expert_bias: Learned per-head bias b, shape (L,). Included as an input so
+                PyTorch tracks it as a computation graph node needing a gradient.
+            routing_freqs: Realized routing frequency f_l per head, shape (L,). Computed
+                from the discrete TopK selection — not differentiable.
+
+        Returns:
+            Scalar loss equal to sum_l |f_l - 1/L|.
+        """
+        L = expert_bias.shape[0]
+        # imbalance = f_l - 1/L for each head: positive means overloaded, negative means
+        # underloaded. Saved for backward where sign(imbalance) determines the direction
+        # of the bias-correction update.
+        imbalance = routing_freqs - 1.0 / L
+        ctx.save_for_backward(imbalance)
+        return imbalance.abs().sum()
+
+    @staticmethod
+    def backward(
+        ctx: torch.autograd.function.FunctionCtx,
+        grad_output: torch.Tensor,
+    ) -> tuple[torch.Tensor, None]:
+        """Emit the DeepSeek-style bias-correction gradient.
+
+        Args:
+            ctx: Autograd context carrying imbalance saved in forward.
+            grad_output: Incoming gradient L_grad (scalar). Any rescaling of the loss
+                by the training loop arrives here and is propagated to grad_b, so the
+                correction magnitude is proportional to the loss weight chosen by the
+                consumer.
+
+        Returns:
+            Gradient for expert_bias: L_grad * sign(f_l - 1/L), shape (L,).
+            None for routing_freqs: no gradient is defined for the discrete routing
+            frequency.
+        """
+        (imbalance,) = ctx.saved_tensors
+        grad_expert_bias = grad_output * imbalance.sign()
+        return grad_expert_bias, None

__attention__mosrah.py

@@ -1,140 +1,140 @@
-"""Full MoSRAH sparse path for SHRAM.
-
-This module coordinates the routed sparse attention path used inside the SHRAM
-hybrid attention layer. The underlying mechanics already live in verified
-subunits. The responsibility here is to connect those subunits without
-corrupting their bridge contracts.
-
-In particular, this path must preserve three architectural distinctions:
-
-- selected head indices are not routing probabilities
-- packed position semantics are chosen before BEA, not inside it
-- weighted reduction must consume the router's unbiased renormalized
-  probabilities after token-choice order has been restored
-"""
-
-import torch
-from torch import nn
-
-from .__cache__mosrah_cache import MoSRAHCache
-from .configuration import ShramConfig
-from .__attention__bottlenecked_ensemble_attention import BottleneckedEnsembleAttention
-from .__attention__expert_packing import (
-    pack_experts,
-    setup_packing,
-    unpack_experts,
-)
-from .__attention__router import MoSRAHRouter
-from .__attention__positions_converter import SparseMoSRAHPositions
-
-
-class MoSRAHLayer(nn.Module):
-    """Full routed sparse attention path for SHRAM.
-
-    The MoSRAH path consumes model-space hidden states together with
-    authoritative per-token positions and returns the model-space sparse-path
-    contribution, the router's load-balance loss, and the router's MaxVio
-    routing-imbalance scalar.
-    """
-
-    def __init__(self, config: ShramConfig) -> None:
-        super().__init__()
-        self.num_experts = config.num_mosrah_heads
-
-        self.router = MoSRAHRouter(config)
-        self.positions = SparseMoSRAHPositions(config)
-        self.bea = BottleneckedEnsembleAttention(config)
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        position_ids: torch.Tensor,
-        active_mask: torch.Tensor,
-        cache: MoSRAHCache | None,
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        """Run the full MoSRAH sparse path.
-
-        Args:
-            hidden_states: Model-space hidden states x of shape (B, N, d).
-            position_ids: Authoritative per-token positions of shape (B, N).
-            active_mask: Current-chunk active mask of shape (B, N), where True
-                means the token is semantically live. Forwarded to the router
-                so dead tokens are excluded from routing statistics, and to
-                pack_experts so dead outer tokens do not become semantically
-                active packed entries.
-            cache: Optional layer-local MoSRAH cache. Pass None for uncached
-                execution and the layer-local cache instance for cached execution.
-
-        Returns:
-            sparse_output: Model-space sparse-path output of shape (B, N, d).
-            load_balance_loss: Scalar router load-balance loss.
-            max_vio: Detached scalar routing-imbalance summary. Passed through
-                unchanged from the router; see MoSRAHRouter for semantics.
-        """
-
-        # -------------------------------------------------------------------
-        # The first transition moves from model-space token-choice input into
-        # the packed expert-choice sparse-attention state. Routing decides both
-        # which experts each token uses and which unbiased probabilities must be
-        # reserved for the final reduction. The active mask is forwarded to the
-        # router so dead tokens are excluded from routing statistics, and to
-        # pack_experts so outer liveness is faithfully carried into the packed
-        # frame. Packing returns both the unpacking mask (slot occupancy, always
-        # B*N*K True entries) and the packed active mask (live slots only);
-        # active_mask is rebound to the packed form after this point.
-        # -------------------------------------------------------------------
-        selected_heads, routing_probs, load_balance_loss, max_vio = self.router(
-            hidden_states, active_mask
-        )
-
-        flattened_selected_heads, permutation, inverse_permutation = setup_packing(
-            selected_heads
-        )
-        packed_hidden_states, packed_positions, unpacking_mask, active_mask = pack_experts(
-            hidden_states=hidden_states,
-            position_ids=position_ids,
-            selected_heads=selected_heads,
-            num_experts=self.num_experts,
-            flattened_selected_heads=flattened_selected_heads,
-            permutation=permutation,
-            outer_active_mask=active_mask,
-        )
-
-        # -------------------------------------------------------------------
-        # Sparse attention runs entirely in the packed expert-choice frame, so
-        # the RoPE position semantics must also be chosen in that frame. The
-        # position layer therefore decides whether BEA should see packed
-        # original-token positions or packed local-slot positions. BEA then
-        # consumes that packed position tensor together with the packed hidden
-        # states and the layer-local sparse cache, which it owns directly.
-        # -------------------------------------------------------------------
-        bea_positions = self.positions(
-            packed_positions=packed_positions,
-            cache=cache,
-        )
-        packed_outputs = self.bea(
-            packed_embeddings=packed_hidden_states,
-            position_ids=bea_positions,
-            active_mask=active_mask,
-            cache=cache,
-        )
-
-        # -------------------------------------------------------------------
-        # The final transition restores token-choice meaning and only then
-        # collapses the K routed copies back into model space. This ordering is
-        # required because routing_probs live in token-choice space, whereas BEA
-        # returns expert-choice packed outputs. The reduction must therefore
-        # happen after unpacking, and it must use the router's unbiased
-        # renormalized probabilities rather than any biased selection scores.
-        # -------------------------------------------------------------------
-        token_choice_outputs = unpack_experts(
-            expert_outputs=packed_outputs,
-            selected_heads=selected_heads,
-            unpacking_mask=unpacking_mask,
-            inverse_permutation=inverse_permutation,
-        )
-        final_output = (
-            token_choice_outputs * routing_probs.unsqueeze(-1)
-        ).sum(dim=2)
-
-        return final_output, load_balance_loss, max_vio
+"""Full MoSRAH sparse path for SHRAM.
+
+This module coordinates the routed sparse attention path used inside the SHRAM
+hybrid attention layer. The underlying mechanics already live in verified
+subunits. The responsibility here is to connect those subunits without
+corrupting their bridge contracts.
+
+In particular, this path must preserve three architectural distinctions:
+
+- selected head indices are not routing probabilities
+- packed position semantics are chosen before BEA, not inside it
+- weighted reduction must consume the router's unbiased renormalized
+  probabilities after token-choice order has been restored
+"""
+
+import torch
+from torch import nn
+
+from .__cache__mosrah_cache import MoSRAHCache
+from .configuration import ShramConfig
+from .__attention__bottlenecked_ensemble_attention import BottleneckedEnsembleAttention
+from .__attention__expert_packing import (
+    pack_experts,
+    setup_packing,
+    unpack_experts,
+)
+from .__attention__router import MoSRAHRouter
+from .__attention__positions_converter import SparseMoSRAHPositions
+
+
+class MoSRAHLayer(nn.Module):
+    """Full routed sparse attention path for SHRAM.
+
+    The MoSRAH path consumes model-space hidden states together with
+    authoritative per-token positions and returns the model-space sparse-path
+    contribution, the router's load-balance loss, and the router's MaxVio
+    routing-imbalance scalar.
+    """
+
+    def __init__(self, config: ShramConfig) -> None:
+        super().__init__()
+        self.num_experts = config.num_mosrah_heads
+
+        self.router = MoSRAHRouter(config)
+        self.positions = SparseMoSRAHPositions(config)
+        self.bea = BottleneckedEnsembleAttention(config)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_ids: torch.Tensor,
+        active_mask: torch.Tensor,
+        cache: MoSRAHCache | None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Run the full MoSRAH sparse path.
+
+        Args:
+            hidden_states: Model-space hidden states x of shape (B, N, d).
+            position_ids: Authoritative per-token positions of shape (B, N).
+            active_mask: Current-chunk active mask of shape (B, N), where True
+                means the token is semantically live. Forwarded to the router
+                so dead tokens are excluded from routing statistics, and to
+                pack_experts so dead outer tokens do not become semantically
+                active packed entries.
+            cache: Optional layer-local MoSRAH cache. Pass None for uncached
+                execution and the layer-local cache instance for cached execution.
+
+        Returns:
+            sparse_output: Model-space sparse-path output of shape (B, N, d).
+            load_balance_loss: Scalar router load-balance loss.
+            max_vio: Detached scalar routing-imbalance summary. Passed through
+                unchanged from the router; see MoSRAHRouter for semantics.
+        """
+
+        # -------------------------------------------------------------------
+        # The first transition moves from model-space token-choice input into
+        # the packed expert-choice sparse-attention state. Routing decides both
+        # which experts each token uses and which unbiased probabilities must be
+        # reserved for the final reduction. The active mask is forwarded to the
+        # router so dead tokens are excluded from routing statistics, and to
+        # pack_experts so outer liveness is faithfully carried into the packed
+        # frame. Packing returns both the unpacking mask (slot occupancy, always
+        # B*N*K True entries) and the packed active mask (live slots only);
+        # active_mask is rebound to the packed form after this point.
+        # -------------------------------------------------------------------
+        selected_heads, routing_probs, load_balance_loss, max_vio = self.router(
+            hidden_states, active_mask
+        )
+
+        flattened_selected_heads, permutation, inverse_permutation = setup_packing(
+            selected_heads
+        )
+        packed_hidden_states, packed_positions, unpacking_mask, active_mask = pack_experts(
+            hidden_states=hidden_states,
+            position_ids=position_ids,
+            selected_heads=selected_heads,
+            num_experts=self.num_experts,
+            flattened_selected_heads=flattened_selected_heads,
+            permutation=permutation,
+            outer_active_mask=active_mask,
+        )
+
+        # -------------------------------------------------------------------
+        # Sparse attention runs entirely in the packed expert-choice frame, so
+        # the RoPE position semantics must also be chosen in that frame. The
+        # position layer therefore decides whether BEA should see packed
+        # original-token positions or packed local-slot positions. BEA then
+        # consumes that packed position tensor together with the packed hidden
+        # states and the layer-local sparse cache, which it owns directly.
+        # -------------------------------------------------------------------
+        bea_positions = self.positions(
+            packed_positions=packed_positions,
+            cache=cache,
+        )
+        packed_outputs = self.bea(
+            packed_embeddings=packed_hidden_states,
+            position_ids=bea_positions,
+            active_mask=active_mask,
+            cache=cache,
+        )
+
+        # -------------------------------------------------------------------
+        # The final transition restores token-choice meaning and only then
+        # collapses the K routed copies back into model space. This ordering is
+        # required because routing_probs live in token-choice space, whereas BEA
+        # returns expert-choice packed outputs. The reduction must therefore
+        # happen after unpacking, and it must use the router's unbiased
+        # renormalized probabilities rather than any biased selection scores.
+        # -------------------------------------------------------------------
+        token_choice_outputs = unpack_experts(
+            expert_outputs=packed_outputs,
+            selected_heads=selected_heads,
+            unpacking_mask=unpacking_mask,
+            inverse_permutation=inverse_permutation,
+        )
+        final_output = (
+            token_choice_outputs * routing_probs.unsqueeze(-1)
+        ).sum(dim=2)
+
+        return final_output, load_balance_loss, max_vio

__attention__positions_converter.py

@@ -1,105 +1,105 @@
-"""Position computation for the MoSRAH sparse path.
-
-This layer computes the packed position tensor P consumed by BEA.
-
-- In main-sequence mode, P is the packed original-token position tensor from the
-  packing path.
-- In semantic-sequence mode, P is a per-expert local sequence over the packed
-  expert-choice layout, optionally offset by the current sparse-cache occupancies
-  during cached inference.
-"""
-
-import torch
-from torch import nn
-
-from .configuration import ShramConfig
-from .__cache__mosrah_cache import MoSRAHCache
-
-
-class SparseMoSRAHPositions(nn.Module):
-    """Compute the packed RoPE position tensor for the MoSRAH sparse path.
-
-    This layer operates in the packed expert-choice frame used by BEA. The input
-    packed_positions tensor is always the packed original-token position tensor
-    produced by the packing path. The configured rope_mode determines whether that
-    tensor is forwarded directly or replaced by a semantic local-slot sequence.
-    """
-
-    def __init__(self, config: ShramConfig) -> None:
-        super().__init__()
-        self.rope_mode = config.rope_mode
-
-    def forward(
-        self,
-        packed_positions: torch.Tensor,
-        cache: MoSRAHCache | None,
-    ) -> torch.Tensor:
-        """Compute the packed position tensor P consumed by BEA.
-
-        Args:
-            packed_positions: Packed original-token positions J' of shape (B, L, T).
-            cache: Optional layer-local MoSRAH cache. When present in semantic-sequence
-                mode, the current per-head occupancies offset the local packed sequence.
-
-        Returns:
-            Packed position tensor P of shape (B, L, T).
-        """
-        if self.rope_mode == "main_sequence":
-            return self._main_sequence_positions(packed_positions)
-
-        if self.rope_mode == "semantic_sequence":
-            return self._semantic_sequence_positions(packed_positions, cache)
-
-        raise NotImplementedError(
-            f"Unsupported MoSRAH rope_mode '{self.rope_mode}'."
-        )
-
-    def _main_sequence_positions(
-        self,
-        packed_positions: torch.Tensor,
-    ) -> torch.Tensor:
-        """Forward packed original-token positions unchanged."""
-        return packed_positions
-
-    def _semantic_sequence_positions(
-        self,
-        packed_positions: torch.Tensor,
-        cache: MoSRAHCache | None,
-    ) -> torch.Tensor:
-        """Compute semantic-sequence packed positions in expert-choice space.
-
-        Without a sparse cache, semantic positions are the local packed sequence
-        0, 1, 2, ... over the expert-local T dimension. With a sparse cache, that
-        same local sequence is offset by the current per-(batch, expert) occupancies
-        returned by get_heads_lengths().
-        """
-        batch_size, num_experts, packed_length = packed_positions.shape
-
-        # -------------------------------------------------------------------
-        # Construct the local packed sequence 0, 1, 2, ... over the expert-local
-        # sequence dimension T. This is then broadcast across batch and experts.
-        # -------------------------------------------------------------------
-        local_positions = torch.arange(
-            packed_length,
-            device=packed_positions.device,
-            dtype=packed_positions.dtype,
-        ).view(1, 1, packed_length).expand(
-            batch_size,
-            num_experts,
-            packed_length,
-        )
-
-        # -------------------------------------------------------------------
-        # In cached semantic-sequence mode, positions continue from the current
-        # sparse-cache occupancies rather than restarting at zero for the local
-        # chunk.
-        # -------------------------------------------------------------------
-        if cache is None:
-            return local_positions
-
-        cached_lengths = cache.get_heads_lengths().to(
-            device=packed_positions.device,
-            dtype=packed_positions.dtype,
-        ).unsqueeze(-1)
-
+"""Position computation for the MoSRAH sparse path.
+
+This layer computes the packed position tensor P consumed by BEA.
+
+- In main-sequence mode, P is the packed original-token position tensor from the
+  packing path.
+- In semantic-sequence mode, P is a per-expert local sequence over the packed
+  expert-choice layout, optionally offset by the current sparse-cache occupancies
+  during cached inference.
+"""
+
+import torch
+from torch import nn
+
+from .configuration import ShramConfig
+from .__cache__mosrah_cache import MoSRAHCache
+
+
+class SparseMoSRAHPositions(nn.Module):
+    """Compute the packed RoPE position tensor for the MoSRAH sparse path.
+
+    This layer operates in the packed expert-choice frame used by BEA. The input
+    packed_positions tensor is always the packed original-token position tensor
+    produced by the packing path. The configured rope_mode determines whether that
+    tensor is forwarded directly or replaced by a semantic local-slot sequence.
+    """
+
+    def __init__(self, config: ShramConfig) -> None:
+        super().__init__()
+        self.rope_mode = config.rope_mode
+
+    def forward(
+        self,
+        packed_positions: torch.Tensor,
+        cache: MoSRAHCache | None,
+    ) -> torch.Tensor:
+        """Compute the packed position tensor P consumed by BEA.
+
+        Args:
+            packed_positions: Packed original-token positions J' of shape (B, L, T).
+            cache: Optional layer-local MoSRAH cache. When present in semantic-sequence
+                mode, the current per-head occupancies offset the local packed sequence.
+
+        Returns:
+            Packed position tensor P of shape (B, L, T).
+        """
+        if self.rope_mode == "main_sequence":
+            return self._main_sequence_positions(packed_positions)
+
+        if self.rope_mode == "semantic_sequence":
+            return self._semantic_sequence_positions(packed_positions, cache)
+
+        raise NotImplementedError(
+            f"Unsupported MoSRAH rope_mode '{self.rope_mode}'."
+        )
+
+    def _main_sequence_positions(
+        self,
+        packed_positions: torch.Tensor,
+    ) -> torch.Tensor:
+        """Forward packed original-token positions unchanged."""
+        return packed_positions
+
+    def _semantic_sequence_positions(
+        self,
+        packed_positions: torch.Tensor,
+        cache: MoSRAHCache | None,
+    ) -> torch.Tensor:
+        """Compute semantic-sequence packed positions in expert-choice space.
+
+        Without a sparse cache, semantic positions are the local packed sequence
+        0, 1, 2, ... over the expert-local T dimension. With a sparse cache, that
+        same local sequence is offset by the current per-(batch, expert) occupancies
+        returned by get_heads_lengths().
+        """
+        batch_size, num_experts, packed_length = packed_positions.shape
+
+        # -------------------------------------------------------------------
+        # Construct the local packed sequence 0, 1, 2, ... over the expert-local
+        # sequence dimension T. This is then broadcast across batch and experts.
+        # -------------------------------------------------------------------
+        local_positions = torch.arange(
+            packed_length,
+            device=packed_positions.device,
+            dtype=packed_positions.dtype,
+        ).view(1, 1, packed_length).expand(
+            batch_size,
+            num_experts,
+            packed_length,
+        )
+
+        # -------------------------------------------------------------------
+        # In cached semantic-sequence mode, positions continue from the current
+        # sparse-cache occupancies rather than restarting at zero for the local
+        # chunk.
+        # -------------------------------------------------------------------
+        if cache is None:
+            return local_positions
+
+        cached_lengths = cache.get_heads_lengths().to(
+            device=packed_positions.device,
+            dtype=packed_positions.dtype,
+        ).unsqueeze(-1)
+
         return local_positions + cached_lengths

__attention__router.py

@@ -1,162 +1,162 @@
-"""Token-choice router for the MoSRAH sparse attention path.
-
-This module implements the routing mechanism described in Appendix A.Routing of the
-paper. Given an input hidden state x, the router produces two outputs used downstream:
-
-  - selected_heads (I): which K of the L available expert heads each token routes to,
-    determined by TopK over biased routing scores.
-  - routing_probs (P): the weights used for the weighted output reduction, gathered from
-    *unbiased* routing scores at the selected indices and renormalized. The learned expert
-    bias b must not influence P.
-
-This separation is architecturally critical: expert_bias drives selection (and thus load
-balancing) but does not corrupt the gradient path from the output through routing_probs
-back to the routing projection weights.
-
-The router also computes and returns the load balance loss via the LoadBalanceLoss custom
-autograd operator (see load_balance_loss.py). This loss is a scalar that the training
-loop can weight and add to the language modeling loss.
-
-The router additionally computes and returns MaxVio, a detached scalar summarising
-routing imbalance for the current forward pass:
-
-    MaxVio = L · max_l(f_l − 1/L)
-
-where f_l is the realised routing frequency of head l and 1/L is the perfectly balanced
-target. MaxVio is a monitoring quantity only; it never contributes gradients.
-
-Paper ref: Appendix A.Routing, Appendix A.Load Balancing, §MaxVio.
-"""
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-from .configuration import ShramConfig
-from .__attention__load_balance_loss import LoadBalanceLoss
-
-
-class MoSRAHRouter(nn.Module):
-    """Token-choice router for MoSRAH sparse attention.
-
-    Each input token independently selects K of the L available expert heads. Selection
-    is driven by biased routing scores to enable load balancing, but the routing
-    probabilities used for output reduction are computed from unbiased scores so that
-    the expert bias does not interfere with the gradient path to the router weights.
-
-    The routing projection W_r has no bias term — the paper specifies xW_r with no
-    additional projection bias. The only bias-like parameter is expert_bias (b), which
-    has an entirely separate role and update mechanism.
-
-    Args:
-        config: Model configuration. Must expose ``hidden_size``, ``num_mosrah_heads``
-            (L), and ``num_selected_heads`` (K).
-    """
-
-    def __init__(self, config: ShramConfig) -> None:
-        super().__init__()
-        self.num_mosrah_heads = config.num_mosrah_heads
-        self.num_selected_heads = config.num_selected_heads
-
-        # W_r: routing projection, no bias (paper specifies xW_r, no additional term).
-        self.routing_projection = nn.Linear(
-            config.hidden_size, config.num_mosrah_heads, bias=False
-        )
-
-        # b: learned per-head bias for load balancing. Initialized to zero so that all
-        # heads start with equal selection probability. Updated by the main optimizer
-        # via the LoadBalanceLoss custom backward.
-        self.expert_bias = nn.Parameter(torch.zeros(config.num_mosrah_heads))
-
-    def forward(
-        self, x: torch.Tensor, active_mask: torch.Tensor
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-        """Route input tokens to K expert heads each and compute routing probabilities.
-
-        Args:
-            x: Input hidden states of shape (batch, seq_len, hidden_size).
-            active_mask: Current-chunk active mask of shape (batch, seq_len), where
-                True means the token is semantically live. Dead tokens do not
-                contribute to routing frequencies, load_balance_loss, or max_vio.
-
-        Returns:
-            selected_heads: Head indices I of shape (batch, seq_len, num_selected_heads).
-                Each token's K selected head indices, determined by TopK on biased scores.
-            routing_probs: Routing probabilities P of shape (batch, seq_len,
-                num_selected_heads). Gathered from unbiased scores at selected_heads
-                indices and renormalized to sum to 1 per token.
-            load_balance_loss: Scalar load balance imbalance loss for this forward pass.
-                Training loop scales this by a weight and adds it to the main loss.
-            max_vio: Detached scalar routing-imbalance summary for this forward pass.
-                Equal to L · max_l(f_l − 1/L). Zero means perfect balance. Not a loss;
-                never contributes gradients.
-        """
-        B, N, _ = x.shape
-        L = self.num_mosrah_heads
-        K = self.num_selected_heads
-
-        # Unbiased routing scores R = Softmax(xW_r). These are the scores used to
-        # compute routing_probs — expert_bias must not influence them.
-        logits = self.routing_projection(x)                    # (B, N, L)
-        routing_scores = F.softmax(logits, dim=-1)             # R, (B, N, L)
-
-        # Biased routing scores R̂ = Softmax(xW_r + b). Used only for TopK head
-        # selection. expert_bias is added to logits before softmax so that the bias
-        # shifts selection probability without rescaling the unbiased distribution.
-        biased_routing_scores = F.softmax(                     # R̂, (B, N, L)
-            logits + self.expert_bias, dim=-1
-        )
-
-        # selected_heads I = TopK(R̂): K head indices per token, shape (B, N, K).
-        selected_heads = biased_routing_scores.topk(K, dim=-1).indices
-
-        # Routing probabilities P: gathered from unbiased R at selected_heads indices,
-        # then renormalized so they sum to 1 per token. Gathering from routing_scores
-        # (not biased_routing_scores) is the invariant that keeps the gradient path from
-        # the output back to the router weights free of expert_bias influence.
-        gathered = routing_scores.gather(dim=-1, index=selected_heads)   # V, (B, N, K)
-        routing_probs = gathered / gathered.sum(dim=-1, keepdim=True)    # P, (B, N, K)
-
-        # Routing frequency f_l: fraction of active (batch, token, head_slot) triples
-        # assigned to each head. Dead tokens are excluded by zeroing their rows in the
-        # assignment mask before reduction. Normalization uses the active assignment
-        # count so frequencies remain properly scaled regardless of how many tokens
-        # are live in this chunk.
-        assignment_mask = torch.zeros(B, N, L, device=x.device, dtype=x.dtype)
-        assignment_mask.scatter_(-1, selected_heads, 1.0)
-        active_assignments = assignment_mask * active_mask.unsqueeze(-1)
-        num_active_assignments = active_mask.sum() * K
-        routing_freqs = active_assignments.sum(dim=(0, 1)) / num_active_assignments  # f, (L,)
-
-        # Load balance loss via custom autograd. expert_bias is an input so PyTorch
-        # registers it as a graph node; the custom backward writes the DeepSeek-style
-        # correction gradient to expert_bias.grad for the optimizer to consume.
-        load_balance_loss = LoadBalanceLoss.apply(self.expert_bias, routing_freqs)
-
-        # MaxVio is a detached monitoring scalar derived from routing_freqs. It must
-        # not contribute gradients under any circumstance, so it is detached at the
-        # point of computation rather than left to callers to detach.
-        max_vio = self._compute_max_vio(routing_freqs, L)
-
-        return selected_heads, routing_probs, load_balance_loss, max_vio
-
-    @staticmethod
-    def _compute_max_vio(routing_freqs: torch.Tensor, num_heads: int) -> torch.Tensor:
-        """Compute the MaxVio routing-imbalance scalar.
-
-        MaxVio = L · max_l(f_l − 1/L), where f_l is the realised routing frequency of
-        head l and 1/L is the perfectly balanced target. A value of zero indicates
-        perfect balance; a value of 1 means the most overloaded head received exactly
-        double its fair share.
-
-        The result is detached from the autograd graph — MaxVio is a monitoring scalar
-        and must never contribute gradients to any parameter.
-
-        Args:
-            routing_freqs: Per-head routing frequencies of shape (L,). Sums to 1.
-            num_heads: Total number of MoSRAH heads L.
-
-        Returns:
-            Detached scalar MaxVio tensor.
-        """
-        return (num_heads * (routing_freqs - 1.0 / num_heads).max()).detach()
+"""Token-choice router for the MoSRAH sparse attention path.
+
+This module implements the routing mechanism described in Appendix A.Routing of the
+paper. Given an input hidden state x, the router produces two outputs used downstream:
+
+  - selected_heads (I): which K of the L available expert heads each token routes to,
+    determined by TopK over biased routing scores.
+  - routing_probs (P): the weights used for the weighted output reduction, gathered from
+    *unbiased* routing scores at the selected indices and renormalized. The learned expert
+    bias b must not influence P.
+
+This separation is architecturally critical: expert_bias drives selection (and thus load
+balancing) but does not corrupt the gradient path from the output through routing_probs
+back to the routing projection weights.
+
+The router also computes and returns the load balance loss via the LoadBalanceLoss custom
+autograd operator (see load_balance_loss.py). This loss is a scalar that the training
+loop can weight and add to the language modeling loss.
+
+The router additionally computes and returns MaxVio, a detached scalar summarising
+routing imbalance for the current forward pass:
+
+    MaxVio = L · max_l(f_l − 1/L)
+
+where f_l is the realised routing frequency of head l and 1/L is the perfectly balanced
+target. MaxVio is a monitoring quantity only; it never contributes gradients.
+
+Paper ref: Appendix A.Routing, Appendix A.Load Balancing, §MaxVio.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .configuration import ShramConfig
+from .__attention__load_balance_loss import LoadBalanceLoss
+
+
+class MoSRAHRouter(nn.Module):
+    """Token-choice router for MoSRAH sparse attention.
+
+    Each input token independently selects K of the L available expert heads. Selection
+    is driven by biased routing scores to enable load balancing, but the routing
+    probabilities used for output reduction are computed from unbiased scores so that
+    the expert bias does not interfere with the gradient path to the router weights.
+
+    The routing projection W_r has no bias term — the paper specifies xW_r with no
+    additional projection bias. The only bias-like parameter is expert_bias (b), which
+    has an entirely separate role and update mechanism.
+
+    Args:
+        config: Model configuration. Must expose ``hidden_size``, ``num_mosrah_heads``
+            (L), and ``num_selected_heads`` (K).
+    """
+
+    def __init__(self, config: ShramConfig) -> None:
+        super().__init__()
+        self.num_mosrah_heads = config.num_mosrah_heads
+        self.num_selected_heads = config.num_selected_heads
+
+        # W_r: routing projection, no bias (paper specifies xW_r, no additional term).
+        self.routing_projection = nn.Linear(
+            config.hidden_size, config.num_mosrah_heads, bias=False
+        )
+
+        # b: learned per-head bias for load balancing. Initialized to zero so that all
+        # heads start with equal selection probability. Updated by the main optimizer
+        # via the LoadBalanceLoss custom backward.
+        self.expert_bias = nn.Parameter(torch.zeros(config.num_mosrah_heads))
+
+    def forward(
+        self, x: torch.Tensor, active_mask: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Route input tokens to K expert heads each and compute routing probabilities.
+
+        Args:
+            x: Input hidden states of shape (batch, seq_len, hidden_size).
+            active_mask: Current-chunk active mask of shape (batch, seq_len), where
+                True means the token is semantically live. Dead tokens do not
+                contribute to routing frequencies, load_balance_loss, or max_vio.
+
+        Returns:
+            selected_heads: Head indices I of shape (batch, seq_len, num_selected_heads).
+                Each token's K selected head indices, determined by TopK on biased scores.
+            routing_probs: Routing probabilities P of shape (batch, seq_len,
+                num_selected_heads). Gathered from unbiased scores at selected_heads
+                indices and renormalized to sum to 1 per token.
+            load_balance_loss: Scalar load balance imbalance loss for this forward pass.
+                Training loop scales this by a weight and adds it to the main loss.
+            max_vio: Detached scalar routing-imbalance summary for this forward pass.
+                Equal to L · max_l(f_l − 1/L). Zero means perfect balance. Not a loss;
+                never contributes gradients.
+        """
+        B, N, _ = x.shape
+        L = self.num_mosrah_heads
+        K = self.num_selected_heads
+
+        # Unbiased routing scores R = Softmax(xW_r). These are the scores used to
+        # compute routing_probs — expert_bias must not influence them.
+        logits = self.routing_projection(x)                    # (B, N, L)
+        routing_scores = F.softmax(logits, dim=-1)             # R, (B, N, L)
+
+        # Biased routing scores R̂ = Softmax(xW_r + b). Used only for TopK head
+        # selection. expert_bias is added to logits before softmax so that the bias
+        # shifts selection probability without rescaling the unbiased distribution.
+        biased_routing_scores = F.softmax(                     # R̂, (B, N, L)
+            logits + self.expert_bias, dim=-1
+        )
+
+        # selected_heads I = TopK(R̂): K head indices per token, shape (B, N, K).
+        selected_heads = biased_routing_scores.topk(K, dim=-1).indices
+
+        # Routing probabilities P: gathered from unbiased R at selected_heads indices,
+        # then renormalized so they sum to 1 per token. Gathering from routing_scores
+        # (not biased_routing_scores) is the invariant that keeps the gradient path from
+        # the output back to the router weights free of expert_bias influence.
+        gathered = routing_scores.gather(dim=-1, index=selected_heads)   # V, (B, N, K)
+        routing_probs = gathered / gathered.sum(dim=-1, keepdim=True)    # P, (B, N, K)
+
+        # Routing frequency f_l: fraction of active (batch, token, head_slot) triples
+        # assigned to each head. Dead tokens are excluded by zeroing their rows in the
+        # assignment mask before reduction. Normalization uses the active assignment
+        # count so frequencies remain properly scaled regardless of how many tokens
+        # are live in this chunk.
+        assignment_mask = torch.zeros(B, N, L, device=x.device, dtype=x.dtype)
+        assignment_mask.scatter_(-1, selected_heads, 1.0)
+        active_assignments = assignment_mask * active_mask.unsqueeze(-1)
+        num_active_assignments = active_mask.sum() * K
+        routing_freqs = active_assignments.sum(dim=(0, 1)) / num_active_assignments  # f, (L,)
+
+        # Load balance loss via custom autograd. expert_bias is an input so PyTorch
+        # registers it as a graph node; the custom backward writes the DeepSeek-style
+        # correction gradient to expert_bias.grad for the optimizer to consume.
+        load_balance_loss = LoadBalanceLoss.apply(self.expert_bias, routing_freqs)
+
+        # MaxVio is a detached monitoring scalar derived from routing_freqs. It must
+        # not contribute gradients under any circumstance, so it is detached at the
+        # point of computation rather than left to callers to detach.
+        max_vio = self._compute_max_vio(routing_freqs, L)
+
+        return selected_heads, routing_probs, load_balance_loss, max_vio
+
+    @staticmethod
+    def _compute_max_vio(routing_freqs: torch.Tensor, num_heads: int) -> torch.Tensor:
+        """Compute the MaxVio routing-imbalance scalar.
+
+        MaxVio = L · max_l(f_l − 1/L), where f_l is the realised routing frequency of
+        head l and 1/L is the perfectly balanced target. A value of zero indicates
+        perfect balance; a value of 1 means the most overloaded head received exactly
+        double its fair share.
+
+        The result is detached from the autograd graph — MaxVio is a monitoring scalar
+        and must never contribute gradients to any parameter.
+
+        Args:
+            routing_freqs: Per-head routing frequencies of shape (L,). Sums to 1.
+            num_heads: Total number of MoSRAH heads L.
+
+        Returns:
+            Detached scalar MaxVio tensor.
+        """
+        return (num_heads * (routing_freqs - 1.0 / num_heads).max()).detach()

__attention__shram.py

@@ -1,116 +1,116 @@
-"""SHRAM hybrid attention layer.
-
-This module implements the hybrid attention construction H(x) = h_l(x) + h_s(x)
-used at one decoder attention slot in SHRAM.
-
-The local sliding-window path and the MoSRAH sparse path are already verified
-independently. The responsibility here is therefore not to introduce new
-attention logic, but to preserve the bridge contracts between them: both paths
-must consume the same input hidden state, each path must receive the sub-cache
-it actually owns, the two model-space outputs must be summed directly, and the
-sparse-path load-balance loss must remain visible to the caller.
-"""
-
-import torch
-from torch import nn
-
-from .__cache__shram_layer_cache import ShramLayerCache
-from .configuration import ShramConfig
-from .__attention__sliding_window_attention import SlidingWindowAttention
-from .__attention__mosrah import MoSRAHLayer
-
-
-class SHRAMHybridLayer(nn.Module):
-    """Hybrid attention layer H(x) = h_l(x) + h_s(x) for one decoder slot.
-
-    The local path preserves nearby-token behavior through sliding-window causal
-    attention. The sparse path is the theorem-facing MoSRAH routed attention
-    path. Both operate over the same model-space hidden state and return
-    model-space outputs, so the hybrid composition is a direct sum in model
-    space.
-    """
-
-    def __init__(self, config: ShramConfig) -> None:
-        super().__init__()
-        self.local_attention = SlidingWindowAttention(config)
-        self.sparse_attention = MoSRAHLayer(config)
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        position_ids: torch.Tensor,
-        active_mask: torch.Tensor,
-        cache: ShramLayerCache | None,
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        """Apply the SHRAM hybrid attention layer.
-
-        Args:
-            hidden_states: Input hidden states of shape (B, N, d).
-            position_ids: Authoritative token positions of shape (B, N).
-            active_mask: Current-chunk active mask of shape (B, N), where True
-                means the token is semantically live. Forwarded unchanged to
-                both the local path and the sparse path.
-            cache: Optional per-layer SHRAM cache. When provided, the owned
-                sliding-window and MoSRAH sub-caches are dispatched directly to
-                their corresponding attention paths.
-
-        Returns:
-            hybrid_output: Model-space hybrid attention output of shape (B, N, d).
-            load_balance_loss: Scalar sparse-path load-balance loss.
-            max_vio: Detached scalar routing-imbalance summary. Passed through
-                unchanged from MoSRAHLayer; see MoSRAHRouter for semantics.
-        """
-        # ------------------------------------------------
-        # It is not possible, due to how bea constructs its block mask,
-        # for the model to process a sequence that does not start at zero
-        # without a cache to track the per-head offsets
-        # ------------------------------------------------
-
-        if cache is None and torch.any(position_ids[:, 0] != 0):
-            raise ValueError(
-                "Uncached SHRAMHybridLayer does not support nonzero starting positions. "
-                "Either provide a matching ShramLayerCache populated by the prefix for "
-                "continued decoding, or rebase the uncached sequence to start at 0."
-            )
-
-        # -------------------------------------------------------------------
-        # The hybrid layer's first responsibility is cache dispatch. The layer
-        # cache already owns the concrete sub-cache objects required by each
-        # path, so this unit should forward those exact references rather than
-        # reinterpret cache ownership or invent a composite update protocol here.
-        # -------------------------------------------------------------------
-        if cache is None:
-            sliding_window_cache = None
-            mosrah_cache = None
-        else:
-            sliding_window_cache = cache.sliding_window_cache
-            mosrah_cache = cache.mosrah_cache
-
-        # -------------------------------------------------------------------
-        # Both attention paths must see the same model-space hidden state for
-        # the current decoder layer. The local path preserves short-range
-        # structure, while the sparse path provides the routed long-range
-        # contribution and emits the load-balance signal used by training.
-        # -------------------------------------------------------------------
-        local_output = self.local_attention(
-            x=hidden_states,
-            position_ids=position_ids,
-            active_mask=active_mask,
-            cache=sliding_window_cache,
-        )
-        sparse_output, load_balance_loss, max_vio = self.sparse_attention(
-            hidden_states=hidden_states,
-            position_ids=position_ids,
-            active_mask=active_mask,
-            cache=mosrah_cache,
-        )
-
-        # -------------------------------------------------------------------
-        # The composition rule is intentionally simple at this boundary. Both
-        # sublayers already return model-space tensors of matching shape, so the
-        # correct hybrid behavior is their direct sum with no additional mixing
-        # logic introduced here.
-        # -------------------------------------------------------------------
-        hybrid_output = local_output + sparse_output
-
+"""SHRAM hybrid attention layer.
+
+This module implements the hybrid attention construction H(x) = h_l(x) + h_s(x)
+used at one decoder attention slot in SHRAM.
+
+The local sliding-window path and the MoSRAH sparse path are already verified
+independently. The responsibility here is therefore not to introduce new
+attention logic, but to preserve the bridge contracts between them: both paths
+must consume the same input hidden state, each path must receive the sub-cache
+it actually owns, the two model-space outputs must be summed directly, and the
+sparse-path load-balance loss must remain visible to the caller.
+"""
+
+import torch
+from torch import nn
+
+from .__cache__shram_layer_cache import ShramLayerCache
+from .configuration import ShramConfig
+from .__attention__sliding_window_attention import SlidingWindowAttention
+from .__attention__mosrah import MoSRAHLayer
+
+
+class SHRAMHybridLayer(nn.Module):
+    """Hybrid attention layer H(x) = h_l(x) + h_s(x) for one decoder slot.
+
+    The local path preserves nearby-token behavior through sliding-window causal
+    attention. The sparse path is the theorem-facing MoSRAH routed attention
+    path. Both operate over the same model-space hidden state and return
+    model-space outputs, so the hybrid composition is a direct sum in model
+    space.
+    """
+
+    def __init__(self, config: ShramConfig) -> None:
+        super().__init__()
+        self.local_attention = SlidingWindowAttention(config)
+        self.sparse_attention = MoSRAHLayer(config)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_ids: torch.Tensor,
+        active_mask: torch.Tensor,
+        cache: ShramLayerCache | None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Apply the SHRAM hybrid attention layer.
+
+        Args:
+            hidden_states: Input hidden states of shape (B, N, d).
+            position_ids: Authoritative token positions of shape (B, N).
+            active_mask: Current-chunk active mask of shape (B, N), where True
+                means the token is semantically live. Forwarded unchanged to
+                both the local path and the sparse path.
+            cache: Optional per-layer SHRAM cache. When provided, the owned
+                sliding-window and MoSRAH sub-caches are dispatched directly to
+                their corresponding attention paths.
+
+        Returns:
+            hybrid_output: Model-space hybrid attention output of shape (B, N, d).
+            load_balance_loss: Scalar sparse-path load-balance loss.
+            max_vio: Detached scalar routing-imbalance summary. Passed through
+                unchanged from MoSRAHLayer; see MoSRAHRouter for semantics.
+        """
+        # ------------------------------------------------
+        # It is not possible, due to how bea constructs its block mask,
+        # for the model to process a sequence that does not start at zero
+        # without a cache to track the per-head offsets
+        # ------------------------------------------------
+
+        if cache is None and torch.any(position_ids[:, 0] != 0):
+            raise ValueError(
+                "Uncached SHRAMHybridLayer does not support nonzero starting positions. "
+                "Either provide a matching ShramLayerCache populated by the prefix for "
+                "continued decoding, or rebase the uncached sequence to start at 0."
+            )
+
+        # -------------------------------------------------------------------
+        # The hybrid layer's first responsibility is cache dispatch. The layer
+        # cache already owns the concrete sub-cache objects required by each
+        # path, so this unit should forward those exact references rather than
+        # reinterpret cache ownership or invent a composite update protocol here.
+        # -------------------------------------------------------------------
+        if cache is None:
+            sliding_window_cache = None
+            mosrah_cache = None
+        else:
+            sliding_window_cache = cache.sliding_window_cache
+            mosrah_cache = cache.mosrah_cache
+
+        # -------------------------------------------------------------------
+        # Both attention paths must see the same model-space hidden state for
+        # the current decoder layer. The local path preserves short-range
+        # structure, while the sparse path provides the routed long-range
+        # contribution and emits the load-balance signal used by training.
+        # -------------------------------------------------------------------
+        local_output = self.local_attention(
+            x=hidden_states,
+            position_ids=position_ids,
+            active_mask=active_mask,
+            cache=sliding_window_cache,
+        )
+        sparse_output, load_balance_loss, max_vio = self.sparse_attention(
+            hidden_states=hidden_states,
+            position_ids=position_ids,
+            active_mask=active_mask,
+            cache=mosrah_cache,
+        )
+
+        # -------------------------------------------------------------------
+        # The composition rule is intentionally simple at this boundary. Both
+        # sublayers already return model-space tensors of matching shape, so the
+        # correct hybrid behavior is their direct sum with no additional mixing
+        # logic introduced here.
+        # -------------------------------------------------------------------
+        hybrid_output = local_output + sparse_output
+
         return hybrid_output, load_balance_loss, max_vio

__attention__sliding_window_attention.py

@@ -1,233 +1,233 @@
-# src/shram/model/attention/sliding_window_attention.py
-
-"""Local sliding-window attention path for SHRAM.
-
-This file defines `SlidingWindowAttention`, the local short-range attention path
-used inside the SHRAM hybrid layer.
-
-In the masked-continuation variant, the local cache no longer returns a
-semantically dense visible frame. Instead, `LocalSlidingWindowLayerCache`
-returns:
-
-- the retained local window memory concatenated with the current chunk
-- an aligned active mask over that returned frame
-
-This module consumes that returned frame directly and constructs effective local
-causal/window visibility from the mask. It does not own cache retention policy;
-it owns only local attention semantics.
-"""
-
-import math
-from typing import Any
-
-import torch
-import torch.nn as nn
-from torch.nn.attention.flex_attention import create_block_mask, flex_attention
-
-from .__cache__sliding_window_cache import LocalSlidingWindowLayerCache
-from .configuration import ShramConfig
-from .rope import RotaryEmbedding
-
-
-class SlidingWindowAttention(nn.Module):
-    """Causal local sliding-window attention for one SHRAM layer.
-
-    Args:
-        config: SHRAM config. Must expose `hidden_size`,
-            `num_sliding_window_heads`, `head_dim`, `window_size`,
-            `attention_dropout`, and `local_rope_theta`.
-
-    Raises:
-        NotImplementedError: If `attention_dropout != 0.0`.
-    """
-
-    def __init__(self, config: ShramConfig) -> None:
-        super().__init__()
-
-        self.hidden_size = config.hidden_size
-        self.num_heads = config.num_sliding_window_heads
-        self.head_dim = config.head_dim
-        self.window_size = config.window_size
-        self.attention_dropout = config.attention_dropout
-
-        if self.attention_dropout != 0.0:
-            raise NotImplementedError(
-                "SlidingWindowAttention currently supports only "
-                "attention_dropout == 0.0."
-            )
-
-        self.inner_dim = self.num_heads * self.head_dim
-
-        # Standard MHA projections for the local path.
-        self.q_proj = nn.Linear(self.hidden_size, self.inner_dim, bias=False)
-        self.k_proj = nn.Linear(self.hidden_size, self.inner_dim, bias=False)
-        self.v_proj = nn.Linear(self.hidden_size, self.inner_dim, bias=False)
-        self.o_proj = nn.Linear(self.inner_dim, self.hidden_size, bias=False)
-
-        # The local path always uses default-mode RoPE with its own theta.
-        self.rope = RotaryEmbedding(
-            mode="default",
-            head_dim=self.head_dim,
-            theta=config.local_rope_theta,
-        )
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        position_ids: torch.Tensor,
-        active_mask: torch.Tensor,
-        cache: LocalSlidingWindowLayerCache | None = None,
-    ) -> torch.Tensor:
-        """Apply local causal sliding-window attention.
-
-        Args:
-            x: Input tensor of shape `(B, N, hidden_size)`.
-            position_ids: Position tensor of shape `(B, N)`.
-            active_mask: Current-chunk active mask of shape `(B, N)`, where
-                `True` means active.
-            cache: Optional `LocalSlidingWindowLayerCache`.
-
-        Returns:
-            Output tensor of shape `(B, N, hidden_size)`.
-        """
-        batch_size, query_len, _ = x.shape
-
-        self._validate_position_shape(x, position_ids)
-        self._validate_active_mask_shape(x, active_mask)
-
-        # (B, N, H*D) -> (B, H, N, D)
-        q = self.q_proj(x).view(
-            batch_size,
-            query_len,
-            self.num_heads,
-            self.head_dim,
-        ).transpose(1, 2)
-        k = self.k_proj(x).view(
-            batch_size,
-            query_len,
-            self.num_heads,
-            self.head_dim,
-        ).transpose(1, 2)
-        v = self.v_proj(x).view(
-            batch_size,
-            query_len,
-            self.num_heads,
-            self.head_dim,
-        ).transpose(1, 2)
-
-        q, k, attention_scaling = self.rope(q, k, position_ids)
-
-        # The cache returns the current-step visible local frame, not merely the
-        # retained next-step cache buffer.
-        if cache is not None:
-            k_full, v_full, full_active_mask = cache.update(k, v, active_mask)
-        else:
-            k_full, v_full, full_active_mask = k, v, active_mask
-
-        block_mask = self._make_block_mask(
-            active_mask=full_active_mask,
-            batch_size=batch_size,
-            num_heads=self.num_heads,
-            query_len=query_len,
-            kv_len=k_full.shape[-2],
-            window_size=self.window_size,
-            device=x.device,
-        )
-
-        attn_output = flex_attention(
-            q,
-            k_full,
-            v_full,
-            block_mask=block_mask,
-            scale=attention_scaling / math.sqrt(self.head_dim),
-        )
-
-        # (B, H, N, D) -> (B, N, H*D) -> (B, N, hidden_size)
-        attn_output = (
-            attn_output.transpose(1, 2)
-            .contiguous()
-            .view(batch_size, query_len, self.inner_dim)
-        )
-
-        return self.o_proj(attn_output)
-
-    def _validate_position_shape(
-        self,
-        x: torch.Tensor,
-        position_ids: torch.Tensor,
-    ) -> None:
-        """Validate the position tensor shape expected by local RoPE."""
-        if position_ids.shape != x.shape[:2]:
-            raise ValueError(
-                f"position_ids must have shape {tuple(x.shape[:2])}, "
-                f"got {tuple(position_ids.shape)}."
-            )
-
-    def _validate_active_mask_shape(
-        self,
-        x: torch.Tensor,
-        active_mask: torch.Tensor,
-    ) -> None:
-        """Validate the current-chunk active-mask contract."""
-        if active_mask.shape != x.shape[:2]:
-            raise ValueError(
-                f"active_mask must have shape {tuple(x.shape[:2])}, "
-                f"got {tuple(active_mask.shape)}."
-            )
-        if active_mask.dtype != torch.bool:
-            raise ValueError(
-                f"active_mask must have dtype torch.bool, got {active_mask.dtype}."
-            )
-
-    def _make_block_mask(
-        self,
-        active_mask: torch.Tensor,
-        batch_size: int,
-        num_heads: int,
-        query_len: int,
-        kv_len: int,
-        window_size: int,
-        device: torch.device,
-    ) -> Any:
-        """Create the FlexAttention block mask for masked local continuation.
-
-        The returned local frame is chronological in raw buffer order, but dead
-        positions may remain inside it. Effective local order is therefore
-        recovered from the active mask itself by taking a cumulative count over
-        active positions.
-
-        Queries still occupy the tail of the returned frame, so raw buffer order
-        is used to locate query rows. Semantic active-token positions are then
-        used to decide causality and sliding-window distance.
-        """
-        query_offset = kv_len - query_len
-        semantic_positions = active_mask.long().cumsum(dim=-1) - 1
-
-        def sliding_window_mask(
-            batch_idx: torch.Tensor,
-            head_idx: torch.Tensor,
-            q_idx: torch.Tensor,
-            kv_idx: torch.Tensor,
-        ) -> torch.Tensor:
-
-            q_abs = query_offset + q_idx
-
-            query_is_active = active_mask[batch_idx, q_abs]
-            key_is_active = active_mask[batch_idx, kv_idx]
-
-            q_sem = semantic_positions[batch_idx, q_abs]
-            k_sem = semantic_positions[batch_idx, kv_idx]
-
-            is_causal = k_sem <= q_sem
-            in_window = (q_sem - k_sem) < window_size
-
-            return query_is_active & key_is_active & is_causal & in_window
-
-        return create_block_mask(
-            sliding_window_mask,
-            B=batch_size,
-            H=num_heads,
-            Q_LEN=query_len,
-            KV_LEN=kv_len,
-            device=device,
+# src/shram/model/attention/sliding_window_attention.py
+
+"""Local sliding-window attention path for SHRAM.
+
+This file defines `SlidingWindowAttention`, the local short-range attention path
+used inside the SHRAM hybrid layer.
+
+In the masked-continuation variant, the local cache no longer returns a
+semantically dense visible frame. Instead, `LocalSlidingWindowLayerCache`
+returns:
+
+- the retained local window memory concatenated with the current chunk
+- an aligned active mask over that returned frame
+
+This module consumes that returned frame directly and constructs effective local
+causal/window visibility from the mask. It does not own cache retention policy;
+it owns only local attention semantics.
+"""
+
+import math
+from typing import Any
+
+import torch
+import torch.nn as nn
+from torch.nn.attention.flex_attention import create_block_mask, flex_attention
+
+from .__cache__sliding_window_cache import LocalSlidingWindowLayerCache
+from .configuration import ShramConfig
+from .rope import RotaryEmbedding
+
+
+class SlidingWindowAttention(nn.Module):
+    """Causal local sliding-window attention for one SHRAM layer.
+
+    Args:
+        config: SHRAM config. Must expose `hidden_size`,
+            `num_sliding_window_heads`, `head_dim`, `window_size`,
+            `attention_dropout`, and `local_rope_theta`.
+
+    Raises:
+        NotImplementedError: If `attention_dropout != 0.0`.
+    """
+
+    def __init__(self, config: ShramConfig) -> None:
+        super().__init__()
+
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_sliding_window_heads
+        self.head_dim = config.head_dim
+        self.window_size = config.window_size
+        self.attention_dropout = config.attention_dropout
+
+        if self.attention_dropout != 0.0:
+            raise NotImplementedError(
+                "SlidingWindowAttention currently supports only "
+                "attention_dropout == 0.0."
+            )
+
+        self.inner_dim = self.num_heads * self.head_dim
+
+        # Standard MHA projections for the local path.
+        self.q_proj = nn.Linear(self.hidden_size, self.inner_dim, bias=False)
+        self.k_proj = nn.Linear(self.hidden_size, self.inner_dim, bias=False)
+        self.v_proj = nn.Linear(self.hidden_size, self.inner_dim, bias=False)
+        self.o_proj = nn.Linear(self.inner_dim, self.hidden_size, bias=False)
+
+        # The local path always uses default-mode RoPE with its own theta.
+        self.rope = RotaryEmbedding(
+            mode="default",
+            head_dim=self.head_dim,
+            theta=config.local_rope_theta,
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        position_ids: torch.Tensor,
+        active_mask: torch.Tensor,
+        cache: LocalSlidingWindowLayerCache | None = None,
+    ) -> torch.Tensor:
+        """Apply local causal sliding-window attention.
+
+        Args:
+            x: Input tensor of shape `(B, N, hidden_size)`.
+            position_ids: Position tensor of shape `(B, N)`.
+            active_mask: Current-chunk active mask of shape `(B, N)`, where
+                `True` means active.
+            cache: Optional `LocalSlidingWindowLayerCache`.
+
+        Returns:
+            Output tensor of shape `(B, N, hidden_size)`.
+        """
+        batch_size, query_len, _ = x.shape
+
+        self._validate_position_shape(x, position_ids)
+        self._validate_active_mask_shape(x, active_mask)
+
+        # (B, N, H*D) -> (B, H, N, D)
+        q = self.q_proj(x).view(
+            batch_size,
+            query_len,
+            self.num_heads,
+            self.head_dim,
+        ).transpose(1, 2)
+        k = self.k_proj(x).view(
+            batch_size,
+            query_len,
+            self.num_heads,
+            self.head_dim,
+        ).transpose(1, 2)
+        v = self.v_proj(x).view(
+            batch_size,
+            query_len,
+            self.num_heads,
+            self.head_dim,
+        ).transpose(1, 2)
+
+        q, k, attention_scaling = self.rope(q, k, position_ids)
+
+        # The cache returns the current-step visible local frame, not merely the
+        # retained next-step cache buffer.
+        if cache is not None:
+            k_full, v_full, full_active_mask = cache.update(k, v, active_mask)
+        else:
+            k_full, v_full, full_active_mask = k, v, active_mask
+
+        block_mask = self._make_block_mask(
+            active_mask=full_active_mask,
+            batch_size=batch_size,
+            num_heads=self.num_heads,
+            query_len=query_len,
+            kv_len=k_full.shape[-2],
+            window_size=self.window_size,
+            device=x.device,
+        )
+
+        attn_output = flex_attention(
+            q,
+            k_full,
+            v_full,
+            block_mask=block_mask,
+            scale=attention_scaling / math.sqrt(self.head_dim),
+        )
+
+        # (B, H, N, D) -> (B, N, H*D) -> (B, N, hidden_size)
+        attn_output = (
+            attn_output.transpose(1, 2)
+            .contiguous()
+            .view(batch_size, query_len, self.inner_dim)
+        )
+
+        return self.o_proj(attn_output)
+
+    def _validate_position_shape(
+        self,
+        x: torch.Tensor,
+        position_ids: torch.Tensor,
+    ) -> None:
+        """Validate the position tensor shape expected by local RoPE."""
+        if position_ids.shape != x.shape[:2]:
+            raise ValueError(
+                f"position_ids must have shape {tuple(x.shape[:2])}, "
+                f"got {tuple(position_ids.shape)}."
+            )
+
+    def _validate_active_mask_shape(
+        self,
+        x: torch.Tensor,
+        active_mask: torch.Tensor,
+    ) -> None:
+        """Validate the current-chunk active-mask contract."""
+        if active_mask.shape != x.shape[:2]:
+            raise ValueError(
+                f"active_mask must have shape {tuple(x.shape[:2])}, "
+                f"got {tuple(active_mask.shape)}."
+            )
+        if active_mask.dtype != torch.bool:
+            raise ValueError(
+                f"active_mask must have dtype torch.bool, got {active_mask.dtype}."
+            )
+
+    def _make_block_mask(
+        self,
+        active_mask: torch.Tensor,
+        batch_size: int,
+        num_heads: int,
+        query_len: int,
+        kv_len: int,
+        window_size: int,
+        device: torch.device,
+    ) -> Any:
+        """Create the FlexAttention block mask for masked local continuation.
+
+        The returned local frame is chronological in raw buffer order, but dead
+        positions may remain inside it. Effective local order is therefore
+        recovered from the active mask itself by taking a cumulative count over
+        active positions.
+
+        Queries still occupy the tail of the returned frame, so raw buffer order
+        is used to locate query rows. Semantic active-token positions are then
+        used to decide causality and sliding-window distance.
+        """
+        query_offset = kv_len - query_len
+        semantic_positions = active_mask.long().cumsum(dim=-1) - 1
+
+        def sliding_window_mask(
+            batch_idx: torch.Tensor,
+            head_idx: torch.Tensor,
+            q_idx: torch.Tensor,
+            kv_idx: torch.Tensor,
+        ) -> torch.Tensor:
+
+            q_abs = query_offset + q_idx
+
+            query_is_active = active_mask[batch_idx, q_abs]
+            key_is_active = active_mask[batch_idx, kv_idx]
+
+            q_sem = semantic_positions[batch_idx, q_abs]
+            k_sem = semantic_positions[batch_idx, kv_idx]
+
+            is_causal = k_sem <= q_sem
+            in_window = (q_sem - k_sem) < window_size
+
+            return query_is_active & key_is_active & is_causal & in_window
+
+        return create_block_mask(
+            sliding_window_mask,
+            B=batch_size,
+            H=num_heads,
+            Q_LEN=query_len,
+            KV_LEN=kv_len,
+            device=device,
         )

__cache__mosrah_cache.py

@@ -1,359 +1,359 @@
-"""MoSRAH sparse KV cache — single-layer implementation.
-
-MoSRAH routes each token to K of L available expert heads, so its KV cache is indexed
-by head rather than by sequence position. The routing is dynamic and produces a ragged
-distribution of token counts across (batch, head) slots — different batch items may
-route different numbers of tokens to the same head, and different heads accumulate at
-different rates. DynamicCache cannot represent this correctly: it concatenates along
-the sequence dimension and assumes uniform token counts across the batch. MoSRAHCache
-therefore uses a custom buffer design.
-
-Keys and values are stored in the CacheLayerMixin-standard self.keys and self.values
-attributes as (B, L, T, u) tensors, where B is batch size, L is the number of expert
-heads (num_mosrah_heads), T is the current buffer capacity, and u is the bottlenecked
-head embedding width (head_dim). A (B, L) integer count tensor _counts tracks the
-valid occupancy of each (batch, head) slot. Buffer capacity is exposed as the
-buffer_capacity property and is derived directly from self.keys rather than tracked
-as a separate variable.
-
-The primary interface is update(key_states, value_states, active_mask), which accepts
-expert-choice layout, stores only active entries in causal order, and returns the full
-accumulated (keys, values, active_mask) for immediate use by BEA. The returned
-active_mask identifies valid cached positions; everything beyond each slot's count is
-junk data that downstream attention must exclude.
-
-BEA applies RoPE and calls update() with post-RoPE keys (K̃). The occupancy counts
-exposed by get_heads_lengths() must be read before update() if the caller needs the
-pre-update occupancy for position computation (Unit 10.A). update() increments counts
-in-place and the pre-update values are not recoverable afterward.
-
-All buffers are allocated at construction time. MoSRAHCache is constructed by
-ShramLayerCache, which has access to batch size, device, and all model config parameters
-needed to fully specify the storage layout upfront.
-"""
-
-import torch
-from transformers.cache_utils import CacheLayerMixin
-
-
-class MoSRAHCache(CacheLayerMixin):
-    """KV cache for the MoSRAH sparse attention path — single decoder layer.
-
-    Subclasses CacheLayerMixin to satisfy the HuggingFace per-layer cache role.
-    Stores keys and values in the mixin-standard self.keys and self.values attributes
-    using a custom (B, L, T, u) layout rather than delegating to DynamicCache,
-    which cannot represent MoSRAH's ragged per-(batch, head) token counts correctly.
-
-    All storage is allocated at construction time and is_initialized is True
-    immediately. The caller (ShramLayerCache) provides batch size, device, and model
-    config parameters so no lazy allocation is needed.
-
-    Input is expected in expert-choice layout: (B, L, T, u) key/value tensors with a
-    (B, L, T) boolean active_mask. Only positions where active_mask is True are written.
-    This matches the packed representation produced by expert packing in the MoSRAH
-    forward pass, where BEA has already applied RoPE before calling update().
-
-    Args:
-        num_mosrah_heads: Total number of MoSRAH expert heads (L). Determines the
-            second dimension of all storage tensors.
-        head_dim: Bottlenecked head embedding width (u). Determines the fourth
-            dimension of all storage tensors.
-        batch_size: Number of sequences in the batch. Determines the first dimension
-            of all storage tensors.
-        device: Device on which to allocate all tensors. Should match the model device.
-        initial_buffer_size: Initial sequence capacity per (batch, head) slot. Doubled
-            when any slot overflows. Defaults to 64 to avoid repeated reallocation
-            during prompt processing.
-    """
-
-    is_compileable = False
-    is_sliding = False
-
-    def __init__(
-        self,
-        num_mosrah_heads: int,
-        head_dim: int,
-        batch_size: int,
-        device: torch.device,
-        initial_buffer_size: int = 64,
-    ) -> None:
-        super().__init__()
-        self.num_mosrah_heads = num_mosrah_heads
-        self.head_dim = head_dim
-        self.batch_size = batch_size
-        self.device = device
-
-        # Allocate primary storage into the mixin-standard self.keys / self.values so
-        # that inherited methods (offload, prefetch) operate on real tensors. _counts
-        # tracks valid occupancy per (batch, head) slot.
-        self.keys: torch.Tensor = torch.zeros(
-            batch_size, num_mosrah_heads, initial_buffer_size, head_dim, device=device
-        )
-        self.values: torch.Tensor = torch.zeros(
-            batch_size, num_mosrah_heads, initial_buffer_size, head_dim, device=device
-        )
-        self._counts: torch.Tensor = torch.zeros(
-            batch_size, num_mosrah_heads, dtype=torch.long, device=device
-        )
-
-        # Storage is fully allocated at construction — the cache is initialized.
-        self.is_initialized = True
-
-    # ---------------------------------------------------------------------------
-    # Properties
-    # ---------------------------------------------------------------------------
-
-    @property
-    def buffer_capacity(self) -> int:
-        """Current number of slots allocated per (batch, head) pair.
-
-        Derived directly from self.keys rather than tracked separately, so it is
-        always consistent with the actual buffer after expansion.
-        """
-        return self.keys.shape[2]
-
-    # ---------------------------------------------------------------------------
-    # Primary API
-    # ---------------------------------------------------------------------------
-
-    def update(  # type: ignore[override]
-        self,
-        key_states: torch.Tensor,
-        value_states: torch.Tensor,
-        active_mask: torch.Tensor,
-        cache_kwargs: dict | None = None,
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        """Scatter active key/value states into the buffer and return the full cache state.
-
-        Accepts expert-choice layout: key_states and value_states are (B, L, T, u);
-        active_mask is (B, L, T) bool with True marking real tokens. Only active
-        positions are written; inactive positions are ignored.
-
-        Uses a cumsum construction to derive the absolute buffer position for each
-        active token without any Python loops. For a given (batch, head) slot,
-        positions are assigned in the order tokens appear along the T dimension,
-        preserving causal ordering.
-
-        Returns the full accumulated (keys, values, active_mask) across the cached
-        sparse sequence. The returned active_mask is True exactly for slots t <
-        counts[b, l]; everything beyond is junk data that BEA must exclude.
-
-        Note: get_heads_lengths() must be called before update() if the caller needs
-        the pre-update occupancy for position computation (Unit 10.A). update()
-        increments counts in-place and the pre-update values are not recoverable.
-
-        Args:
-            key_states: Shape (B, L, T, u) — post-RoPE key vectors in expert-choice layout.
-            value_states: Shape (B, L, T, u) — value vectors in expert-choice layout.
-            active_mask: Shape (B, L, T) bool — True for real tokens, False for padding.
-            cache_kwargs: Unused; present to satisfy the CacheLayerMixin signature.
-
-        Returns:
-            Tuple of (keys, values, active_mask):
-              keys: (B, L, T, u) float — full key buffer including junk slots.
-              values: (B, L, T, u) float — full value buffer including junk slots.
-              active_mask: (B, L, T) bool — True iff slot (b, l, t) has been written.
-        """
-        incoming_delta = active_mask.long().sum(dim=2)  # (B, L)
-
-        if (self._counts + incoming_delta).max().item() > self.buffer_capacity:
-            self._expand()
-
-        # Cumulative count of active positions along T for each (b, l) slot. Entry
-        # [b, l, t] is the 1-based rank of position t among all active positions in
-        # that slot. Subtract 1 for a zero-based within-slot index. Inactive positions
-        # produce a negative value, which is excluded by the mask gate below.
-        within_slot = active_mask.long().cumsum(dim=2) - 1  # (B, L, T)
-
-        # Add the pre-update count to get the absolute buffer position for each
-        # active token.
-        abs_pos = within_slot + self._counts.unsqueeze(-1)  # (B, L, T)
-
-        # Scatter key and value vectors at all active positions.
-        b_idx, l_idx, t_idx = torch.where(active_mask)
-        self.keys[b_idx, l_idx, abs_pos[b_idx, l_idx, t_idx]] = (
-            key_states[b_idx, l_idx, t_idx]
-        )
-        self.values[b_idx, l_idx, abs_pos[b_idx, l_idx, t_idx]] = (
-            value_states[b_idx, l_idx, t_idx]
-        )
-
-        self._counts += incoming_delta
-
-        return self.keys, self.values, self._make_active_mask()
-
-    def get_heads_lengths(self) -> torch.Tensor:
-        """Return the per-(batch, head) token count for this layer.
-
-        This is the authoritative occupancy tensor consumed by BEA for attention
-        masking and by position computation (Unit 10.A) for semantic-sequence
-        position computation.
-
-        Note: in the MoSRAH forward pass, this must be called before update() if the
-        caller needs the pre-update occupancy. update() increments these counts in-place.
-
-        Returns:
-            Integer tensor of shape (B, L) where entry [b, h] is the number of valid
-            tokens stored in the (b, h) slot. Zero for slots with no writes yet.
-        """
-        return self._counts
-
-    # ---------------------------------------------------------------------------
-    # CacheLayerMixin — overridden coordination methods
-    # ---------------------------------------------------------------------------
-
-    def reset(self) -> None:
-        """Clear all cached key and value tensors.
-
-        Zeroes self.keys, self.values, and _counts in place. Storage remains allocated
-        and is_initialized remains True — only the contents are cleared.
-        """
-        self.keys.zero_()
-        self.values.zero_()
-        self._counts.zero_()
-
-    def reorder_cache(self, beam_idx: torch.LongTensor) -> None:
-        """Reorder the batch dimension of all cached tensors for beam search.
-
-        Applied atomically across self.keys, self.values, and _counts. Beam search
-        must reorder all three together or the occupancy counts and buffer contents
-        will correspond to different beam hypotheses.
-
-        Overrides the parent because the parent's implementation calls get_seq_length(),
-        which is not supported for this cache.
-
-        Args:
-            beam_idx: Permutation indices of shape (batch,) produced by the beam
-                search algorithm.
-        """
-        self.keys = self.keys[beam_idx]
-        self.values = self.values[beam_idx]
-        self._counts = self._counts[beam_idx]
-
-    def batch_repeat_interleave(self, repeats: int) -> None:
-        """Expand the batch dimension by repeating each entry repeats times.
-
-        Used at beam search initialisation to expand the cache from batch size B to
-        B * repeats, matching the expanded beam candidate batch. Applied atomically
-        across keys, values, and _counts; batch_size is updated to reflect the new size.
-
-        Args:
-            repeats: Number of times to repeat each batch entry.
-        """
-        self.keys = self.keys.repeat_interleave(repeats, dim=0)
-        self.values = self.values.repeat_interleave(repeats, dim=0)
-        self._counts = self._counts.repeat_interleave(repeats, dim=0)
-        self.batch_size = self.batch_size * repeats
-
-    def batch_select_indices(self, indices: torch.Tensor) -> None:
-        """Select a subset of batch entries by index.
-
-        Used in contrastive search to retain only the selected candidate entries.
-        Applied atomically across keys, values, and _counts; batch_size is updated
-        to reflect the number of retained entries.
-
-        Args:
-            indices: 1-D integer tensor of batch indices to retain.
-        """
-        self.keys = self.keys[indices]
-        self.values = self.values[indices]
-        self._counts = self._counts[indices]
-        self.batch_size = indices.shape[0]
-
-    def offload(self) -> None:
-        """Offload all cached tensors to CPU.
-
-        Extends the parent to also offload _counts, which the parent does not know
-        about. All three tensors are moved atomically so device state remains consistent.
-        """
-        super().offload()
-        self._counts = self._counts.to("cpu", non_blocking=True)
-
-    def prefetch(self) -> None:
-        """Move all cached tensors back to the model device ahead of time.
-
-        Extends the parent to also prefetch _counts, which the parent does not know
-        about. _counts is synced to self.keys.device after the parent moves keys and
-        values, so all three remain consistent.
-        """
-        super().prefetch()
-        if self._counts.device != self.keys.device:
-            self._counts = self._counts.to(self.keys.device, non_blocking=True)
-
-    def lazy_initialization(  # type: ignore[override]
-        self, key_states: torch.Tensor, value_states: torch.Tensor
-    ) -> None:
-        """No-op — storage is fully allocated at construction time."""
-        pass
-
-    # ---------------------------------------------------------------------------
-    # CacheLayerMixin — unsupported abstract methods
-    # ---------------------------------------------------------------------------
-
-    def get_seq_length(self) -> int:  # type: ignore[override]
-        """Not supported — no single sequence length represents this cache's state.
-
-        MoSRAH heads accumulate independently; (batch, head) slots have different
-        lengths depending on routing history. There is no meaningful scalar summary.
-        Use get_heads_lengths() for per-head occupancy.
-        """
-        raise NotImplementedError(
-            "MoSRAHCache has no single sequence length. "
-            "Use get_heads_lengths() for per-head occupancy."
-        )
-
-    def get_max_cache_shape(self) -> int:  # type: ignore[override]
-        """Not supported — MoSRAHCache is dynamic and unbounded."""
-        raise NotImplementedError(
-            "MoSRAHCache is unbounded; get_max_cache_shape() is not supported."
-        )
-
-    def get_mask_sizes(  # type: ignore[override]
-        self,
-        cache_position: torch.Tensor,
-    ) -> tuple[int, int]:
-        """Not supported — MoSRAHCache does not participate in HF mask construction."""
-        raise NotImplementedError(
-            "MoSRAHCache does not support get_mask_sizes()."
-        )
-
-    # ---------------------------------------------------------------------------
-    # Internal helpers
-    # ---------------------------------------------------------------------------
-
-    def _make_active_mask(self) -> torch.Tensor:
-        """Construct the (B, L, T) active mask from current counts.
-
-        Returns True at position [b, l, t] iff t < _counts[b, l], i.e. the slot
-        has been written. Positions at or beyond the count are junk and must be
-        excluded by downstream attention.
-        """
-        cap = self.buffer_capacity
-        return (
-            torch.arange(cap, device=self.keys.device)
-            .expand(self.batch_size, self.num_mosrah_heads, cap)
-            < self._counts.unsqueeze(-1)
-        )
-
-    def _expand(self) -> None:
-        """Double the buffer capacity, preserving existing data.
-
-        Called by update() when an incoming batch of tokens would cause any
-        (batch, head) slot to exceed the current buffer capacity. All existing
-        key and value data is copied into the low half of the new buffer; the
-        high half is zero-initialised and will be filled by subsequent writes.
-        After reassignment, buffer_capacity reflects the new size automatically.
-        """
-        old_cap = self.buffer_capacity
-        new_cap = old_cap * 2
-        dev = self.keys.device
-        new_keys = torch.zeros(
-            self.batch_size, self.num_mosrah_heads, new_cap, self.head_dim, device=dev
-        )
-        new_values = torch.zeros(
-            self.batch_size, self.num_mosrah_heads, new_cap, self.head_dim, device=dev
-        )
-        new_keys[:, :, :old_cap, :] = self.keys
-        new_values[:, :, :old_cap, :] = self.values
-        self.keys = new_keys
-        self.values = new_values
+"""MoSRAH sparse KV cache — single-layer implementation.
+
+MoSRAH routes each token to K of L available expert heads, so its KV cache is indexed
+by head rather than by sequence position. The routing is dynamic and produces a ragged
+distribution of token counts across (batch, head) slots — different batch items may
+route different numbers of tokens to the same head, and different heads accumulate at
+different rates. DynamicCache cannot represent this correctly: it concatenates along
+the sequence dimension and assumes uniform token counts across the batch. MoSRAHCache
+therefore uses a custom buffer design.
+
+Keys and values are stored in the CacheLayerMixin-standard self.keys and self.values
+attributes as (B, L, T, u) tensors, where B is batch size, L is the number of expert
+heads (num_mosrah_heads), T is the current buffer capacity, and u is the bottlenecked
+head embedding width (head_dim). A (B, L) integer count tensor _counts tracks the
+valid occupancy of each (batch, head) slot. Buffer capacity is exposed as the
+buffer_capacity property and is derived directly from self.keys rather than tracked
+as a separate variable.
+
+The primary interface is update(key_states, value_states, active_mask), which accepts
+expert-choice layout, stores only active entries in causal order, and returns the full
+accumulated (keys, values, active_mask) for immediate use by BEA. The returned
+active_mask identifies valid cached positions; everything beyond each slot's count is
+junk data that downstream attention must exclude.
+
+BEA applies RoPE and calls update() with post-RoPE keys (K̃). The occupancy counts
+exposed by get_heads_lengths() must be read before update() if the caller needs the
+pre-update occupancy for position computation (Unit 10.A). update() increments counts
+in-place and the pre-update values are not recoverable afterward.
+
+All buffers are allocated at construction time. MoSRAHCache is constructed by
+ShramLayerCache, which has access to batch size, device, and all model config parameters
+needed to fully specify the storage layout upfront.
+"""
+
+import torch
+from transformers.cache_utils import CacheLayerMixin
+
+
+class MoSRAHCache(CacheLayerMixin):
+    """KV cache for the MoSRAH sparse attention path — single decoder layer.
+
+    Subclasses CacheLayerMixin to satisfy the HuggingFace per-layer cache role.
+    Stores keys and values in the mixin-standard self.keys and self.values attributes
+    using a custom (B, L, T, u) layout rather than delegating to DynamicCache,
+    which cannot represent MoSRAH's ragged per-(batch, head) token counts correctly.
+
+    All storage is allocated at construction time and is_initialized is True
+    immediately. The caller (ShramLayerCache) provides batch size, device, and model
+    config parameters so no lazy allocation is needed.
+
+    Input is expected in expert-choice layout: (B, L, T, u) key/value tensors with a
+    (B, L, T) boolean active_mask. Only positions where active_mask is True are written.
+    This matches the packed representation produced by expert packing in the MoSRAH
+    forward pass, where BEA has already applied RoPE before calling update().
+
+    Args:
+        num_mosrah_heads: Total number of MoSRAH expert heads (L). Determines the
+            second dimension of all storage tensors.
+        head_dim: Bottlenecked head embedding width (u). Determines the fourth
+            dimension of all storage tensors.
+        batch_size: Number of sequences in the batch. Determines the first dimension
+            of all storage tensors.
+        device: Device on which to allocate all tensors. Should match the model device.
+        initial_buffer_size: Initial sequence capacity per (batch, head) slot. Doubled
+            when any slot overflows. Defaults to 64 to avoid repeated reallocation
+            during prompt processing.
+    """
+
+    is_compileable = False
+    is_sliding = False
+
+    def __init__(
+        self,
+        num_mosrah_heads: int,
+        head_dim: int,
+        batch_size: int,
+        device: torch.device,
+        initial_buffer_size: int = 64,
+    ) -> None:
+        super().__init__()
+        self.num_mosrah_heads = num_mosrah_heads
+        self.head_dim = head_dim
+        self.batch_size = batch_size
+        self.device = device
+
+        # Allocate primary storage into the mixin-standard self.keys / self.values so
+        # that inherited methods (offload, prefetch) operate on real tensors. _counts
+        # tracks valid occupancy per (batch, head) slot.
+        self.keys: torch.Tensor = torch.zeros(
+            batch_size, num_mosrah_heads, initial_buffer_size, head_dim, device=device
+        )
+        self.values: torch.Tensor = torch.zeros(
+            batch_size, num_mosrah_heads, initial_buffer_size, head_dim, device=device
+        )
+        self._counts: torch.Tensor = torch.zeros(
+            batch_size, num_mosrah_heads, dtype=torch.long, device=device
+        )
+
+        # Storage is fully allocated at construction — the cache is initialized.
+        self.is_initialized = True
+
+    # ---------------------------------------------------------------------------
+    # Properties
+    # ---------------------------------------------------------------------------
+
+    @property
+    def buffer_capacity(self) -> int:
+        """Current number of slots allocated per (batch, head) pair.
+
+        Derived directly from self.keys rather than tracked separately, so it is
+        always consistent with the actual buffer after expansion.
+        """
+        return self.keys.shape[2]
+
+    # ---------------------------------------------------------------------------
+    # Primary API
+    # ---------------------------------------------------------------------------
+
+    def update(  # type: ignore[override]
+        self,
+        key_states: torch.Tensor,
+        value_states: torch.Tensor,
+        active_mask: torch.Tensor,
+        cache_kwargs: dict | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Scatter active key/value states into the buffer and return the full cache state.
+
+        Accepts expert-choice layout: key_states and value_states are (B, L, T, u);
+        active_mask is (B, L, T) bool with True marking real tokens. Only active
+        positions are written; inactive positions are ignored.
+
+        Uses a cumsum construction to derive the absolute buffer position for each
+        active token without any Python loops. For a given (batch, head) slot,
+        positions are assigned in the order tokens appear along the T dimension,
+        preserving causal ordering.
+
+        Returns the full accumulated (keys, values, active_mask) across the cached
+        sparse sequence. The returned active_mask is True exactly for slots t <
+        counts[b, l]; everything beyond is junk data that BEA must exclude.
+
+        Note: get_heads_lengths() must be called before update() if the caller needs
+        the pre-update occupancy for position computation (Unit 10.A). update()
+        increments counts in-place and the pre-update values are not recoverable.
+
+        Args:
+            key_states: Shape (B, L, T, u) — post-RoPE key vectors in expert-choice layout.
+            value_states: Shape (B, L, T, u) — value vectors in expert-choice layout.
+            active_mask: Shape (B, L, T) bool — True for real tokens, False for padding.
+            cache_kwargs: Unused; present to satisfy the CacheLayerMixin signature.
+
+        Returns:
+            Tuple of (keys, values, active_mask):
+              keys: (B, L, T, u) float — full key buffer including junk slots.
+              values: (B, L, T, u) float — full value buffer including junk slots.
+              active_mask: (B, L, T) bool — True iff slot (b, l, t) has been written.
+        """
+        incoming_delta = active_mask.long().sum(dim=2)  # (B, L)
+
+        if (self._counts + incoming_delta).max().item() > self.buffer_capacity:
+            self._expand()
+
+        # Cumulative count of active positions along T for each (b, l) slot. Entry
+        # [b, l, t] is the 1-based rank of position t among all active positions in
+        # that slot. Subtract 1 for a zero-based within-slot index. Inactive positions
+        # produce a negative value, which is excluded by the mask gate below.
+        within_slot = active_mask.long().cumsum(dim=2) - 1  # (B, L, T)
+
+        # Add the pre-update count to get the absolute buffer position for each
+        # active token.
+        abs_pos = within_slot + self._counts.unsqueeze(-1)  # (B, L, T)
+
+        # Scatter key and value vectors at all active positions.
+        b_idx, l_idx, t_idx = torch.where(active_mask)
+        self.keys[b_idx, l_idx, abs_pos[b_idx, l_idx, t_idx]] = (
+            key_states[b_idx, l_idx, t_idx]
+        )
+        self.values[b_idx, l_idx, abs_pos[b_idx, l_idx, t_idx]] = (
+            value_states[b_idx, l_idx, t_idx]
+        )
+
+        self._counts += incoming_delta
+
+        return self.keys, self.values, self._make_active_mask()
+
+    def get_heads_lengths(self) -> torch.Tensor:
+        """Return the per-(batch, head) token count for this layer.
+
+        This is the authoritative occupancy tensor consumed by BEA for attention
+        masking and by position computation (Unit 10.A) for semantic-sequence
+        position computation.
+
+        Note: in the MoSRAH forward pass, this must be called before update() if the
+        caller needs the pre-update occupancy. update() increments these counts in-place.
+
+        Returns:
+            Integer tensor of shape (B, L) where entry [b, h] is the number of valid
+            tokens stored in the (b, h) slot. Zero for slots with no writes yet.
+        """
+        return self._counts
+
+    # ---------------------------------------------------------------------------
+    # CacheLayerMixin — overridden coordination methods
+    # ---------------------------------------------------------------------------
+
+    def reset(self) -> None:
+        """Clear all cached key and value tensors.
+
+        Zeroes self.keys, self.values, and _counts in place. Storage remains allocated
+        and is_initialized remains True — only the contents are cleared.
+        """
+        self.keys.zero_()
+        self.values.zero_()
+        self._counts.zero_()
+
+    def reorder_cache(self, beam_idx: torch.LongTensor) -> None:
+        """Reorder the batch dimension of all cached tensors for beam search.
+
+        Applied atomically across self.keys, self.values, and _counts. Beam search
+        must reorder all three together or the occupancy counts and buffer contents
+        will correspond to different beam hypotheses.
+
+        Overrides the parent because the parent's implementation calls get_seq_length(),
+        which is not supported for this cache.
+
+        Args:
+            beam_idx: Permutation indices of shape (batch,) produced by the beam
+                search algorithm.
+        """
+        self.keys = self.keys[beam_idx]
+        self.values = self.values[beam_idx]
+        self._counts = self._counts[beam_idx]
+
+    def batch_repeat_interleave(self, repeats: int) -> None:
+        """Expand the batch dimension by repeating each entry repeats times.
+
+        Used at beam search initialisation to expand the cache from batch size B to
+        B * repeats, matching the expanded beam candidate batch. Applied atomically
+        across keys, values, and _counts; batch_size is updated to reflect the new size.
+
+        Args:
+            repeats: Number of times to repeat each batch entry.
+        """
+        self.keys = self.keys.repeat_interleave(repeats, dim=0)
+        self.values = self.values.repeat_interleave(repeats, dim=0)
+        self._counts = self._counts.repeat_interleave(repeats, dim=0)
+        self.batch_size = self.batch_size * repeats
+
+    def batch_select_indices(self, indices: torch.Tensor) -> None:
+        """Select a subset of batch entries by index.
+
+        Used in contrastive search to retain only the selected candidate entries.
+        Applied atomically across keys, values, and _counts; batch_size is updated
+        to reflect the number of retained entries.
+
+        Args:
+            indices: 1-D integer tensor of batch indices to retain.
+        """
+        self.keys = self.keys[indices]
+        self.values = self.values[indices]
+        self._counts = self._counts[indices]
+        self.batch_size = indices.shape[0]
+
+    def offload(self) -> None:
+        """Offload all cached tensors to CPU.
+
+        Extends the parent to also offload _counts, which the parent does not know
+        about. All three tensors are moved atomically so device state remains consistent.
+        """
+        super().offload()
+        self._counts = self._counts.to("cpu", non_blocking=True)
+
+    def prefetch(self) -> None:
+        """Move all cached tensors back to the model device ahead of time.
+
+        Extends the parent to also prefetch _counts, which the parent does not know
+        about. _counts is synced to self.keys.device after the parent moves keys and
+        values, so all three remain consistent.
+        """
+        super().prefetch()
+        if self._counts.device != self.keys.device:
+            self._counts = self._counts.to(self.keys.device, non_blocking=True)
+
+    def lazy_initialization(  # type: ignore[override]
+        self, key_states: torch.Tensor, value_states: torch.Tensor
+    ) -> None:
+        """No-op — storage is fully allocated at construction time."""
+        pass
+
+    # ---------------------------------------------------------------------------
+    # CacheLayerMixin — unsupported abstract methods
+    # ---------------------------------------------------------------------------
+
+    def get_seq_length(self) -> int:  # type: ignore[override]
+        """Not supported — no single sequence length represents this cache's state.
+
+        MoSRAH heads accumulate independently; (batch, head) slots have different
+        lengths depending on routing history. There is no meaningful scalar summary.
+        Use get_heads_lengths() for per-head occupancy.
+        """
+        raise NotImplementedError(
+            "MoSRAHCache has no single sequence length. "
+            "Use get_heads_lengths() for per-head occupancy."
+        )
+
+    def get_max_cache_shape(self) -> int:  # type: ignore[override]
+        """Not supported — MoSRAHCache is dynamic and unbounded."""
+        raise NotImplementedError(
+            "MoSRAHCache is unbounded; get_max_cache_shape() is not supported."
+        )
+
+    def get_mask_sizes(  # type: ignore[override]
+        self,
+        cache_position: torch.Tensor,
+    ) -> tuple[int, int]:
+        """Not supported — MoSRAHCache does not participate in HF mask construction."""
+        raise NotImplementedError(
+            "MoSRAHCache does not support get_mask_sizes()."
+        )
+
+    # ---------------------------------------------------------------------------
+    # Internal helpers
+    # ---------------------------------------------------------------------------
+
+    def _make_active_mask(self) -> torch.Tensor:
+        """Construct the (B, L, T) active mask from current counts.
+
+        Returns True at position [b, l, t] iff t < _counts[b, l], i.e. the slot
+        has been written. Positions at or beyond the count are junk and must be
+        excluded by downstream attention.
+        """
+        cap = self.buffer_capacity
+        return (
+            torch.arange(cap, device=self.keys.device)
+            .expand(self.batch_size, self.num_mosrah_heads, cap)
+            < self._counts.unsqueeze(-1)
+        )
+
+    def _expand(self) -> None:
+        """Double the buffer capacity, preserving existing data.
+
+        Called by update() when an incoming batch of tokens would cause any
+        (batch, head) slot to exceed the current buffer capacity. All existing
+        key and value data is copied into the low half of the new buffer; the
+        high half is zero-initialised and will be filled by subsequent writes.
+        After reassignment, buffer_capacity reflects the new size automatically.
+        """
+        old_cap = self.buffer_capacity
+        new_cap = old_cap * 2
+        dev = self.keys.device
+        new_keys = torch.zeros(
+            self.batch_size, self.num_mosrah_heads, new_cap, self.head_dim, device=dev
+        )
+        new_values = torch.zeros(
+            self.batch_size, self.num_mosrah_heads, new_cap, self.head_dim, device=dev
+        )
+        new_keys[:, :, :old_cap, :] = self.keys
+        new_values[:, :, :old_cap, :] = self.values
+        self.keys = new_keys
+        self.values = new_values

__cache__shram_cache.py

@@ -1,141 +1,141 @@
-"""SHRAM top-level cache — model-wide owner for the full SHRAM decoder stack.
-
-The HuggingFace Cache protocol expects a single top-level Cache object that owns one
-CacheLayerMixin per decoder layer. The actual SHRAM caching responsibilities live one level
-lower in ShramLayerCache — each of which owns a LocalSlidingWindowLayerCache and a MoSRAHCache.
-ShramCache bridges those two levels: it constructs one ShramLayerCache per decoder layer,
-presents them through the Cache interface, and transparently forwards model-wide operations
-across all of them.
-
-ShramCache does not define a composite update() interface. The two attention paths inside each
-SHRAM layer have different update semantics, and neither the layer-level boundary (Unit 6.B)
-nor the model-level boundary here can meaningfully unify them. Callers must reach down to the
-relevant sub-cache directly. ShramCache's role is ownership, construction, and model-wide
-coordination of the layer caches — not routing attention inputs.
-
-Sequence length is reported by delegating to the local sliding-window sub-cache of the
-specified layer, which tracks the cumulative count of token positions processed. This is
-what HuggingFace generation reads through get_seq_length().
-"""
-
-import torch
-from transformers.cache_utils import Cache
-
-from .__cache__shram_layer_cache import ShramLayerCache
-
-
-class ShramCache(Cache):
-    """Top-level cache for the full SHRAM model.
-
-    Owns one ShramLayerCache per decoder layer. Satisfies the HuggingFace top-level Cache
-    role and transparently forwards reset, reorder, and sequence-length queries across all
-    owned layer caches.
-
-    No composite update() interface is provided. The two attention paths inside each SHRAM
-    layer have materially different update semantics; callers must update sub-caches directly
-    via cache.layers[layer_idx].sliding_window_cache or cache.layers[layer_idx].mosrah_cache.
-
-    Args:
-        num_hidden_layers: Number of SHRAM decoder layers. Determines how many
-            ShramLayerCache objects are constructed.
-        sliding_window: Token window size passed to each layer's LocalSlidingWindowLayerCache.
-        num_local_heads: Number of local attention heads per layer.
-        local_head_dim: Per-head embedding width for the local path.
-        num_mosrah_heads: Total number of MoSRAH expert heads (L) per layer.
-        mosrah_head_dim: Bottlenecked head embedding width (u) for the MoSRAH path.
-        batch_size: Number of sequences in the batch.
-        device: Device on which to allocate cache tensors.
-        initial_buffer_size: Initial per-(batch, head) capacity for each MoSRAHCache.
-            Doubled when any slot overflows. Defaults to 64 to avoid repeated reallocation
-            during prompt processing.
-    """
-
-    def __init__(
-        self,
-        num_hidden_layers: int,
-        sliding_window: int,
-        num_local_heads: int,
-        local_head_dim: int,
-        num_mosrah_heads: int,
-        mosrah_head_dim: int,
-        batch_size: int,
-        device: torch.device,
-        initial_buffer_size: int = 64,
-    ) -> None:
-        layers = [
-            ShramLayerCache(
-                sliding_window=sliding_window,
-                num_local_heads=num_local_heads,
-                local_head_dim=local_head_dim,
-                num_mosrah_heads=num_mosrah_heads,
-                mosrah_head_dim=mosrah_head_dim,
-                batch_size=batch_size,
-                device=device,
-                initial_buffer_size=initial_buffer_size,
-            )
-            for _ in range(num_hidden_layers)
-        ]
-        super().__init__(layers=layers)
-
-    # ---------------------------------------------------------------------------
-    # Cache — composite-meaningful methods
-    # ---------------------------------------------------------------------------
-    #
-    # reset(): Inherited. Iterates all layer caches and calls reset() on each.
-    #
-    # reorder_cache(beam_idx): Inherited. Iterates all layer caches and reorders each.
-    #
-    # is_initialized: Inherited property. True iff all layer caches are initialized.
-    #   Since ShramLayerCache.is_initialized is True from construction, this is True
-    #   immediately after ShramCache.__init__ returns.
-
-    def get_seq_length(self, layer_idx: int = 0) -> int:  # type: ignore[override]
-        """Return the cumulative sequence length for the specified layer.
-
-        Delegates to the layer cache at layer_idx, which in turn delegates to the
-        local sliding-window sub-cache. That sub-cache is authoritative for sequence
-        progress: it sees every token presented to the layer and accumulates a truthful
-        total count. Defaults to layer 0, which is sufficient for HuggingFace generation.
-        """
-        return self.layers[layer_idx].get_seq_length()
-
-    # ---------------------------------------------------------------------------
-    # Cache — unsupported methods
-    # ---------------------------------------------------------------------------
-
-    def update(  # type: ignore[override]
-        self,
-        key_states: torch.Tensor,
-        value_states: torch.Tensor,
-        layer_idx: int,
-        cache_kwargs: dict | None = None,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        """Not supported — ShramCache has no composite update interface.
-
-        The two attention paths inside each SHRAM layer have different update semantics.
-        Callers must update sub-caches directly:
-          cache.layers[layer_idx].sliding_window_cache.update(key_states, value_states)
-          cache.layers[layer_idx].mosrah_cache.update(key_states, value_states, active_mask)
-        """
-        raise NotImplementedError(
-            "ShramCache has no composite update interface. "
-            "Update sliding_window_cache or mosrah_cache on the relevant layer directly."
-        )
-
-    def crop(self, max_length: int) -> None:
-        """Not supported — ShramCache layers do not implement crop()."""
-        raise NotImplementedError("ShramCache does not support crop().")
-
-    @property
-    def max_batch_size(self) -> int:
-        """Not supported — ShramCache does not track a uniform batch size across layers."""
-        raise NotImplementedError("ShramCache does not expose max_batch_size.")
-
-    @property
-    def max_cache_len(self) -> int:
-        """Not supported — ShramCache has no single maximum cache length.
-
-        The sliding-window side is bounded by sliding_window; the MoSRAH side is unbounded.
-        No truthful scalar maximum represents the composite.
-        """
-        raise NotImplementedError("ShramCache does not expose max_cache_len.")
+"""SHRAM top-level cache — model-wide owner for the full SHRAM decoder stack.
+
+The HuggingFace Cache protocol expects a single top-level Cache object that owns one
+CacheLayerMixin per decoder layer. The actual SHRAM caching responsibilities live one level
+lower in ShramLayerCache — each of which owns a LocalSlidingWindowLayerCache and a MoSRAHCache.
+ShramCache bridges those two levels: it constructs one ShramLayerCache per decoder layer,
+presents them through the Cache interface, and transparently forwards model-wide operations
+across all of them.
+
+ShramCache does not define a composite update() interface. The two attention paths inside each
+SHRAM layer have different update semantics, and neither the layer-level boundary (Unit 6.B)
+nor the model-level boundary here can meaningfully unify them. Callers must reach down to the
+relevant sub-cache directly. ShramCache's role is ownership, construction, and model-wide
+coordination of the layer caches — not routing attention inputs.
+
+Sequence length is reported by delegating to the local sliding-window sub-cache of the
+specified layer, which tracks the cumulative count of token positions processed. This is
+what HuggingFace generation reads through get_seq_length().
+"""
+
+import torch
+from transformers.cache_utils import Cache
+
+from .__cache__shram_layer_cache import ShramLayerCache
+
+
+class ShramCache(Cache):
+    """Top-level cache for the full SHRAM model.
+
+    Owns one ShramLayerCache per decoder layer. Satisfies the HuggingFace top-level Cache
+    role and transparently forwards reset, reorder, and sequence-length queries across all
+    owned layer caches.
+
+    No composite update() interface is provided. The two attention paths inside each SHRAM
+    layer have materially different update semantics; callers must update sub-caches directly
+    via cache.layers[layer_idx].sliding_window_cache or cache.layers[layer_idx].mosrah_cache.
+
+    Args:
+        num_hidden_layers: Number of SHRAM decoder layers. Determines how many
+            ShramLayerCache objects are constructed.
+        sliding_window: Token window size passed to each layer's LocalSlidingWindowLayerCache.
+        num_local_heads: Number of local attention heads per layer.
+        local_head_dim: Per-head embedding width for the local path.
+        num_mosrah_heads: Total number of MoSRAH expert heads (L) per layer.
+        mosrah_head_dim: Bottlenecked head embedding width (u) for the MoSRAH path.
+        batch_size: Number of sequences in the batch.
+        device: Device on which to allocate cache tensors.
+        initial_buffer_size: Initial per-(batch, head) capacity for each MoSRAHCache.
+            Doubled when any slot overflows. Defaults to 64 to avoid repeated reallocation
+            during prompt processing.
+    """
+
+    def __init__(
+        self,
+        num_hidden_layers: int,
+        sliding_window: int,
+        num_local_heads: int,
+        local_head_dim: int,
+        num_mosrah_heads: int,
+        mosrah_head_dim: int,
+        batch_size: int,
+        device: torch.device,
+        initial_buffer_size: int = 64,
+    ) -> None:
+        layers = [
+            ShramLayerCache(
+                sliding_window=sliding_window,
+                num_local_heads=num_local_heads,
+                local_head_dim=local_head_dim,
+                num_mosrah_heads=num_mosrah_heads,
+                mosrah_head_dim=mosrah_head_dim,
+                batch_size=batch_size,
+                device=device,
+                initial_buffer_size=initial_buffer_size,
+            )
+            for _ in range(num_hidden_layers)
+        ]
+        super().__init__(layers=layers)
+
+    # ---------------------------------------------------------------------------
+    # Cache — composite-meaningful methods
+    # ---------------------------------------------------------------------------
+    #
+    # reset(): Inherited. Iterates all layer caches and calls reset() on each.
+    #
+    # reorder_cache(beam_idx): Inherited. Iterates all layer caches and reorders each.
+    #
+    # is_initialized: Inherited property. True iff all layer caches are initialized.
+    #   Since ShramLayerCache.is_initialized is True from construction, this is True
+    #   immediately after ShramCache.__init__ returns.
+
+    def get_seq_length(self, layer_idx: int = 0) -> int:  # type: ignore[override]
+        """Return the cumulative sequence length for the specified layer.
+
+        Delegates to the layer cache at layer_idx, which in turn delegates to the
+        local sliding-window sub-cache. That sub-cache is authoritative for sequence
+        progress: it sees every token presented to the layer and accumulates a truthful
+        total count. Defaults to layer 0, which is sufficient for HuggingFace generation.
+        """
+        return self.layers[layer_idx].get_seq_length()
+
+    # ---------------------------------------------------------------------------
+    # Cache — unsupported methods
+    # ---------------------------------------------------------------------------
+
+    def update(  # type: ignore[override]
+        self,
+        key_states: torch.Tensor,
+        value_states: torch.Tensor,
+        layer_idx: int,
+        cache_kwargs: dict | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Not supported — ShramCache has no composite update interface.
+
+        The two attention paths inside each SHRAM layer have different update semantics.
+        Callers must update sub-caches directly:
+          cache.layers[layer_idx].sliding_window_cache.update(key_states, value_states)
+          cache.layers[layer_idx].mosrah_cache.update(key_states, value_states, active_mask)
+        """
+        raise NotImplementedError(
+            "ShramCache has no composite update interface. "
+            "Update sliding_window_cache or mosrah_cache on the relevant layer directly."
+        )
+
+    def crop(self, max_length: int) -> None:
+        """Not supported — ShramCache layers do not implement crop()."""
+        raise NotImplementedError("ShramCache does not support crop().")
+
+    @property
+    def max_batch_size(self) -> int:
+        """Not supported — ShramCache does not track a uniform batch size across layers."""
+        raise NotImplementedError("ShramCache does not expose max_batch_size.")
+
+    @property
+    def max_cache_len(self) -> int:
+        """Not supported — ShramCache has no single maximum cache length.
+
+        The sliding-window side is bounded by sliding_window; the MoSRAH side is unbounded.
+        No truthful scalar maximum represents the composite.
+        """
+        raise NotImplementedError("ShramCache does not expose max_cache_len.")

__cache__shram_layer_cache.py

@@ -1,233 +1,233 @@
-"""SHRAM per-layer cache — composite owner for one SHRAM decoder layer.
-
-A SHRAM decoder layer contains two distinct attention pathways at one attention slot: the
-local sliding-window path and the MoSRAH sparse path. Each path has its own cache with
-different semantics and a different downstream consumer. ShramLayerCache owns both, satisfies
-the HuggingFace per-layer cache role, and exposes each sub-cache directly so its attention
-path can interact with it without indirection.
-
-ShramLayerCache does not define a composite update() interface. The two paths have materially
-different update semantics — the local side uses chunk-local key/value/mask concatenation
-while the MoSRAH side uses expert-choice scatter with an active mask — and merging these
-behind a single update() would hide those differences behind a misleading abstraction. Instead,
-each attention path calls update() on the sub-cache it owns. ShramLayerCache acts as the
-ownership, coordination, and reset/reorder boundary for one decoder layer.
-
-Sequence length at this boundary is reported by delegating to the local sliding-window
-sub-cache, which tracks the cumulative count of token positions processed. This is the
-quantity HuggingFace generation reads through get_seq_length().
-"""
-
-import torch
-from transformers.cache_utils import CacheLayerMixin
-
-from .__cache__mosrah_cache import MoSRAHCache
-from .__cache__sliding_window_cache import LocalSlidingWindowLayerCache
-
-
-class ShramLayerCache(CacheLayerMixin):
-    """Cache subsystem for one SHRAM decoder layer.
-
-    Owns and coordinates two sub-caches:
-      - sliding_window_cache: LocalSlidingWindowLayerCache for the local sliding-window path.
-      - mosrah_cache: MoSRAHCache for the MoSRAH sparse attention path.
-
-    Satisfies the HuggingFace per-layer cache role (CacheLayerMixin). The two sub-caches are
-    exposed directly for their downstream attention paths — no composite update() interface is
-    provided, because the two paths have materially different update semantics.
-
-    Sequence length is reported by delegating to the local sliding-window sub-cache, which
-    tracks the cumulative count of token positions processed across all update() calls.
-
-    Args:
-        sliding_window: Number of tokens retained by the local sliding-window cache.
-        num_local_heads: Number of local attention heads.
-        local_head_dim: Per-head embedding width for the local path.
-        num_mosrah_heads: Total number of MoSRAH expert heads (L).
-        mosrah_head_dim: Bottlenecked head embedding width (u) for the MoSRAH path.
-        batch_size: Number of sequences in the batch.
-        device: Device on which to allocate cache tensors.
-        initial_buffer_size: Initial per-(batch, head) capacity for MoSRAHCache. Doubled
-            when any slot overflows. Defaults to 64 to avoid repeated reallocation during
-            prompt processing.
-    """
-
-    is_compileable = False
-    is_sliding = False
-
-    def __init__(
-        self,
-        sliding_window: int,
-        num_local_heads: int,
-        local_head_dim: int,
-        num_mosrah_heads: int,
-        mosrah_head_dim: int,
-        batch_size: int,
-        device: torch.device,
-        initial_buffer_size: int = 64,
-    ) -> None:
-        super().__init__()
-        self.sliding_window_cache = LocalSlidingWindowLayerCache(
-            sliding_window=sliding_window,
-            num_heads=num_local_heads,
-            head_dim=local_head_dim,
-            batch_size=batch_size,
-            device=device,
-        )
-        self.mosrah_cache = MoSRAHCache(
-            num_mosrah_heads=num_mosrah_heads,
-            head_dim=mosrah_head_dim,
-            batch_size=batch_size,
-            device=device,
-            initial_buffer_size=initial_buffer_size,
-        )
-
-    # ---------------------------------------------------------------------------
-    # Properties
-    # ---------------------------------------------------------------------------
-
-    @property
-    def is_initialized(self) -> bool:
-        """True iff both sub-caches have allocated their storage.
-
-        Both LocalSlidingWindowLayerCache and MoSRAHCache pre-allocate at construction,
-        so this is True immediately after ShramLayerCache.__init__ returns.
-        """
-        return self.sliding_window_cache.is_initialized and self.mosrah_cache.is_initialized
-
-    @is_initialized.setter
-    def is_initialized(self, value: bool) -> None:
-        # CacheLayerMixin.__init__ assigns self.is_initialized = False as an instance
-        # attribute. Since property is a data descriptor it takes precedence, but Python
-        # still routes the assignment through __set__. Absorb it silently — state is
-        # derived from sub-caches, not stored here.
-        pass
-
-    # ---------------------------------------------------------------------------
-    # CacheLayerMixin — composite-meaningful methods
-    # ---------------------------------------------------------------------------
-
-    def get_seq_length(self) -> int:  # type: ignore[override]
-        """Return the cumulative sequence length from the local sliding-window path.
-
-        The local path is authoritative for sequence progress: it sees every token
-        presented to this layer and accumulates a truthful total. Delegates to
-        sliding_window_cache.get_seq_length().
-        """
-        return self.sliding_window_cache.get_seq_length()
-
-    def reset(self) -> None:
-        """Clear both sub-caches.
-
-        Delegates reset to each sub-cache. Both are cleared atomically so the sliding-window
-        state and MoSRAH sparse state remain consistent.
-        """
-        self.sliding_window_cache.reset()
-        self.mosrah_cache.reset()
-
-    def reorder_cache(self, beam_idx: torch.LongTensor) -> None:
-        """Reorder the batch dimension of both sub-caches for beam search.
-
-        Delegates to each sub-cache. Both are reordered atomically so the sliding-window
-        and MoSRAH state correspond to the same beam hypotheses after reordering.
-
-        Args:
-            beam_idx: Permutation indices of shape (batch,) produced by beam search.
-        """
-        self.sliding_window_cache.reorder_cache(beam_idx)
-        self.mosrah_cache.reorder_cache(beam_idx)
-
-    def batch_repeat_interleave(self, repeats: int) -> None:
-        """Expand the batch dimension of both sub-caches for beam search initialisation.
-
-        Delegates atomically to each sub-cache. Both must be expanded together so the
-        sliding-window and MoSRAH state correspond to the same beam candidates.
-
-        Args:
-            repeats: Number of times to repeat each batch entry.
-        """
-        self.sliding_window_cache.batch_repeat_interleave(repeats)
-        self.mosrah_cache.batch_repeat_interleave(repeats)
-
-    def batch_select_indices(self, indices: torch.Tensor) -> None:
-        """Select a subset of batch entries in both sub-caches for contrastive search.
-
-        Delegates atomically to each sub-cache. Both must be trimmed together so the
-        sliding-window and MoSRAH state remain consistent.
-
-        Args:
-            indices: 1-D integer tensor of batch indices to retain.
-        """
-        self.sliding_window_cache.batch_select_indices(indices)
-        self.mosrah_cache.batch_select_indices(indices)
-
-    def offload(self) -> None:
-        """Offload both sub-caches to CPU.
-
-        Delegates to each sub-cache's offload method. Does not call super() — ShramLayerCache
-        does not own self.keys/self.values directly; all cached data lives in the sub-caches.
-        """
-        self.sliding_window_cache.offload()
-        self.mosrah_cache.offload()
-
-    def prefetch(self) -> None:
-        """Move both sub-caches back to their model device ahead of time.
-
-        Delegates to each sub-cache's prefetch method. Does not call super() — ShramLayerCache
-        does not own self.keys/self.values directly; all cached data lives in the sub-caches.
-        """
-        self.sliding_window_cache.prefetch()
-        self.mosrah_cache.prefetch()
-
-    def lazy_initialization(  # type: ignore[override]
-        self, key_states: torch.Tensor, value_states: torch.Tensor
-    ) -> None:
-        """No-op — both sub-caches handle their own initialization."""
-        pass
-
-    # ---------------------------------------------------------------------------
-    # CacheLayerMixin — unsupported abstract methods
-    # ---------------------------------------------------------------------------
-
-    def update(  # type: ignore[override]
-        self,
-        key_states: torch.Tensor,
-        value_states: torch.Tensor,
-        cache_kwargs: dict | None = None,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        """Not supported — ShramLayerCache has no composite update interface.
-
-        The two sub-caches have materially different update semantics: the sliding-window
-        side uses standard key/value concatenation while the MoSRAH side uses expert-choice
-        scatter with an active mask. Callers must update each sub-cache directly via
-        sliding_window_cache.update() or mosrah_cache.update().
-        """
-        raise NotImplementedError(
-            "ShramLayerCache has no composite update interface. "
-            "Update sliding_window_cache or mosrah_cache directly."
-        )
-
-    def get_max_cache_shape(self) -> int:  # type: ignore[override]
-        """Not supported — the composite cache has no single maximum shape.
-
-        The sliding-window cache is bounded by sliding_window; the MoSRAH cache is
-        unbounded. No truthful scalar maximum represents the composite.
-        """
-        raise NotImplementedError(
-            "ShramLayerCache has no single maximum cache shape. "
-            "Query sliding_window_cache or mosrah_cache directly."
-        )
-
-    def get_mask_sizes(  # type: ignore[override]
-        self,
-        cache_position: torch.Tensor,
-    ) -> tuple[int, int]:
-        """Not supported — ShramLayerCache does not participate in HF mask construction.
-
-        The two sub-caches have different mask semantics and their respective attention
-        paths handle masking directly.
-        """
-        raise NotImplementedError(
-            "ShramLayerCache does not support get_mask_sizes(). "
-            "Query sliding_window_cache or mosrah_cache directly."
-        )
+"""SHRAM per-layer cache — composite owner for one SHRAM decoder layer.
+
+A SHRAM decoder layer contains two distinct attention pathways at one attention slot: the
+local sliding-window path and the MoSRAH sparse path. Each path has its own cache with
+different semantics and a different downstream consumer. ShramLayerCache owns both, satisfies
+the HuggingFace per-layer cache role, and exposes each sub-cache directly so its attention
+path can interact with it without indirection.
+
+ShramLayerCache does not define a composite update() interface. The two paths have materially
+different update semantics — the local side uses chunk-local key/value/mask concatenation
+while the MoSRAH side uses expert-choice scatter with an active mask — and merging these
+behind a single update() would hide those differences behind a misleading abstraction. Instead,
+each attention path calls update() on the sub-cache it owns. ShramLayerCache acts as the
+ownership, coordination, and reset/reorder boundary for one decoder layer.
+
+Sequence length at this boundary is reported by delegating to the local sliding-window
+sub-cache, which tracks the cumulative count of token positions processed. This is the
+quantity HuggingFace generation reads through get_seq_length().
+"""
+
+import torch
+from transformers.cache_utils import CacheLayerMixin
+
+from .__cache__mosrah_cache import MoSRAHCache
+from .__cache__sliding_window_cache import LocalSlidingWindowLayerCache
+
+
+class ShramLayerCache(CacheLayerMixin):
+    """Cache subsystem for one SHRAM decoder layer.
+
+    Owns and coordinates two sub-caches:
+      - sliding_window_cache: LocalSlidingWindowLayerCache for the local sliding-window path.
+      - mosrah_cache: MoSRAHCache for the MoSRAH sparse attention path.
+
+    Satisfies the HuggingFace per-layer cache role (CacheLayerMixin). The two sub-caches are
+    exposed directly for their downstream attention paths — no composite update() interface is
+    provided, because the two paths have materially different update semantics.
+
+    Sequence length is reported by delegating to the local sliding-window sub-cache, which
+    tracks the cumulative count of token positions processed across all update() calls.
+
+    Args:
+        sliding_window: Number of tokens retained by the local sliding-window cache.
+        num_local_heads: Number of local attention heads.
+        local_head_dim: Per-head embedding width for the local path.
+        num_mosrah_heads: Total number of MoSRAH expert heads (L).
+        mosrah_head_dim: Bottlenecked head embedding width (u) for the MoSRAH path.
+        batch_size: Number of sequences in the batch.
+        device: Device on which to allocate cache tensors.
+        initial_buffer_size: Initial per-(batch, head) capacity for MoSRAHCache. Doubled
+            when any slot overflows. Defaults to 64 to avoid repeated reallocation during
+            prompt processing.
+    """
+
+    is_compileable = False
+    is_sliding = False
+
+    def __init__(
+        self,
+        sliding_window: int,
+        num_local_heads: int,
+        local_head_dim: int,
+        num_mosrah_heads: int,
+        mosrah_head_dim: int,
+        batch_size: int,
+        device: torch.device,
+        initial_buffer_size: int = 64,
+    ) -> None:
+        super().__init__()
+        self.sliding_window_cache = LocalSlidingWindowLayerCache(
+            sliding_window=sliding_window,
+            num_heads=num_local_heads,
+            head_dim=local_head_dim,
+            batch_size=batch_size,
+            device=device,
+        )
+        self.mosrah_cache = MoSRAHCache(
+            num_mosrah_heads=num_mosrah_heads,
+            head_dim=mosrah_head_dim,
+            batch_size=batch_size,
+            device=device,
+            initial_buffer_size=initial_buffer_size,
+        )
+
+    # ---------------------------------------------------------------------------
+    # Properties
+    # ---------------------------------------------------------------------------
+
+    @property
+    def is_initialized(self) -> bool:
+        """True iff both sub-caches have allocated their storage.
+
+        Both LocalSlidingWindowLayerCache and MoSRAHCache pre-allocate at construction,
+        so this is True immediately after ShramLayerCache.__init__ returns.
+        """
+        return self.sliding_window_cache.is_initialized and self.mosrah_cache.is_initialized
+
+    @is_initialized.setter
+    def is_initialized(self, value: bool) -> None:
+        # CacheLayerMixin.__init__ assigns self.is_initialized = False as an instance
+        # attribute. Since property is a data descriptor it takes precedence, but Python
+        # still routes the assignment through __set__. Absorb it silently — state is
+        # derived from sub-caches, not stored here.
+        pass
+
+    # ---------------------------------------------------------------------------
+    # CacheLayerMixin — composite-meaningful methods
+    # ---------------------------------------------------------------------------
+
+    def get_seq_length(self) -> int:  # type: ignore[override]
+        """Return the cumulative sequence length from the local sliding-window path.
+
+        The local path is authoritative for sequence progress: it sees every token
+        presented to this layer and accumulates a truthful total. Delegates to
+        sliding_window_cache.get_seq_length().
+        """
+        return self.sliding_window_cache.get_seq_length()
+
+    def reset(self) -> None:
+        """Clear both sub-caches.
+
+        Delegates reset to each sub-cache. Both are cleared atomically so the sliding-window
+        state and MoSRAH sparse state remain consistent.
+        """
+        self.sliding_window_cache.reset()
+        self.mosrah_cache.reset()
+
+    def reorder_cache(self, beam_idx: torch.LongTensor) -> None:
+        """Reorder the batch dimension of both sub-caches for beam search.
+
+        Delegates to each sub-cache. Both are reordered atomically so the sliding-window
+        and MoSRAH state correspond to the same beam hypotheses after reordering.
+
+        Args:
+            beam_idx: Permutation indices of shape (batch,) produced by beam search.
+        """
+        self.sliding_window_cache.reorder_cache(beam_idx)
+        self.mosrah_cache.reorder_cache(beam_idx)
+
+    def batch_repeat_interleave(self, repeats: int) -> None:
+        """Expand the batch dimension of both sub-caches for beam search initialisation.
+
+        Delegates atomically to each sub-cache. Both must be expanded together so the
+        sliding-window and MoSRAH state correspond to the same beam candidates.
+
+        Args:
+            repeats: Number of times to repeat each batch entry.
+        """
+        self.sliding_window_cache.batch_repeat_interleave(repeats)
+        self.mosrah_cache.batch_repeat_interleave(repeats)
+
+    def batch_select_indices(self, indices: torch.Tensor) -> None:
+        """Select a subset of batch entries in both sub-caches for contrastive search.
+
+        Delegates atomically to each sub-cache. Both must be trimmed together so the
+        sliding-window and MoSRAH state remain consistent.
+
+        Args:
+            indices: 1-D integer tensor of batch indices to retain.
+        """
+        self.sliding_window_cache.batch_select_indices(indices)
+        self.mosrah_cache.batch_select_indices(indices)
+
+    def offload(self) -> None:
+        """Offload both sub-caches to CPU.
+
+        Delegates to each sub-cache's offload method. Does not call super() — ShramLayerCache
+        does not own self.keys/self.values directly; all cached data lives in the sub-caches.
+        """
+        self.sliding_window_cache.offload()
+        self.mosrah_cache.offload()
+
+    def prefetch(self) -> None:
+        """Move both sub-caches back to their model device ahead of time.
+
+        Delegates to each sub-cache's prefetch method. Does not call super() — ShramLayerCache
+        does not own self.keys/self.values directly; all cached data lives in the sub-caches.
+        """
+        self.sliding_window_cache.prefetch()
+        self.mosrah_cache.prefetch()
+
+    def lazy_initialization(  # type: ignore[override]
+        self, key_states: torch.Tensor, value_states: torch.Tensor
+    ) -> None:
+        """No-op — both sub-caches handle their own initialization."""
+        pass
+
+    # ---------------------------------------------------------------------------
+    # CacheLayerMixin — unsupported abstract methods
+    # ---------------------------------------------------------------------------
+
+    def update(  # type: ignore[override]
+        self,
+        key_states: torch.Tensor,
+        value_states: torch.Tensor,
+        cache_kwargs: dict | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Not supported — ShramLayerCache has no composite update interface.
+
+        The two sub-caches have materially different update semantics: the sliding-window
+        side uses standard key/value concatenation while the MoSRAH side uses expert-choice
+        scatter with an active mask. Callers must update each sub-cache directly via
+        sliding_window_cache.update() or mosrah_cache.update().
+        """
+        raise NotImplementedError(
+            "ShramLayerCache has no composite update interface. "
+            "Update sliding_window_cache or mosrah_cache directly."
+        )
+
+    def get_max_cache_shape(self) -> int:  # type: ignore[override]
+        """Not supported — the composite cache has no single maximum shape.
+
+        The sliding-window cache is bounded by sliding_window; the MoSRAH cache is
+        unbounded. No truthful scalar maximum represents the composite.
+        """
+        raise NotImplementedError(
+            "ShramLayerCache has no single maximum cache shape. "
+            "Query sliding_window_cache or mosrah_cache directly."
+        )
+
+    def get_mask_sizes(  # type: ignore[override]
+        self,
+        cache_position: torch.Tensor,
+    ) -> tuple[int, int]:
+        """Not supported — ShramLayerCache does not participate in HF mask construction.
+
+        The two sub-caches have different mask semantics and their respective attention
+        paths handle masking directly.
+        """
+        raise NotImplementedError(
+            "ShramLayerCache does not support get_mask_sizes(). "
+            "Query sliding_window_cache or mosrah_cache directly."
+        )

__cache__sliding_window_cache.py

@@ -1,289 +1,289 @@
-# src/shram/model/cache/sliding_window_cache.py
-
-"""Local sliding-window cache for the SHRAM local attention path.
-
-This file defines `LocalSlidingWindowLayerCache`, the local sub-cache owned by
-`ShramLayerCache` and consumed by `SlidingWindowAttention`.
-
-Its job is narrow:
-
-- accept the current chunk's local key/value tensors and active mask
-- return the current-step local frame consumed by local attention
-- separately retain the next-step sliding-window cache state
-
-It does not decide local causal visibility. That is owned by
-`SlidingWindowAttention`, which consumes the returned key/value/mask frame and
-constructs the effective local attention mask from it.
-"""
-
-import torch
-from transformers.cache_utils import CacheLayerMixin
-
-
-class LocalSlidingWindowLayerCache(CacheLayerMixin):
-    """Fixed-width local cache for one SHRAM decoder layer.
-
-    The cache keeps a retained local sliding-window buffer and an aligned active
-    mask. On update, it returns the current-step local frame formed by
-    concatenating retained cache state with the new chunk, then remembers only
-    the last `sliding_window` positions for the next step.
-
-    Dead positions are allowed to remain in both the returned frame and the
-    retained cache. Correctness is carried by the aligned active mask.
-
-    Args:
-        sliding_window: Width of the retained local sliding-window buffer.
-        num_heads: Number of local attention heads.
-        head_dim: Per-head embedding width for the local path.
-        batch_size: Number of sequences in the batch.
-        device: Device on which to allocate cache storage.
-    """
-
-    is_compileable = False
-    is_sliding = True
-
-    def __init__(
-        self,
-        sliding_window: int,
-        num_heads: int,
-        head_dim: int,
-        batch_size: int,
-        device: torch.device,
-    ) -> None:
-        super().__init__()
-
-        if sliding_window < 1:
-            raise ValueError(
-                f"sliding_window must be >= 1, got {sliding_window}."
-            )
-        if num_heads < 1:
-            raise ValueError(f"num_heads must be >= 1, got {num_heads}.")
-        if head_dim < 1:
-            raise ValueError(f"head_dim must be >= 1, got {head_dim}.")
-        if batch_size < 1:
-            raise ValueError(f"batch_size must be >= 1, got {batch_size}.")
-
-        self.sliding_window = sliding_window
-        self.num_heads = num_heads
-        self.head_dim = head_dim
-        self.batch_size = batch_size
-        self.device = device
-
-        # Retained next-step local cache state. Storage is fixed-width from the
-        # start; semantic validity is carried by `active_mask`.
-        self.keys = torch.zeros(
-            batch_size,
-            num_heads,
-            sliding_window,
-            head_dim,
-            device=device,
-        )
-        self.values = torch.zeros(
-            batch_size,
-            num_heads,
-            sliding_window,
-            head_dim,
-            device=device,
-        )
-        self.active_mask = torch.zeros(
-            batch_size,
-            sliding_window,
-            dtype=torch.bool,
-            device=device,
-        )
-
-        self.is_initialized = True
-
-        # Cumulative count of all token positions presented through update() for
-        # this cache instance. This is the quantity HuggingFace generation reads
-        # through get_seq_length() to track how far along the sequence we are.
-        self._total_processed: int = 0
-
-    def update(  # type: ignore[override]
-        self,
-        key_states: torch.Tensor,
-        value_states: torch.Tensor,
-        active_mask: torch.Tensor,
-        cache_kwargs: dict | None = None,
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        """Return the current-step local frame and retain the next-step window.
-
-        Args:
-            key_states: Shape `(B, H, T_new, D)` local key vectors for the
-                current chunk.
-            value_states: Shape `(B, H, T_new, D)` local value vectors for the
-                current chunk.
-            active_mask: Shape `(B, T_new)` bool. `True` means the
-                corresponding token position in the current chunk is active.
-            cache_kwargs: Present only to satisfy the `CacheLayerMixin`
-                interface. Unused by this cache.
-
-        Returns:
-            Tuple of:
-              - visible_keys: `(B, H, sliding_window + T_new, D)`
-              - visible_values: `(B, H, sliding_window + T_new, D)`
-              - visible_active_mask: `(B, sliding_window + T_new)`
-
-            These are the tensors the local attention path should consume
-            directly for the current step.
-        """
-        self._ensure_state_compatibility(
-            key_states=key_states,
-            value_states=value_states,
-        )
-
-        # The current-step local frame is just retained cache state followed by
-        # the current chunk in chronological order.
-        composite_keys, composite_values, composite_mask = self._make_composite_frame(
-            key_states=key_states,
-            value_states=value_states,
-            active_mask=active_mask,
-        )
-
-        # The cache remembers only the last raw sliding-window positions of that
-        # composite frame for the next step. Dead positions are allowed to
-        # survive; downstream local attention will ignore them using the mask.
-        self._retain_next_window(
-            composite_keys=composite_keys,
-            composite_values=composite_values,
-            composite_mask=composite_mask,
-        )
-
-        self._total_processed += key_states.shape[2]
-
-        return composite_keys, composite_values, composite_mask
-
-    def _ensure_state_compatibility(
-        self,
-        key_states: torch.Tensor,
-        value_states: torch.Tensor,
-    ) -> None:
-        """Keep retained cache buffers compatible with the incoming update tensors.
-
-        The cache is allocated eagerly for simplicity. If later updates arrive on
-        a different device or in a different floating dtype, move the retained
-        state to match while preserving its contents.
-        """
-        if self.keys.dtype != key_states.dtype or self.keys.device != key_states.device:
-            self.keys = self.keys.to(
-                device=key_states.device,
-                dtype=key_states.dtype,
-            )
-
-        if (
-            self.values.dtype != value_states.dtype
-            or self.values.device != value_states.device
-        ):
-            self.values = self.values.to(
-                device=value_states.device,
-                dtype=value_states.dtype,
-            )
-
-        if self.active_mask.device != key_states.device:
-            self.active_mask = self.active_mask.to(
-                key_states.device,
-                non_blocking=True,
-            )
-
-    def _make_composite_frame(
-        self,
-        key_states: torch.Tensor,
-        value_states: torch.Tensor,
-        active_mask: torch.Tensor,
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        """Build the current-step local frame in chronological order."""
-        return (
-            torch.cat([self.keys, key_states], dim=-2),
-            torch.cat([self.values, value_states], dim=-2),
-            torch.cat([self.active_mask, active_mask], dim=-1),
-        )
-
-    def _retain_next_window(
-        self,
-        composite_keys: torch.Tensor,
-        composite_values: torch.Tensor,
-        composite_mask: torch.Tensor,
-    ) -> None:
-        """Remember the next-step retained local state.
-
-        This is a raw positional trim to the last `sliding_window` positions, not
-        a semantic live-token trim.
-        """
-        self.keys = composite_keys[:, :, -self.sliding_window :, :]
-        self.values = composite_values[:, :, -self.sliding_window :, :]
-        self.active_mask = composite_mask[:, -self.sliding_window :]
-
-    def get_seq_length(self) -> int:
-        """Return the cumulative number of token positions processed by this cache.
-
-        This is the total count of token positions presented across all update()
-        calls since construction or the last reset(). It is the quantity HuggingFace
-        generation reads to track sequence progress and is not the same as active-token
-        count or current window occupancy.
-        """
-        return self._total_processed
-
-    def get_max_cache_shape(self) -> int:
-        return self.sliding_window
-
-    def get_mask_sizes(  # type: ignore[override]
-        self,
-        cache_position: torch.Tensor,
-    ) -> tuple[int, int]:
-        raise NotImplementedError(
-            "LocalSlidingWindowLayerCache does not support get_mask_sizes()."
-        )
-
-    def reset(self) -> None:
-        """Restore fresh-cache behavior."""
-        self.keys.zero_()
-        self.values.zero_()
-        self.active_mask.zero_()
-        self._total_processed = 0
-
-    def reorder_cache(self, beam_idx: torch.LongTensor) -> None:
-        """Reorder the batch dimension for beam search."""
-        self.keys = self.keys[beam_idx]
-        self.values = self.values[beam_idx]
-        self.active_mask = self.active_mask[beam_idx]
-
-    def batch_repeat_interleave(self, repeats: int) -> None:
-        """Expand the batch dimension for beam-search initialisation."""
-        self.keys = self.keys.repeat_interleave(repeats, dim=0)
-        self.values = self.values.repeat_interleave(repeats, dim=0)
-        self.active_mask = self.active_mask.repeat_interleave(repeats, dim=0)
-        self.batch_size = self.batch_size * repeats
-
-    def batch_select_indices(self, indices: torch.Tensor) -> None:
-        """Select a subset of batch entries for contrastive search."""
-        self.keys = self.keys[indices]
-        self.values = self.values[indices]
-        self.active_mask = self.active_mask[indices]
-        self.batch_size = int(indices.shape[0])
-
-    def offload(self) -> None:
-        """Offload cache tensors to CPU."""
-        super().offload()
-        self.active_mask = self.active_mask.to("cpu", non_blocking=True)
-
-    def prefetch(self) -> None:
-        """Move cache tensors back to the model device ahead of time."""
-        super().prefetch()
-        if self.active_mask.device != self.keys.device:
-            self.active_mask = self.active_mask.to(
-                self.keys.device,
-                non_blocking=True,
-            )
-
-    def crop(self, max_length: int) -> None:
-        raise NotImplementedError(
-            "LocalSlidingWindowLayerCache does not support crop()."
-        )
-
-    def lazy_initialization(
-        self,
-        key_states: torch.Tensor,
-        value_states: torch.Tensor,
-    ) -> None:
-        """No-op — this cache allocates its fixed buffers at construction time."""
+# src/shram/model/cache/sliding_window_cache.py
+
+"""Local sliding-window cache for the SHRAM local attention path.
+
+This file defines `LocalSlidingWindowLayerCache`, the local sub-cache owned by
+`ShramLayerCache` and consumed by `SlidingWindowAttention`.
+
+Its job is narrow:
+
+- accept the current chunk's local key/value tensors and active mask
+- return the current-step local frame consumed by local attention
+- separately retain the next-step sliding-window cache state
+
+It does not decide local causal visibility. That is owned by
+`SlidingWindowAttention`, which consumes the returned key/value/mask frame and
+constructs the effective local attention mask from it.
+"""
+
+import torch
+from transformers.cache_utils import CacheLayerMixin
+
+
+class LocalSlidingWindowLayerCache(CacheLayerMixin):
+    """Fixed-width local cache for one SHRAM decoder layer.
+
+    The cache keeps a retained local sliding-window buffer and an aligned active
+    mask. On update, it returns the current-step local frame formed by
+    concatenating retained cache state with the new chunk, then remembers only
+    the last `sliding_window` positions for the next step.
+
+    Dead positions are allowed to remain in both the returned frame and the
+    retained cache. Correctness is carried by the aligned active mask.
+
+    Args:
+        sliding_window: Width of the retained local sliding-window buffer.
+        num_heads: Number of local attention heads.
+        head_dim: Per-head embedding width for the local path.
+        batch_size: Number of sequences in the batch.
+        device: Device on which to allocate cache storage.
+    """
+
+    is_compileable = False
+    is_sliding = True
+
+    def __init__(
+        self,
+        sliding_window: int,
+        num_heads: int,
+        head_dim: int,
+        batch_size: int,
+        device: torch.device,
+    ) -> None:
+        super().__init__()
+
+        if sliding_window < 1:
+            raise ValueError(
+                f"sliding_window must be >= 1, got {sliding_window}."
+            )
+        if num_heads < 1:
+            raise ValueError(f"num_heads must be >= 1, got {num_heads}.")
+        if head_dim < 1:
+            raise ValueError(f"head_dim must be >= 1, got {head_dim}.")
+        if batch_size < 1:
+            raise ValueError(f"batch_size must be >= 1, got {batch_size}.")
+
+        self.sliding_window = sliding_window
+        self.num_heads = num_heads
+        self.head_dim = head_dim
+        self.batch_size = batch_size
+        self.device = device
+
+        # Retained next-step local cache state. Storage is fixed-width from the
+        # start; semantic validity is carried by `active_mask`.
+        self.keys = torch.zeros(
+            batch_size,
+            num_heads,
+            sliding_window,
+            head_dim,
+            device=device,
+        )
+        self.values = torch.zeros(
+            batch_size,
+            num_heads,
+            sliding_window,
+            head_dim,
+            device=device,
+        )
+        self.active_mask = torch.zeros(
+            batch_size,
+            sliding_window,
+            dtype=torch.bool,
+            device=device,
+        )
+
+        self.is_initialized = True
+
+        # Cumulative count of all token positions presented through update() for
+        # this cache instance. This is the quantity HuggingFace generation reads
+        # through get_seq_length() to track how far along the sequence we are.
+        self._total_processed: int = 0
+
+    def update(  # type: ignore[override]
+        self,
+        key_states: torch.Tensor,
+        value_states: torch.Tensor,
+        active_mask: torch.Tensor,
+        cache_kwargs: dict | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Return the current-step local frame and retain the next-step window.
+
+        Args:
+            key_states: Shape `(B, H, T_new, D)` local key vectors for the
+                current chunk.
+            value_states: Shape `(B, H, T_new, D)` local value vectors for the
+                current chunk.
+            active_mask: Shape `(B, T_new)` bool. `True` means the
+                corresponding token position in the current chunk is active.
+            cache_kwargs: Present only to satisfy the `CacheLayerMixin`
+                interface. Unused by this cache.
+
+        Returns:
+            Tuple of:
+              - visible_keys: `(B, H, sliding_window + T_new, D)`
+              - visible_values: `(B, H, sliding_window + T_new, D)`
+              - visible_active_mask: `(B, sliding_window + T_new)`
+
+            These are the tensors the local attention path should consume
+            directly for the current step.
+        """
+        self._ensure_state_compatibility(
+            key_states=key_states,
+            value_states=value_states,
+        )
+
+        # The current-step local frame is just retained cache state followed by
+        # the current chunk in chronological order.
+        composite_keys, composite_values, composite_mask = self._make_composite_frame(
+            key_states=key_states,
+            value_states=value_states,
+            active_mask=active_mask,
+        )
+
+        # The cache remembers only the last raw sliding-window positions of that
+        # composite frame for the next step. Dead positions are allowed to
+        # survive; downstream local attention will ignore them using the mask.
+        self._retain_next_window(
+            composite_keys=composite_keys,
+            composite_values=composite_values,
+            composite_mask=composite_mask,
+        )
+
+        self._total_processed += key_states.shape[2]
+
+        return composite_keys, composite_values, composite_mask
+
+    def _ensure_state_compatibility(
+        self,
+        key_states: torch.Tensor,
+        value_states: torch.Tensor,
+    ) -> None:
+        """Keep retained cache buffers compatible with the incoming update tensors.
+
+        The cache is allocated eagerly for simplicity. If later updates arrive on
+        a different device or in a different floating dtype, move the retained
+        state to match while preserving its contents.
+        """
+        if self.keys.dtype != key_states.dtype or self.keys.device != key_states.device:
+            self.keys = self.keys.to(
+                device=key_states.device,
+                dtype=key_states.dtype,
+            )
+
+        if (
+            self.values.dtype != value_states.dtype
+            or self.values.device != value_states.device
+        ):
+            self.values = self.values.to(
+                device=value_states.device,
+                dtype=value_states.dtype,
+            )
+
+        if self.active_mask.device != key_states.device:
+            self.active_mask = self.active_mask.to(
+                key_states.device,
+                non_blocking=True,
+            )
+
+    def _make_composite_frame(
+        self,
+        key_states: torch.Tensor,
+        value_states: torch.Tensor,
+        active_mask: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Build the current-step local frame in chronological order."""
+        return (
+            torch.cat([self.keys, key_states], dim=-2),
+            torch.cat([self.values, value_states], dim=-2),
+            torch.cat([self.active_mask, active_mask], dim=-1),
+        )
+
+    def _retain_next_window(
+        self,
+        composite_keys: torch.Tensor,
+        composite_values: torch.Tensor,
+        composite_mask: torch.Tensor,
+    ) -> None:
+        """Remember the next-step retained local state.
+
+        This is a raw positional trim to the last `sliding_window` positions, not
+        a semantic live-token trim.
+        """
+        self.keys = composite_keys[:, :, -self.sliding_window :, :]
+        self.values = composite_values[:, :, -self.sliding_window :, :]
+        self.active_mask = composite_mask[:, -self.sliding_window :]
+
+    def get_seq_length(self) -> int:
+        """Return the cumulative number of token positions processed by this cache.
+
+        This is the total count of token positions presented across all update()
+        calls since construction or the last reset(). It is the quantity HuggingFace
+        generation reads to track sequence progress and is not the same as active-token
+        count or current window occupancy.
+        """
+        return self._total_processed
+
+    def get_max_cache_shape(self) -> int:
+        return self.sliding_window
+
+    def get_mask_sizes(  # type: ignore[override]
+        self,
+        cache_position: torch.Tensor,
+    ) -> tuple[int, int]:
+        raise NotImplementedError(
+            "LocalSlidingWindowLayerCache does not support get_mask_sizes()."
+        )
+
+    def reset(self) -> None:
+        """Restore fresh-cache behavior."""
+        self.keys.zero_()
+        self.values.zero_()
+        self.active_mask.zero_()
+        self._total_processed = 0
+
+    def reorder_cache(self, beam_idx: torch.LongTensor) -> None:
+        """Reorder the batch dimension for beam search."""
+        self.keys = self.keys[beam_idx]
+        self.values = self.values[beam_idx]
+        self.active_mask = self.active_mask[beam_idx]
+
+    def batch_repeat_interleave(self, repeats: int) -> None:
+        """Expand the batch dimension for beam-search initialisation."""
+        self.keys = self.keys.repeat_interleave(repeats, dim=0)
+        self.values = self.values.repeat_interleave(repeats, dim=0)
+        self.active_mask = self.active_mask.repeat_interleave(repeats, dim=0)
+        self.batch_size = self.batch_size * repeats
+
+    def batch_select_indices(self, indices: torch.Tensor) -> None:
+        """Select a subset of batch entries for contrastive search."""
+        self.keys = self.keys[indices]
+        self.values = self.values[indices]
+        self.active_mask = self.active_mask[indices]
+        self.batch_size = int(indices.shape[0])
+
+    def offload(self) -> None:
+        """Offload cache tensors to CPU."""
+        super().offload()
+        self.active_mask = self.active_mask.to("cpu", non_blocking=True)
+
+    def prefetch(self) -> None:
+        """Move cache tensors back to the model device ahead of time."""
+        super().prefetch()
+        if self.active_mask.device != self.keys.device:
+            self.active_mask = self.active_mask.to(
+                self.keys.device,
+                non_blocking=True,
+            )
+
+    def crop(self, max_length: int) -> None:
+        raise NotImplementedError(
+            "LocalSlidingWindowLayerCache does not support crop()."
+        )
+
+    def lazy_initialization(
+        self,
+        key_states: torch.Tensor,
+        value_states: torch.Tensor,
+    ) -> None:
+        """No-op — this cache allocates its fixed buffers at construction time."""
         return

__cache__slow_mosrah_cache.py

@@ -1,321 +1,321 @@
-"""Unvectorized reference implementation of the MoSRAH sparse KV cache.
-
-This module exists solely as a correctness oracle. SlowMoSRAHCache implements the same
-interface and storage layout as MoSRAHCache but uses an explicit Python loop over
-(b, l, t) triples in update(). The loop is obviously correct by inspection: each active
-position's key and value are written to the next available slot for that (batch, head)
-pair, in the order positions appear along the T dimension, which directly enforces
-causal ordering without any index arithmetic to verify.
-
-SlowMoSRAHCache is never instantiated in the model path. Its role is to provide a
-trusted ground truth against which the vectorized MoSRAHCache.update() is validated in
-Unit 6.A tests, and as a reference for the Unit 10.A position decoder. Because the
-vectorized implementation is validated by asserting exact agreement with this one on all
-test inputs, the correctness of SlowMoSRAHCache is load-bearing: its own test suite
-(test_slow_mosrah_cache.py) must establish it is trustworthy before it can be used as
-an oracle.
-"""
-
-import torch
-from transformers.cache_utils import CacheLayerMixin
-
-
-class SlowMoSRAHCache(CacheLayerMixin):
-    """Unvectorized reference implementation of the MoSRAH KV cache.
-
-    Identical storage layout to MoSRAHCache: (B, L, T, u) tensors in the
-    mixin-standard self.keys and self.values attributes, plus a (B, L) _counts tensor,
-    with the same constructor signature and the same CacheLayerMixin protocol methods.
-    The sole difference is update(), which uses an explicit Python loop over (b, l, t)
-    triples rather than vectorized index arithmetic.
-
-    This class is not used in the model path. It exists so that MoSRAHCache.update()
-    can be validated by asserting exact agreement with this implementation on all test
-    inputs. See module docstring for the trust chain this enables.
-
-    Args:
-        num_mosrah_heads: Total number of MoSRAH expert heads (L). Determines the
-            second dimension of all storage tensors.
-        head_dim: Bottlenecked head embedding width (u). Determines the fourth
-            dimension of all storage tensors.
-        batch_size: Number of sequences in the batch. Determines the first dimension
-            of all storage tensors.
-        device: Device on which to allocate all tensors. Should match the model device.
-        initial_buffer_size: Initial sequence capacity per (batch, head) slot. Doubled
-            when any slot overflows. Defaults to 64 to avoid repeated reallocation
-            during prompt processing.
-    """
-
-    is_compileable = False
-    is_sliding = False
-
-    def __init__(
-        self,
-        num_mosrah_heads: int,
-        head_dim: int,
-        batch_size: int,
-        device: torch.device,
-        initial_buffer_size: int = 64,
-    ) -> None:
-        super().__init__()
-        self.num_mosrah_heads = num_mosrah_heads
-        self.head_dim = head_dim
-        self.batch_size = batch_size
-        self.device = device
-
-        # Allocate primary storage into the mixin-standard self.keys / self.values so
-        # that inherited methods (offload, prefetch) operate on real tensors. _counts
-        # tracks valid occupancy per (batch, head) slot.
-        self.keys: torch.Tensor = torch.zeros(
-            batch_size, num_mosrah_heads, initial_buffer_size, head_dim, device=device
-        )
-        self.values: torch.Tensor = torch.zeros(
-            batch_size, num_mosrah_heads, initial_buffer_size, head_dim, device=device
-        )
-        self._counts: torch.Tensor = torch.zeros(
-            batch_size, num_mosrah_heads, dtype=torch.long, device=device
-        )
-
-        # Storage is fully allocated at construction — the cache is initialized.
-        self.is_initialized = True
-
-    # ---------------------------------------------------------------------------
-    # Properties
-    # ---------------------------------------------------------------------------
-
-    @property
-    def buffer_capacity(self) -> int:
-        """Current number of slots allocated per (batch, head) pair.
-
-        Derived directly from self.keys rather than tracked separately, so it is
-        always consistent with the actual buffer after expansion.
-        """
-        return self.keys.shape[2]
-
-    # ---------------------------------------------------------------------------
-    # Primary API
-    # ---------------------------------------------------------------------------
-
-    def update(  # type: ignore[override]
-        self,
-        key_states: torch.Tensor,
-        value_states: torch.Tensor,
-        active_mask: torch.Tensor,
-        cache_kwargs: dict | None = None,
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        """Scatter active key/value states using an explicit loop; return full cache state.
-
-        Iterates over every (b, l, t) triple. For each position where active_mask is
-        True, the key and value are written to the next available slot for that
-        (batch, head) pair and the count is incremented. Causal ordering is guaranteed
-        because the t dimension is traversed from 0 to T-1 and counts are updated
-        immediately after each write.
-
-        Buffer expansion (doubling buffer_capacity) is triggered before any writes if
-        the incoming tokens would cause any slot to overflow the current capacity.
-
-        Args:
-            key_states: Shape (B, L, T, u) — post-RoPE key vectors in expert-choice layout.
-            value_states: Shape (B, L, T, u) — value vectors in expert-choice layout.
-            active_mask: Shape (B, L, T) bool — True for real tokens, False for padding.
-            cache_kwargs: Unused; present to satisfy the CacheLayerMixin signature.
-
-        Returns:
-            Tuple of (keys, values, active_mask):
-              keys: (B, L, T, u) float — full key buffer including junk slots.
-              values: (B, L, T, u) float — full value buffer including junk slots.
-              active_mask: (B, L, T) bool — True iff slot (b, l, t) has been written.
-        """
-        B, L, T = active_mask.shape
-
-        # Expansion check uses the total active tokens per slot, same as the
-        # vectorized implementation, so both expand under identical conditions.
-        incoming_delta = active_mask.long().sum(dim=2)  # (B, L)
-        if (self._counts + incoming_delta).max().item() > self.buffer_capacity:
-            self._expand()
-
-        # Write each active position into the next available slot for its (batch, head)
-        # pair. Iterating t from 0 to T-1 preserves causal ordering within each slot.
-        for b in range(B):
-            for l in range(L):
-                for t in range(T):
-                    if active_mask[b, l, t]:
-                        pos = self._counts[b, l].item()
-                        self.keys[b, l, pos, :] = key_states[b, l, t, :]
-                        self.values[b, l, pos, :] = value_states[b, l, t, :]
-                        self._counts[b, l] += 1
-
-        return self.keys, self.values, self._make_active_mask()
-
-    def get_heads_lengths(self) -> torch.Tensor:
-        """Return the per-(batch, head) token count for this layer.
-
-        This is the authoritative occupancy tensor consumed by BEA for attention
-        masking and by position computation (Unit 10.A) for semantic-sequence
-        position computation.
-
-        Returns:
-            Integer tensor of shape (B, L) where entry [b, h] is the number of valid
-            tokens stored in the (b, h) slot. Zero for slots with no writes yet.
-        """
-        return self._counts
-
-    # ---------------------------------------------------------------------------
-    # CacheLayerMixin — overridden coordination methods
-    # ---------------------------------------------------------------------------
-
-    def reset(self) -> None:
-        """Clear all cached key and value tensors.
-
-        Zeroes self.keys, self.values, and _counts in place. Storage remains allocated
-        and is_initialized remains True — only the contents are cleared.
-        """
-        self.keys.zero_()
-        self.values.zero_()
-        self._counts.zero_()
-
-    def reorder_cache(self, beam_idx: torch.LongTensor) -> None:
-        """Reorder the batch dimension of all cached tensors for beam search.
-
-        Applied atomically across self.keys, self.values, and _counts. Beam search
-        must reorder all three together or the occupancy counts and buffer contents
-        will correspond to different beam hypotheses.
-
-        Overrides the parent because the parent's implementation calls get_seq_length(),
-        which is not supported for this cache.
-
-        Args:
-            beam_idx: Permutation indices of shape (batch,) produced by the beam
-                search algorithm.
-        """
-        self.keys = self.keys[beam_idx]
-        self.values = self.values[beam_idx]
-        self._counts = self._counts[beam_idx]
-
-    def batch_repeat_interleave(self, repeats: int) -> None:
-        """Expand the batch dimension by repeating each entry repeats times.
-
-        Used at beam search initialisation to expand the cache from batch size B to
-        B * repeats, matching the expanded beam candidate batch. Applied atomically
-        across keys, values, and _counts; batch_size is updated to reflect the new size.
-
-        Args:
-            repeats: Number of times to repeat each batch entry.
-        """
-        self.keys = self.keys.repeat_interleave(repeats, dim=0)
-        self.values = self.values.repeat_interleave(repeats, dim=0)
-        self._counts = self._counts.repeat_interleave(repeats, dim=0)
-        self.batch_size = self.batch_size * repeats
-
-    def batch_select_indices(self, indices: torch.Tensor) -> None:
-        """Select a subset of batch entries by index.
-
-        Used in contrastive search to retain only the selected candidate entries.
-        Applied atomically across keys, values, and _counts; batch_size is updated
-        to reflect the number of retained entries.
-
-        Args:
-            indices: 1-D integer tensor of batch indices to retain.
-        """
-        self.keys = self.keys[indices]
-        self.values = self.values[indices]
-        self._counts = self._counts[indices]
-        self.batch_size = indices.shape[0]
-
-    def offload(self) -> None:
-        """Offload all cached tensors to CPU.
-
-        Extends the parent to also offload _counts, which the parent does not know
-        about. All three tensors are moved atomically so device state remains consistent.
-        """
-        super().offload()
-        self._counts = self._counts.to("cpu", non_blocking=True)
-
-    def prefetch(self) -> None:
-        """Move all cached tensors back to the model device ahead of time.
-
-        Extends the parent to also prefetch _counts, which the parent does not know
-        about. _counts is synced to self.keys.device after the parent moves keys and
-        values, so all three remain consistent.
-        """
-        super().prefetch()
-        if self._counts.device != self.keys.device:
-            self._counts = self._counts.to(self.keys.device, non_blocking=True)
-
-    def lazy_initialization(  # type: ignore[override]
-        self, key_states: torch.Tensor, value_states: torch.Tensor
-    ) -> None:
-        """No-op — storage is fully allocated at construction time."""
-        pass
-
-    # ---------------------------------------------------------------------------
-    # CacheLayerMixin — unsupported abstract methods
-    # ---------------------------------------------------------------------------
-
-    def get_seq_length(self) -> int:  # type: ignore[override]
-        """Not supported — no single sequence length represents this cache's state.
-
-        MoSRAH heads accumulate independently; (batch, head) slots have different
-        lengths depending on routing history. There is no meaningful scalar summary.
-        Use get_heads_lengths() for per-head occupancy.
-        """
-        raise NotImplementedError(
-            "SlowMoSRAHCache has no single sequence length. "
-            "Use get_heads_lengths() for per-head occupancy."
-        )
-
-    def get_max_cache_shape(self) -> int:  # type: ignore[override]
-        """Not supported — SlowMoSRAHCache is dynamic and unbounded."""
-        raise NotImplementedError(
-            "SlowMoSRAHCache is unbounded; get_max_cache_shape() is not supported."
-        )
-
-    def get_mask_sizes(  # type: ignore[override]
-        self,
-        cache_position: torch.Tensor,
-    ) -> tuple[int, int]:
-        """Not supported — SlowMoSRAHCache does not participate in HF mask construction."""
-        raise NotImplementedError(
-            "SlowMoSRAHCache does not support get_mask_sizes()."
-        )
-
-    # ---------------------------------------------------------------------------
-    # Internal helpers
-    # ---------------------------------------------------------------------------
-
-    def _make_active_mask(self) -> torch.Tensor:
-        """Construct the (B, L, T) active mask from current counts.
-
-        Returns True at position [b, l, t] iff t < _counts[b, l], i.e. the slot
-        has been written. Positions at or beyond the count are junk and must be
-        excluded by downstream attention.
-        """
-        cap = self.buffer_capacity
-        return (
-            torch.arange(cap, device=self.keys.device)
-            .expand(self.batch_size, self.num_mosrah_heads, cap)
-            < self._counts.unsqueeze(-1)
-        )
-
-    def _expand(self) -> None:
-        """Double the buffer capacity, preserving existing data.
-
-        Called by update() when an incoming batch of tokens would cause any
-        (batch, head) slot to exceed the current buffer capacity. All existing
-        key and value data is copied into the low half of the new buffer; the
-        high half is zero-initialised and will be filled by subsequent writes.
-        After reassignment, buffer_capacity reflects the new size automatically.
-        """
-        old_cap = self.buffer_capacity
-        new_cap = old_cap * 2
-        dev = self.keys.device
-        new_keys = torch.zeros(
-            self.batch_size, self.num_mosrah_heads, new_cap, self.head_dim, device=dev
-        )
-        new_values = torch.zeros(
-            self.batch_size, self.num_mosrah_heads, new_cap, self.head_dim, device=dev
-        )
-        new_keys[:, :, :old_cap, :] = self.keys
-        new_values[:, :, :old_cap, :] = self.values
-        self.keys = new_keys
-        self.values = new_values
+"""Unvectorized reference implementation of the MoSRAH sparse KV cache.
+
+This module exists solely as a correctness oracle. SlowMoSRAHCache implements the same
+interface and storage layout as MoSRAHCache but uses an explicit Python loop over
+(b, l, t) triples in update(). The loop is obviously correct by inspection: each active
+position's key and value are written to the next available slot for that (batch, head)
+pair, in the order positions appear along the T dimension, which directly enforces
+causal ordering without any index arithmetic to verify.
+
+SlowMoSRAHCache is never instantiated in the model path. Its role is to provide a
+trusted ground truth against which the vectorized MoSRAHCache.update() is validated in
+Unit 6.A tests, and as a reference for the Unit 10.A position decoder. Because the
+vectorized implementation is validated by asserting exact agreement with this one on all
+test inputs, the correctness of SlowMoSRAHCache is load-bearing: its own test suite
+(test_slow_mosrah_cache.py) must establish it is trustworthy before it can be used as
+an oracle.
+"""
+
+import torch
+from transformers.cache_utils import CacheLayerMixin
+
+
+class SlowMoSRAHCache(CacheLayerMixin):
+    """Unvectorized reference implementation of the MoSRAH KV cache.
+
+    Identical storage layout to MoSRAHCache: (B, L, T, u) tensors in the
+    mixin-standard self.keys and self.values attributes, plus a (B, L) _counts tensor,
+    with the same constructor signature and the same CacheLayerMixin protocol methods.
+    The sole difference is update(), which uses an explicit Python loop over (b, l, t)
+    triples rather than vectorized index arithmetic.
+
+    This class is not used in the model path. It exists so that MoSRAHCache.update()
+    can be validated by asserting exact agreement with this implementation on all test
+    inputs. See module docstring for the trust chain this enables.
+
+    Args:
+        num_mosrah_heads: Total number of MoSRAH expert heads (L). Determines the
+            second dimension of all storage tensors.
+        head_dim: Bottlenecked head embedding width (u). Determines the fourth
+            dimension of all storage tensors.
+        batch_size: Number of sequences in the batch. Determines the first dimension
+            of all storage tensors.
+        device: Device on which to allocate all tensors. Should match the model device.
+        initial_buffer_size: Initial sequence capacity per (batch, head) slot. Doubled
+            when any slot overflows. Defaults to 64 to avoid repeated reallocation
+            during prompt processing.
+    """
+
+    is_compileable = False
+    is_sliding = False
+
+    def __init__(
+        self,
+        num_mosrah_heads: int,
+        head_dim: int,
+        batch_size: int,
+        device: torch.device,
+        initial_buffer_size: int = 64,
+    ) -> None:
+        super().__init__()
+        self.num_mosrah_heads = num_mosrah_heads
+        self.head_dim = head_dim
+        self.batch_size = batch_size
+        self.device = device
+
+        # Allocate primary storage into the mixin-standard self.keys / self.values so
+        # that inherited methods (offload, prefetch) operate on real tensors. _counts
+        # tracks valid occupancy per (batch, head) slot.
+        self.keys: torch.Tensor = torch.zeros(
+            batch_size, num_mosrah_heads, initial_buffer_size, head_dim, device=device
+        )
+        self.values: torch.Tensor = torch.zeros(
+            batch_size, num_mosrah_heads, initial_buffer_size, head_dim, device=device
+        )
+        self._counts: torch.Tensor = torch.zeros(
+            batch_size, num_mosrah_heads, dtype=torch.long, device=device
+        )
+
+        # Storage is fully allocated at construction — the cache is initialized.
+        self.is_initialized = True
+
+    # ---------------------------------------------------------------------------
+    # Properties
+    # ---------------------------------------------------------------------------
+
+    @property
+    def buffer_capacity(self) -> int:
+        """Current number of slots allocated per (batch, head) pair.
+
+        Derived directly from self.keys rather than tracked separately, so it is
+        always consistent with the actual buffer after expansion.
+        """
+        return self.keys.shape[2]
+
+    # ---------------------------------------------------------------------------
+    # Primary API
+    # ---------------------------------------------------------------------------
+
+    def update(  # type: ignore[override]
+        self,
+        key_states: torch.Tensor,
+        value_states: torch.Tensor,
+        active_mask: torch.Tensor,
+        cache_kwargs: dict | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Scatter active key/value states using an explicit loop; return full cache state.
+
+        Iterates over every (b, l, t) triple. For each position where active_mask is
+        True, the key and value are written to the next available slot for that
+        (batch, head) pair and the count is incremented. Causal ordering is guaranteed
+        because the t dimension is traversed from 0 to T-1 and counts are updated
+        immediately after each write.
+
+        Buffer expansion (doubling buffer_capacity) is triggered before any writes if
+        the incoming tokens would cause any slot to overflow the current capacity.
+
+        Args:
+            key_states: Shape (B, L, T, u) — post-RoPE key vectors in expert-choice layout.
+            value_states: Shape (B, L, T, u) — value vectors in expert-choice layout.
+            active_mask: Shape (B, L, T) bool — True for real tokens, False for padding.
+            cache_kwargs: Unused; present to satisfy the CacheLayerMixin signature.
+
+        Returns:
+            Tuple of (keys, values, active_mask):
+              keys: (B, L, T, u) float — full key buffer including junk slots.
+              values: (B, L, T, u) float — full value buffer including junk slots.
+              active_mask: (B, L, T) bool — True iff slot (b, l, t) has been written.
+        """
+        B, L, T = active_mask.shape
+
+        # Expansion check uses the total active tokens per slot, same as the
+        # vectorized implementation, so both expand under identical conditions.
+        incoming_delta = active_mask.long().sum(dim=2)  # (B, L)
+        if (self._counts + incoming_delta).max().item() > self.buffer_capacity:
+            self._expand()
+
+        # Write each active position into the next available slot for its (batch, head)
+        # pair. Iterating t from 0 to T-1 preserves causal ordering within each slot.
+        for b in range(B):
+            for l in range(L):
+                for t in range(T):
+                    if active_mask[b, l, t]:
+                        pos = self._counts[b, l].item()
+                        self.keys[b, l, pos, :] = key_states[b, l, t, :]
+                        self.values[b, l, pos, :] = value_states[b, l, t, :]
+                        self._counts[b, l] += 1
+
+        return self.keys, self.values, self._make_active_mask()
+
+    def get_heads_lengths(self) -> torch.Tensor:
+        """Return the per-(batch, head) token count for this layer.
+
+        This is the authoritative occupancy tensor consumed by BEA for attention
+        masking and by position computation (Unit 10.A) for semantic-sequence
+        position computation.
+
+        Returns:
+            Integer tensor of shape (B, L) where entry [b, h] is the number of valid
+            tokens stored in the (b, h) slot. Zero for slots with no writes yet.
+        """
+        return self._counts
+
+    # ---------------------------------------------------------------------------
+    # CacheLayerMixin — overridden coordination methods
+    # ---------------------------------------------------------------------------
+
+    def reset(self) -> None:
+        """Clear all cached key and value tensors.
+
+        Zeroes self.keys, self.values, and _counts in place. Storage remains allocated
+        and is_initialized remains True — only the contents are cleared.
+        """
+        self.keys.zero_()
+        self.values.zero_()
+        self._counts.zero_()
+
+    def reorder_cache(self, beam_idx: torch.LongTensor) -> None:
+        """Reorder the batch dimension of all cached tensors for beam search.
+
+        Applied atomically across self.keys, self.values, and _counts. Beam search
+        must reorder all three together or the occupancy counts and buffer contents
+        will correspond to different beam hypotheses.
+
+        Overrides the parent because the parent's implementation calls get_seq_length(),
+        which is not supported for this cache.
+
+        Args:
+            beam_idx: Permutation indices of shape (batch,) produced by the beam
+                search algorithm.
+        """
+        self.keys = self.keys[beam_idx]
+        self.values = self.values[beam_idx]
+        self._counts = self._counts[beam_idx]
+
+    def batch_repeat_interleave(self, repeats: int) -> None:
+        """Expand the batch dimension by repeating each entry repeats times.
+
+        Used at beam search initialisation to expand the cache from batch size B to
+        B * repeats, matching the expanded beam candidate batch. Applied atomically
+        across keys, values, and _counts; batch_size is updated to reflect the new size.
+
+        Args:
+            repeats: Number of times to repeat each batch entry.
+        """
+        self.keys = self.keys.repeat_interleave(repeats, dim=0)
+        self.values = self.values.repeat_interleave(repeats, dim=0)
+        self._counts = self._counts.repeat_interleave(repeats, dim=0)
+        self.batch_size = self.batch_size * repeats
+
+    def batch_select_indices(self, indices: torch.Tensor) -> None:
+        """Select a subset of batch entries by index.
+
+        Used in contrastive search to retain only the selected candidate entries.
+        Applied atomically across keys, values, and _counts; batch_size is updated
+        to reflect the number of retained entries.
+
+        Args:
+            indices: 1-D integer tensor of batch indices to retain.
+        """
+        self.keys = self.keys[indices]
+        self.values = self.values[indices]
+        self._counts = self._counts[indices]
+        self.batch_size = indices.shape[0]
+
+    def offload(self) -> None:
+        """Offload all cached tensors to CPU.
+
+        Extends the parent to also offload _counts, which the parent does not know
+        about. All three tensors are moved atomically so device state remains consistent.
+        """
+        super().offload()
+        self._counts = self._counts.to("cpu", non_blocking=True)
+
+    def prefetch(self) -> None:
+        """Move all cached tensors back to the model device ahead of time.
+
+        Extends the parent to also prefetch _counts, which the parent does not know
+        about. _counts is synced to self.keys.device after the parent moves keys and
+        values, so all three remain consistent.
+        """
+        super().prefetch()
+        if self._counts.device != self.keys.device:
+            self._counts = self._counts.to(self.keys.device, non_blocking=True)
+
+    def lazy_initialization(  # type: ignore[override]
+        self, key_states: torch.Tensor, value_states: torch.Tensor
+    ) -> None:
+        """No-op — storage is fully allocated at construction time."""
+        pass
+
+    # ---------------------------------------------------------------------------
+    # CacheLayerMixin — unsupported abstract methods
+    # ---------------------------------------------------------------------------
+
+    def get_seq_length(self) -> int:  # type: ignore[override]
+        """Not supported — no single sequence length represents this cache's state.
+
+        MoSRAH heads accumulate independently; (batch, head) slots have different
+        lengths depending on routing history. There is no meaningful scalar summary.
+        Use get_heads_lengths() for per-head occupancy.
+        """
+        raise NotImplementedError(
+            "SlowMoSRAHCache has no single sequence length. "
+            "Use get_heads_lengths() for per-head occupancy."
+        )
+
+    def get_max_cache_shape(self) -> int:  # type: ignore[override]
+        """Not supported — SlowMoSRAHCache is dynamic and unbounded."""
+        raise NotImplementedError(
+            "SlowMoSRAHCache is unbounded; get_max_cache_shape() is not supported."
+        )
+
+    def get_mask_sizes(  # type: ignore[override]
+        self,
+        cache_position: torch.Tensor,
+    ) -> tuple[int, int]:
+        """Not supported — SlowMoSRAHCache does not participate in HF mask construction."""
+        raise NotImplementedError(
+            "SlowMoSRAHCache does not support get_mask_sizes()."
+        )
+
+    # ---------------------------------------------------------------------------
+    # Internal helpers
+    # ---------------------------------------------------------------------------
+
+    def _make_active_mask(self) -> torch.Tensor:
+        """Construct the (B, L, T) active mask from current counts.
+
+        Returns True at position [b, l, t] iff t < _counts[b, l], i.e. the slot
+        has been written. Positions at or beyond the count are junk and must be
+        excluded by downstream attention.
+        """
+        cap = self.buffer_capacity
+        return (
+            torch.arange(cap, device=self.keys.device)
+            .expand(self.batch_size, self.num_mosrah_heads, cap)
+            < self._counts.unsqueeze(-1)
+        )
+
+    def _expand(self) -> None:
+        """Double the buffer capacity, preserving existing data.
+
+        Called by update() when an incoming batch of tokens would cause any
+        (batch, head) slot to exceed the current buffer capacity. All existing
+        key and value data is copied into the low half of the new buffer; the
+        high half is zero-initialised and will be filled by subsequent writes.
+        After reassignment, buffer_capacity reflects the new size automatically.
+        """
+        old_cap = self.buffer_capacity
+        new_cap = old_cap * 2
+        dev = self.keys.device
+        new_keys = torch.zeros(
+            self.batch_size, self.num_mosrah_heads, new_cap, self.head_dim, device=dev
+        )
+        new_values = torch.zeros(
+            self.batch_size, self.num_mosrah_heads, new_cap, self.head_dim, device=dev
+        )
+        new_keys[:, :, :old_cap, :] = self.keys
+        new_values[:, :, :old_cap, :] = self.values
+        self.keys = new_keys
+        self.values = new_values

__init__.py

@@ -1,21 +1,21 @@
-from .configuration import ShramConfig
-from .decoder_layer import DecoderLayer
-from .huggingface import ShramForCausalLM
-from .__attention__load_balance_loss import LoadBalanceLoss
-from .mlp import SwiGLUMLP
-from .model import ShramModel
-from .rope import RotaryEmbedding
-from .__attention__router import MoSRAHRouter
-from .__cache__mosrah_cache import MoSRAHCache
-
-__all__ = [
-    "DecoderLayer",
-    "LoadBalanceLoss",
-    "MoSRAHCache",
-    "MoSRAHRouter",
-    "ShramConfig",
-    "ShramForCausalLM",
-    "ShramModel",
-    "RotaryEmbedding",
-    "SwiGLUMLP",
-]
+from .configuration import ShramConfig
+from .decoder_layer import DecoderLayer
+from .huggingface import ShramForCausalLM
+from .__attention__load_balance_loss import LoadBalanceLoss
+from .mlp import SwiGLUMLP
+from .model import ShramModel
+from .rope import RotaryEmbedding
+from .__attention__router import MoSRAHRouter
+from .__cache__mosrah_cache import MoSRAHCache
+
+__all__ = [
+    "DecoderLayer",
+    "LoadBalanceLoss",
+    "MoSRAHCache",
+    "MoSRAHRouter",
+    "ShramConfig",
+    "ShramForCausalLM",
+    "ShramModel",
+    "RotaryEmbedding",
+    "SwiGLUMLP",
+]

config.json

@@ -1,28 +1,28 @@
-{
-  "alpha": 1.0,
-  "attention_dropout": 0.0,
-  "auto_map": {
-    "AutoConfig": "configuration.ShramConfig",
-    "AutoModelForCausalLM": "huggingface.ShramForCausalLM"
-  },
-  "beta": 32.0,
-  "head_dim": 16,
-  "hidden_size": 512,
-  "inference_sequence_length": 1024,
-  "intermediate_size": 1366,
-  "local_rope_theta": 10000.0,
-  "model_type": "shram",
-  "mosrah_rope_theta": 10000.0,
-  "num_hidden_layers": 12,
-  "num_mosrah_heads": 16,
-  "num_selected_heads": 16,
-  "num_sliding_window_heads": 16,
-  "rms_norm_eps": 1e-05,
-  "rope_mode": "main_sequence",
-  "tie_word_embeddings": false,
-  "training_sequence_length": 1024,
-  "transformers_version": "5.3.0",
-  "use_cache": true,
-  "vocab_size": 50277,
-  "window_size": 128
-}
+{
+  "alpha": 1.0,
+  "attention_dropout": 0.0,
+  "auto_map": {
+    "AutoConfig": "configuration.ShramConfig",
+    "AutoModelForCausalLM": "huggingface.ShramForCausalLM"
+  },
+  "beta": 32.0,
+  "head_dim": 16,
+  "hidden_size": 512,
+  "inference_sequence_length": 1024,
+  "intermediate_size": 1366,
+  "local_rope_theta": 10000.0,
+  "model_type": "shram",
+  "mosrah_rope_theta": 10000.0,
+  "num_hidden_layers": 12,
+  "num_mosrah_heads": 16,
+  "num_selected_heads": 16,
+  "num_sliding_window_heads": 16,
+  "rms_norm_eps": 1e-05,
+  "rope_mode": "main_sequence",
+  "tie_word_embeddings": false,
+  "training_sequence_length": 1024,
+  "transformers_version": "5.7.0",
+  "use_cache": true,
+  "vocab_size": 50277,
+  "window_size": 128
+}

configuration.py

@@ -1,199 +1,178 @@
-"""Configuration for the SHRAM transformer.
-
-All architectural parameters that vary across model scales or are meaningful research
-variables are expressed here. Architectural constants (no bias in linear layers,
-SwiGLU activation with SiLU gate) are implemented in the relevant modules and
-documented at the point of use — they are not config parameters because they do not
-vary and changing them produces a different architecture, not a different scale.
-
-RoPE configuration is owned entirely by this config. Each attention path reads its
-parameters directly and constructs its own RotaryEmbedding instance explicitly — no
-HuggingFace rope infrastructure is used. See Unit 5.A design decisions in plan.md.
-"""
-
-from transformers import PretrainedConfig
-
-
-class ShramConfig(PretrainedConfig):
-    """Configuration class for the SHRAM decoder-only transformer.
-
-    SHRAM (Sparse Hybrid Token Routed Attention Mixture) replaces every standard
-    attention layer with a hybrid layer H(x) = h_l(x) + h_s(x), where h_l is a
-    local sliding-window causal attention path and h_s is the MoSRAH sparse routed
-    path. All other components follow the Llama 3 baseline.
-
-    This config is the single source of truth for every architectural dimension of the
-    model. Nothing in the architecture may use a literal number that belongs here.
-
-    Two independent RoPE configurations exist — one per attention path:
-
-    - h_l always uses standard RoPE with ``local_rope_theta``.
-    - BEA always uses YaRN with ``mosrah_rope_theta``, ``training_sequence_length``,
-      ``inference_sequence_length``, ``alpha``, and ``beta``. When
-      ``inference_sequence_length == training_sequence_length`` the YaRN scale factor
-      ``s = 1`` and YaRN reduces exactly to standard RoPE — this is the default state
-      and the correct setting for experiments that do not require context extension.
-
-    Registered with HuggingFace AutoClass via ``auto_map``. Instantiate from the Hub::
-
-        config = AutoConfig.from_pretrained(
-            "your-namespace/advanced-transformers-lib",
-            trust_remote_code=True,
-            num_hidden_layers=12,
-        )
-        model = AutoModelForCausalLM.from_config(config)
-
-    Args:
-        vocab_size: Vocabulary size. Controls the embedding table and output logits
-            dimension. Must match the tokenizer.
-        embedding_width: Model width ``d``. The dimension of the residual stream.
-        mlp_width: FFN hidden dimension.
-        num_decoder_layers: Number of transformer blocks stacked in sequence.
-        num_sliding_window_heads: Number of heads in the local sliding-window path h_l.
-        num_mosrah_heads: Total MoSRAH expert heads available ``L``.
-        num_selected_heads: MoSRAH heads each token selects ``K``.
-        head_dim: Per-head dimension, shared by both attention paths. Must be even
-            (RoPE rotates dimensions in pairs). Paper uses 16.
-        window_size: Sliding window size for h_l. Paper uses 128.
-        rope_mode: RoPE position encoding mode for BEA. ``"main_sequence"`` supplies
-            original sequence positions; ``"semantic_sequence"`` supplies local slot
-            indices. Both are required; experimentally correct mode is undetermined
-            (paper §4). Default ``"main_sequence"``.
-        rms_norm_eps: Epsilon for RMSNorm layers.
-        local_rope_theta: RoPE base frequency ``b`` for the local attention path h_l.
-            Paper uses b=10000.
-        mosrah_rope_theta: RoPE base frequency ``b`` for the BEA path. Paper uses
-            b=10000.
-        training_sequence_length: Context length ``C_train`` the model was or will be
-            trained at. Used to compute the YaRN scale factor for BEA.
-        inference_sequence_length: Context length ``C_target`` the model must support
-            at inference. When equal to ``training_sequence_length``, scale ``s=1``
-            and YaRN reduces to standard RoPE.
-        alpha: YaRN ramp lower boundary α (paper §A.2). Frequency dimensions with
-            ``r(d) < alpha`` are fully interpolated by scale s. Paper value: 1.0.
-        beta: YaRN ramp upper boundary β (paper §A.2). Frequency dimensions with
-            ``r(d) > beta`` are left unscaled. Paper value: 32.0.
-        attention_dropout: Dropout probability on attention weights. Default 0.0.
-        use_cache: Whether to return past_key_values for KV caching.
-        output_hidden_states: Whether to return hidden states after each layer.
-        tie_word_embeddings: Whether input embedding and LM head share weights.
-    """
-
-    model_type = "shram"
-
-    auto_map = {
-        "AutoConfig": "configuration.ShramConfig",
-        "AutoModelForCausalLM": "huggingface.ShramForCausalLM",
-    }
-
-    def __init__(
-        self,
-        vocab_size: int = 50277,
-        embedding_width: int = 512,
-        mlp_width: int = 1366,
-        num_decoder_layers: int = 12,
-        num_sliding_window_heads: int = 16,
-        num_mosrah_heads: int = 16,
-        num_selected_heads: int = 16,
-        head_dim: int = 16,
-        window_size: int = 128,
-        rope_mode: str = "main_sequence",
-        rms_norm_eps: float = 1e-5,
-        local_rope_theta: float = 10000.0,
-        mosrah_rope_theta: float = 10000.0,
-        training_sequence_length: int = 1024,
-        alpha: float = 1.0,
-        beta: float = 32.0,
-        attention_dropout: float = 0.0,
-        use_cache: bool = True,
-        output_hidden_states: bool = False,
-        tie_word_embeddings: bool = False,
-        **kwargs,
-    ):
-        if head_dim % 2 != 0:
-            raise ValueError(
-                f"head_dim must be even (RoPE rotates dimensions in pairs). "
-                f"Got head_dim={head_dim}."
-            )
-
-        if rope_mode not in {"main_sequence", "semantic_sequence"}:
-            raise ValueError(
-                f"rope_mode must be 'main_sequence' or 'semantic_sequence', "
-                f"got '{rope_mode}'."
-            )
-
-        if training_sequence_length <= 0:
-            raise ValueError(
-                f"training_sequence_length must be positive, "
-                f"got {training_sequence_length}."
-            )
-
-        # inference_sequence_length is not a constructor parameter. It defaults to
-        # training_sequence_length (scale=1.0, standard RoPE). If a saved config
-        # carries the field through kwargs (after set_inference_context() was called
-        # before saving), restore it here with validation.
-        saved_inference_length = kwargs.pop("inference_sequence_length", training_sequence_length)
-        if saved_inference_length <= 0:
-            raise ValueError(
-                f"inference_sequence_length must be positive, "
-                f"got {saved_inference_length}."
-            )
-
-        self.vocab_size = vocab_size
-        self.hidden_size = embedding_width
-        self.intermediate_size = mlp_width
-        self.num_hidden_layers = num_decoder_layers
-        self.num_sliding_window_heads = num_sliding_window_heads
-        self.num_mosrah_heads = num_mosrah_heads
-        self.num_selected_heads = num_selected_heads
-        self.head_dim = head_dim
-        self.window_size = window_size
-        self.rope_mode = rope_mode
-        self.rms_norm_eps = rms_norm_eps
-        self.local_rope_theta = local_rope_theta
-        self.mosrah_rope_theta = mosrah_rope_theta
-        self.training_sequence_length = training_sequence_length
-        self.inference_sequence_length = saved_inference_length
-        self.alpha = alpha
-        self.beta = beta
-        self.attention_dropout = attention_dropout
-        self.use_cache = use_cache
-
-        super().__init__(
-            tie_word_embeddings=tie_word_embeddings,
-            output_hidden_states=output_hidden_states,
-            **kwargs,
-        )
-
-        # Promote auto_map to an instance attribute so PretrainedConfig.to_dict()
-        # serialises it into config.json.
-        self.auto_map = type(self).auto_map
-
-    @property
-    def scale(self) -> float:
-        """YaRN context extension scale factor s = inference_sequence_length / training_sequence_length.
-
-        When scale == 1.0, YaRN reduces exactly to standard RoPE — all frequency
-        adjustments cancel and A_rope = 1. This is the default state.
-        """
-        return self.inference_sequence_length / self.training_sequence_length
-
-    def set_inference_context(self, inference_sequence_length: int) -> None:
-        """Set the inference context length for YaRN context extension.
-
-        This is the only supported way to set inference_sequence_length. At construction
-        the inference context defaults to training_sequence_length (scale=1.0, standard
-        RoPE). Call this method to configure a longer inference context, which causes
-        YaRN to interpolate frequencies and extend the effective context window.
-
-        Args:
-            inference_sequence_length: Target inference context length. Must be positive.
-                Values equal to training_sequence_length produce scale=1.0 (standard RoPE).
-                Values greater than training_sequence_length enable YaRN extrapolation.
-        """
-        if inference_sequence_length <= 0:
-            raise ValueError(
-                f"inference_sequence_length must be positive, "
-                f"got {inference_sequence_length}."
-            )
-        self.inference_sequence_length = inference_sequence_length
+"""Configuration for the SHRAM transformer.
+
+All architectural parameters that vary across model scales or are meaningful research
+variables are expressed here. Architectural constants (no bias in linear layers,
+SwiGLU activation with SiLU gate) are implemented in the relevant modules and
+documented at the point of use — they are not config parameters because they do not
+vary and changing them produces a different architecture, not a different scale.
+
+RoPE configuration is owned entirely by this config. Each attention path reads its
+parameters directly and constructs its own RotaryEmbedding instance explicitly — no
+HuggingFace rope infrastructure is used. See Unit 5.A design decisions in plan.md.
+"""
+
+from transformers import PretrainedConfig
+
+
+class ShramConfig(PretrainedConfig):
+    """Configuration class for the SHRAM decoder-only transformer.
+
+    SHRAM (Sparse Hybrid Token Routed Attention Mixture) replaces every standard
+    attention layer with a hybrid layer H(x) = h_l(x) + h_s(x), where h_l is a
+    local sliding-window causal attention path and h_s is the MoSRAH sparse routed
+    path. All other components follow the Llama 3 baseline.
+
+    This config is the single source of truth for every architectural dimension of the
+    model. Nothing in the architecture may use a literal number that belongs here.
+
+    Two independent RoPE configurations exist — one per attention path:
+
+    - h_l always uses standard RoPE with ``local_rope_theta``.
+    - BEA always uses YaRN with ``mosrah_rope_theta``, ``training_sequence_length``,
+      ``inference_sequence_length``, ``alpha``, and ``beta``. When
+      ``inference_sequence_length == training_sequence_length`` the YaRN scale factor
+      ``s = 1`` and YaRN reduces exactly to standard RoPE — this is the default state
+      and the correct setting for experiments that do not require context extension.
+
+    Registered with HuggingFace AutoClass via ``auto_map``. Instantiate from the Hub::
+
+        config = AutoConfig.from_pretrained(
+            "your-namespace/advanced-transformers-lib",
+            trust_remote_code=True,
+            num_hidden_layers=12,
+        )
+        model = AutoModelForCausalLM.from_config(config)
+
+    Args:
+        vocab_size: Vocabulary size. Controls the embedding table and output logits
+            dimension. Must match the tokenizer.
+        embedding_width: Model width ``d``. The dimension of the residual stream.
+        mlp_width: FFN hidden dimension.
+        num_decoder_layers: Number of transformer blocks stacked in sequence.
+        num_sliding_window_heads: Number of heads in the local sliding-window path h_l.
+        num_mosrah_heads: Total MoSRAH expert heads available ``L``.
+        num_selected_heads: MoSRAH heads each token selects ``K``.
+        head_dim: Per-head dimension, shared by both attention paths. Must be even
+            (RoPE rotates dimensions in pairs). Paper uses 16.
+        window_size: Sliding window size for h_l. Paper uses 128.
+        rope_mode: RoPE position encoding mode for BEA. ``"main_sequence"`` supplies
+            original sequence positions; ``"semantic_sequence"`` supplies local slot
+            indices. Both are required; experimentally correct mode is undetermined
+            (paper §4). Default ``"main_sequence"``.
+        rms_norm_eps: Epsilon for RMSNorm layers.
+        local_rope_theta: RoPE base frequency ``b`` for the local attention path h_l.
+            Paper uses b=10000.
+        mosrah_rope_theta: RoPE base frequency ``b`` for the BEA path. Paper uses
+            b=10000.
+        training_sequence_length: Context length ``C_train`` the model was or will be
+            trained at. Used to compute the YaRN scale factor for BEA.
+        inference_sequence_length: Context length ``C_target`` the model must support
+            at inference. Optional; defaults to ``training_sequence_length`` so that
+            ``scale=1`` and YaRN reduces to standard RoPE unless explicitly extended.
+        alpha: YaRN ramp lower boundary α (paper §A.2). Frequency dimensions with
+            ``r(d) < alpha`` are fully interpolated by scale s. Paper value: 1.0.
+        beta: YaRN ramp upper boundary β (paper §A.2). Frequency dimensions with
+            ``r(d) > beta`` are left unscaled. Paper value: 32.0.
+        attention_dropout: Dropout probability on attention weights. Default 0.0.
+        use_cache: Whether to return past_key_values for KV caching.
+        output_hidden_states: Whether to return hidden states after each layer.
+        tie_word_embeddings: Whether input embedding and LM head share weights.
+    """
+
+    model_type = "shram"
+
+    auto_map = {
+        "AutoConfig": "configuration.ShramConfig",
+        "AutoModelForCausalLM": "huggingface.ShramForCausalLM",
+    }
+
+    def __init__(
+        self,
+        vocab_size: int = 50277,
+        embedding_width: int = 512,
+        mlp_width: int = 1366,
+        num_decoder_layers: int = 12,
+        num_sliding_window_heads: int = 16,
+        num_mosrah_heads: int = 16,
+        num_selected_heads: int = 16,
+        head_dim: int = 16,
+        window_size: int = 128,
+        rope_mode: str = "main_sequence",
+        rms_norm_eps: float = 1e-5,
+        local_rope_theta: float = 10000.0,
+        mosrah_rope_theta: float = 10000.0,
+        training_sequence_length: int = 1024,
+        inference_sequence_length: int | None = None,
+        alpha: float = 1.0,
+        beta: float = 32.0,
+        attention_dropout: float = 0.0,
+        use_cache: bool = True,
+        output_hidden_states: bool = False,
+        tie_word_embeddings: bool = False,
+        **kwargs,
+    ):
+        if head_dim % 2 != 0:
+            raise ValueError(
+                f"head_dim must be even (RoPE rotates dimensions in pairs). "
+                f"Got head_dim={head_dim}."
+            )
+
+        if rope_mode not in {"main_sequence", "semantic_sequence"}:
+            raise ValueError(
+                f"rope_mode must be 'main_sequence' or 'semantic_sequence', "
+                f"got '{rope_mode}'."
+            )
+
+        if training_sequence_length <= 0:
+            raise ValueError(
+                f"training_sequence_length must be positive, "
+                f"got {training_sequence_length}."
+            )
+
+        if inference_sequence_length is None:
+            inference_sequence_length = training_sequence_length
+        if inference_sequence_length <= 0:
+            raise ValueError(
+                f"inference_sequence_length must be positive, "
+                f"got {inference_sequence_length}."
+            )
+
+        self.vocab_size = vocab_size
+        self.hidden_size = embedding_width
+        self.intermediate_size = mlp_width
+        self.num_hidden_layers = num_decoder_layers
+        self.num_sliding_window_heads = num_sliding_window_heads
+        self.num_mosrah_heads = num_mosrah_heads
+        self.num_selected_heads = num_selected_heads
+        self.head_dim = head_dim
+        self.window_size = window_size
+        self.rope_mode = rope_mode
+        self.rms_norm_eps = rms_norm_eps
+        self.local_rope_theta = local_rope_theta
+        self.mosrah_rope_theta = mosrah_rope_theta
+        self.training_sequence_length = training_sequence_length
+        self.inference_sequence_length = inference_sequence_length
+        self.alpha = alpha
+        self.beta = beta
+        self.attention_dropout = attention_dropout
+        self.use_cache = use_cache
+
+        super().__init__(
+            tie_word_embeddings=tie_word_embeddings,
+            output_hidden_states=output_hidden_states,
+            **kwargs,
+        )
+
+        # Promote auto_map to an instance attribute so PretrainedConfig.to_dict()
+        # serialises it into config.json.
+        self.auto_map = type(self).auto_map
+
+    @property
+    def scale(self) -> float:
+        """YaRN context extension scale factor s = inference_sequence_length / training_sequence_length.
+
+        When scale == 1.0, YaRN reduces exactly to standard RoPE — all frequency
+        adjustments cancel and A_rope = 1. This is the default state.
+        """
+        return self.inference_sequence_length / self.training_sequence_length
+

decoder_layer.py

@@ -1,87 +1,87 @@
-"""Decoder layer — a single transformer block.
-
-Each block applies pre-norm hybrid attention followed by pre-norm MLP, with
-residual connections around both sublayers:
-
-    normed_attn = RMSNorm(x)
-    attn_out, load_balance_loss, max_vio = SHRAMHybridLayer(normed_attn, ...)
-    h = x + attn_out
-
-    normed_mlp = RMSNorm(h)
-    mlp_out = SwiGLUMLP(normed_mlp)
-    out = h + mlp_out
-
-Pre-norm keeps the residual stream unnormalised. Gradients flow more cleanly
-through unnormalised residuals at depth, and each sublayer receives a stable,
-normalised view of the signal.
-
-Two independent RMSNorm instances are used — one before attention, one before
-MLP. They learn different scalings because they precede layers with different
-dynamic ranges. Sharing them would be wrong.
-
-torch.nn.RMSNorm is used directly (available from PyTorch 2.4+). It omits mean
-subtraction, is faster than LayerNorm, and proved more stable at scale.
-"""
-
-import torch
-import torch.nn as nn
-
-from .__attention__shram import SHRAMHybridLayer
-from .__cache__shram_layer_cache import ShramLayerCache
-from .configuration import ShramConfig
-from .mlp import SwiGLUMLP
-
-
-class DecoderLayer(nn.Module):
-    """A single pre-norm SHRAM decoder block.
-
-    Composes SHRAMHybridLayer and SwiGLUMLP with residual connections and
-    independent RMSNorm instances on each sublayer input.
-
-    Args:
-        config: SHRAM config. Must expose ``hidden_size`` and ``rms_norm_eps``
-            in addition to the fields required by SHRAMHybridLayer and
-            SwiGLUMLP.
-    """
-
-    def __init__(self, config: ShramConfig) -> None:
-        super().__init__()
-        self.attn_norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        self.mlp_norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        self.attention = SHRAMHybridLayer(config)
-        self.mlp = SwiGLUMLP(config)
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        position_ids: torch.Tensor,
-        active_mask: torch.Tensor,
-        cache: ShramLayerCache | None = None,
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        """Apply one decoder block to the input.
-
-        Args:
-            x: Input of shape (batch, seq_len, hidden_size).
-            position_ids: Authoritative positions of shape (batch, seq_len).
-            active_mask: Current-chunk active mask of shape (batch, seq_len),
-                where True means the token is semantically live. Forwarded
-                unchanged to the hybrid attention layer.
-            cache: Optional per-layer SHRAM cache passed through to the hybrid
-                attention layer unchanged.
-
-        Returns:
-            output: Tensor of shape (batch, seq_len, hidden_size).
-            load_balance_loss: Scalar sparse-path load-balance loss propagated
-                from SHRAMHybridLayer.
-            max_vio: Detached scalar routing-imbalance summary. Passed through
-                unchanged from SHRAMHybridLayer; see MoSRAHRouter for semantics.
-        """
-        attn_out, load_balance_loss, max_vio = self.attention(
-            hidden_states=self.attn_norm(x),
-            position_ids=position_ids,
-            active_mask=active_mask,
-            cache=cache,
-        )
-        hidden_states = x + attn_out
-        output = hidden_states + self.mlp(self.mlp_norm(hidden_states))
+"""Decoder layer — a single transformer block.
+
+Each block applies pre-norm hybrid attention followed by pre-norm MLP, with
+residual connections around both sublayers:
+
+    normed_attn = RMSNorm(x)
+    attn_out, load_balance_loss, max_vio = SHRAMHybridLayer(normed_attn, ...)
+    h = x + attn_out
+
+    normed_mlp = RMSNorm(h)
+    mlp_out = SwiGLUMLP(normed_mlp)
+    out = h + mlp_out
+
+Pre-norm keeps the residual stream unnormalised. Gradients flow more cleanly
+through unnormalised residuals at depth, and each sublayer receives a stable,
+normalised view of the signal.
+
+Two independent RMSNorm instances are used — one before attention, one before
+MLP. They learn different scalings because they precede layers with different
+dynamic ranges. Sharing them would be wrong.
+
+torch.nn.RMSNorm is used directly (available from PyTorch 2.4+). It omits mean
+subtraction, is faster than LayerNorm, and proved more stable at scale.
+"""
+
+import torch
+import torch.nn as nn
+
+from .__attention__shram import SHRAMHybridLayer
+from .__cache__shram_layer_cache import ShramLayerCache
+from .configuration import ShramConfig
+from .mlp import SwiGLUMLP
+
+
+class DecoderLayer(nn.Module):
+    """A single pre-norm SHRAM decoder block.
+
+    Composes SHRAMHybridLayer and SwiGLUMLP with residual connections and
+    independent RMSNorm instances on each sublayer input.
+
+    Args:
+        config: SHRAM config. Must expose ``hidden_size`` and ``rms_norm_eps``
+            in addition to the fields required by SHRAMHybridLayer and
+            SwiGLUMLP.
+    """
+
+    def __init__(self, config: ShramConfig) -> None:
+        super().__init__()
+        self.attn_norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.mlp_norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.attention = SHRAMHybridLayer(config)
+        self.mlp = SwiGLUMLP(config)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        position_ids: torch.Tensor,
+        active_mask: torch.Tensor,
+        cache: ShramLayerCache | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Apply one decoder block to the input.
+
+        Args:
+            x: Input of shape (batch, seq_len, hidden_size).
+            position_ids: Authoritative positions of shape (batch, seq_len).
+            active_mask: Current-chunk active mask of shape (batch, seq_len),
+                where True means the token is semantically live. Forwarded
+                unchanged to the hybrid attention layer.
+            cache: Optional per-layer SHRAM cache passed through to the hybrid
+                attention layer unchanged.
+
+        Returns:
+            output: Tensor of shape (batch, seq_len, hidden_size).
+            load_balance_loss: Scalar sparse-path load-balance loss propagated
+                from SHRAMHybridLayer.
+            max_vio: Detached scalar routing-imbalance summary. Passed through
+                unchanged from SHRAMHybridLayer; see MoSRAHRouter for semantics.
+        """
+        attn_out, load_balance_loss, max_vio = self.attention(
+            hidden_states=self.attn_norm(x),
+            position_ids=position_ids,
+            active_mask=active_mask,
+            cache=cache,
+        )
+        hidden_states = x + attn_out
+        output = hidden_states + self.mlp(self.mlp_norm(hidden_states))
         return output, load_balance_loss, max_vio

huggingface.py

@@ -1,518 +1,518 @@
-"""HuggingFace causal-LM wrapper for SHRAM.
-
-ShramForCausalLM is the HuggingFace-facing language-model boundary for SHRAM.
-It owns token embedding lookup, LM-head projection, wrapper-level next-token
-cross-entropy loss, config-controlled tied embeddings, and generation/cache
-orchestration at the wrapper boundary.
-
-The backbone remains a pure transformer stack. ShramModel accepts pre-embedded
-hidden states together with current position IDs, a current active mask, and an
-optional ShramCache. It has no knowledge of token IDs, vocabulary projection,
-or causal-LM loss.
-
-HuggingFace generation reaches this wrapper with two different tensor
-conventions:
-
-- ``position_ids`` is a current-step tensor. GenerationMixin updates the total
-  sequence state between steps, then slices position-bearing tensors back down
-  before calling ``forward()``.
-- ``attention_mask`` is a full 2D mask over the total sequence so far. This
-  wrapper slices its recent chunk to produce the current semantic liveness mask
-  expected by the backbone.
-
-Generation-created caches are handled in ``_prepare_cache_for_generation``.
-That hook ensures HuggingFace generation uses ShramCache rather than a generic
-dynamic cache. The direct ``forward()`` path does not silently create caches;
-when ``use_cache=True`` it expects a truthful ShramCache to have been supplied.
-"""
-
-from dataclasses import dataclass
-from typing import Any
-
-import torch
-import torch.nn as nn
-from transformers import GenerationMixin, PreTrainedModel
-from transformers.cache_utils import Cache
-from transformers.generation.configuration_utils import GenerationMode
-from transformers.modeling_outputs import CausalLMOutputWithPast
-
-from .__cache__shram_cache import ShramCache
-from .configuration import ShramConfig
-from .model import ShramModel
-
-
-@dataclass
-class ShramCausalLMOutput(CausalLMOutputWithPast):
-    """SHRAM causal-LM wrapper output.
-
-    This subclasses HuggingFace's standard ``CausalLMOutputWithPast``.
-    Dataclass inheritance is sufficient here: all standard causal-LM fields and
-    ModelOutput behavior are inherited from the parent, and this subclass adds
-    only the SHRAM-specific wrapper outputs.
-    """
-
-    ce_loss: torch.FloatTensor | None = None
-    load_balance_loss: torch.FloatTensor | None = None
-    max_vio: torch.FloatTensor | None = None
-
-
-class ShramForCausalLM(PreTrainedModel, GenerationMixin):
-    """HuggingFace-facing causal language model wrapper for SHRAM.
-
-    Owns token embeddings, LM-head projection, wrapper-level shifted CE loss,
-    tied embedding configuration, and generation/cache boundary behavior.
-    Delegates all transformer computation to ``ShramModel``.
-
-    Args:
-        config: SHRAM model configuration.
-    """
-
-    config_class = ShramConfig
-    base_model_prefix = "model"
-    _no_split_modules = ["DecoderLayer"]
-    supports_gradient_checkpointing = True
-
-    def __init__(self, config: ShramConfig) -> None:
-        super().__init__(config)
-        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
-        self.model = ShramModel(config)
-        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
-        self._configure_tied_embeddings()
-        self.post_init()
-
-    def _configure_tied_embeddings(self) -> None:
-        """Apply config-controlled tied embedding behavior on this instance."""
-        if self.config.tie_word_embeddings:
-            self.lm_head.weight = self.embed_tokens.weight
-            self._tied_weights_keys = {
-                "lm_head.weight": "embed_tokens.weight",
-            }
-        else:
-            self._tied_weights_keys = {}
-
-    def get_input_embeddings(self) -> nn.Embedding:
-        """Return the token embedding matrix."""
-        return self.embed_tokens
-
-    def set_input_embeddings(self, value: nn.Embedding) -> None:
-        """Replace the token embedding matrix."""
-        self.embed_tokens = value
-        self._configure_tied_embeddings()
-
-    def get_output_embeddings(self) -> nn.Linear:
-        """Return the LM head."""
-        return self.lm_head
-
-    def set_output_embeddings(self, value: nn.Linear) -> None:
-        """Replace the LM head."""
-        self.lm_head = value
-        self._configure_tied_embeddings()
-
-    def _build_shram_cache(
-        self,
-        batch_size: int,
-        device: torch.device,
-    ) -> ShramCache:
-        """Construct a fresh top-level SHRAM cache."""
-        return ShramCache(
-            num_hidden_layers=self.config.num_hidden_layers,
-            sliding_window=self.config.window_size,
-            num_local_heads=self.config.num_sliding_window_heads,
-            local_head_dim=self.config.head_dim,
-            num_mosrah_heads=self.config.num_mosrah_heads,
-            mosrah_head_dim=self.config.head_dim,
-            batch_size=batch_size,
-            device=device,
-        )
-
-    def _validate_generation_cache_request(
-        self,
-        generation_config: Any,
-        model_kwargs: dict[str, Any],
-        generation_mode: GenerationMode,
-    ) -> None:
-        """Validate SHRAM's generation-side cache policy."""
-        if generation_mode in {
-            GenerationMode.ASSISTED_GENERATION,
-            GenerationMode.CONTRASTIVE_SEARCH,
-        }:
-            raise NotImplementedError(
-                "ShramForCausalLM does not currently support assisted generation "
-                "or contrastive search because ShramCache does not support crop()."
-            )
-
-        user_defined_cache = model_kwargs.get("past_key_values")
-        if user_defined_cache is not None:
-            if generation_config.cache_implementation is not None:
-                raise ValueError(
-                    "Passing both `cache_implementation` and `past_key_values` "
-                    "is unsupported. Please use only one."
-                )
-            if isinstance(user_defined_cache, tuple):
-                raise ValueError(
-                    "Passing a tuple of `past_key_values` is not supported. "
-                    "Please use a `ShramCache` instance."
-                )
-            if not isinstance(user_defined_cache, ShramCache):
-                raise TypeError(
-                    "ShramForCausalLM requires `past_key_values` to be a "
-                    "`ShramCache` instance."
-                )
-
-        if (
-            user_defined_cache is None
-            and generation_config.use_cache
-            and generation_config.cache_implementation is not None
-        ):
-            raise ValueError(
-                "ShramForCausalLM does not support `cache_implementation`. "
-                "Generation-created caches must be `ShramCache` objects."
-            )
-
-    def _prepare_cache_for_generation(
-        self,
-        generation_config: Any,
-        model_kwargs: dict[str, Any],
-        generation_mode: GenerationMode,
-        batch_size: int,
-        max_cache_length: int,
-    ) -> None:
-        """Ensure HuggingFace generation uses ShramCache.
-
-        This is the SHRAM-specific generation hook. The rest of the default
-        generation plumbing is kept intact as much as possible.
-
-        Args:
-            generation_config: Active generation configuration.
-            model_kwargs: Generation kwargs, updated in place.
-            generation_mode: HuggingFace generation mode.
-            batch_size: Effective generation batch size.
-            max_cache_length: Requested cache length. Accepted but unused here.
-        """
-        self._validate_generation_cache_request(
-            generation_config=generation_config,
-            model_kwargs=model_kwargs,
-            generation_mode=generation_mode,
-        )
-
-        if model_kwargs.get("past_key_values") is not None:
-            return
-
-        if not generation_config.use_cache:
-            return
-
-        num_repeats = max(
-            generation_config.num_beams or 1,
-            generation_config.num_return_sequences or 1,
-        )
-        model_kwargs["past_key_values"] = self._build_shram_cache(
-            batch_size=batch_size*num_repeats,
-            device=self.embed_tokens.weight.device,
-        )
-
-    def _reorder_cache(
-        self,
-        past_key_values: Cache,
-        beam_idx: torch.Tensor,
-    ) -> Cache:
-        """Reorder the cache in place for beam search."""
-        past_key_values.reorder_cache(beam_idx)
-        return past_key_values
-
-    def _validate_input_ids(self, input_ids: torch.Tensor) -> None:
-        """Validate token IDs at the wrapper boundary."""
-        if input_ids.ndim != 2:
-            raise ValueError("input_ids must have shape (batch, seq_len).")
-        if input_ids.shape[1] == 0:
-            raise ValueError("input_ids sequence length must be nonzero.")
-        if input_ids.dtype != torch.long:
-            raise TypeError("input_ids must be an long int tensor.")
-
-    def _validate_attention_mask(
-        self,
-        input_ids: torch.Tensor,
-        attention_mask: torch.Tensor | None,
-    ) -> None:
-        """Validate the full-sequence attention mask."""
-        if attention_mask is None:
-            return
-        if attention_mask.ndim != 2:
-            raise ValueError("attention_mask must have shape (batch, total_seq_len).")
-        if attention_mask.shape[0] != input_ids.shape[0]:
-            raise ValueError("attention_mask batch dimension must match input_ids.")
-        if attention_mask.shape[1] < input_ids.shape[1]:
-            raise ValueError(
-                "attention_mask must be at least as long as the current input_ids chunk."
-            )
-
-    def _validate_position_ids(
-        self,
-        input_ids: torch.Tensor,
-        position_ids: torch.Tensor | None,
-    ) -> None:
-        """Validate current-step position IDs."""
-        if position_ids is None:
-            return
-        if position_ids.ndim != 2:
-            raise ValueError("position_ids must have shape (batch, seq_len).")
-        if position_ids.shape != input_ids.shape:
-            raise ValueError(
-                "position_ids must match the current input_ids shape exactly."
-            )
-        if input_ids.dtype != torch.long:
-            raise TypeError("position_ids must be an long tensor.")
-
-    def _validate_labels(
-        self,
-        input_ids: torch.Tensor,
-        labels: torch.Tensor | None,
-    ) -> None:
-        """Validate label shape at the wrapper boundary."""
-        if labels is None:
-            return
-        if labels.ndim != 2:
-            raise ValueError("labels must have shape (batch, seq_len).")
-        if labels.shape != input_ids.shape:
-            raise ValueError("labels must have the same shape as input_ids.")
-        if input_ids.dtype != torch.long:
-            raise TypeError("labels must be a long tensor.")
-
-    def _validate_cache_inputs(
-        self,
-        use_cache: bool,
-        past_key_values: Cache | None,
-    ) -> None:
-        """Validate cache policy for direct wrapper calls."""
-        if use_cache:
-            if past_key_values is None:
-                raise ValueError(
-                    "use_cache=True requires an explicit ShramCache. During "
-                    "generate(), HuggingFace should supply this through "
-                    "_prepare_cache_for_generation()."
-                )
-            if not isinstance(past_key_values, ShramCache):
-                raise TypeError(
-                    "past_key_values must be a ShramCache when use_cache=True."
-                )
-            return
-
-        if past_key_values is not None:
-            raise ValueError("past_key_values was provided while use_cache=False.")
-
-    def _validate_position_sources(
-        self,
-        use_cache: bool,
-        attention_mask: torch.Tensor | None,
-        position_ids: torch.Tensor | None,
-    ) -> None:
-        """Validate that cached forward has a truthful source of positions."""
-        if use_cache and attention_mask is None and position_ids is None:
-            raise ValueError(
-                "Cached forward requires either position_ids or attention_mask."
-            )
-
-    def _validate_hf_boundary(
-        self,
-        output_attentions: bool | None,
-        return_dict: bool | None,
-        inputs_embeds: torch.Tensor | None,
-        cache_position: torch.Tensor | None,
-        extra_kwargs: dict[str, Any],
-    ) -> None:
-        """Validate unsupported HuggingFace-facing wrapper inputs."""
-        if output_attentions:
-            raise NotImplementedError(
-                "ShramForCausalLM does not expose output_attentions."
-            )
-        if return_dict is False:
-            raise ValueError(
-                "return_dict=False is not supported. "
-                "ShramForCausalLM always returns ShramCausalLMOutput."
-            )
-        if inputs_embeds is not None:
-            raise ValueError(
-                "inputs_embeds is not supported at the SHRAM wrapper boundary. "
-                "Pass input_ids instead."
-            )
-        if extra_kwargs:
-            unsupported = ", ".join(sorted(extra_kwargs))
-            raise TypeError(
-                f"Unsupported forward kwargs for ShramForCausalLM: {unsupported}"
-            )
-
-    def _standardize_full_attention_mask(
-        self,
-        input_ids: torch.Tensor,
-        attention_mask: torch.Tensor | None,
-    ) -> torch.BoolTensor:
-        """Return a concrete full-sequence boolean attention mask."""
-        if attention_mask is None:
-            return torch.ones_like(input_ids, dtype=torch.bool)
-        return attention_mask.to(dtype=torch.bool)
-
-    def _resolve_current_position_ids(
-        self,
-        input_ids: torch.Tensor,
-        position_ids: torch.Tensor | None,
-        full_attention_mask: torch.BoolTensor,
-    ) -> torch.LongTensor:
-        """Resolve concrete current-step position IDs for the backbone."""
-        if position_ids is not None:
-            return position_ids.to(dtype=torch.long)
-
-        full_position_ids = full_attention_mask.to(dtype=torch.long).cumsum(dim=-1) - 1
-        full_position_ids = full_position_ids.masked_fill(~full_attention_mask, 0)
-        current_length = input_ids.shape[1]
-        return full_position_ids[:, -current_length:]
-
-    def forward(
-        self,
-        input_ids: torch.Tensor,
-        attention_mask: torch.Tensor | None = None,
-        position_ids: torch.Tensor | None = None,
-        past_key_values: Cache | None = None,
-        use_cache: bool | None = None,
-        output_hidden_states: bool | None = None,
-        labels: torch.Tensor | None = None,
-        return_dict: bool | None = None,
-        ce_weight: float = 1.0,
-        load_balance_weight: float = 0.01,
-        **kwargs: Any,
-    ) -> ShramCausalLMOutput:
-        """Run the SHRAM causal language model wrapper.
-
-        Args:
-            input_ids: Current token IDs of shape ``(batch, seq_len)``.
-            attention_mask: Optional full 2D mask of shape
-                ``(batch, total_seq_len)``. The wrapper slices its recent chunk
-                to produce the current semantic liveness mask expected by the
-                backbone.
-            position_ids: Optional current-step position IDs of shape
-                ``(batch, seq_len)``. In ordinary HuggingFace generation this is
-                already the current-step tensor when it reaches ``forward()``.
-            past_key_values: Optional SHRAM cache. Required when
-                ``use_cache=True``.
-            use_cache: Whether to use and return a cache. Defaults to
-                ``config.use_cache``.
-            output_hidden_states: Whether to return backbone hidden states.
-                Defaults to ``config.output_hidden_states``.
-            labels: Optional target token IDs of shape ``(batch, seq_len)``.
-            return_dict: Must be ``True`` or ``None``.
-            ce_weight: Weight applied to the cross-entropy loss when combining with
-                the load-balance loss. Default 1.0.
-            load_balance_weight: Weight applied to the load-balance auxiliary loss.
-                Default 0.01, matching the paper's recommendation.
-            **kwargs: Unsupported HuggingFace kwargs fail explicitly.
-
-        Returns:
-            ``ShramCausalLMOutput`` with:
-            - ``logits`` of shape ``(batch, seq_len, vocab_size)``,
-            - ``loss`` = ``ce_weight * ce_loss + load_balance_weight * load_balance_loss``
-              when labels are provided (``None`` otherwise),
-            - ``ce_loss`` — raw unweighted cross-entropy loss for logging,
-            - ``past_key_values`` as the active ``ShramCache`` or ``None``,
-            - ``hidden_states`` when requested,
-            - ``load_balance_loss`` — raw unweighted load-balance loss from the backbone,
-            - detached ``max_vio`` from the backbone.
-        """
-        use_cache = use_cache if use_cache is not None else self.config.use_cache
-        output_hidden_states = (
-            output_hidden_states
-            if output_hidden_states is not None
-            else self.config.output_hidden_states
-        )
-
-        inputs_embeds = kwargs.pop("inputs_embeds", None)
-        output_attentions = kwargs.pop("output_attentions", None)
-        cache_position = kwargs.pop("cache_position", None)
-
-        # ------------------------------------------------------------------
-        # Validation zone.
-        #
-        # The wrapper boundary is where HuggingFace-facing inputs are judged
-        # for truthfulness before any internal work begins. These checks are
-        # intentionally front-loaded so the core logic below can assume one
-        # coherent interpretation of the call rather than defensively checking
-        # shapes, cache policy, or unsupported HF knobs at the point of use.
-        # This keeps the main sequence readable while ensuring invalid states
-        # fail before they can silently contaminate backbone execution.
-        # ------------------------------------------------------------------
-        self._validate_input_ids(input_ids)
-        self._validate_attention_mask(input_ids, attention_mask)
-        self._validate_position_ids(input_ids, position_ids)
-        self._validate_labels(input_ids, labels)
-        self._validate_cache_inputs(use_cache, past_key_values)
-        self._validate_position_sources(use_cache, attention_mask, position_ids)
-        self._validate_hf_boundary(
-            output_attentions=output_attentions,
-            return_dict=return_dict,
-            inputs_embeds=inputs_embeds,
-            cache_position=cache_position,
-            extra_kwargs=kwargs,
-        )
-
-        # ------------------------------------------------------------------
-        # Standardization zone.
-        #
-        # HuggingFace and SHRAM use different boundary conventions: generation
-        # carries a full-sequence 2D attention mask, while the SHRAM backbone
-        # wants a current-step active mask and concrete current position IDs.
-        # This zone collapses those wrapper-facing conventions into one valid
-        # backbone-facing state. After this point the core no longer reasons
-        # about optional or ambiguous input forms; it works only with concrete
-        # tensors whose semantics are already fixed.
-        # ------------------------------------------------------------------
-        full_attention_mask: torch.BoolTensor = self._standardize_full_attention_mask(
-            input_ids=input_ids,
-            attention_mask=attention_mask,
-        )
-        current_length: int = input_ids.shape[1]
-        current_active_mask: torch.BoolTensor = full_attention_mask[:, -current_length:]
-        current_position_ids: torch.LongTensor = self._resolve_current_position_ids(
-            input_ids=input_ids,
-            position_ids=position_ids,
-            full_attention_mask=full_attention_mask,
-        )
-        shram_cache: ShramCache | None = past_key_values if use_cache else None
-
-        # ------------------------------------------------------------------
-        # Core wrapper responsibilities.
-        #
-        # The wrapper's primary job is kept visible here: convert token IDs to
-        # embeddings, delegate transformer computation to ShramModel, project
-        # hidden states back to vocabulary logits, optionally compute the
-        # wrapper-level shifted next-token loss, and return the HuggingFace-
-        # facing output object. The backbone remains responsible only for
-        # transformer semantics; token/vocabulary/loss concerns stay here.
-        # ------------------------------------------------------------------
-        token_embeddings: torch.FloatTensor = self.embed_tokens(input_ids)
-        backbone_outputs = self.model(
-            inputs_embeds=token_embeddings,
-            position_ids=current_position_ids,
-            active_mask=current_active_mask,
-            cache=shram_cache,
-            output_hidden_states=output_hidden_states,
-        )
-
-        logits: torch.FloatTensor = self.lm_head(backbone_outputs["last_hidden_state"])
-
-        ce_loss: torch.FloatTensor | None = None
-        loss: torch.FloatTensor | None = None
-        if labels is not None:
-            shift_logits = logits[:, :-1, :].contiguous()
-            shift_labels = labels[:, 1:].contiguous()
-            ce_loss = nn.functional.cross_entropy(
-                shift_logits.view(-1, self.config.vocab_size),
-                shift_labels.view(-1),
-            )
-            loss = ce_weight * ce_loss + load_balance_weight * backbone_outputs["load_balance_loss"]
-
-        return ShramCausalLMOutput(
-            loss=loss,
-            ce_loss=ce_loss,
-            logits=logits,
-            past_key_values=backbone_outputs["past_key_values"],
-            hidden_states=backbone_outputs["hidden_states"],
-            load_balance_loss=backbone_outputs["load_balance_loss"],
-            max_vio=backbone_outputs["max_vio"],
+"""HuggingFace causal-LM wrapper for SHRAM.
+
+ShramForCausalLM is the HuggingFace-facing language-model boundary for SHRAM.
+It owns token embedding lookup, LM-head projection, wrapper-level next-token
+cross-entropy loss, config-controlled tied embeddings, and generation/cache
+orchestration at the wrapper boundary.
+
+The backbone remains a pure transformer stack. ShramModel accepts pre-embedded
+hidden states together with current position IDs, a current active mask, and an
+optional ShramCache. It has no knowledge of token IDs, vocabulary projection,
+or causal-LM loss.
+
+HuggingFace generation reaches this wrapper with two different tensor
+conventions:
+
+- ``position_ids`` is a current-step tensor. GenerationMixin updates the total
+  sequence state between steps, then slices position-bearing tensors back down
+  before calling ``forward()``.
+- ``attention_mask`` is a full 2D mask over the total sequence so far. This
+  wrapper slices its recent chunk to produce the current semantic liveness mask
+  expected by the backbone.
+
+Generation-created caches are handled in ``_prepare_cache_for_generation``.
+That hook ensures HuggingFace generation uses ShramCache rather than a generic
+dynamic cache. The direct ``forward()`` path does not silently create caches;
+when ``use_cache=True`` it expects a truthful ShramCache to have been supplied.
+"""
+
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+import torch.nn as nn
+from transformers import GenerationMixin, PreTrainedModel
+from transformers.cache_utils import Cache
+from transformers.generation.configuration_utils import GenerationMode
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+from .__cache__shram_cache import ShramCache
+from .configuration import ShramConfig
+from .model import ShramModel
+
+
+@dataclass
+class ShramCausalLMOutput(CausalLMOutputWithPast):
+    """SHRAM causal-LM wrapper output.
+
+    This subclasses HuggingFace's standard ``CausalLMOutputWithPast``.
+    Dataclass inheritance is sufficient here: all standard causal-LM fields and
+    ModelOutput behavior are inherited from the parent, and this subclass adds
+    only the SHRAM-specific wrapper outputs.
+    """
+
+    ce_loss: torch.FloatTensor | None = None
+    load_balance_loss: torch.FloatTensor | None = None
+    max_vio: torch.FloatTensor | None = None
+
+
+class ShramForCausalLM(PreTrainedModel, GenerationMixin):
+    """HuggingFace-facing causal language model wrapper for SHRAM.
+
+    Owns token embeddings, LM-head projection, wrapper-level shifted CE loss,
+    tied embedding configuration, and generation/cache boundary behavior.
+    Delegates all transformer computation to ``ShramModel``.
+
+    Args:
+        config: SHRAM model configuration.
+    """
+
+    config_class = ShramConfig
+    base_model_prefix = "model"
+    _no_split_modules = ["DecoderLayer"]
+    supports_gradient_checkpointing = True
+
+    def __init__(self, config: ShramConfig) -> None:
+        super().__init__(config)
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.model = ShramModel(config)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self._configure_tied_embeddings()
+        self.post_init()
+
+    def _configure_tied_embeddings(self) -> None:
+        """Apply config-controlled tied embedding behavior on this instance."""
+        if self.config.tie_word_embeddings:
+            self.lm_head.weight = self.embed_tokens.weight
+            self._tied_weights_keys = {
+                "lm_head.weight": "embed_tokens.weight",
+            }
+        else:
+            self._tied_weights_keys = {}
+
+    def get_input_embeddings(self) -> nn.Embedding:
+        """Return the token embedding matrix."""
+        return self.embed_tokens
+
+    def set_input_embeddings(self, value: nn.Embedding) -> None:
+        """Replace the token embedding matrix."""
+        self.embed_tokens = value
+        self._configure_tied_embeddings()
+
+    def get_output_embeddings(self) -> nn.Linear:
+        """Return the LM head."""
+        return self.lm_head
+
+    def set_output_embeddings(self, value: nn.Linear) -> None:
+        """Replace the LM head."""
+        self.lm_head = value
+        self._configure_tied_embeddings()
+
+    def _build_shram_cache(
+        self,
+        batch_size: int,
+        device: torch.device,
+    ) -> ShramCache:
+        """Construct a fresh top-level SHRAM cache."""
+        return ShramCache(
+            num_hidden_layers=self.config.num_hidden_layers,
+            sliding_window=self.config.window_size,
+            num_local_heads=self.config.num_sliding_window_heads,
+            local_head_dim=self.config.head_dim,
+            num_mosrah_heads=self.config.num_mosrah_heads,
+            mosrah_head_dim=self.config.head_dim,
+            batch_size=batch_size,
+            device=device,
+        )
+
+    def _validate_generation_cache_request(
+        self,
+        generation_config: Any,
+        model_kwargs: dict[str, Any],
+        generation_mode: GenerationMode,
+    ) -> None:
+        """Validate SHRAM's generation-side cache policy."""
+        if generation_mode in {
+            GenerationMode.ASSISTED_GENERATION,
+            GenerationMode.CONTRASTIVE_SEARCH,
+        }:
+            raise NotImplementedError(
+                "ShramForCausalLM does not currently support assisted generation "
+                "or contrastive search because ShramCache does not support crop()."
+            )
+
+        user_defined_cache = model_kwargs.get("past_key_values")
+        if user_defined_cache is not None:
+            if generation_config.cache_implementation is not None:
+                raise ValueError(
+                    "Passing both `cache_implementation` and `past_key_values` "
+                    "is unsupported. Please use only one."
+                )
+            if isinstance(user_defined_cache, tuple):
+                raise ValueError(
+                    "Passing a tuple of `past_key_values` is not supported. "
+                    "Please use a `ShramCache` instance."
+                )
+            if not isinstance(user_defined_cache, ShramCache):
+                raise TypeError(
+                    "ShramForCausalLM requires `past_key_values` to be a "
+                    "`ShramCache` instance."
+                )
+
+        if (
+            user_defined_cache is None
+            and generation_config.use_cache
+            and generation_config.cache_implementation is not None
+        ):
+            raise ValueError(
+                "ShramForCausalLM does not support `cache_implementation`. "
+                "Generation-created caches must be `ShramCache` objects."
+            )
+
+    def _prepare_cache_for_generation(
+        self,
+        generation_config: Any,
+        model_kwargs: dict[str, Any],
+        generation_mode: GenerationMode,
+        batch_size: int,
+        max_cache_length: int,
+    ) -> None:
+        """Ensure HuggingFace generation uses ShramCache.
+
+        This is the SHRAM-specific generation hook. The rest of the default
+        generation plumbing is kept intact as much as possible.
+
+        Args:
+            generation_config: Active generation configuration.
+            model_kwargs: Generation kwargs, updated in place.
+            generation_mode: HuggingFace generation mode.
+            batch_size: Effective generation batch size.
+            max_cache_length: Requested cache length. Accepted but unused here.
+        """
+        self._validate_generation_cache_request(
+            generation_config=generation_config,
+            model_kwargs=model_kwargs,
+            generation_mode=generation_mode,
+        )
+
+        if model_kwargs.get("past_key_values") is not None:
+            return
+
+        if not generation_config.use_cache:
+            return
+
+        num_repeats = max(
+            generation_config.num_beams or 1,
+            generation_config.num_return_sequences or 1,
+        )
+        model_kwargs["past_key_values"] = self._build_shram_cache(
+            batch_size=batch_size*num_repeats,
+            device=self.embed_tokens.weight.device,
+        )
+
+    def _reorder_cache(
+        self,
+        past_key_values: Cache,
+        beam_idx: torch.Tensor,
+    ) -> Cache:
+        """Reorder the cache in place for beam search."""
+        past_key_values.reorder_cache(beam_idx)
+        return past_key_values
+
+    def _validate_input_ids(self, input_ids: torch.Tensor) -> None:
+        """Validate token IDs at the wrapper boundary."""
+        if input_ids.ndim != 2:
+            raise ValueError("input_ids must have shape (batch, seq_len).")
+        if input_ids.shape[1] == 0:
+            raise ValueError("input_ids sequence length must be nonzero.")
+        if input_ids.dtype != torch.long:
+            raise TypeError("input_ids must be an long int tensor.")
+
+    def _validate_attention_mask(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None,
+    ) -> None:
+        """Validate the full-sequence attention mask."""
+        if attention_mask is None:
+            return
+        if attention_mask.ndim != 2:
+            raise ValueError("attention_mask must have shape (batch, total_seq_len).")
+        if attention_mask.shape[0] != input_ids.shape[0]:
+            raise ValueError("attention_mask batch dimension must match input_ids.")
+        if attention_mask.shape[1] < input_ids.shape[1]:
+            raise ValueError(
+                "attention_mask must be at least as long as the current input_ids chunk."
+            )
+
+    def _validate_position_ids(
+        self,
+        input_ids: torch.Tensor,
+        position_ids: torch.Tensor | None,
+    ) -> None:
+        """Validate current-step position IDs."""
+        if position_ids is None:
+            return
+        if position_ids.ndim != 2:
+            raise ValueError("position_ids must have shape (batch, seq_len).")
+        if position_ids.shape != input_ids.shape:
+            raise ValueError(
+                "position_ids must match the current input_ids shape exactly."
+            )
+        if input_ids.dtype != torch.long:
+            raise TypeError("position_ids must be an long tensor.")
+
+    def _validate_labels(
+        self,
+        input_ids: torch.Tensor,
+        labels: torch.Tensor | None,
+    ) -> None:
+        """Validate label shape at the wrapper boundary."""
+        if labels is None:
+            return
+        if labels.ndim != 2:
+            raise ValueError("labels must have shape (batch, seq_len).")
+        if labels.shape != input_ids.shape:
+            raise ValueError("labels must have the same shape as input_ids.")
+        if input_ids.dtype != torch.long:
+            raise TypeError("labels must be a long tensor.")
+
+    def _validate_cache_inputs(
+        self,
+        use_cache: bool,
+        past_key_values: Cache | None,
+    ) -> None:
+        """Validate cache policy for direct wrapper calls."""
+        if use_cache:
+            if past_key_values is None:
+                raise ValueError(
+                    "use_cache=True requires an explicit ShramCache. During "
+                    "generate(), HuggingFace should supply this through "
+                    "_prepare_cache_for_generation()."
+                )
+            if not isinstance(past_key_values, ShramCache):
+                raise TypeError(
+                    "past_key_values must be a ShramCache when use_cache=True."
+                )
+            return
+
+        if past_key_values is not None:
+            raise ValueError("past_key_values was provided while use_cache=False.")
+
+    def _validate_position_sources(
+        self,
+        use_cache: bool,
+        attention_mask: torch.Tensor | None,
+        position_ids: torch.Tensor | None,
+    ) -> None:
+        """Validate that cached forward has a truthful source of positions."""
+        if use_cache and attention_mask is None and position_ids is None:
+            raise ValueError(
+                "Cached forward requires either position_ids or attention_mask."
+            )
+
+    def _validate_hf_boundary(
+        self,
+        output_attentions: bool | None,
+        return_dict: bool | None,
+        inputs_embeds: torch.Tensor | None,
+        cache_position: torch.Tensor | None,
+        extra_kwargs: dict[str, Any],
+    ) -> None:
+        """Validate unsupported HuggingFace-facing wrapper inputs."""
+        if output_attentions:
+            raise NotImplementedError(
+                "ShramForCausalLM does not expose output_attentions."
+            )
+        if return_dict is False:
+            raise ValueError(
+                "return_dict=False is not supported. "
+                "ShramForCausalLM always returns ShramCausalLMOutput."
+            )
+        if inputs_embeds is not None:
+            raise ValueError(
+                "inputs_embeds is not supported at the SHRAM wrapper boundary. "
+                "Pass input_ids instead."
+            )
+        if extra_kwargs:
+            unsupported = ", ".join(sorted(extra_kwargs))
+            raise TypeError(
+                f"Unsupported forward kwargs for ShramForCausalLM: {unsupported}"
+            )
+
+    def _standardize_full_attention_mask(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None,
+    ) -> torch.BoolTensor:
+        """Return a concrete full-sequence boolean attention mask."""
+        if attention_mask is None:
+            return torch.ones_like(input_ids, dtype=torch.bool)
+        return attention_mask.to(dtype=torch.bool)
+
+    def _resolve_current_position_ids(
+        self,
+        input_ids: torch.Tensor,
+        position_ids: torch.Tensor | None,
+        full_attention_mask: torch.BoolTensor,
+    ) -> torch.LongTensor:
+        """Resolve concrete current-step position IDs for the backbone."""
+        if position_ids is not None:
+            return position_ids.to(dtype=torch.long)
+
+        full_position_ids = full_attention_mask.to(dtype=torch.long).cumsum(dim=-1) - 1
+        full_position_ids = full_position_ids.masked_fill(~full_attention_mask, 0)
+        current_length = input_ids.shape[1]
+        return full_position_ids[:, -current_length:]
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.Tensor | None = None,
+        past_key_values: Cache | None = None,
+        use_cache: bool | None = None,
+        output_hidden_states: bool | None = None,
+        labels: torch.Tensor | None = None,
+        return_dict: bool | None = None,
+        ce_weight: float = 1.0,
+        load_balance_weight: float = 0.01,
+        **kwargs: Any,
+    ) -> ShramCausalLMOutput:
+        """Run the SHRAM causal language model wrapper.
+
+        Args:
+            input_ids: Current token IDs of shape ``(batch, seq_len)``.
+            attention_mask: Optional full 2D mask of shape
+                ``(batch, total_seq_len)``. The wrapper slices its recent chunk
+                to produce the current semantic liveness mask expected by the
+                backbone.
+            position_ids: Optional current-step position IDs of shape
+                ``(batch, seq_len)``. In ordinary HuggingFace generation this is
+                already the current-step tensor when it reaches ``forward()``.
+            past_key_values: Optional SHRAM cache. Required when
+                ``use_cache=True``.
+            use_cache: Whether to use and return a cache. Defaults to
+                ``config.use_cache``.
+            output_hidden_states: Whether to return backbone hidden states.
+                Defaults to ``config.output_hidden_states``.
+            labels: Optional target token IDs of shape ``(batch, seq_len)``.
+            return_dict: Must be ``True`` or ``None``.
+            ce_weight: Weight applied to the cross-entropy loss when combining with
+                the load-balance loss. Default 1.0.
+            load_balance_weight: Weight applied to the load-balance auxiliary loss.
+                Default 0.01, matching the paper's recommendation.
+            **kwargs: Unsupported HuggingFace kwargs fail explicitly.
+
+        Returns:
+            ``ShramCausalLMOutput`` with:
+            - ``logits`` of shape ``(batch, seq_len, vocab_size)``,
+            - ``loss`` = ``ce_weight * ce_loss + load_balance_weight * load_balance_loss``
+              when labels are provided (``None`` otherwise),
+            - ``ce_loss`` — raw unweighted cross-entropy loss for logging,
+            - ``past_key_values`` as the active ``ShramCache`` or ``None``,
+            - ``hidden_states`` when requested,
+            - ``load_balance_loss`` — raw unweighted load-balance loss from the backbone,
+            - detached ``max_vio`` from the backbone.
+        """
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+        output_hidden_states = (
+            output_hidden_states
+            if output_hidden_states is not None
+            else self.config.output_hidden_states
+        )
+
+        inputs_embeds = kwargs.pop("inputs_embeds", None)
+        output_attentions = kwargs.pop("output_attentions", None)
+        cache_position = kwargs.pop("cache_position", None)
+
+        # ------------------------------------------------------------------
+        # Validation zone.
+        #
+        # The wrapper boundary is where HuggingFace-facing inputs are judged
+        # for truthfulness before any internal work begins. These checks are
+        # intentionally front-loaded so the core logic below can assume one
+        # coherent interpretation of the call rather than defensively checking
+        # shapes, cache policy, or unsupported HF knobs at the point of use.
+        # This keeps the main sequence readable while ensuring invalid states
+        # fail before they can silently contaminate backbone execution.
+        # ------------------------------------------------------------------
+        self._validate_input_ids(input_ids)
+        self._validate_attention_mask(input_ids, attention_mask)
+        self._validate_position_ids(input_ids, position_ids)
+        self._validate_labels(input_ids, labels)
+        self._validate_cache_inputs(use_cache, past_key_values)
+        self._validate_position_sources(use_cache, attention_mask, position_ids)
+        self._validate_hf_boundary(
+            output_attentions=output_attentions,
+            return_dict=return_dict,
+            inputs_embeds=inputs_embeds,
+            cache_position=cache_position,
+            extra_kwargs=kwargs,
+        )
+
+        # ------------------------------------------------------------------
+        # Standardization zone.
+        #
+        # HuggingFace and SHRAM use different boundary conventions: generation
+        # carries a full-sequence 2D attention mask, while the SHRAM backbone
+        # wants a current-step active mask and concrete current position IDs.
+        # This zone collapses those wrapper-facing conventions into one valid
+        # backbone-facing state. After this point the core no longer reasons
+        # about optional or ambiguous input forms; it works only with concrete
+        # tensors whose semantics are already fixed.
+        # ------------------------------------------------------------------
+        full_attention_mask: torch.BoolTensor = self._standardize_full_attention_mask(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+        )
+        current_length: int = input_ids.shape[1]
+        current_active_mask: torch.BoolTensor = full_attention_mask[:, -current_length:]
+        current_position_ids: torch.LongTensor = self._resolve_current_position_ids(
+            input_ids=input_ids,
+            position_ids=position_ids,
+            full_attention_mask=full_attention_mask,
+        )
+        shram_cache: ShramCache | None = past_key_values if use_cache else None
+
+        # ------------------------------------------------------------------
+        # Core wrapper responsibilities.
+        #
+        # The wrapper's primary job is kept visible here: convert token IDs to
+        # embeddings, delegate transformer computation to ShramModel, project
+        # hidden states back to vocabulary logits, optionally compute the
+        # wrapper-level shifted next-token loss, and return the HuggingFace-
+        # facing output object. The backbone remains responsible only for
+        # transformer semantics; token/vocabulary/loss concerns stay here.
+        # ------------------------------------------------------------------
+        token_embeddings: torch.FloatTensor = self.embed_tokens(input_ids)
+        backbone_outputs = self.model(
+            inputs_embeds=token_embeddings,
+            position_ids=current_position_ids,
+            active_mask=current_active_mask,
+            cache=shram_cache,
+            output_hidden_states=output_hidden_states,
+        )
+
+        logits: torch.FloatTensor = self.lm_head(backbone_outputs["last_hidden_state"])
+
+        ce_loss: torch.FloatTensor | None = None
+        loss: torch.FloatTensor | None = None
+        if labels is not None:
+            shift_logits = logits[:, :-1, :].contiguous()
+            shift_labels = labels[:, 1:].contiguous()
+            ce_loss = nn.functional.cross_entropy(
+                shift_logits.view(-1, self.config.vocab_size),
+                shift_labels.view(-1),
+            )
+            loss = ce_weight * ce_loss + load_balance_weight * backbone_outputs["load_balance_loss"]
+
+        return ShramCausalLMOutput(
+            loss=loss,
+            ce_loss=ce_loss,
+            logits=logits,
+            past_key_values=backbone_outputs["past_key_values"],
+            hidden_states=backbone_outputs["hidden_states"],
+            load_balance_loss=backbone_outputs["load_balance_loss"],
+            max_vio=backbone_outputs["max_vio"],
         )

mlp.py

@@ -1,52 +1,52 @@
-"""SwiGLU feed-forward sublayer.
-
-SwiGLU is a gated linear unit variant that multiplies a SiLU-gated projection
-element-wise against a separate up-projection:
-
-    output = W_down(SiLU(W_gate(x)) ⊙ W_up(x))
-
-The gating mechanism gives the network more expressive control over which features
-to propagate than a plain two-matrix FFN. It requires three weight matrices instead
-of two, which is why intermediate_size in Llama 3 is set lower than the 4× multiplier
-typical of two-matrix FFNs — the total parameter count remains comparable.
-
-SiLU is used as the gate activation because Llama 3 committed to SwiGLU specifically
-— a fixed architectural choice.
-"""
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from transformers import PretrainedConfig
-
-
-class SwiGLUMLP(nn.Module):
-    """SwiGLU feed-forward sublayer.
-
-    Implements the three-matrix SwiGLU FFN used in Llama 3:
-
-        output = W_down(SiLU(W_gate(x)) ⊙ W_up(x))
-
-    No bias on any projection. SiLU as the gate activation is an architectural
-    constant — it is what defines SwiGLU specifically.
-
-    Args:
-        config: Model config. Must expose ``hidden_size`` and ``intermediate_size``.
-    """
-
-    def __init__(self, config: PretrainedConfig) -> None:
-        super().__init__()
-        self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
-        self.up_proj   = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
-        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """Apply the SwiGLU feed-forward transformation.
-
-        Args:
-            x: Input tensor of shape (batch, seq_len, hidden_size).
-
-        Returns:
-            Output tensor of shape (batch, seq_len, hidden_size).
-        """
-        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
+"""SwiGLU feed-forward sublayer.
+
+SwiGLU is a gated linear unit variant that multiplies a SiLU-gated projection
+element-wise against a separate up-projection:
+
+    output = W_down(SiLU(W_gate(x)) ⊙ W_up(x))
+
+The gating mechanism gives the network more expressive control over which features
+to propagate than a plain two-matrix FFN. It requires three weight matrices instead
+of two, which is why intermediate_size in Llama 3 is set lower than the 4× multiplier
+typical of two-matrix FFNs — the total parameter count remains comparable.
+
+SiLU is used as the gate activation because Llama 3 committed to SwiGLU specifically
+— a fixed architectural choice.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PretrainedConfig
+
+
+class SwiGLUMLP(nn.Module):
+    """SwiGLU feed-forward sublayer.
+
+    Implements the three-matrix SwiGLU FFN used in Llama 3:
+
+        output = W_down(SiLU(W_gate(x)) ⊙ W_up(x))
+
+    No bias on any projection. SiLU as the gate activation is an architectural
+    constant — it is what defines SwiGLU specifically.
+
+    Args:
+        config: Model config. Must expose ``hidden_size`` and ``intermediate_size``.
+    """
+
+    def __init__(self, config: PretrainedConfig) -> None:
+        super().__init__()
+        self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.up_proj   = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Apply the SwiGLU feed-forward transformation.
+
+        Args:
+            x: Input tensor of shape (batch, seq_len, hidden_size).
+
+        Returns:
+            Output tensor of shape (batch, seq_len, hidden_size).
+        """
+        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))

model.py

@@ -1,132 +1,132 @@
-"""Transformer backbone for Shram.
-
-ShramModel is a pure PyTorch module: a sequence of DecoderLayer blocks followed
-by a final RMSNorm. It accepts pre-embedded hidden states and returns contextual
-representations. It has no knowledge of tokens, vocabulary, generation, or the
-HuggingFace causal-LM wrapper contract.
-
-Keeping the embedding out of the backbone is the correct convention and makes
-the backbone genuinely modality-agnostic. The token interface — embedding lookup,
-LM head, weight tying, and generation-facing naming conventions — belongs on the
-task wrapper (ShramForCausalLM), which is the only class that knows this
-backbone is being used for language modelling.
-
-The final RMSNorm is necessary because the decoder stack uses pre-norm throughout:
-each sublayer normalises its own input, leaving the residual stream itself
-unnormalised. After many layers of accumulated residuals, that stream arrives at
-the top with uncontrolled magnitude. The final norm brings it to a well-scaled
-state before any projection. Without it, the LM head would receive signals of
-arbitrary scale.
-
-Caching is caller-managed. If a ShramCache is provided, ShramModel threads the
-corresponding per-layer ShramLayerCache into each DecoderLayer and returns the
-same top-level ShramCache object in the output dict. If None is provided, no
-caching occurs.
-
-Returns a plain dict with keys:
-- "last_hidden_state": normed backbone output, shape (batch, seq_len, hidden_size)
-- "past_key_values": the ShramCache object passed in, or None
-- "hidden_states": tuple of per-layer activations if output_hidden_states=True, else None
-- "load_balance_loss": scalar sum of per-layer SHRAM load-balance losses
-- "max_vio": detached scalar maximum routing-imbalance across all decoder layers
-"""
-
-import torch
-import torch.nn as nn
-
-from .__cache__shram_cache import ShramCache
-from .configuration import ShramConfig
-from .decoder_layer import DecoderLayer
-
-
-class ShramModel(nn.Module):
-    """Pure transformer backbone: decoder stack and final normalisation.
-
-    Accepts pre-embedded hidden states of shape (batch, seq_len, hidden_size)
-    and returns contextual representations of the same shape. No token embedding,
-    vocabulary projection, or causal-LM lifecycle concerns.
-
-    RoPE is applied inside each attention layer. Positional information is
-    encoded in the relationship between Q and K, not added to the residual
-    stream, so the backbone is agnostic to how positions are represented.
-
-    Args:
-        config: Model configuration. Must be a ``ShramConfig`` instance.
-    """
-
-    def __init__(self, config: ShramConfig) -> None:
-        super().__init__()
-        self.config = config
-        self.layers = nn.ModuleList(
-            [DecoderLayer(config) for _ in range(config.num_hidden_layers)]
-        )
-        self.norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-
-    def forward(
-        self,
-        inputs_embeds: torch.Tensor,
-        position_ids: torch.Tensor,
-        active_mask: torch.Tensor,
-        cache: ShramCache | None = None,
-        output_hidden_states: bool = False,
-    ) -> dict:
-        """Run the transformer stack over a batch of pre-embedded sequences.
-
-        Args:
-            inputs_embeds: Pre-embedded input of shape (batch, seq_len, hidden_size).
-            position_ids: Absolute positions of shape (batch, seq_len). Required.
-                Must be provided explicitly by the caller — this module does not
-                infer positions from cache state.
-            active_mask: Current-chunk active mask of shape (batch, seq_len),
-                where True means the token is semantically live. Forwarded
-                unchanged to every decoder layer.
-            cache: Optional top-level ShramCache. When provided, each DecoderLayer
-                receives its own layer-local cache via ``cache.layers[layer_idx]``.
-                The top-level cache object is updated in place and returned unchanged.
-            output_hidden_states: When True, the output dict includes a tuple of
-                per-layer hidden states: (inputs_embeds, layer_0_out, ..., layer_N_out),
-                collected before the final norm.
-
-        Returns:
-            Plain dict with keys:
-            - ``"last_hidden_state"``: normed backbone output,
-              shape (batch, seq_len, hidden_size).
-            - ``"past_key_values"``: the cache object passed in, or None.
-            - ``"hidden_states"``: tuple of per-layer activations (including
-              inputs_embeds as position 0) if ``output_hidden_states`` is True,
-              else None. Collected before the final norm so each entry reflects the
-              unnormalised residual stream at that depth.
-            - ``"load_balance_loss"``: scalar sum of per-layer SHRAM
-              load-balance losses.
-            - ``"max_vio"``: detached scalar maximum routing-imbalance across
-              all decoder layers. Zero means perfectly balanced routing across
-              every layer; higher values identify the worst-case head imbalance.
-        """
-        hidden_states = inputs_embeds
-        all_hidden_states = (hidden_states,) if output_hidden_states else None
-        total_load_balance_loss = inputs_embeds.new_zeros(())
-        max_vio = inputs_embeds.new_zeros(())
-
-        for layer_idx, layer in enumerate(self.layers):
-            layer_cache = None if cache is None else cache.layers[layer_idx]
-            hidden_states, layer_load_balance_loss, layer_max_vio = layer(
-                hidden_states,
-                position_ids,
-                active_mask,
-                cache=layer_cache,
-            )
-            total_load_balance_loss = total_load_balance_loss + layer_load_balance_loss
-            max_vio = torch.maximum(max_vio, layer_max_vio)
-
-            if output_hidden_states:
-                all_hidden_states = all_hidden_states + (hidden_states,)
-
-        hidden_states = self.norm(hidden_states)
-
-        return {
-            "last_hidden_state": hidden_states,
-            "past_key_values": cache,
-            "hidden_states": all_hidden_states,
-            "load_balance_loss": total_load_balance_loss,
-            "max_vio": max_vio,
+"""Transformer backbone for Shram.
+
+ShramModel is a pure PyTorch module: a sequence of DecoderLayer blocks followed
+by a final RMSNorm. It accepts pre-embedded hidden states and returns contextual
+representations. It has no knowledge of tokens, vocabulary, generation, or the
+HuggingFace causal-LM wrapper contract.
+
+Keeping the embedding out of the backbone is the correct convention and makes
+the backbone genuinely modality-agnostic. The token interface — embedding lookup,
+LM head, weight tying, and generation-facing naming conventions — belongs on the
+task wrapper (ShramForCausalLM), which is the only class that knows this
+backbone is being used for language modelling.
+
+The final RMSNorm is necessary because the decoder stack uses pre-norm throughout:
+each sublayer normalises its own input, leaving the residual stream itself
+unnormalised. After many layers of accumulated residuals, that stream arrives at
+the top with uncontrolled magnitude. The final norm brings it to a well-scaled
+state before any projection. Without it, the LM head would receive signals of
+arbitrary scale.
+
+Caching is caller-managed. If a ShramCache is provided, ShramModel threads the
+corresponding per-layer ShramLayerCache into each DecoderLayer and returns the
+same top-level ShramCache object in the output dict. If None is provided, no
+caching occurs.
+
+Returns a plain dict with keys:
+- "last_hidden_state": normed backbone output, shape (batch, seq_len, hidden_size)
+- "past_key_values": the ShramCache object passed in, or None
+- "hidden_states": tuple of per-layer activations if output_hidden_states=True, else None
+- "load_balance_loss": scalar sum of per-layer SHRAM load-balance losses
+- "max_vio": detached scalar maximum routing-imbalance across all decoder layers
+"""
+
+import torch
+import torch.nn as nn
+
+from .__cache__shram_cache import ShramCache
+from .configuration import ShramConfig
+from .decoder_layer import DecoderLayer
+
+
+class ShramModel(nn.Module):
+    """Pure transformer backbone: decoder stack and final normalisation.
+
+    Accepts pre-embedded hidden states of shape (batch, seq_len, hidden_size)
+    and returns contextual representations of the same shape. No token embedding,
+    vocabulary projection, or causal-LM lifecycle concerns.
+
+    RoPE is applied inside each attention layer. Positional information is
+    encoded in the relationship between Q and K, not added to the residual
+    stream, so the backbone is agnostic to how positions are represented.
+
+    Args:
+        config: Model configuration. Must be a ``ShramConfig`` instance.
+    """
+
+    def __init__(self, config: ShramConfig) -> None:
+        super().__init__()
+        self.config = config
+        self.layers = nn.ModuleList(
+            [DecoderLayer(config) for _ in range(config.num_hidden_layers)]
+        )
+        self.norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def forward(
+        self,
+        inputs_embeds: torch.Tensor,
+        position_ids: torch.Tensor,
+        active_mask: torch.Tensor,
+        cache: ShramCache | None = None,
+        output_hidden_states: bool = False,
+    ) -> dict:
+        """Run the transformer stack over a batch of pre-embedded sequences.
+
+        Args:
+            inputs_embeds: Pre-embedded input of shape (batch, seq_len, hidden_size).
+            position_ids: Absolute positions of shape (batch, seq_len). Required.
+                Must be provided explicitly by the caller — this module does not
+                infer positions from cache state.
+            active_mask: Current-chunk active mask of shape (batch, seq_len),
+                where True means the token is semantically live. Forwarded
+                unchanged to every decoder layer.
+            cache: Optional top-level ShramCache. When provided, each DecoderLayer
+                receives its own layer-local cache via ``cache.layers[layer_idx]``.
+                The top-level cache object is updated in place and returned unchanged.
+            output_hidden_states: When True, the output dict includes a tuple of
+                per-layer hidden states: (inputs_embeds, layer_0_out, ..., layer_N_out),
+                collected before the final norm.
+
+        Returns:
+            Plain dict with keys:
+            - ``"last_hidden_state"``: normed backbone output,
+              shape (batch, seq_len, hidden_size).
+            - ``"past_key_values"``: the cache object passed in, or None.
+            - ``"hidden_states"``: tuple of per-layer activations (including
+              inputs_embeds as position 0) if ``output_hidden_states`` is True,
+              else None. Collected before the final norm so each entry reflects the
+              unnormalised residual stream at that depth.
+            - ``"load_balance_loss"``: scalar sum of per-layer SHRAM
+              load-balance losses.
+            - ``"max_vio"``: detached scalar maximum routing-imbalance across
+              all decoder layers. Zero means perfectly balanced routing across
+              every layer; higher values identify the worst-case head imbalance.
+        """
+        hidden_states = inputs_embeds
+        all_hidden_states = (hidden_states,) if output_hidden_states else None
+        total_load_balance_loss = inputs_embeds.new_zeros(())
+        max_vio = inputs_embeds.new_zeros(())
+
+        for layer_idx, layer in enumerate(self.layers):
+            layer_cache = None if cache is None else cache.layers[layer_idx]
+            hidden_states, layer_load_balance_loss, layer_max_vio = layer(
+                hidden_states,
+                position_ids,
+                active_mask,
+                cache=layer_cache,
+            )
+            total_load_balance_loss = total_load_balance_loss + layer_load_balance_loss
+            max_vio = torch.maximum(max_vio, layer_max_vio)
+
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (hidden_states,)
+
+        hidden_states = self.norm(hidden_states)
+
+        return {
+            "last_hidden_state": hidden_states,
+            "past_key_values": cache,
+            "hidden_states": all_hidden_states,
+            "load_balance_loss": total_load_balance_loss,
+            "max_vio": max_vio,
         }

rope.py

@@ -1,291 +1,291 @@
-"""Rotary Position Embeddings (RoPE).
-
-RoPE encodes position in the *relationship* between query and key vectors. When the
-attention dot product Q·Kᵀ is computed, the per-position rotations cancel to produce
-a score that depends only on the relative distance — not on absolute positions.
-
-Two modes are supported:
-
-  default  Standard RoPE with base frequency b. Each dimension pair d is assigned
-           frequency θ_d = b^{-2d/u} where u is the head dimension. The attention
-           scaling A_rope = 1.
-
-  yarn     YaRN frequency interpolation for long-context extrapolation (Peng et al.,
-           "YaRN: Efficient Context Window Extension of Large Language Models", 2023,
-           §A.2). Three frequency regimes:
-             - Low-frequency dimensions (r < α): fully interpolated by scale s.
-               These dimensions have long wavelengths relative to the training window
-               and must be compressed to avoid out-of-distribution positions.
-             - High-frequency dimensions (r > β): left unchanged. Short-wavelength
-               dimensions already encode relative position accurately at any scale.
-             - Intermediate dimensions (α ≤ r ≤ β): linearly blended via ramp γ(r).
-           Returns A_rope = (0.1·ln(s)+1)². When s = 1, YaRN reduces exactly to
-           standard RoPE.
-
-Each attention path (h_l and BEA) constructs its own RotaryEmbedding with explicit
-parameters — no shared instance, no config reading. See Unit 5.A design decisions.
-
-Cache sharing: all instances with identical parameters share one cos/sin table via a
-class-level registry. The first instance that needs a particular (parameters, seq_len,
-device, dtype) combination builds the table; all subsequent instances reference it
-directly. This avoids redundant builds across the num_hidden_layers instances that
-share the same parametrisation.
-"""
-
-import math
-
-import torch
-import torch.nn as nn
-
-
-# ---------------------------------------------------------------------------
-# Rotation helper
-# ---------------------------------------------------------------------------
-
-def _rotate_half(x: torch.Tensor) -> torch.Tensor:
-    """Apply the 90° rotation used in the RoPE update formula.
-
-    Splits the last dimension into two halves [x1, x2] and returns [-x2, x1].
-    Combined with ``x * cos + rotate_half(x) * sin``, this implements a 2D rotation
-    on each consecutive pair of dimensions, matching the block-diagonal operator
-    R^u_{Θ,p} in the paper.
-    """
-    d = x.shape[-1] // 2
-    x1, x2 = x[..., :d], x[..., d:]
-    return torch.cat([-x2, x1], dim=-1)
-
-
-# ---------------------------------------------------------------------------
-# RotaryEmbedding
-# ---------------------------------------------------------------------------
-
-class RotaryEmbedding(nn.Module):
-    """Rotary Position Embeddings with explicit mode and parameter control.
-
-    Each caller constructs its own instance with the exact parameters it needs.
-    h_l always uses ``mode="default"``; BEA always uses ``mode="yarn"``. No
-    config object is read inside this module.
-
-    The cos/sin cache is built lazily on the first forward call and extended
-    automatically when a longer sequence is encountered. Instances with identical
-    parameters share one cache via the class-level ``_cache`` registry,
-    avoiding redundant computation across decoder layers.
-
-    Args:
-        mode: ``"default"`` for standard RoPE; ``"yarn"`` for YaRN extrapolation.
-        head_dim: Per-head embedding dimension ``u``. Must be even.
-        theta: Base frequency ``b`` in θ_d = b^{-2d/u}.
-        initial_seq_length: ``C_train`` — context length the model was trained at.
-            Required for ``mode="yarn"``.
-        dilation: Scale factor ``s = C_target / C_train`` — how much the context
-            window is extended beyond training length. Required for ``mode="yarn"``.
-            When ``dilation=1.0``, YaRN reduces to standard RoPE.
-        alpha: YaRN ramp lower boundary α. Dimensions with r(d) < α are fully
-            interpolated. Required for ``mode="yarn"``.
-        beta: YaRN ramp upper boundary β. Dimensions with r(d) > β are left
-            unchanged. Required for ``mode="yarn"``.
-        device: Optional device for initial buffer placement.
-
-    Raises:
-        NotImplementedError: If ``mode`` is not ``"default"`` or ``"yarn"``.
-        ValueError: If ``mode="yarn"`` and any of ``initial_seq_length``,
-            ``dilation``, ``alpha``, ``beta`` are absent.
-    """
-
-    # Maps (freq_key, seq_len, device_str, dtype_str) → (cos_table, sin_table).
-    # Shared across all RotaryEmbedding instances in the process. Keys include device
-    # and dtype so that tables built on different devices or in different precisions
-    # are stored independently.
-    _cache: dict = {}
-
-    def __init__(
-        self,
-        mode: str,
-        head_dim: int,
-        theta: float,
-        initial_seq_length: int | None = None,
-        dilation: float | None = None,
-        alpha: float | None = None,
-        beta: float | None = None,
-        device: torch.device | None = None,
-    ) -> None:
-        super().__init__()
-
-        self._validate_mode(mode)
-        self._validate_yarn_params(mode, initial_seq_length, dilation, alpha, beta)
-        self.mode = mode
-
-        # Compute per-dimension rotation frequencies θ_d (default) or θ_d' (yarn).
-        # d_index ranges over 0, 2, 4, ..., head_dim-2 — one index per dimension pair,
-        # so rotation_freqs has head_dim/2 entries.
-        d_index = torch.arange(0, head_dim, 2, dtype=torch.float32, device=device)
-        base_freqs = 1.0 / (theta ** (d_index / head_dim))  # θ_d = b^{-2d/u}
-
-        if mode == "default":
-            rotation_freqs = base_freqs
-            self.attention_scaling: float = 1.0
-
-        else:  # yarn
-            s = dilation
-
-            # r(d) = C_train · θ_d / (2π) — normalized frequency used by the ramp
-            # function to classify each dimension into one of three regimes.
-            normalized_freqs = initial_seq_length * base_freqs / (2.0 * math.pi)
-
-            # γ(r) ramp: 0 for r < α (fully interpolate), 1 for r > β (unchanged),
-            # linear blend between α and β.
-            blend_weights = ((normalized_freqs - alpha) / (beta - alpha)).clamp(0.0, 1.0)
-
-            # θ_d' = (1 − γ) · θ_d / s + γ · θ_d
-            rotation_freqs = (1.0 - blend_weights) * (base_freqs / s) + blend_weights * base_freqs
-
-            # A_rope = (0.1 · ln(s) + 1)² — attention logit scaling returned to caller.
-            self.attention_scaling = (0.1 * math.log(s) + 1.0) ** 2
-
-        # freq_key uniquely identifies the parameter set that produced rotation_freqs.
-        # Used as the primary component of the cache registry key.
-        if mode == "default":
-            self._freq_key: tuple = ("default", head_dim, float(theta))
-        else:
-            self._freq_key = (
-                "yarn", head_dim, float(theta),
-                int(initial_seq_length), float(dilation),
-                float(alpha), float(beta),
-            )
-
-        # rotation_freqs is a non-persistent buffer so it moves with the model across
-        # devices via .to() / .cuda() without appearing in saved checkpoints.
-        # It is stored per-instance rather than in the shared cache because it is
-        # small (head_dim/2 floats) — negligible cost compared to the cos/sin tables
-        # it is used to build. The meaningful sharing win is on those tables.
-        self.register_buffer("rotation_freqs", rotation_freqs, persistent=False)
-
-        # Cache tensors are plain instance attributes (not registered buffers) so that
-        # sharing across identically-parametrised instances survives .to() calls.
-        # Registered buffers are copied on device move; plain attributes are aliased,
-        # preserving the shared-tensor identity that the cache design depends on.
-        self._cos_cached: torch.Tensor | None = None
-        self._sin_cached: torch.Tensor | None = None
-
-    # ---------------------------------------------------------------------------
-    # Validation helpers
-    # ---------------------------------------------------------------------------
-
-    @staticmethod
-    def _validate_mode(mode: str) -> None:
-        """Raise NotImplementedError if mode is not a supported value."""
-        if mode not in {"default", "yarn"}:
-            raise NotImplementedError(
-                f"RoPE mode '{mode}' is not supported. Supported modes: 'default', 'yarn'."
-            )
-
-    @staticmethod
-    def _validate_yarn_params(
-        mode: str,
-        initial_seq_length: int | None,
-        dilation: float | None,
-        alpha: float | None,
-        beta: float | None,
-    ) -> None:
-        """Raise ValueError if mode='yarn' and any required parameter is absent."""
-        if mode != "yarn":
-            return
-        missing = [
-            name for name, val in [
-                ("initial_seq_length", initial_seq_length),
-                ("dilation", dilation),
-                ("alpha", alpha),
-                ("beta", beta),
-            ]
-            if val is None
-        ]
-        if missing:
-            raise ValueError(f"mode='yarn' requires {missing}.")
-
-    # ---------------------------------------------------------------------------
-    # Cache management
-    # ---------------------------------------------------------------------------
-
-    def _extend_cache(self, seq_len: int, device: torch.device, dtype: torch.dtype) -> None:
-        """Build the cos/sin table to cover positions [0, seq_len).
-
-        Checks the class-level registry first. If a table already exists for this
-        exact (parameters, seq_len, device, dtype) combination it is reused directly;
-        otherwise it is computed and stored. The instance attributes are pointed at
-        the registry entry so that all layers sharing the same parametrisation
-        reference the same tensor.
-        """
-        cache_key = (self._freq_key, seq_len, str(device), str(dtype))
-
-        if cache_key not in RotaryEmbedding._cache:
-            positions = torch.arange(seq_len, device=device, dtype=torch.float32)
-            # outer product → (seq_len, head_dim // 2); duplicate to (seq_len, head_dim)
-            freqs = torch.outer(
-                positions,
-                self.rotation_freqs.to(device=device, dtype=torch.float32),
-            )
-            angle_embedding = torch.cat((freqs, freqs), dim=-1)
-            RotaryEmbedding._cache[cache_key] = (
-                angle_embedding.cos().to(dtype),
-                angle_embedding.sin().to(dtype),
-            )
-
-        self._cos_cached, self._sin_cached = RotaryEmbedding._cache[cache_key]
-
-    def forward(
-        self,
-        q: torch.Tensor,
-        k: torch.Tensor,
-        position_ids: torch.Tensor,
-    ) -> tuple[torch.Tensor, torch.Tensor, float]:
-        """Apply rotary embeddings to query and key tensors.
-
-        The cos/sin cache is extended lazily when position_ids reference positions
-        beyond its current length, or when the device or dtype has changed.
-
-        ``position_ids`` may be any integer tensor shape. Its values are valid
-        position indices into the cos/sin cache:
-
-        - h_l (standard causal): position_ids (B, N), q/k (B, H, N, head_dim).
-        - BEA (packed):          position_ids (B, L, T), q/k (B, L, T, head_dim).
-
-        When q/k have head dimensions absent from position_ids, broadcast dimensions
-        are inserted automatically at dim 1.
-
-        Args:
-            q: Query tensor of shape (batch, [heads,] *pos_dims, head_dim).
-            k: Key tensor of shape (batch, [heads,] *pos_dims, head_dim).
-            position_ids: Integer positions of shape (batch, *pos_dims).
-
-        Returns:
-            Tuple of (q_rotated, k_rotated, attention_scaling). attention_scaling is
-            1.0 for default mode; YaRN returns (0.1·ln(s)+1)² which the caller must
-            apply to attention logits before softmax.
-        """
-        seq_len = int(position_ids.max().item()) + 1
-
-        # The cache is valid when it exists, covers all positions referenced by
-        # position_ids, and matches q's dtype and device. Each condition is named
-        # separately so the rebuild trigger is readable rather than a compound predicate.
-        cache_missing = self._cos_cached is None
-        cache_too_short = not cache_missing and seq_len > self._cos_cached.shape[0]
-        wrong_dtype = not cache_missing and self._cos_cached.dtype != q.dtype
-        wrong_device = not cache_missing and self._cos_cached.device != q.device
-
-        if cache_missing or cache_too_short or wrong_dtype or wrong_device:
-            self._extend_cache(seq_len, device=q.device, dtype=q.dtype)
-
-        cos = self._cos_cached[position_ids]
-        sin = self._sin_cached[position_ids]
-
-        # Insert broadcast dimensions for any head axes present in q/k but absent
-        # from position_ids. Standard: pos (B,N) → cos (B,N,D), q (B,H,N,D) → unsqueeze once.
-        # BEA: pos (B,L,T) → cos (B,L,T,D), q (B,L,T,D) → no unsqueeze needed.
-        while cos.ndim < q.ndim:
-            cos = cos.unsqueeze(1)
-            sin = sin.unsqueeze(1)
-
-        q_rotated = q * cos + _rotate_half(q) * sin
-        k_rotated = k * cos + _rotate_half(k) * sin
-
-        return q_rotated, k_rotated, self.attention_scaling
+"""Rotary Position Embeddings (RoPE).
+
+RoPE encodes position in the *relationship* between query and key vectors. When the
+attention dot product Q·Kᵀ is computed, the per-position rotations cancel to produce
+a score that depends only on the relative distance — not on absolute positions.
+
+Two modes are supported:
+
+  default  Standard RoPE with base frequency b. Each dimension pair d is assigned
+           frequency θ_d = b^{-2d/u} where u is the head dimension. The attention
+           scaling A_rope = 1.
+
+  yarn     YaRN frequency interpolation for long-context extrapolation (Peng et al.,
+           "YaRN: Efficient Context Window Extension of Large Language Models", 2023,
+           §A.2). Three frequency regimes:
+             - Low-frequency dimensions (r < α): fully interpolated by scale s.
+               These dimensions have long wavelengths relative to the training window
+               and must be compressed to avoid out-of-distribution positions.
+             - High-frequency dimensions (r > β): left unchanged. Short-wavelength
+               dimensions already encode relative position accurately at any scale.
+             - Intermediate dimensions (α ≤ r ≤ β): linearly blended via ramp γ(r).
+           Returns A_rope = (0.1·ln(s)+1)². When s = 1, YaRN reduces exactly to
+           standard RoPE.
+
+Each attention path (h_l and BEA) constructs its own RotaryEmbedding with explicit
+parameters — no shared instance, no config reading. See Unit 5.A design decisions.
+
+Cache sharing: all instances with identical parameters share one cos/sin table via a
+class-level registry. The first instance that needs a particular (parameters, seq_len,
+device, dtype) combination builds the table; all subsequent instances reference it
+directly. This avoids redundant builds across the num_hidden_layers instances that
+share the same parametrisation.
+"""
+
+import math
+
+import torch
+import torch.nn as nn
+
+
+# ---------------------------------------------------------------------------
+# Rotation helper
+# ---------------------------------------------------------------------------
+
+def _rotate_half(x: torch.Tensor) -> torch.Tensor:
+    """Apply the 90° rotation used in the RoPE update formula.
+
+    Splits the last dimension into two halves [x1, x2] and returns [-x2, x1].
+    Combined with ``x * cos + rotate_half(x) * sin``, this implements a 2D rotation
+    on each consecutive pair of dimensions, matching the block-diagonal operator
+    R^u_{Θ,p} in the paper.
+    """
+    d = x.shape[-1] // 2
+    x1, x2 = x[..., :d], x[..., d:]
+    return torch.cat([-x2, x1], dim=-1)
+
+
+# ---------------------------------------------------------------------------
+# RotaryEmbedding
+# ---------------------------------------------------------------------------
+
+class RotaryEmbedding(nn.Module):
+    """Rotary Position Embeddings with explicit mode and parameter control.
+
+    Each caller constructs its own instance with the exact parameters it needs.
+    h_l always uses ``mode="default"``; BEA always uses ``mode="yarn"``. No
+    config object is read inside this module.
+
+    The cos/sin cache is built lazily on the first forward call and extended
+    automatically when a longer sequence is encountered. Instances with identical
+    parameters share one cache via the class-level ``_cache`` registry,
+    avoiding redundant computation across decoder layers.
+
+    Args:
+        mode: ``"default"`` for standard RoPE; ``"yarn"`` for YaRN extrapolation.
+        head_dim: Per-head embedding dimension ``u``. Must be even.
+        theta: Base frequency ``b`` in θ_d = b^{-2d/u}.
+        initial_seq_length: ``C_train`` — context length the model was trained at.
+            Required for ``mode="yarn"``.
+        dilation: Scale factor ``s = C_target / C_train`` — how much the context
+            window is extended beyond training length. Required for ``mode="yarn"``.
+            When ``dilation=1.0``, YaRN reduces to standard RoPE.
+        alpha: YaRN ramp lower boundary α. Dimensions with r(d) < α are fully
+            interpolated. Required for ``mode="yarn"``.
+        beta: YaRN ramp upper boundary β. Dimensions with r(d) > β are left
+            unchanged. Required for ``mode="yarn"``.
+        device: Optional device for initial buffer placement.
+
+    Raises:
+        NotImplementedError: If ``mode`` is not ``"default"`` or ``"yarn"``.
+        ValueError: If ``mode="yarn"`` and any of ``initial_seq_length``,
+            ``dilation``, ``alpha``, ``beta`` are absent.
+    """
+
+    # Maps (freq_key, seq_len, device_str, dtype_str) → (cos_table, sin_table).
+    # Shared across all RotaryEmbedding instances in the process. Keys include device
+    # and dtype so that tables built on different devices or in different precisions
+    # are stored independently.
+    _cache: dict = {}
+
+    def __init__(
+        self,
+        mode: str,
+        head_dim: int,
+        theta: float,
+        initial_seq_length: int | None = None,
+        dilation: float | None = None,
+        alpha: float | None = None,
+        beta: float | None = None,
+        device: torch.device | None = None,
+    ) -> None:
+        super().__init__()
+
+        self._validate_mode(mode)
+        self._validate_yarn_params(mode, initial_seq_length, dilation, alpha, beta)
+        self.mode = mode
+
+        # Compute per-dimension rotation frequencies θ_d (default) or θ_d' (yarn).
+        # d_index ranges over 0, 2, 4, ..., head_dim-2 — one index per dimension pair,
+        # so rotation_freqs has head_dim/2 entries.
+        d_index = torch.arange(0, head_dim, 2, dtype=torch.float32, device=device)
+        base_freqs = 1.0 / (theta ** (d_index / head_dim))  # θ_d = b^{-2d/u}
+
+        if mode == "default":
+            rotation_freqs = base_freqs
+            self.attention_scaling: float = 1.0
+
+        else:  # yarn
+            s = dilation
+
+            # r(d) = C_train · θ_d / (2π) — normalized frequency used by the ramp
+            # function to classify each dimension into one of three regimes.
+            normalized_freqs = initial_seq_length * base_freqs / (2.0 * math.pi)
+
+            # γ(r) ramp: 0 for r < α (fully interpolate), 1 for r > β (unchanged),
+            # linear blend between α and β.
+            blend_weights = ((normalized_freqs - alpha) / (beta - alpha)).clamp(0.0, 1.0)
+
+            # θ_d' = (1 − γ) · θ_d / s + γ · θ_d
+            rotation_freqs = (1.0 - blend_weights) * (base_freqs / s) + blend_weights * base_freqs
+
+            # A_rope = (0.1 · ln(s) + 1)² — attention logit scaling returned to caller.
+            self.attention_scaling = (0.1 * math.log(s) + 1.0) ** 2
+
+        # freq_key uniquely identifies the parameter set that produced rotation_freqs.
+        # Used as the primary component of the cache registry key.
+        if mode == "default":
+            self._freq_key: tuple = ("default", head_dim, float(theta))
+        else:
+            self._freq_key = (
+                "yarn", head_dim, float(theta),
+                int(initial_seq_length), float(dilation),
+                float(alpha), float(beta),
+            )
+
+        # rotation_freqs is a non-persistent buffer so it moves with the model across
+        # devices via .to() / .cuda() without appearing in saved checkpoints.
+        # It is stored per-instance rather than in the shared cache because it is
+        # small (head_dim/2 floats) — negligible cost compared to the cos/sin tables
+        # it is used to build. The meaningful sharing win is on those tables.
+        self.register_buffer("rotation_freqs", rotation_freqs, persistent=False)
+
+        # Cache tensors are plain instance attributes (not registered buffers) so that
+        # sharing across identically-parametrised instances survives .to() calls.
+        # Registered buffers are copied on device move; plain attributes are aliased,
+        # preserving the shared-tensor identity that the cache design depends on.
+        self._cos_cached: torch.Tensor | None = None
+        self._sin_cached: torch.Tensor | None = None
+
+    # ---------------------------------------------------------------------------
+    # Validation helpers
+    # ---------------------------------------------------------------------------
+
+    @staticmethod
+    def _validate_mode(mode: str) -> None:
+        """Raise NotImplementedError if mode is not a supported value."""
+        if mode not in {"default", "yarn"}:
+            raise NotImplementedError(
+                f"RoPE mode '{mode}' is not supported. Supported modes: 'default', 'yarn'."
+            )
+
+    @staticmethod
+    def _validate_yarn_params(
+        mode: str,
+        initial_seq_length: int | None,
+        dilation: float | None,
+        alpha: float | None,
+        beta: float | None,
+    ) -> None:
+        """Raise ValueError if mode='yarn' and any required parameter is absent."""
+        if mode != "yarn":
+            return
+        missing = [
+            name for name, val in [
+                ("initial_seq_length", initial_seq_length),
+                ("dilation", dilation),
+                ("alpha", alpha),
+                ("beta", beta),
+            ]
+            if val is None
+        ]
+        if missing:
+            raise ValueError(f"mode='yarn' requires {missing}.")
+
+    # ---------------------------------------------------------------------------
+    # Cache management
+    # ---------------------------------------------------------------------------
+
+    def _extend_cache(self, seq_len: int, device: torch.device, dtype: torch.dtype) -> None:
+        """Build the cos/sin table to cover positions [0, seq_len).
+
+        Checks the class-level registry first. If a table already exists for this
+        exact (parameters, seq_len, device, dtype) combination it is reused directly;
+        otherwise it is computed and stored. The instance attributes are pointed at
+        the registry entry so that all layers sharing the same parametrisation
+        reference the same tensor.
+        """
+        cache_key = (self._freq_key, seq_len, str(device), str(dtype))
+
+        if cache_key not in RotaryEmbedding._cache:
+            positions = torch.arange(seq_len, device=device, dtype=torch.float32)
+            # outer product → (seq_len, head_dim // 2); duplicate to (seq_len, head_dim)
+            freqs = torch.outer(
+                positions,
+                self.rotation_freqs.to(device=device, dtype=torch.float32),
+            )
+            angle_embedding = torch.cat((freqs, freqs), dim=-1)
+            RotaryEmbedding._cache[cache_key] = (
+                angle_embedding.cos().to(dtype),
+                angle_embedding.sin().to(dtype),
+            )
+
+        self._cos_cached, self._sin_cached = RotaryEmbedding._cache[cache_key]
+
+    def forward(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        position_ids: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor, float]:
+        """Apply rotary embeddings to query and key tensors.
+
+        The cos/sin cache is extended lazily when position_ids reference positions
+        beyond its current length, or when the device or dtype has changed.
+
+        ``position_ids`` may be any integer tensor shape. Its values are valid
+        position indices into the cos/sin cache:
+
+        - h_l (standard causal): position_ids (B, N), q/k (B, H, N, head_dim).
+        - BEA (packed):          position_ids (B, L, T), q/k (B, L, T, head_dim).
+
+        When q/k have head dimensions absent from position_ids, broadcast dimensions
+        are inserted automatically at dim 1.
+
+        Args:
+            q: Query tensor of shape (batch, [heads,] *pos_dims, head_dim).
+            k: Key tensor of shape (batch, [heads,] *pos_dims, head_dim).
+            position_ids: Integer positions of shape (batch, *pos_dims).
+
+        Returns:
+            Tuple of (q_rotated, k_rotated, attention_scaling). attention_scaling is
+            1.0 for default mode; YaRN returns (0.1·ln(s)+1)² which the caller must
+            apply to attention logits before softmax.
+        """
+        seq_len = int(position_ids.max().item()) + 1
+
+        # The cache is valid when it exists, covers all positions referenced by
+        # position_ids, and matches q's dtype and device. Each condition is named
+        # separately so the rebuild trigger is readable rather than a compound predicate.
+        cache_missing = self._cos_cached is None
+        cache_too_short = not cache_missing and seq_len > self._cos_cached.shape[0]
+        wrong_dtype = not cache_missing and self._cos_cached.dtype != q.dtype
+        wrong_device = not cache_missing and self._cos_cached.device != q.device
+
+        if cache_missing or cache_too_short or wrong_dtype or wrong_device:
+            self._extend_cache(seq_len, device=q.device, dtype=q.dtype)
+
+        cos = self._cos_cached[position_ids]
+        sin = self._sin_cached[position_ids]
+
+        # Insert broadcast dimensions for any head axes present in q/k but absent
+        # from position_ids. Standard: pos (B,N) → cos (B,N,D), q (B,H,N,D) → unsqueeze once.
+        # BEA: pos (B,L,T) → cos (B,L,T,D), q (B,L,T,D) → no unsqueeze needed.
+        while cos.ndim < q.ndim:
+            cos = cos.unsqueeze(1)
+            sin = sin.unsqueeze(1)
+
+        q_rotated = q * cos + _rotate_half(q) * sin
+        k_rotated = k * cos + _rotate_half(k) * sin
+
+        return q_rotated, k_rotated, self.attention_scaling

tokenizer_config.json

@@ -1,13 +1,14 @@
-{
-  "add_prefix_space": false,
-  "backend": "tokenizers",
-  "bos_token": "<|endoftext|>",
-  "eos_token": "<|endoftext|>",
-  "errors": "replace",
-  "is_local": false,
-  "model_max_length": 1000000000000000019884624838656,
-  "pad_token": "<|padding|>",
-  "tokenizer_class": "GPTNeoXTokenizerFast",
-  "trim_offsets": true,
-  "unk_token": "<|endoftext|>"
+{
+  "add_prefix_space": false,
+  "backend": "tokenizers",
+  "bos_token": "<|endoftext|>",
+  "eos_token": "<|endoftext|>",
+  "errors": "replace",
+  "is_local": false,
+  "local_files_only": false,
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "<|padding|>",
+  "tokenizer_class": "GPTNeoXTokenizerFast",
+  "trim_offsets": true,
+  "unk_token": "<|endoftext|>"
 }