Skip to content

Commit 9e3e7b9

Browse files
committed
Add TODO for custom metal kernel
1 parent e4018c7 commit 9e3e7b9

File tree

1 file changed

+3
-5
lines changed

1 file changed

+3
-5
lines changed

Libraries/MLXVLM/Models/Gemma3n.swift

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1474,12 +1474,10 @@ private func gemma3nAttentionWithCacheUpdate(
14741474
private func bicubicInterpolate(
14751475
_ x: MLXArray, to targetSize: (Int, Int), alignCorners: Bool = false
14761476
) -> MLXArray {
1477-
// TODO: This implementation uses nested loops and sequential MLX operations, which is much slower
1477+
// This implementation uses nested loops and sequential MLX operations, which is much slower
14781478
// than the Python version that uses mx.fast.metal_kernel() for parallel GPU computation.
1479-
// MLX Swift currently doesn't have custom Metal kernel creation capabilities like Python's
1480-
// mx.fast.metal_kernel(). Consider optimizing with vectorized MLX operations or requesting
1481-
// custom kernel support from the MLX Swift team for better performance.
1482-
1479+
// TODO: Port the custom metal kernel from Python to Swift using `MLXFast.metalKernel`.
1480+
//
14831481
// Input: NHWC format [batch, height, width, channels]
14841482
// Output: NHWC format [batch, target_height, target_width, channels]
14851483

0 commit comments

Comments
 (0)