From 7cf0772826f72613c2ee8fb7ba5d79d254021b00 Mon Sep 17 00:00:00 2001
From: nan <zjn32202153@gmail.com>
Date: Thu, 3 Sep 2026 00:26:18 +0800
Subject: [PATCH] Use the FFN drop path in ViT residuals

---
 test/eval_core/test_vit_block.py  | 62 +++++++++++++++++++++++++++++++
 worldfoundry/core/nn/vit_block.py | 11 ++++--
 2 files changed, 69 insertions(+), 4 deletions(-)
 create mode 100644 test/eval_core/test_vit_block.py

diff --git a/test/eval_core/test_vit_block.py b/test/eval_core/test_vit_block.py
new file mode 100644
index 000000000..d1ab142ba
--- /dev/null
+++ b/test/eval_core/test_vit_block.py
@@ -0,0 +1,62 @@
+"""Focused tests for the shared ViT residual scheduling helper."""
+
+from __future__ import annotations
+
+import torch
+from torch import Tensor, nn
+
+from worldfoundry.core.nn.vit_block import apply_prenorm_transformer_residuals
+
+
+class _RecordingScale(nn.Module):
+    """Deterministic residual-path stand-in that records its invocations."""
+
+    def __init__(self, scale: float) -> None:
+        super().__init__()
+        self.scale = scale
+        self.calls = 0
+
+    def forward(self, value: Tensor) -> Tensor:
+        self.calls += 1
+        return value * self.scale
+
+
+def test_moderate_stochastic_depth_uses_ffn_drop_path() -> None:
+    """Attention and FFN use their respective paths in the per-sample branch."""
+
+    x = torch.zeros(2, 3)
+    attention_path = _RecordingScale(2.0)
+    ffn_path = _RecordingScale(3.0)
+
+    result = apply_prenorm_transformer_residuals(
+        x,
+        attn_residual=lambda value: torch.ones_like(value),
+        ffn_residual=lambda value: torch.ones_like(value),
+        sample_drop_ratio=0.05,
+        drop_path1=attention_path,
+        drop_path2=ffn_path,
+        training=True,
+    )
+
+    torch.testing.assert_close(result, torch.full_like(x, 5.0))
+    assert attention_path.calls == 1
+    assert ffn_path.calls == 1
+
+
+def test_missing_ffn_drop_path_preserves_legacy_fallback() -> None:
+    """Callers that omit the optional second path retain the old behavior."""
+
+    x = torch.zeros(1, 2)
+    shared_path = _RecordingScale(4.0)
+
+    result = apply_prenorm_transformer_residuals(
+        x,
+        attn_residual=lambda value: torch.ones_like(value),
+        ffn_residual=lambda value: torch.ones_like(value),
+        sample_drop_ratio=0.05,
+        drop_path1=shared_path,
+        training=True,
+    )
+
+    torch.testing.assert_close(result, torch.full_like(x, 8.0))
+    assert shared_path.calls == 2
diff --git a/worldfoundry/core/nn/vit_block.py b/worldfoundry/core/nn/vit_block.py
index e7175de91..3573b6637 100644
--- a/worldfoundry/core/nn/vit_block.py
+++ b/worldfoundry/core/nn/vit_block.py
@@ -27,6 +27,10 @@ def apply_prenorm_transformer_residuals(
 
     attn_kwargs = dict(attn_residual_kwargs or {})
     sd_pos = stochastic_depth_pos
+    # Keep the two residual branches independently configurable.  ``drop_path2``
+    # is optional for callers of this helper that predate the second path, so
+    # retain the historical behavior by falling back to ``drop_path1``.
+    ffn_drop_path = drop_path1 if drop_path2 is None else drop_path2
 
     if training and sample_drop_ratio > 0.1:
         x = drop_add_residual_stochastic_depth(
@@ -41,12 +45,11 @@ def apply_prenorm_transformer_residuals(
             sample_drop_ratio=sample_drop_ratio,
         )
     elif training and sample_drop_ratio > 0.0:
-        drop_path = drop_path1
         if sd_pos is not None or attn_kwargs:
-            x = x + drop_path(attn_residual(x, pos=sd_pos, **attn_kwargs))
+            x = x + drop_path1(attn_residual(x, pos=sd_pos, **attn_kwargs))
         else:
-            x = x + drop_path(attn_residual(x))
-        x = x + drop_path(ffn_residual(x))  # FIXME: drop_path2
+            x = x + drop_path1(attn_residual(x))
+        x = x + ffn_drop_path(ffn_residual(x))
     else:
         if sd_pos is not None or attn_kwargs:
             x = x + attn_residual(x, pos=sd_pos, **attn_kwargs)