From 55fa80d07a9f2e9c51632328665d00f2cbdf7b02 Mon Sep 17 00:00:00 2001 From: Daniel Hiltgen Date: Thu, 16 Apr 2026 13:07:19 -0700 Subject: [PATCH] mlx: additional gemma4 cache fixes (#15607) Harden additional corner cases --- x/mlxrunner/cache/cache.go | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/x/mlxrunner/cache/cache.go b/x/mlxrunner/cache/cache.go index a513b5717..f9202d90f 100644 --- a/x/mlxrunner/cache/cache.go +++ b/x/mlxrunner/cache/cache.go @@ -337,9 +337,10 @@ func (c *RotatingKVCache) State() []*mlx.Array { if c.keys == nil || c.values == nil { return nil } + liveLen := min(c.offset, c.keys.Dim(2)) return []*mlx.Array{ - c.keys.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(0, c.offset), mlx.Slice()), - c.values.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(0, c.offset), mlx.Slice()), + c.keys.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(0, liveLen), mlx.Slice()), + c.values.Slice(mlx.Slice(), mlx.Slice(), mlx.Slice(0, liveLen), mlx.Slice()), } }