Skip to content

Fix matmul autotune key stride factors overstating real alignment - #488

Open
RAV64 wants to merge 1 commit into
tracel-ai:mainfrom
RAV64:fix/autotune-key-real-stride-alignment
Open

Fix matmul autotune key stride factors overstating real alignment#488
RAV64 wants to merge 1 commit into
tracel-ai:mainfrom
RAV64:fix/autotune-key-real-stride-alignment

Conversation

@RAV64

@RAV64 RAV64 commented Aug 13, 2026

Copy link
Copy Markdown

Addresses tracel-ai/burn#5352 (one of two layers; see Scope).

Problem

MatmulAutotuneKey::from_parts computes lhs_stride_factor / rhs_stride_factor from 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) = 128 keys 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:

  1. The bucket is first tuned on an aligned representative; an async-copy kernel wins and is cached.
  2. A later cache hit from an under-aligned bucket member re-runs setup, validation fails — after autotune has committed to the winner.
  3. The caller sees Unable to launch matmul because the config is invalid at a point where errors are no longer recoverable: .expect("Should run when selected by autotune.") in cubecl-runtime's LocalTuner::execute, or matmul(...).unwrap() in burn-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 backward dW = x^T @ dy matmul 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: the tma group already refuses its strategies unless both stride factors clear the 16-byte threshold, and the specialized cyclic strategies (which read through AsyncPartialCyclicLoading) 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 transposed dW = x^T @ dy shape of a Linear(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) — branch validate/keyfix-on-burn-pin — burn compiles cleanly with std,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 16 cubek-matmul unit tests pass on both cubek main and the pinned base.

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.
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