Update modeling_neollm.py

modeling_neollm.py

@@ -12,6 +12,10 @@ factorised attention log-priors (Litman & Guo, 2026), optional StackMemory
 (Zhang et al., NeurIPS 2025), and optional HOLA-derived episodic memory for
 the FlashAttention-2 IHA seq-expand path (Cui, 2026).
 
+The FlashAttention training paths deliberately use dense right-padded
+attention with a shared dtype and compilation contract across IHA and non-IHA.
+See ``NeoLLMModel.forward`` for its padding/packing constraints.
+
 Attention stack (orthogonal, all active simultaneously when enabled):
     1. Gated Attention (use_gated_attention implicit via q_proj gate chunk):
        applies a head-specific elementwise sigmoid gate to the concatenated
@@ -105,7 +109,11 @@ except ImportError:
 
 from transformers.generation import GenerationMixin
 from transformers.masking_utils import create_causal_mask
-from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
+from transformers.modeling_flash_attention_utils import (
+    FlashAttentionKwargs,
+    _flash_attention_forward,
+    flash_attn_supports_top_left_mask,
+)
 from transformers.modeling_layers import GradientCheckpointingLayer
 from transformers.modeling_outputs import (
     BaseModelOutputWithPast,
@@ -121,6 +129,7 @@ from transformers import AutoConfig, AutoModel, AutoModelForCausalLM
 
 torch._dynamo.config.capture_scalar_outputs = True
 logger = logging.get_logger(__name__)
+_NEOLLM_FA_USE_TOP_LEFT_MASK = flash_attn_supports_top_left_mask()
 
 # ── Optional flash_attn direct import (required only for IHA seq-expand, P>1) ──
 # Attempted once at module load time so the symbols are available when
@@ -1351,6 +1360,93 @@ def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
     return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
 
 
+def _flash_target_dtype(
+    query: torch.Tensor,
+    module: nn.Module,
+) -> Optional[torch.dtype]:
+    """Return the half dtype required by FlashAttention for FP32 Q/K states.
+
+    PyTorch RMSNorm intentionally returns FP32 under BF16 autocast.  IHA's
+    pseudo-head einsums cast those states back through autocast, while the
+    standard path has no equivalent operation.  Resolve the target here so
+    both routes enter every FlashAttention backend with the same dtype instead
+    of relying on a backend-specific compatibility cast.
+    """
+    if query.dtype != torch.float32:
+        return None
+    device_type = query.device.type
+    if torch.is_autocast_enabled(device_type):
+        return torch.get_autocast_dtype(device_type)
+    if hasattr(module.config, "_is_quantized"):
+        return module.config.dtype
+    return next(
+        layer.weight.dtype
+        for layer in module.modules()
+        if isinstance(layer, nn.Linear)
+    )
+
+
+def neollm_flash_attention_forward(
+    module: nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: Optional[torch.Tensor],
+    dropout: float = 0.0,
+    scaling: Optional[float] = None,
+    sliding_window: Optional[int] = None,
+    softcap: Optional[float] = None,
+    is_causal: Optional[bool] = None,
+    s_aux: Optional[torch.Tensor] = None,
+    **kwargs: Unpack[TransformersKwargs],
+):
+    """Compile-stable HF FlashAttention bridge for NeoLLM.
+
+    This mirrors the Transformers FA2/FA3 integration but deliberately does
+    not forward ``module.layer_idx``.  That integer is not consumed by the
+    current FA2/FA3 implementation (it is discarded through ``**kwargs``), yet
+    Dynamo guards its value and otherwise compiles one copy per decoder layer.
+    Keeping the bridge local preserves the public integer attribute for model
+    tooling, checkpoints, and diagnostics without making it part of the
+    attention graph signature.
+    """
+    if kwargs.get("output_attentions", False):
+        logger.warning_once(
+            "Flash Attention does not support `output_attentions=True`. "
+            "Please set attention to `eager` to request attention weights."
+        )
+
+    seq_len = query.shape[2]
+    if any(dim == 0 for dim in query.shape):
+        raise ValueError(
+            "FlashAttention does not support a query tensor with a zero dimension."
+        )
+
+    query = query.transpose(1, 2)
+    key = key.transpose(1, 2)
+    value = value.transpose(1, 2)
+    is_causal = is_causal if is_causal is not None else module.is_causal
+
+    attn_output = _flash_attention_forward(
+        query,
+        key,
+        value,
+        attention_mask,
+        query_length=seq_len,
+        is_causal=is_causal,
+        dropout=dropout,
+        softmax_scale=scaling,
+        sliding_window=sliding_window,
+        softcap=softcap,
+        use_top_left_mask=_NEOLLM_FA_USE_TOP_LEFT_MASK,
+        target_dtype=None,  # NeoLLMAttention normalizes Q/K/V before dispatch.
+        attn_implementation=module.config._attn_implementation,
+        s_aux=s_aux.to(query.dtype) if s_aux is not None else None,
+        **kwargs,
+    )
+    return attn_output, None
+
+
 def causal_first_difference(x: torch.Tensor) -> torch.Tensor:
     return x - F.pad(x[..., :-1, :], (0, 0, 1, 0))
 
@@ -1795,7 +1891,12 @@ def affine_scaled_flash_attention_forward(
     """
     # ── Term 1: standard flash / sdpa output ─────────────────────────────
     # dropout=0.0: we apply dropout to the combined output below instead.
-    attn_fn = ALL_ATTENTION_FUNCTIONS[module.config._attn_implementation]
+    attn_impl = module.config._attn_implementation
+    attn_fn = (
+        neollm_flash_attention_forward
+        if attn_impl in {"flash_attention_2", "flash_attention_3"}
+        else ALL_ATTENTION_FUNCTIONS[attn_impl]
+    )
     flash_out, _ = attn_fn(
         module,
         query,
@@ -4010,6 +4111,18 @@ class NeoLLMAttention(nn.Module):
         # was designed for the head-expand layout and would mismatch here.
         # For P=1 or no IHA: the existing branch logic in _apply_mea_head_mixing
         # is unchanged and handles both the pseudo-slot-aware and normal cases.
+        # Standard FlashAttention layers use the dense right-padded training
+        # contract even inside a mixed IHA schedule.  Keep ``attention_mask``
+        # itself untouched for IHA packing detection and auxiliary operators;
+        # only the backend argument is normalized to ``None``.
+        backend_attention_mask = attention_mask
+        if (
+            self.training
+            and self.config._attn_implementation
+            in {"flash_attention_2", "flash_attention_3"}
+        ):
+            backend_attention_mask = None
+
         if _iha_seq_expand:
             k, v = self._apply_mea_head_mixing(k, v, _force_standard=True)
         else:
@@ -4033,6 +4146,23 @@ class NeoLLMAttention(nn.Module):
                 position_ids=position_ids,
             )
 
+        # Normalize FlashAttention inputs for every topology, not only the
+        # non-IHA fallback.  RMSNorm produces FP32 under BF16 autocast; IHA's
+        # pseudo-head einsums happened to cast it back implicitly, whereas the
+        # standard route reached the Transformers wrapper in FP32.  An explicit
+        # dispatch-boundary cast gives IHA and non-IHA identical dtype behavior,
+        # avoids hidden backend casts/warnings, and keeps FP32 work where it is
+        # intentional (normalization, RoPE/REPO and momentum calculations).
+        if self.config._attn_implementation in {
+            "flash_attention_2",
+            "flash_attention_3",
+        }:
+            flash_dtype = _flash_target_dtype(q, self)
+            if flash_dtype is not None:
+                q = q.to(flash_dtype)
+                k = k.to(flash_dtype)
+                v = v.to(flash_dtype)
+
         # ── IHA: local sliding-window mask for non-seq-expand path (P=1) ─────
         # For P>1 (seq-expand): the sliding window is expressed as window_size
         # in flash_attn_func and handled inside _iha_seq_expand_flash_forward.
@@ -4164,7 +4294,7 @@ class NeoLLMAttention(nn.Module):
                     q,
                     k,
                     v,
-                    attention_mask,
+                    backend_attention_mask,
                     scaling=self.scaling,
                     alpha=alpha,
                     beta=beta,
@@ -4174,6 +4304,11 @@ class NeoLLMAttention(nn.Module):
         else:
             if self.config._attn_implementation == "eager":
                 attn_fn = eager_attention_forward
+            elif self.config._attn_implementation in {
+                "flash_attention_2",
+                "flash_attention_3",
+            }:
+                attn_fn = neollm_flash_attention_forward
             else:
                 attn_fn = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
             backend_kwargs = kwargs
@@ -4187,7 +4322,7 @@ class NeoLLMAttention(nn.Module):
                 q,
                 k,
                 v,
-                attention_mask,
+                backend_attention_mask,
                 dropout=0.0 if not self.training else self.attention_dropout,
                 scaling=self.scaling,
                 **backend_kwargs,
@@ -4414,6 +4549,21 @@ class NeoLLMDecoderLayer(GradientCheckpointingLayer):
         self.siamese_attn_x_scale_init = float(
             getattr(config, "siamese_attn_x_scale_init", 1.0)
         )
+        siamese_stream_scale = (
+            1.0
+            if not self.siamese_depth_scaling
+            else 1.0 / math.sqrt(2.0 * float(layer_idx + 1))
+        )
+        # A tensor buffer keeps the depth-dependent value out of Python control
+        # flow during forward.  Computing it from ``self.layer_idx`` inside a
+        # Dynamo resume frame specialized that frame once per decoder layer.
+        # Non-persistent preserves checkpoint compatibility: the value is fully
+        # determined by layer_idx and the config on every construction.
+        self.register_buffer(
+            "_siamese_stream_scale_value",
+            torch.tensor(siamese_stream_scale, dtype=torch.float32),
+            persistent=False,
+        )
         # Controls only the first pre-attention normalisation applied directly
         # to the embedding stream. Defaults to True for checkpoint/config
         # backward compatibility.  When False, layer 0 does not instantiate
@@ -4710,9 +4860,10 @@ class NeoLLMDecoderLayer(GradientCheckpointingLayer):
             return delta + lr_delta + residual
 
     def _siamese_stream_scale(self, ref: torch.Tensor) -> torch.Tensor:
-        if not self.siamese_depth_scaling:
-            return ref.new_tensor(1.0)
-        return ref.new_tensor(1.0 / math.sqrt(2.0 * float(self.layer_idx + 1)))
+        return self._siamese_stream_scale_value.to(
+            device=ref.device,
+            dtype=ref.dtype,
+        )
 
     def forward_siamesenorm(
         self,
@@ -5647,13 +5798,54 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
                     0, inputs_embeds.shape[1], device=inputs_embeds.device
                 ).unsqueeze(0)
 
-        causal_mask = create_causal_mask(
-            config=self.config,
-            inputs_embeds=inputs_embeds,
-            attention_mask=attention_mask,
-            past_key_values=None,
-            position_ids=position_ids,
-        )
+        # ── Compile-stable FlashAttention training mask contract ─────────────
+        # The Transformers FA2/FA3 wrapper treats a 2-D padding mask as a
+        # variable-length problem: it branches between ``None`` and a tensor,
+        # calls ``nonzero`` to unpad it, and specializes on both the resulting
+        # token count and ``layer_idx``.  With batches containing different
+        # amounts of padding, Dynamo therefore recompiles layer by layer and
+        # retains several compiled graphs, which can *increase* VRAM precisely
+        # when IHA is disabled.  IHA's seq-expand kernel does not use that
+        # wrapper, so it did not expose the same behavior.
+        #
+        # Training in this project has a strict dense, causal, right-padding
+        # contract (no packed sequences and no interior holes).  Under that
+        # contract, valid queries occur before every padded key and causal
+        # attention cannot see those future padded positions.  Outputs for
+        # padded queries are discarded by the masked objectives.  Passing
+        # ``None`` to standard decoder attention backends is thus equivalent on
+        # trainable tokens while keeping one fixed dense kernel topology.
+        #
+        # IMPORTANT: the original ``attention_mask`` remains intact above and
+        # outside the decoder: it still builds ``position_ids`` and is still
+        # consumed by CCE/TWEO/NITP/temporal/NextLat and label masking.  Do not
+        # broaden this branch to packed/left-padded data.  Validate that kind
+        # of input at the data-pipeline boundary rather than with a tensor-value
+        # check here, because a runtime ``all()``/``nonzero()`` would recreate
+        # the graph breaks this branch is intended to remove.  Evaluation,
+        # generation and non-FlashAttention backends retain standard mask
+        # construction below.  IHA receives the original 2-D mask only as a
+        # static dense/padded sentinel for its packing detector; its direct
+        # kernel intentionally does not unpad fixed-length right-padded batches.
+        use_dense_right_padded_flash_train = (
+            self.training
+            and self.config._attn_implementation
+            in {"flash_attention_2", "flash_attention_3"}
+        )
+        if use_dense_right_padded_flash_train:
+            causal_mask = (
+                attention_mask
+                if bool(getattr(self.config, "use_iha", False))
+                else None
+            )
+        else:
+            causal_mask = create_causal_mask(
+                config=self.config,
+                inputs_embeds=inputs_embeds,
+                attention_mask=attention_mask,
+                past_key_values=None,
+                position_ids=position_ids,
+            )
 
         hidden_states = inputs_embeds
         use_siamesenorm = bool(getattr(self.config, "use_siamesenorm", False))