Repository navigation
Conversation
jeffbolznv
left a comment
There was a problem hiding this comment.
IIUC this allows ggml-alloc to reuse the scratch memory for other tensors, which is nice. Seems theoretically possible this could be worse in some cases, but is probably a win on average.
| } | ||
|
|
||
| bool ggml_vk_scratch_available(const ggml_backend_vk_context * ctx, const ggml_tensor * dst) { | ||
| if (dst->view_src != nullptr) { |
There was a problem hiding this comment.
This is fine for this optional FA optimization, but it seems like if we have required scratch space for other ops then we wouldn't be able to use this kind of scratch space for operations that have a view for the destination.
| } | ||
| // meta backend (-sm tensor) reserves for the full tensor, slices may not fit | ||
| ggml_backend_vk_buffer_context * buf_ctx = (ggml_backend_vk_buffer_context *)dst->buffer->context; | ||
| return vk_tensor_offset(dst) + ggml_vk_scratch_offset(ctx->device, ggml_nbytes(dst)) + scratch <= buf_ctx->dev_buffer->size; |
There was a problem hiding this comment.
codex reports an issue here:
• [P1] The availability check validates the whole buffer, not the tensor’s allocation. ggml-vulkan-scratch.cpp:93 (ggml/src/ggml-vulkan/ggml-vulkan-scratch.cpp:93) can accept
scratch extending into another tensor’s slot in tensor-split mode. The allocation-size safeguard makes this concrete: a full tensor can exceed the limit and reserve no
scratch, while its smaller slices enable scratch. Using the actual helpers, I reproduced a full tensor reserving 512 KiB, a half-head slice requiring 1.25 MiB, and
scratch_available() returning true because the backing buffer was large enough. This needs bounds from the actual reserved slot, or a fallback when that capacity cannot be
established. Additional barriers cannot fix overwriting another live allocation.
There was a problem hiding this comment.
Yes, it is not covering all cases. I haven't found a way to do that, since I can't see the actual allocation size. It does block an actual crash I was able to produce though, so the check does help. As far as I understood it, an actual fix would have to be in the meta backend.
I haven't seen a way it would fail in actual use, have you?
There was a problem hiding this comment.
I haven't run the change yet, just reviewed it. I agree this check is necessary, I was just thinking about whether we could eventually port all scratch memory usage to work like this, and I don't think we can due to this limitation. If the dst is usually not a view for FA, then it should be fine for this change.
Overview
Add Vulkan scratch buffer system, use for Flash Attention dequant_kv to make its memory use visible to GGML (to resolve #30577). Other uses of internal buffers in the backend can be integrated into this system as follow-up work.
There seem to be some potential issues with meta-backend handling of larger compute buffers in combination with ggml-alloc multi-buffers, in specific worst case scenarios (non-GQA model with huge kv cache) it can cause large chunks of memory to be allocated, but unused. But in regular use this should not happen. I hope eventually we can improve the behaviour of meta backend and allocator for these cases.
This should also fix #30048.
Requirements