fix: add Q/K norm handling to BitNetAttention.forward

Address Copilot reviewer comment: BitNetAttention.forward was missing
the optional Q/K normalization path (q_norm/k_norm) that the base
Attention class supports. Added the full q_norm/k_norm handling
(both full and per-head modes) between Q/K/V projections and RoPE,
matching the base Attention.forward logic.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>

Author: justinchuby

src/mobius/models/bitnet.py

@@ -61,6 +61,19 @@ def forward(
         key_states = self.k_proj(op, hidden_states)
         value_states = self.v_proj(op, hidden_states)
 
+        # Optional Q/K normalization (inherited from base Attention)
+        if self.q_norm is not None and self.k_norm is not None:
+            if self._qk_norm_full:
+                query_states = self.q_norm(op, query_states)
+                key_states = self.k_norm(op, key_states)
+            else:
+                query_states = op.Reshape(query_states, [0, 0, -1, self.head_dim])
+                key_states = op.Reshape(key_states, [0, 0, -1, self.head_dim])
+                query_states = self.q_norm(op, query_states)
+                key_states = self.k_norm(op, key_states)
+                query_states = op.Reshape(query_states, [0, 0, -1])
+                key_states = op.Reshape(key_states, [0, 0, -1])
+
         # RoPE
         if position_embeddings is not None:
             query_states = apply_rotary_pos_emb(