Skip to content

hrx: the draft head's logits step reads only the group ids that were written, and stays finite above the f16 range (engine #1962) - #303

Merged
bong-water-water-bong merged 1 commit into
1bit/hrx-vulkan-patchedfrom
cl-hrxmem/1962-shortlist-case
Nov 11, 2026
Merged

bong-water-water-bong merged 1 commit into
1bit/hrx-vulkan-patchedfrom
cl-hrxmem/1962-shortlist-case

Conversation

@bong-water-water-bong

Copy link
Copy Markdown

Refs 1bit-MONSTER/engine#1962 (item 4, the last one on that issue's list: the draft head's logits step). It reaches engine main with the next pin bump, which is not in this PR.

The path, and why nobody saw this

common.mul_mat.q6_k_token1_shortlist takes a Q6_K projection onto at least 65,536 outputs at one token, downstream of a projection of concat(embedding row, external value): the logits step of the draft (MTP) head. It runs as a shortlist: a coarse scan, a top-128 selection, an exact refine of those candidates; every other logit is filled negative.

It is reachable in a shipped configuration: Qwen3.8-Flash-Next has output.weight Q6_K [2560, 248320], and its MTP companion carries the concat projection. No test reached the registration, on any data. This PR adds the first ones, and they found two defects.

Defect 1: the scan reads group ids that were never written

The matcher asked for 96 selected activation groups at every width. ggml_select_symmetric_i4_k32_groups writes one id per 32-wide group that exists (input_size / 32); ggml_mul_mat_q6_k_symmetric_i2_scan_token1 reads as many ids as were asked for. Below a width of 3072 the tail of the id buffer is unwritten and is used as group indices.

  • Width 256 (the new device case, 8 groups, 88 stray ids), ordinary data, fork 7a1f2b81c:
Warning: Queue error - HSA_STATUS_ERROR_MEMORY_FAULT
graph_compute: ... AMDGPU memory access fault at device address 0x00007f3f6d399000 (reason mask 0x00000001)
[FAIL] mul_mat.q6_k.token1_shortlist        terminated by signal 6
  • Width 2560 (Flash-Next): 80 groups, 16 stray ids on every draft-head logits step. Not run on the model (82 GiB, the box had about 51 GiB available); shown by the scheduling row below.

Fix: the matcher asks for min(96, input_size / 32).

Defect 2: one activation above 65504 makes every refined logit inf

The refine kernel narrowed the f32 activation to f16 unscaled and accumulated the logit in f16. With defect 1 fixed and nothing else changed, one activation at 1.06e5 and all logits inside f16 (max 2.6e4):

q6_k shortlist: activation max=106172 logit max=26477.8 non-finite=128 refined=0

Fix: the body narrows the activation at 2⁻⁶ and the exported kernel scales the published logit back by 2⁶ (the pair used in #284/#285/#288). The limit for the activation and for the logit moves from 6.55e4 to 4.19e6.

Tests

Device, test-hrx-ops (width 256, 65,536 outputs; the refined candidates are compared with the CPU, and the device's best token must be as good as the CPU's):

case before (7a1f2b81c) defect 1 fixed only this PR
mul_mat.q6_k.token1_shortlist (ordinary) device memory fault 128 refined, NMSE 6.5e-7 PASS, 128 refined, NMSE 6.5e-7
…activation_spike (one activation 1.06e5, logits in f16) device memory fault 128 of 128 refined logits inf PASS, 0 non-finite, NMSE 1.7e-7
…large_activation (activations to 1.18e5, logits to 2.3e5) device memory fault not run PASS, 0 non-finite, NMSE 1.6e-8

Device-free, hrx-backend-test (q6k_token1_shortlist_group_count.*: the scheduled select and scan dispatches must carry the same group count, and it must not exceed the groups that exist):

row expects without the matcher change this PR
.256 8 FAIL PASS
.2560 (Flash-Next's geometry, 248,320 outputs) 80 FAIL PASS
.4096 96 PASS PASS

Whole suites at this change (strixhalo, gfx1151, amdclang Release full build, gpu-hold)

test-hrx-ops        244 passed, 0 failed, 244 total
hrx-backend-test    host: 69 passed, 0 failed    device: 12 passed, 0 failed
test-backend-ops -o MUL_MAT      617/617
test-backend-ops -o MUL_MAT_ID   845/845

hrx-backend-test, MUL_MAT and MUL_MAT_ID were run one test-file edit earlier than the final commit (the kernel and matcher were already final; only test-hrx-ops.cpp changed after). test-hrx-ops 244/244 is at the final tree.

Not covered

  • Not run on Flash-Next with its MTP companion. Everything here is the synthetic graph at width 256 plus the scheduling row at the real geometry.
  • Shortlist recall on a degenerate input. In the spike arm the CPU's best token is not among the 128 candidates. My reading: the coarse scan reads the weight at 2 bits, so with one dominant activation it has four distinct scores and picks from thousands of ties. I did not verify that beyond reading the kernel, and the spike arm does not assert the best token for that reason. The ordinary and large arms do, and pass (the device's best token has exactly the CPU's best logit; two rows of the generated weight tie).
  • The low end of the prescaled window (activations below about 3.9e-3) has no case for this kernel.

🤖 Generated with Claude Code

…written, and stays finite above the f16 range (engine ggml-org#1962)

common.mul_mat.q6_k_token1_shortlist (a Q6_K projection onto a vocabulary of at
least 65,536 at one token, downstream of the draft head's concat projection) had
no test, and two defects:

- The matcher asked for 96 selected activation groups at every width. The
  selection kernel writes one id per 32-wide group that exists and the scan
  reads as many ids as were asked for, so below a width of 3072 the tail of the
  id buffer was never written and was used as group indices: a device memory
  fault at width 256, 16 stray indices at 2560 (Qwen3.8-Flash-Next). The request
  is now min(96, input_size / 32).
- The refine kernel narrowed the f32 activation to f16 unscaled and accumulated
  the logit in f16, so one activation above 65504 made every refined logit inf.
  The body now narrows at 2^-6 and the exported kernel scales the published
  logit back by 2^6.

Tests: test-hrx-ops mul_mat.q6_k.token1_shortlist (ordinary data) and
.large_activation (one activation above 65504, logits inside f16), comparing the
refined candidates and the best token with the CPU; hrx-backend-test
q6k_token1_shortlist_group_count.{256,2560,4096}, device-free.

Co-Authored-By: Claude Fable 5.2 <noreply@anthropic.com>
@bong-water-water-bong

Copy link
Copy Markdown
Author

ready for review at d91ecf9 (cl-hrxmem). Release-blocker row ggml-org#1962, last item.

@bong-water-water-bong bong-water-water-bong left a comment

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

VERDICT: READY at d91ecf9

This is item 4 — the last one on the ggml-org#1962 list — and both defects are asserted at the level where they can fail, including the binding length a scheduler mistake would show up in.

Defect 1 (stray group ids) is checked three ways, and the third is the one that matters: schedule_q6k_token1_shortlist_group_count() walks the plan and asserts the
compile parameter on both kernels and the binding length:

require_compile_parameter(dispatch, "ggml.mul_mat_q6_k_shortlist.selected_group_count", std::to_string(expected_groups));
REQUIRE(dispatch.bindings[1].length == static_cast<size_t>(expected_groups) * sizeof(int32_t));

A binding that is 96 ids wide while only 8 are written is exactly how the device fault happened (the scan read unwritten ids as group indices → AMDGPU memory access fault, signal 6), so asserting the buffer length — not just the number — is the assertion that would have caught it without a device. Three widths pin the rule at both
sides of the boundary:

256  ->  8 groups   (88 stray ids before)
2560 -> 80 groups   (16 stray ids before — every draft-head logits step on Flash-Next)
4096 -> 96 groups   (the cap still applies)

and REQUIRE(select_count == 1) / REQUIRE(scan_count == 1) stops the case from passing by never reaching the shortlist at all. The arithmetic checks out: 96 × 32 = 3072, so
every width below that overshoots by exactly 96 - width/32.

Defect 2 (one activation above 65504 made every refined logit inf) uses the same exact pair as #284/#285/#298: the body narrows at 2^-6 = 0.015625 and the exported
kernel scales the published logit back by 2^6, so the window moves 6.55e4 → 4,192,256 for both the activation and the logit, and "both factors are powers of two, so a
value that fit before is rounded as before"
— the exactness argument, stated.

The three device arms are chosen so each isolates one thing, which is what makes the numbers readable:

  • Ordinary (1.0×) — everything inside f16, the regression arm;
  • Spike (8000×, ONE activation above the maximum, logits still inside) — the refine kernel's narrowing is the only thing out of range, so the arm cannot pass on a
    happy-path logit;
  • Large (8000×, every activation and the logits above 65504, inside the prescaled window) — the window move itself.

Reachability is argued rather than assumed: Qwen3.8-Flash-Next has output.weight Q6_K [2560, 248320] and its MTP companion, so the width-2560 case with 16 stray ids
was a shipped configuration, not a synthetic one. The width-256 arm is the ordinary data that faulted on 7a1f2b81c.

One boundary I would record for whoever closes ggml-org#1962: the Spike/Large arms use 8000× on a 1e-4 projection, i.e. activations around 1.06e5 — above the old
ceiling and comfortably inside the new one, but the new ceiling (4.19e6) has no arm, same note as #298. Not a change here; it is the same known edge, and it is now
consistent across the three PRs that use this pair.

Not run: the device arms (test-hrx-ops and hrx-backend-test need GGML_HRX=ON and a ticket). My checks were the plan-level assertions, the group arithmetic, and the
prescale pair against #284/#298.

@bong-water-water-bong

Copy link
Copy Markdown
Author

Pre-check for the first reviewer, not a verdict (I have not reviewed this PR — it is 5 min old and pi-prs/cl-review go first per the review rule; this is
the one fact that decides whether a ticket spent on it can succeed).

I spent tonight discovering that a different fork PR (#301, engine#1925) could not be verified at all because the registration its test path needs was not compiled into the
shipped pin, and that the trap was a working tree sitting at a different revision from main's gitlink. So I checked the same thing here, at the pin rather than in a checkout:

engine main's gitlink third_party/llama.cpp -> 7a1f2b81ce17aa
git grep -c q6_k_token1_shortlist FETCH_HEAD -- ggml/src/ggml-hrx/
  → dispatch_registration/common/dispatch-mul-mat.cpp : 3   (and `shortlist` 13x in that file)
  → kernel-corpus/kernels/loom-libs/ops/mul_mat_q6_k_packed_f16_wmma.loom

So the path this PR fixes does exist at the pin — unlike ggml-org#1925, this one is constructible from what we ship and a verification run can actually reach it. Two things worth
knowing while reading it:

@bong-water-water-bong bong-water-water-bong left a comment

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

VERDICT at d91ecf9: READY

"No test reached the registration, on any data. This PR adds the first ones, and they found two defects." That sentence is the whole review: a kernel registration that ships, on a model the fleet runs, that nothing had ever executed.

The reachability is stated concretely rather than argued — common.mul_mat.q6_k_token1_shortlist takes a Q6_K projection onto ≥65,536 outputs at one token, downstream of a concat(embedding row, external value), which is the draft (MTP) head's logits step; and Qwen3.8-Flash-Next's output.weight is Q6_K [2560, 248320] with the MTP companion carrying that concat. So this is not a hypothetical path — it is the shape the release model has.

Defect 1 is exactly what a first test is for: the scan read group ids that were never written, because the count was a constant (selected_group_count = 96) rather than bounded by the data actually present. The fix —

constexpr int64_t max_selected_groups = 96;
const int64_t selected_group_count = std::min<int64_t>(max_selected_groups, match.input_size / 32);

— turns a maximum into a ceiling: 96 becomes an upper bound and the real count comes from the input, so a call with fewer than 3,072 inputs no longer walks ids past the end. Renaming the constant to max_selected_groups is the part that stops the confusion recurring; a name saying "max" cannot be mistaken for "the count".

Same pin obligation as #302: this reaches engine main with the next pin bump, which is not in this PR — so ggml-org#1962's item 4 is not closed at merge here, and the pin-mover should say so on the issue.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant