feat(aero_realtime): make audio tower causality configurable via is_causal

Author: kcz358

src/lmms_engine/models/aero_realtime/aero_realtime_audio_ops.py

@@ -102,7 +102,7 @@ def _varlen_attention_forward(
         cu_seqlens_k=cu_seq_lens,
         max_seqlen_q=max_seqlen,
         max_seqlen_k=max_seqlen,
-        causal=True,
+        causal=bool(getattr(self.config, "is_causal", False)),
         softmax_scale=self.scaling,
         window_size=window_size,
         dropout_p=0.0 if not self.training else self.attention_dropout,

src/lmms_engine/models/aero_realtime/configuration_aero_realtime.py

@@ -73,6 +73,7 @@ def __init__(
         sliding_window: int | None = 750,
         attention_window_left: int | None = None,
         attention_window_right: int = 0,
+        is_causal: bool = False,
         norm_type: str = "rms_norm",
         mlp_type: str = "swiglu",
         conv_padding: str = "causal",
@@ -100,6 +101,7 @@ def __init__(
             attention_window_left = (sliding_window - 1) if sliding_window is not None else -1
         self.attention_window_left = attention_window_left
         self.attention_window_right = attention_window_right
+        self.is_causal = bool(is_causal)
 
         if norm_type not in ("rms_norm", "layer_norm"):
             raise ValueError(f"norm_type must be 'rms_norm' or 'layer_norm', got {norm_type}")

src/lmms_engine/models/aero_realtime/modeling_aero_realtime.py

@@ -952,6 +952,9 @@ def _aero_audio_eager_forward(
     key_rep = repeat_kv(key, module.num_key_value_groups)
     value_rep = repeat_kv(value, module.num_key_value_groups)
 
+    if bool(getattr(module.config, "is_causal", False)):
+        window_right = 0 if window_right < 0 else min(window_right, 0)
+
     local_mask = _aero_audio_build_local_additive_mask(
         q_len=q_len,
         k_len=k_len,
@@ -989,6 +992,9 @@ def _aero_audio_sdpa_forward(
     key_rep = repeat_kv(key, module.num_key_value_groups)
     value_rep = repeat_kv(value, module.num_key_value_groups)
 
+    if bool(getattr(module.config, "is_causal", False)):
+        window_right = 0 if window_right < 0 else min(window_right, 0)
+
     local_mask = _aero_audio_build_local_additive_mask(
         q_len=q_len,
         k_len=k_len,
@@ -1033,7 +1039,7 @@ def _aero_audio_fa_forward(
     v = value.transpose(1, 2)
     bsz, q_len, _, _ = q.shape
 
-    causal = q_len == k.shape[1]
+    causal = bool(getattr(module.config, "is_causal", False)) or (q_len == k.shape[1])
     window = (window_left, window_right)
 
     if attention_mask is None or attention_mask.dim() != 2: