docs: add a usage page on growing arrays - #4019
Open
apocryphx wants to merge 1 commit into
Open
Conversation
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
reviewed
Aug 6, 2026
zcbenz
left a comment
Member
There was a problem hiding this comment.
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 |
Member
There was a problem hiding this comment.
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++: |
Member
There was a problem hiding this comment.
Let's just ignore C++ in the tutorial.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adds a short usage page documenting the preallocate-and-update pattern for arrays that grow a step at a time, and why growing with
concatenatein 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.
BufferCachereuses 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: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
concatenateloop.