[bong-water-water-bong]

Masked KV rows now reach the HRX flash-attention P*V as V = +0. Before this change, identical requests in one process gave different logits.

Problem

Identical greedy requests in one llama-server process (cache_prompt: false) gave different logprobs:

• Qwen3.8-27B UD-Q4_K_XL, no MTP, --ctx-checkpoints 0, /completion top-5 logprobs, 5 requests: up to 0.042 difference between requests, and the top-5 ids changed at step 0.
• Qwen3-0.6B, repeated runs in one context with seq_rm between them: 0.10-0.58 max |dlogit|.
• GLM-4.7-Flash with a ~750-token prompt: up to 13.5 max |dlogit|.

Zeroing only the KV buffers (llama_memory_clear(mem, true)) made every run identical. So did GGML_HRX_DISABLE_DISPATCH=flash_attention.

Cause

Both FA kernels multiply P = 0 for masked keys by the V rows of those keys, in an F16-accumulating WMMA. The result depended on the sign of that V: flipping the sign of one masked V row changed one output element by one F16 ulp. KV rows past the sequence end hold whatever an earlier request or a rejected draft wrote, so the output depended on request history. Accumulating in F32 only made it rarer (3/40 random cases instead of 12/40).

Change

These are in-place edits to the two FA kernel files. Their existing notices and attribution are unchanged.

• ops/flash_attention_decode_split_f32_f16_wmma.loom (1..15 query rows): each 16-key tile is handled one of three ways.
  • Every key visible: direct V fragments, as before.
  • No key visible: a zero V fragment.
  • Mixed: V is staged per subgroup in LDS with the masked rows zeroed.
  • The partial-last-block path also zeroes masked rows in its existing staging.
• ops/flash_attention_f32_f16_wmma.loom (16+ rows): blocks where some valid row masks some key stage V through LDS. Keys that no row of the tile gives a probability are zeroed. The KV tail path does the same.
• New test in its own file, tests/test-hrx-fa-masked-v.cpp (our header), plus a 5-line hook in tests/CMakeLists.txt. It has 27 cases covering the decode split and prefill kernels, with the sequence end inside a block and in the partial last block. In each case the masked K/V rows are rewritten five ways (random, +x, -x, +0, -0) and the outputs must be bit-identical. NMSE against CPU must be <= 5e-4.

Checks

Strixhalo, balanced 85 W, rebased on 388290b:

• test-hrx-fa-masked-v: 27/27, 3 repeats. 388290b without this change: 5/27 fail.

• Random-input sweep, masked V +x vs -x:

  kernel	388290b	this change
  decode split	12/40 cases differ	0/40
  prefill	3/40 cases differ	0/40

• Qwen3-0.6B, in one context:

  • 5 identical runs: bit-identical logits for 40 steps.
  • Runs with a different request in between: bit-identical logits for 40 steps.
  • A 280-token prompt at ubatch 64: bit-identical.

• GLM-4.7-Flash Q4_K_M, ~750-token prompt through the prefill kernel with ggml-hrx: no V transpose for MLA strided V (GLM-4.7-Flash prompts >= 512 tokens: PPL 315,664 -> 5.92) #70's row-major V, 4 runs with a different request in between: bit-identical logits. Before: up to 13.5 max |dlogit|. The first run's tokens are equal across both builds.

• Qwen3.8-27B llama-server, no MTP, --ctx-checkpoints 0, /completion top-5 x5: 5/5 identical, and equal to the first request on 388290b. Before: 5 different results.

• Qwen3.8-27B MTP n3 chat x8: 8/8 identical text.

• test-backend-ops -o MUL_MAT -b HRX0: 322/322, 2 runs.

• iree-test-loom (hrx-system 244cd38) with --sanitizer=access: all 6 prefill cases listed below and all 4 decode-split cases pass.

  • Prefill: selected_row, multirow_selected, multirow_online, query_tail, key_value_tail, pruned_block.
  • Decode split: case, masked_blocks, decode_513, tail.
  • multirow_differential fails SUBRANGE/024 in its test-only helper, both before and after this change.

• With ggml-hrx: attention sinks after FlashAttention (gpt-oss fully on HRX: 1 graph split) #68 (attention sinks) merged on top: test-hrx-attention-sink passes and test-hrx-fa-masked-v is 27/27, 2 runs each.

Cost

llama-bench -fa 1, interleaved, 3 rounds, each starting at <= 60 C, medians, balanced 85 W:

model	test	before	after
Qwen3-0.6B	pp512	20815	19835 (-4.7%)
Qwen3-0.6B	tg128	301.5	301.0
Qwen3.8-27B	pp512	344.7	345.3
Qwen3.8-27B	tg128	10.16	10.49 (within spread)

Decode-split FA dispatch: 9.34 -> 11.27 us. Qwen3-0.6B device time for 64 decode steps: 237.5 -> 241.6 ms.

Not fixed here

With MTP, context checkpoints add a separate request-to-request dependence that alternates every other request. It is filed as an engine issue. --ctx-checkpoints 0 removes it.

🤖 Generated with Claude Code

[bong-water-water-bong]

Review (PR-Agent duty): approve; merge after the #67 placement guard, before #68.

• The bug is real and release-relevant: masked KV rows held stale V from earlier requests or rejected drafts. The f16 WMMA made P=0 x V sign-dependent, so identical requests gave different logits. It reproduces on dense Qwen3-0.6B.
• The fix is in-place in AMD's two FA kernels, with their notice kept. Key tiles that are all visible keep the direct V path, so the common case is unchanged. Fully masked tiles use a zero V fragment. Mixed tiles stage V into the (free) score LDS with masked rows zeroed, behind subgroup barriers.
• The test is in our own file (tests/test-hrx-fa-masked-v.cpp).
• Evidence:
  • regression 27/27 (old 5/27 fail);
  • random sweep 0/40 on both kernels (old 12/40 and 3/40);
  • bit-identical repeats on Qwen3-0.6B and Qwen3.8-27B;
  • GLM-4.7-Flash at ~750 tokens bit-identical on top of ggml-hrx: no V transpose for MLA strided V (GLM-4.7-Flash prompts >= 512 tokens: PPL 315,664 -> 5.92) #70 (was up to 13.5 max |dlogit|);
  • test-backend-ops 322/322;
  • sanitizer clean.
• Cost: pp512 -4.7% on Qwen3-0.6B, the rest within noise.
• Note for ggml-hrx: attention sinks after FlashAttention (gpt-oss fully on HRX: 1 graph split) #68: a one-line tests/CMakeLists.txt conflict.

[bong-water-water-bong]

Rebased onto 35027f7 (after #73) by the orchestrator. The only conflict was tests/CMakeLists.txt: both PRs append a test block, and both are kept. The FA kernel diffs are unchanged. Merging after CI.