[Bugfix] Fix sinks related argument error (#1044)

Signed-off-by: Kyuyeun Kim <kyuyeunk@google.com>

Author: kyuyeunk

tpu_inference/layers/jax/attention_interface.py

@@ -342,14 +342,14 @@ def attention(
     q: jax.Array,
     k: jax.Array,
     v: jax.Array,
-    sinks: jax.Array | None,
     attention_metadata: AttentionMetadata,
     mesh: Mesh,
     head_dim_original: int | None = None,  # before padding,
     attention_chunk_size: int | None = None,
     q_scale: float | None = None,
     k_scale: float | None = None,
     v_scale: float | None = None,
+    sinks: jax.Array | None = None,
 ) -> Tuple[jax.Array, jax.Array]:
     # T: seq_len
     # N: num_heads

tpu_inference/layers/vllm/attention.py

@@ -178,12 +178,12 @@ def _jax_attn_func(
         q,
         k,
         v,
-        sinks,
         attention_metadata,
         mesh,
         q_scale=q_scale,
         k_scale=k_scale,
         v_scale=v_scale,
+        sinks=sinks,
     )
 
     # Convert the shape back to vLLM's convention