chore: cleanup garbage code (#29770)

This commit is contained in:
Xiaoyu Zhang
2026-07-02 16:14:01 +08:00
committed by GitHub
parent 0ae76117ef
commit b276a9acee
55 changed files with 86 additions and 262 deletions
+8 -8
View File
@@ -32,19 +32,19 @@ else:
def rope_pool_fused(
q: "mx.array",
k: "mx.array",
v: "mx.array",
positions: "mx.array",
slots: "mx.array",
k_pool: "mx.array",
v_pool: "mx.array",
q: mx.array,
k: mx.array,
v: mx.array,
positions: mx.array,
slots: mx.array,
k_pool: mx.array,
v_pool: mx.array,
*,
head_dim: int,
num_qo_heads: int,
num_kv_heads: int,
rope_base: float,
) -> tuple["mx.array", "mx.array", "mx.array", "mx.array"]:
) -> tuple[mx.array, mx.array, mx.array, mx.array]:
"""Apply NeoX RoPE to Q/K and scatter K/V into the MLX KV pool.
Args: