Fix matmul autotune key stride factors overstating real alignment - #488
Open
RAV64 wants to merge 1 commit into
Open
Fix matmul autotune key stride factors overstating real alignment#488RAV64 wants to merge 1 commit into
RAV64 wants to merge 1 commit into
Conversation
The stride-alignment factors were computed from the anchored (bucketed) dims, so a problem whose real canonical stride is under the 16-byte async-copy/TMA threshold (e.g. an inner dim of 126 in f32: 504 bytes) shared its autotune bucket with fully aligned problems (126 anchors to 128: 512 bytes). A winner tuned on an aligned representative is then an invalid config for the under-aligned members: selection-time validation (validate_async_copy_with_problem) fails at launch, after autotune has already committed to the kernel. Aligned dims keep the anchored bucketing (no re-tuning churn for runtime-dependent lengths); only real strides under the threshold split out, which they must for soundness.
This was referenced Aug 13, 2026
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
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.
Addresses tracel-ai/burn#5352 (one of two layers; see Scope).
Problem
MatmulAutotuneKey::from_partscomputeslhs_stride_factor/rhs_stride_factorfrom the anchored dims (stride_factor(k_anchored, …)), never from the raw values. The comment above the anchored block explains why: a runtime-dependent dim would re-split every anchored bucket by the raw value's alignment class and re-tune endlessly.But this makes the key overstate alignment: a dim of 126 in f32 has a real canonical row stride of 126 × 4 = 504 bytes (8-byte aligned), while
anchor(126) = 128keys it as 512 bytes (512-byte aligned). The problem then shares an autotune bucket with genuinely aligned problems (e.g. a 128-wide dim).That is unsound for the async-copy/TMA read strategies, which are selection-time validated against the real strides (
validate_async_copy_with_problem: "Async copy requires strides to be aligned to 16 bytes"). The failure sequence:Unable to launch matmul because the config is invalidat a point where errors are no longer recoverable:.expect("Should run when selected by autotune.")incubecl-runtime'sLocalTuner::execute, ormatmul(...).unwrap()inburn-cubecl. On an async device runner the panic is caught and only warn-logged, and the op's registered output buffers are left uninitialized — downstream this surfaces as silent NaN corruption in training (see the linked burn issue).A
Linear(126, N)layer hits this in ordinary training: the backwarddW = x^T @ dymatmul has canonical stride 126 and shares its bucket with the 128-wide hidden-layer gradients.Fix
Keep the anchored bucketing for every dim whose real stride meets the 16-byte reader threshold — the no-churn property the existing tests pin down is untouched for them — but report the real alignment class when it falls below the threshold. The invariant after this change:
factor < 4 ⟺ the real canonical stride is under 16 bytes, so a bucket can never mix members that pass the reader validators with members that fail them. The anchored branch is clamped to the threshold so the invariant also holds when the anchor base is not a power of two (autotune levels 0 and 2).Under-aligned dims are rare (the vast majority of models use aligned widths), so the extra key granularity is confined to problems that genuinely need their own tuning decision.
Scope
This fixes the key's soundness (a bucket can no longer mix members that pass the async-copy/TMA validators with members that fail them — the collision buckets are present in real-workload autotune caches, e.g. a
Linear(126, _)'s problems sharing buckets with 128-wide ones). It is one of two layers needed for the training-corruption issue linked above, and the two only work together. The matmul tuner decides which strategies may run a problem from this key: thetmagroup already refuses its strategies unless both stride factors clear the 16-byte threshold, and the specialized cyclic strategies (which read throughAsyncPartialCyclicLoading) need the same treatment. Neither gate can fire while the key derives its factors from the anchored dims, so this PR is what makes them effective; the companion burn PR (tracel-ai/burn#5370) adds the missing gate. Merging either alone is a no-op for the bug.An earlier attempt at that second layer made the tuner fall back to other candidates after a launch failure (cubecl#1510). It was rejected upstream, correctly — recovery would have to hand the real inputs to another candidate, and autotune must not clone them.
Tests
under_aligned_lengths_split_from_the_aligned_bucket— the regression: 126 (504 B) no longer shares the 128 bucket; distinct under-aligned classes (odd vs 2-aligned) stay distinct.transposed_lhs_shares_bucket_key— extended with the transposeddW = x^T @ dyshape of aLinear(126, _)backward.raw_lengths_within_a_bucket_share_one_key/bucket_maxima_share_the_bucket_key— updated to assert the preserved no-churn property over the 16-byte-aligned members.Validate your PR with burn
Validation details: burn main does not currently compile against cubecl/cubek main (pre-existing quant API drift from cubecl#1479, unrelated to this change), so validation ran against burn's pinned revs, mirroring the maintainers' rev-bump flow. With this fix cherry-picked onto burn's pinned cubek (
611491b) — branchvalidate/keyfix-on-burn-pin— burn compiles cleanly withstd,autodiff,optim,cuda,fusion,autotune, and the training reproduction from the linked issue trains to completion (3/3 fresh-cache runs, finite decreasing losses), identical to the unpatched baseline. The change is API-neutral: key values change for under-aligned dims; the key type and construction are unchanged. All 16cubek-matmulunit tests pass on both cubek main and the pinned base.