Expose mlx_fast_rope_offset_array

[ronaldmannak]

Proposed changes

Add array offset support to MLXFast.RoPE to match the Python MLX API.

Dependencies

• Requires: MLX-C PR #85 the underlying C binding and core support for array-offset RoPE must be merged first.
• Once MLX-C is updated, this PR updates the Swift wrapper and exposes the new capability to Swift clients.

Motivation

The Python mlx.core.fast.rope function accepts offset as either an int or array.
This enables several important use cases:

• Continuous batching: Different sequences in a batch at different positions
• Speculative decoding: Verifying multiple candidate tokens at different positions in parallel
• Sliding window attention: Processing long sequences in chunks starting at different offsets

Changes

• Add MLXFast.RoPE(..., offset: MLXArray, ...) overload
• Add MLXNN.RoPE.callAsFunction(_:offset:) overload
• Add unit tests

Example

// Batch of 3 sequences at positions [50, 20, 0]
let offsets = MLXArray([50, 20, 0])
let result = MLXFast.RoPE(queries, dimensions: 64, traditional: false, base: 10000, scale: 1.0, offset: offsets)

Checklist

Put an x in the boxes that apply.

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (if needed)

[davidkoski]

Requires: https://github.com/ml-explore/mlx-c/pull/85 the underlying C binding and core support for array-offset RoPE must be merged first.

Ahh, missed that one. Likely Ronan will do that a different way as he has some automated tools for keeping the API in sync.

[ronaldmannak]

@davidkoski Ah, that explains the lack of newlines in MLX-c :) Happy to wait until it's exposed using the automated tools. Is there any ETA?

[davidkoski]

Is there any ETA?

Not sure, but I can ping Ronan.