model : fix k-pool indexer softmax gridDim.y overflow - #10
Closed
AIalliAI wants to merge 1 commit into
Closed
Conversation
Reshape the k-pool gate logits before softmax so n_new does not map to gridDim.y, which is capped at 65535 on CUDA. Fixes prompt processing abort at n_kv >= 262144 reported in ggml-org#27773. Assisted-by: Cursor
Author
|
Closing - pursuing local Strix Halo testing instead. |
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.
Overview
Fixes CUDA
SOFT_MAXcrash during k-pool indexer prefill atn_kv >= 262144, reported in ggml-org#27773.Additional information
The k-pool gate softmax builds
[kpool, head_dim, n_new]. CUDA mapsne[2]togridDim.y(65535 cap). At 262k context with kpool=4,n_newhits 65536 and the launch fails.Reshape logits to
[kpool, head_dim * n_new]before softmax, then reshape probs back. Same idea as the block-count reshape inqwen4exp.cppfor the rms_norm gridDim.y cap.Related sibling issue for the same 262144 threshold: ggml-org#27901 (rms_norm in
norm.cu). Fix there is tracked in ggml-org#28175 / ggml-org#28039.CC @timkhronos @ggerganov @CISC @ngxson
Requirements