Fix sorted gather_qmm on ragged K - #4009
Conversation
|
Ran a wider sweep while reviewing this, and the fix turns out to cover more than the description above claims. Posting the numbers since they change how the patch should be judged. Sweep: 264 configs over Result: 42 configs improve, 0 regress, 219 unchanged. The severity on the fp modes is much worse than the affine numbers in the description suggested. With unaligned K, several configs return values consistent with uninitialized memory rather than merely inaccurate ones:
Two things worth noting from that table:
Three configs move in the other direction, all nvfp4: 0.0002 to 0.0012, and two from 0.0000 to 0.0009 and 0.0001. These are deterministic across repeated runs, and they sit well below the roughly 0.05 quantization floor for the fp4 modes, so I do not think they are meaningful. Flagging them rather than leaving them out. Also for whoever picks these up: this PR and #4010 both add a test immediately before |
|
Some context I found afterwards that makes the case better than the description does, including an alternative fix a reviewer may prefer. The NAX dispatch condition already carries a // line 1007, qmm_nax
if (metal::is_nax_available() && transpose && (K % 64 == 0) && ...)
// line 1202, gather_qmm_nax
if (metal::is_nax_available() && transpose && (K % 64 == 0) && ...)
// line 1563, gather_qmm_rhs_nax
if (metal::is_nax_available() && transpose && ...)The third one does not have it, and that is the entry point this pull request fixes. The guard came from 0dbc7e5, "Centralize NAX condition" (#2811), so the invariant is established, one call site just does not honour it. That suggests a second, smaller fix: add The trade is the fast path. Falling back gives up NAX for every unaligned-K shape, and in a quantized MoE that is a normal shape rather than an edge case, so the cost lands on real workloads. This pull request instead makes the kernel correct for those shapes and keeps NAX, which is why I went that way. Happy to switch to the guard, or to add it as a belt-and-braces alongside the kernel fix, if you would rather be conservative here. Worth noting either way that the two are not equivalent: the guard hides the defect, the kernel fix removes it. The wrong-axis bound in |
|
Closing since there are already PRs with similar changes. |
|
@erwinzhang7 There are also other fixes (#3912, #3922) which I do not have much confidence in, I'm not familiar with this code so it will take some before I can verify the correctness of the fixes, but can you check if your PR could be missing anything? |
zcbenz
added
the
await verification
Two bounds defects in the NAX quantized GEMM path made gather_qmm with sorted_indices=True return silently wrong results whenever K % BK != 0. The activation tile in the K tail was bounded with BK rather than the number of K elements actually remaining, so it read past the end of the row into the next one. The quantized block loader zeroed partial tiles using the wrong component of src_tile_dim. src_tile_dim.y bounds the rows and src_tile_dim.x bounds the columns for both reduction dims, but the reduction_dim == 1 branch compared the row index against the column bound. With BROWS = BN and a K remainder of 32, that zeroed output rows 32 to 63 instead of masking the columns past the K edge, so half the output columns silently lost their tail contribution. Both defects only affect the M5 NAX kernels. Earlier hardware takes a different code path.
b7c5341 to
31d685f
Compare
erwinzhang7
changed the title
|
Rebased onto main. I measured all three PRs against each other and found something I did not Since you said you wanted to verify these before merging, I have put all three on this branch That packaging is only for convenience and it is entirely your call. If you would rather What I didBuilt each PR separately on an M5 Max and measured against a dequantize-then-dense-matmul
Each PR is the only one that fixes one of those rows. On #3912It reproduces on M5 and its diagnosis holds exactly. nvfp4 It does not regress anything. nvfp4 sorted This one matters beyond M5. On an M4 Pro, which has no NAX path at all, main is still On #3922Its I did not take the other two parts of that PR, and I want to be explicit about why, because The column term is not independently safe. #3922 widens the partial-tile guard to On the affine path I could not construct any shape where the column term changes a result: The dispatch guard costs about 1.75x. Median of 5 after a global warm-up, without which
Smaller shapes ran at 15-28% spread and I would not read anything into them. The aligned Non-NAX hardware (M4 Pro)Everything above is M5, where NAX exists. I re-ran the whole thing on an M4 Pro, which has no
Three things follow:
Coverage of the rest of the surface
Reproduction is three short scripts that each take a |
|
Yeah, if I could get my own commit in there, it would be nice. This would be my first official commit with Apple. A little notch on my belt <3. |
|
The non-NAX half of this is separable, and it is verifiable without an M5. #3912 touches only fp_quantized.h and quantized.h, nothing in quantized.cpp. On 39d9a8a each defect reproduces and is clean after: the K overread (0.69 rel err, 72% of outputs wrong), the N store race (200 identical runs give 61 distinct results, 1 after), and the dropped K tail (wrong columns are exactly 16..31 mod 32 and equal the reference minus the last 16 K elements). Metal shader validation flags both out of bounds accesses on main and is clean after: invalid load at 2129920 in nvfp4_qmm_t, invalid store at 1064976 in nvfp4_qmm_n. It also flags a read past the buffer in affine today, invalid load at 213056 in affine_qmm_t_splitk at gs=32 N=100, which the same change removes. On not slowing other modes down: I compiled fp_quantized.metal against main and against the branch and diffed every emitted kernel after normalizing metadata ids. Of 306 gs_32 kernels, 253 are byte identical and none gained an instruction. Plus 4515 randomized differential cases and a 289 case misalignment sweep against an fp64 reference, zero failures, and the added tests fail on main for the right reason (17 of 19 subtests) and pass on the branch. Upstream CI on #3912 is 28/28 green. I also reproduced the cpu job locally with MLX_BUILD_METAL=OFF: 807 python tests OK, C++ suite 245/245. Suggestion: land #3912 first for the non-NAX half, then rebase the NAX commits on top as a follow-up. That splits the review into a part you can verify on hardware you have and a part that needs an M5, and each gets its own CI run. No objection to the NAX work landing right behind it. |
|
Dropping the two cherry-picked commits, this PR is back to just my own fix. @kapellirohith I checked and found MLX squash-merges, so all three commits would have collapsed into one commit under my name at merge. The authorship the cherry-pick was meant to preserve would not have survived it. Philip @kapellirohith's sequencing is also better for testing, so: land #3912 first. It is Three things from the measurements above that are worth having in one place, since they are
The coverage matrix above stands whoever ships what. Happy to rebase this PR onto whatever |
31d685f to
d9cbbdc
Compare
erwinzhang7
changed the title
4 participants
Fixes #3887.
mx.gather_qmm(..., sorted_indices=True)returns silently wrong results on M5 wheneverK % BK != 0. Aligned K is unaffected andsorted_indices=Falseis unaffected, so quantized MoE inference degrades quietly instead of failing.Root cause
There are two separate bounds defects in the NAX path, both in the K tail.
1. Activation tile bounded with
BKinstead of the K remainder. Inaffine_gather_qmm_rhs_naxand its fp counterpart:k_remainis already computed higher up and used for the weight tile (tile_w), so the tail simply needs to use it here too.2.
QuantizedBlockLoader::load_safezeroes partial tiles on the wrong axis.src_tile_dim.ybounds the rows of the tile andsrc_tile_dim.xbounds the columns, for both reduction dims. Thereduction_dim == 1branch compares the row index against the column bound:In the transpose case
BROWS = BN, andtile_w = short2(k_remain, tgp_bn). Withk_remain = 32this zeroes tile rows 32 to 63, which are output columns, instead of masking the columns past the K edge. Those rows are exactly one of the two N direction simdgroups (SN = BN / WN = 32), so half the output columns lose their tail contribution entirely.Fixing only the first defect is not sufficient. It narrows the corruption from every output column to the
tn = 32half, which is why the columns withcol % 64 >= 32are the ones left wrong.Evidence
Measured on an M5 Max (40 core,
applegpu_g17s), on 0.32.0 and at 2c46b95. The check comparessorted_indices=Trueagainstsorted_indices=Falseon identical, already sorted inputs. The flag is only a performance hint, so the two must agree and there is no reference implementation ambiguity. E=8, N=256, M=512, group_size=32, fp16.Max absolute difference relative to
mean(abs(unsorted)), and the fraction of output elements that differ:Across the full sweep of 21 configurations, every failure had
K % 64 != 0andsorted_indices=True, and no configuration with those two properties passed. After the change all 21 agree.The relative error shrinks as K grows, which is consistent with a fixed size corrupted tail region while the correct aligned bulk grows around it.
Test
test_gather_qmm_sortedparameterizes over K, but every case usesK = 512, and the affine cases usegroup_size = 64, which forces K to be a multiple of 64. The K tail was therefore never exercised for this kernel.Added
test_gather_qmm_sorted_unaligned_kcovering K in {160, 288, 544} for affine (group_size 32) and mxfp4. On 0.32.0 all six subtests fail, one of them producing5.24887e+29from uninitialized threadgroup memory. All pass with this change.test_quantized.py,test_blas.py,test_nn.pyandtest_ops.pypass (278 tests) on top of the new one.Performance
No measurable change. On a
gather_qmmbenchmark (E=8, N=4096, M=4096, 4 bit, group_size=32, sorted) the run to run variance at a fixed build was larger than the difference between the patched and unpatched builds at the same commit, so I would not claim a delta in either direction.