Skip to content

docs: add a usage page on growing arrays - #4019

Open
apocryphx wants to merge 1 commit into
ml-explore:mainfrom
apocryphx:docs-growing-arrays
Open

docs: add a usage page on growing arrays#4019
apocryphx wants to merge 1 commit into
ml-explore:mainfrom
apocryphx:docs-growing-arrays

Conversation

@apocryphx

Copy link
Copy Markdown
Contributor

Adds a short usage page documenting the preallocate-and-update pattern for arrays that grow a step at a time, and why growing with concatenate in a loop is worse than it looks.

The copying cost is the obvious half. The half worth documenting is that the pattern defeats buffer reuse. BufferCache reuses a pooled buffer only whole — it does not split one to serve a smaller request or combine several to serve a larger one — so a loop whose requests grow each step never matches what it just freed. Every step allocates from the driver, and the freed buffers accumulate at sizes nothing will ask for again.

The cost is CPU-side, so it profiles as GPU idle time between kernels rather than as a slow kernel, which is what makes it hard to find.

mlx-lm sidesteps all of this by preallocating in 256-step chunks, so Python users of it never see the problem. Anyone writing a cache directly, especially in C++, hits it with no warning.

Measured on an M4 Max (mlx 0.32.0), appending one position per step to 20 buffers of [1, 4, N, 512] bfloat16:

context concatenate preallocate + update
512 0.90 ms/step 0.24 ms/step
1024 1.11 ms/step 0.21 ms/step
4096 3.73 ms/step 0.22 ms/step

The preallocated version is flat in context length; the concatenating one is not.

The page also notes that the chunk size should be a multiple of 256, since CUDA's cuDNN attention routing requires the key/value arrays to be slices of a cache whose capacity is a multiple of kv_cache_step (mlx/backend/cuda/scaled_dot_product_attention.cpp). That requirement is currently only discoverable by reading the routing predicate.

Background and fuller measurements in #3886. This documents the existing behavior rather than changing it — I looked at fixing it in the allocator and could not find a version that doesn't pad every allocation, including weights that never recycle, so that seemed like the wrong trade to propose.

Verification: docs build clean with sphinx, no warnings attributable to the new page. The C++ example compiles against the current headers, and both Python forms were run against mlx 0.32.0 and checked to produce results identical to the naive concatenate loop.

Appending with concatenate in a loop copies on the order of n^2 elements, and
also defeats the buffer pool: a monotonically growing request is always
slightly larger than what was just freed, so the pooled buffer is too small
and every step allocates from the driver instead. The cost is CPU-side and
shows up in a profile as GPU idle time rather than as a slow kernel.

Document the preallocate-and-update pattern instead, in Python and C++, and
why the chunk size should be a multiple of 256.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

@zcbenz zcbenz left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for writing this, it is very useful, especially for documenting the behavior of cuDNN kernel.

I wonder if it would make sense to slightly change it to explicitly be a guide on writing performant KV cache, which would attract much more readers.


auto keys = mx::slice(cache, {0, 0, 0}, {1, offset, d});

Why It Matters

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This section is too long and verbose, it should just briefly explain MLX's buffer cache policy, and then list the numbers showing the difference.


cache[:, offset : offset + 1, :] = x

The same pattern in C++:

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Let's just ignore C++ in the tutorial.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants