optimize E8: precompute causal mask buffer, add eval early exit, and enable compile_model in e8.csv

1 files changed, 6 insertions, 2 deletions

diff --git a/src/models/transformer.py b/src/models/transformer.py
index 2b986ce..a128ad9 100644
--- a/src/models/transformer.py
+++ b/src/models/transformer.py
@@ -36,6 +36,11 @@ class TransformerBaseline(PrimeModel):
         self.blocks = nn.ModuleList([CausalBlock(d, cfg.n_heads) for _ in range(cfg.n_layers)])
         self.ln_f = nn.LayerNorm(d)
         self.head = nn.Linear(d, cfg.vocab)
+        self.register_buffer(
+            "causal_triu",
+            torch.triu(torch.ones(256, 256, dtype=torch.bool), diagonal=1),
+            persistent=False,
+        )
 
     def forward(self, x: torch.Tensor, y_in: torch.Tensor) -> dict:
         cfg = self.cfg
@@ -61,8 +66,7 @@ class TransformerBaseline(PrimeModel):
         # causal mask: input positions (j < T_in) fully visible; output positions causal
         # (True = blocked, per torch.nn.MultiheadAttention bool convention)
         blocked = torch.zeros(T, T, dtype=torch.bool, device=device)
-        causal = torch.triu(torch.ones(T, T, dtype=torch.bool, device=device), diagonal=1)
-        blocked[:, T_in:] = causal[:, T_in:]
+        blocked[:, T_in:] = self.causal_triu[:T, T_in:T]
         key_pad = seq == cfg.pad_id                       # (B,T) True = ignore
         h = e
         for blk in self.blocks: