Use explicit causal attention for stable remote inference

modeling_tinyqwen3_novelty.py

@@ -89,6 +89,19 @@ def apply_rope(x, cos, sin):
     return (x * cos.to(dtype=x.dtype)) + (rotate_half(x) * sin.to(dtype=x.dtype))
 
 
+def causal_attention(q, k, v):
+    scores = (q.float() @ k.float().transpose(-2, -1)) / (q.size(-1) ** 0.5)
+    causal_mask = torch.ones(
+        scores.size(-2),
+        scores.size(-1),
+        dtype=torch.bool,
+        device=scores.device,
+    ).triu(1)
+    scores = scores.masked_fill(causal_mask, torch.finfo(scores.dtype).min)
+    probs = F.softmax(scores, dim=-1).to(dtype=v.dtype)
+    return probs @ v
+
+
 class MathNoveltyGate(nn.Module):
     def __init__(self, head_dim, floor=0.05):
         super().__init__()
@@ -141,7 +154,7 @@ class NoveltyGQA(nn.Module):
         k = apply_rope(k, cos, sin)
         k = k.repeat_interleave(self.kv_repeat, dim=1)
         v = v.repeat_interleave(self.kv_repeat, dim=1)
-        heads = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+        heads = causal_attention(q, k, v)
         heads = self.novelty(heads)
         out = heads.transpose(1, 2).contiguous().view(bsz, seq_len, self.dim)
         return self.o_proj(out)