Ensure cache_offset is Python int for 0-d scalar

2026-03-15 12:51:46 +00:00
parent 2be3615994
commit 3e953daf5f
1 changed files with 1 additions and 1 deletions
@@ -185,7 +185,7 @@ def fused_qk_norm_rope(queries, keys, q_norm_weight, k_norm_weight,

    q_flat = queries.reshape(B, H_q * D)
    k_flat = keys.reshape(B, H_kv * D)
-    pos = mx.array(cache_offset, dtype=mx.int32)
+    pos = mx.array(int(cache_offset), dtype=mx.int32)

    n_heads = H_q + H_kv
    results = kern(