Repository navigation
Conversation
The Q length 3-4 quantized-V fallback in ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1 always launched the F16 instance of the ncols=64 kernel. When the call came from the in-place q4_0/q8_0 K/V dispatch (hs 128/256, K and V the same type), that F16 instance converts K and V to F16 into scratch after dst, but ggml_cuda_flash_attn_ext_get_alloc_size reserves no such scratch for the in-place kernels. The conversion then wrote past the dst allocation and FLASH_ATTN_EXT hit an illegal memory access (e.g. hsk=256, nr23=[6,1], kv=1024, nb=3, q4_0/q8_0). Pass type_KV through so the fallback uses the matching in-place instance.
Collaborator
Author
|
Superseded by #355, which merged the same fix with tests. |
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
1 participant
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.
Fixes an illegal memory access in CUDA flash attention with quantized K/V. Stacked on #352 (base
merge/upstream-20261011);prismhas the same bug and this patch applies there unchanged.What
The short-query fallback added in #189 (
ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1, quantized V with 3-4 queries) always launched the F16 instance of the ncols=64 MMA kernel. When the call comes from the in-place q4_0/q8_0 K/V dispatch (head size 128/256, K and V the same type), the F16 instance converts K and V to F16 into scratch afterdst, butggml_cuda_flash_attn_ext_get_alloc_sizereserves no scratch for the in-place kernels. The conversion writes past the allocation.The fix passes
type_KVthrough, so the fallback uses the matching in-place instance (already instantiated). The F16 and mixed-type paths and the #189 routing are unchanged.Validation (H200, sm_90a)
test-backend-ops -o FLASH_ATTN_EXTtest-backend-ops -o MUL_MAT-fa on -ctk q8_0 -ctv q8_0, vs unpatchedMUL_MAT_IDstill fails 4 f16dup=1cases on CUDA, identically on unmodified upstream master; unrelated to this change.