Extend block wrapper parameters

src/exo/worker/engines/image/pipeline/block_wrapper.py

@@ -1,3 +1,5 @@
+from typing import Any
+
 import mlx.core as mx
 
 from exo.worker.engines.image.pipeline.adapter import (
@@ -45,6 +47,7 @@ class JointBlockWrapper:
         mode: BlockWrapperMode,
         patch_start: int | None = None,
         patch_end: int | None = None,
+        **kwargs: Any,
     ) -> tuple[mx.array, mx.array]:
         """Apply the joint block.
 
@@ -58,6 +61,8 @@ class JointBlockWrapper:
             mode: CACHING (populate cache) or PATCHED (use cached K/V)
             patch_start: Start index for patched mode (required if mode=PATCHED)
             patch_end: End index for patched mode (required if mode=PATCHED)
+            **kwargs: Additional model-specific arguments (e.g., encoder_hidden_states_mask,
+                block_idx for Qwen)
 
         Returns:
             Tuple of (encoder_hidden_states, hidden_states)
@@ -73,6 +78,7 @@ class JointBlockWrapper:
             text_seq_len=text_seq_len,
             patch_start=patch_start,
             patch_end=patch_end,
+            **kwargs,
         )