Repository navigation
qwen4exp : use maskless block selection only where the selected-key kernels run - #115
Conversation
|
Tested this on a Strix Halo (gfx1151, Vulkan) and it fixes it for us. With Flash-Next on current master, perplexity at It needs a small rebase after #111: both changed the same line in ggml_type type_k, ggml_type type_v, bool selected_key_attn) {
return selected_key_attn && blk_bias && cparams.causal_attn && n_stream==1 && ratio>1 && hparams.indexer_top_k%ratio==0 &&The rebased commit is on my fork if you want it, and I'm happy to push it to your branch if that saves you a step. Assisted-by: Claude Opus 5.5 |
…ernels run Complete-block selection builds a maskless flash-attention op whose visibility lives only in the selected-cell list (src[5]). Only the selected-key kernels read that list, and they exist only in the HIP backend (RDNA3.5); the gate did not check the device. ggml-cuda aborts on the op when no QSA kernel takes it. The Vulkan backend accepted it and ran its dense kernel, which ignores src[5], so the QSA layers attended to the whole K/V view, unmasked, as soon as the context exceeded indexer_top_k + ratio - 1: - other sequences' cells: with several conversations in one unified cache, a conversation's output depended on another conversation's content (a 2048-token prompt, greedy, changed at the 2nd token next to related text, where it had been 6.6 nats ahead; unchanged next to unrelated text) - no sparsity: every cell instead of the top-k blocks test-backend-ops shows it directly: QSA_PREFILL_MASKLESS fails on Vulkan (0/6, ERR 0.48 to 0.98 against the FP64 reference). - The memory decides once per model whether its QSA layers all run on ROCm (LLAMA_QSA_SELECTED_KEY=0/1 overrides that), and qwen4exp_use_block_selection() requires it. Everywhere else the masked top-k path is used, which the dense kernels compute correctly (QSA_PREFILL and QSA_DECODE pass on Vulkan except their poison=1 cases, which fail on master as well). - Vulkan no longer claims flash attention with a selection list and no mask, so the op can never be computed wrongly there. Assisted-by: Claude Opus 5.5 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Complete-block selection builds a maskless FLASH_ATTN_EXT whose visibility is only in the selection list (src[5]). Vulkan had no kernel that reads the list: halo-box#115 makes it decline the op, and Vulkan then takes the masked per-cell top-k path, which is correct but slow at depth. flash_attn_sel.comp computes the op: one workgroup per query x K/V head, one 64-lane subgroup per query head of that K/V head (GQA 4/8/12, head 256). It sorts the query's list in shared memory and skips invalid and duplicate entries, as the CPU path does, then gathers K and V in blocks of 32 cells. Two lanes score one key, so the online softmax needs two subgroup reductions per block. With fewer than 64 workgroups the rows are split, and flash_attn_split_k_reduce combines the parts. supports_op takes a maskless selected-key op exactly when this kernel does. llama_memory_hybrid_idx asks each QSA layer's device with a probe op instead of checking for ROCm by name. HIP stays accepted; CUDA is not asked, because its supports_op ignores the list. Radeon 890M (gfx1150), Qwen3.8-Flash-Next UD-IQ4_XS: - test-backend-ops: QSA_PREFILL_MASKLESS 6/6 and QSA_DECODE_MASKLESS 25/25 (master: 0/6 and 0/25), NMSE <= 6.9e-13 against the FP64 reference - pp2048 at depth 12000 / 32000: 140.9 / 134.1 t/s (masked path, halo-box#115: 139.0 / 88.7) Assisted-by: Claude Opus 5.5 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2b74a36 to
72f0f12
Compare
|
Thanks for testing this on gfx1151, and for the rebase. I rebased onto 307c50d with the resolution you suggested ( Your perplexity numbers were the clue we had missed. Our ubatch-size check only compares the tokens generated after a prompt, so it could not see look-ahead inside the prompt. On gfx1150, with text the model cannot have memorised (this fork's own August/September docs), at One more thing changed with the new master: our token-level isolation test no longer flips on master, but the leak is still there. With a second set of prompts, master's output depends on the other conversations' content: it diverges at token 9, where the gap is 1.86, so not a near-tie. Instrumentation shows the Vulkan kernel reading all 50176 cells of the view, the residents' cells included. Details are in the description. Assisted-by: Claude Opus 5.5 |
Complete-block selection builds a maskless FLASH_ATTN_EXT whose visibility is only in the selection list (src[5]). Vulkan had no kernel that reads the list: halo-box#115 makes it decline the op, and Vulkan then takes the masked per-cell top-k path, which is correct but slow at depth. flash_attn_sel.comp computes the op: one workgroup per query x K/V head, one 64-lane subgroup per query head of that K/V head (GQA 4/8/12, head 256). It sorts the query's list in shared memory and skips invalid and duplicate entries, as the CPU path does, then gathers K and V in blocks of 32 cells. Two lanes score one key, so the online softmax needs two subgroup reductions per block. With fewer than 64 workgroups the rows are split, and flash_attn_split_k_reduce combines the parts. supports_op takes a maskless selected-key op exactly when this kernel does. llama_memory_hybrid_idx asks each QSA layer's device with a probe op instead of checking for ROCm by name. HIP stays accepted; CUDA is not asked, because its supports_op ignores the list. Radeon 890M (gfx1150), Qwen3.8-Flash-Next UD-IQ4_XS: - test-backend-ops: QSA_PREFILL_MASKLESS 6/6 and QSA_DECODE_MASKLESS 25/25 (master: 0/6 and 0/25), NMSE <= 6.9e-13 against the FP64 reference - pp2048 at depth 12000 / 32000: 140.9 / 134.1 t/s (masked path, halo-box#115: 139.0 / 88.7) Assisted-by: Claude Opus 5.5 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> (cherry picked from commit 3bf5356)
2 participants
Overview
Fixes #130.
Complete-block selection builds a maskless
FLASH_ATTN_EXTop, where the visibility is only in the selection list(
src[5]). Only the selected-key kernels read that list, and those kernels exist only in the HIP backend forRDNA3.5. The gate does not check the device. On Vulkan the dense kernel runs the op, ignores the list, and attends
to every cell of the K/V view: other sequences' cells, the later tokens of the query's own ubatch (look-ahead), and
all cells instead of the top-k blocks.
llama_memory_hybrid_idxdecides once per model whether all QSA layers run on ROCm (selected_key_attn, loggedat load).
LLAMA_QSA_SELECTED_KEY=0/1overrides that decision for diagnosis.qwen4exp_use_block_selection()requiresselected_key_attn. On any other backend the masked top-k path isused, and the dense kernels compute it correctly.
computed wrongly there, whatever the caller does.
When all QSA layers run on ROCm, the gate gives the same result as before. When only some of them do (partial
offload), all of them now use the masked path, because the maskless op must not reach a layer on another backend.
This PR adds one clause to the gate, and #110 turns the gate into a reason classifier. The two PRs touch the same
lines. Whichever merges second, we rebase it: here the clause becomes one more refusal reason ("no selected-key
kernels on this backend").
Measurements
The baseline is not a correct reference past 2051 cells: that is the bug. So speed at depth compares the wrong
(dense, unmasked) path with the correct masked top-k path.
llama-bench -m MODEL -p 2048 -d 0,12000,32000,64000 -b 2048 -ub 512 -n 0 -r 4 -ngl 99 -fa on -ctk f16 -ctv f16 -lm none -lzm on-direct -o jsonlBaseline: (307c50d)
After: (72f0f12)
The fix costs prompt speed at depth, 18% at 32k and 25% at 64k. The baseline was faster there because it skipped
the top-k selection and attended to every cell in one dense pass. The correct path is the masked per-cell top-k,
which builds
[n_kv, strip]tensors that grow with the context. Decode was 4-18% faster after the fix in thellama-server runs below (one conversation, 12.7k-53.4k tokens). #119 adds a Vulkan kernel that reads the selection list, like the HIP
selected-key kernels. With it, prompt speed is 149.8 t/s at 32k and 144.1 t/s at 64k (this PR: 93.3 and 60.8), and
the masked cells are skipped instead of weighted with zero.
Correctness:
test-backend-ops -b Vulkan0, baseline -> change:QSA_PREFILL_MASKLESSQSA_DECODE_MASKLESSQSA_PREFILLpoison=1cases failQSA_DECODEpoison=1cases failTOPK_QSAFLASH_ATTN_EXTThe
poison=1cases put NaN into masked cells. Vulkan fails them before and after this change: its dense kernelstill reads masked cells. The masked path is correct for finite cache contents, but it does not isolate a sequence
from a NaN in another sequence's cells. The selected-key kernels would not read those cells at all. That gap
already exists and is not addressed here.
The four repository corpora (
tests/corpus/correctness-*.txt) go throughllama-server(1 slot, no MTP). Foreach one we compare the prompt token IDs, the complete next-token distribution after prefill (every
vocabulary entry) and 64 greedy steps with their top-20 log-probabilities:
all four corpora, including every log-probability and all 64 decode steps.
wrong path there. The top token after prefill is the same for all four, and so are all 64 greedy tokens. The
distributions move by KL 0.002 (prose, 12.7k), 0.004 (code, 22.8k), 0.098 (structured, 27.1k) and 0.013
(numeric, 53.4k).
Look-ahead.
llama-perplexity -c 4096 -b 2048 -ub 512scores positions 2048-4095 of each chunk, where blockselection runs. The text is one the model cannot have memorised: this fork's own docs, written in August and
September 2026 (
docs/development/DSV4-vulkan-sparse-prefill-progress.md,backend-scheduler.md,DSV4-vulkan-lightning-indexer-progress.md,rocm-mmb-context.md,vulkan-sparse-fa-deferred.md,docs/speculative-prefill.mdanddocs/build-profiling.mdat 307c50d, concatenated in that order, sha256a3b9cd54...; 5 chunks at-c 4096, 10 at-c 2048):A perplexity of 1.14 on text the model has never seen means the baseline sees the tokens it is asked to predict:
the dense kernel attends to the later tokens of the query's own ubatch. Below the block-selection threshold all
three builds agree to the last digit. @Nathanw1014 measured the same on Strix Halo (gfx1151): 1.06 on master,
3.94 with this PR, at
-c 4096(comment below).llama-perplexity, 16384 context,-b 2048 -ub 512, 2 chunks of Python standard-library source (the first400k characters of
/usr/lib/python3.14/*.pyin sorted order, sha256108e868c...): baseline 1.0303 +/- 0.0030,after 1.0664 +/- 0.0046. The model has largely memorised this text, so every value is close to 1, and look-ahead
can lower it only a little. The baseline's lower value is the same look-ahead, not a better result.
Isolation,
-c 65536 -np 4 -kvu, no MTP: slot 0 generates 48 greedy tokens from a 2048-token prompt alone, next tothree 16k-token conversations X (related text) and next to three Y (unrelated text):
identical, X vs Y differ at token 2. The output depends on the other conversations' content.
is the four repository corpora (2048 tokens), X repeats them, Y is Python standard-library source. X vs Y differ
at token 9, where X's top-2 gap is 1.86, so not a near-tie: the output still depends on the other conversations'
content. Instrumentation shows the Vulkan kernel running the maskless op without a mask over all 50176 cells of
the view, the residents' cells included, for the probe's prompt and for every decode step.
give the same 48 tokens with identical log-probabilities. Alone vs next to residents differ at token 26, a
near-tie (top-2 gap 0.04) that X and Y share, so it comes from the larger cache's rounding, not from the other
conversations.
Ubatch size: a 6000-token prompt with
-ub 512and with-ub 256gives the same 32 tokens after the change (maxtop-1 log-prob difference 0.0010). That test compares only the tokens generated after the prompt, so it does not
see look-ahead inside the prompt; the perplexity test above does.
All results above are from master 307c50d after the rebase, except where a revision is named.
Additional information
below.
ROCm, so qwen4exp: GGML_ABORT ("maskless selected-key attention was not taken by a QSA kernel") on KV-cache exhaustion under --kv-unified with concurrent slots #85 stays open.
device itself says whether it takes the maskless op. ggml-cuda would still need to decline the shapes its QSA
kernels do not take (qwen4exp: GGML_ABORT ("maskless selected-key attention was not taken by a QSA kernel") on KV-cache exhaustion under --kv-unified with concurrent slots #85). That needs HIP hardware to verify, so neither PR does it.
llama-server, the memory-logger probe, themulti-slot tests) can be shared if it is useful here.
Requirements
fork's qwen4exp / QSA code on the Vulkan backend. It does not depend on the device. The measurements come from
Strix Point (gfx1150), not Strix Halo.
description, and ran every measurement. Every commit carries
Assisted-by: Claude Opus 5.5.for gfx1150 here, and ROCm's GTT limit cannot hold this model on our machine.