From 6398cd5b78083372639e00a0cbd26d500761b2ed Mon Sep 17 00:00:00 2001 From: Daniel Hiltgen Date: Wed, 13 May 2026 13:37:31 -0700 Subject: [PATCH] mlx: add memory trace logging (#16131) This should help narrow down the root cause of #16030 --- x/mlxrunner/pipeline.go | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/x/mlxrunner/pipeline.go b/x/mlxrunner/pipeline.go index 89afcf6ab..777653c04 100644 --- a/x/mlxrunner/pipeline.go +++ b/x/mlxrunner/pipeline.go @@ -133,6 +133,7 @@ func (r *Runner) TextGenerationPipeline(ctx context.Context, request Request) er processed += n position += n slog.Info("Prompt processing progress", "processed", processed, "total", total) + logutil.TraceContext(ctx, "mlx prompt forward", "processed", processed, "total", total, "tokens", n, "memory", mlx.Memory{}) // Create snapshot if we've reached a pending offset. if snapOffset := session.nextPendingSnapshot(); snapOffset > 0 { @@ -171,6 +172,7 @@ func (r *Runner) TextGenerationPipeline(ctx context.Context, request Request) er } sample = step(mlx.FromValues(tokens[processed:], 1, total-processed)) + logutil.TraceContext(ctx, "mlx decode seed", "tokens", total-processed, "memory", mlx.Memory{}) dec := decoder{ tokenizer: r.Tokenizer, @@ -194,6 +196,9 @@ func (r *Runner) TextGenerationPipeline(ctx context.Context, request Request) er output := int32(sample.Token.Int()) session.outputs = append(session.outputs, output) + if i == 0 { + logutil.TraceContext(ctx, "mlx decode first token", "memory", mlx.Memory{}) + } if r.Tokenizer.IsEOS(output) { final.DoneReason = 0