diff --git a/crates/cubek-matmul/src/strategy/tune_key.rs b/crates/cubek-matmul/src/strategy/tune_key.rs index 4f125af4c..6d6cc6f97 100644 --- a/crates/cubek-matmul/src/strategy/tune_key.rs +++ b/crates/cubek-matmul/src/strategy/tune_key.rs @@ -137,7 +137,7 @@ impl MatmulAutotuneKey { k: k as u32, }); - // The alignment factors below are computed from the *anchored* dims — + // The pow2 factors below are computed from the *anchored* dims — // the same bucketing the `m`/`n`/`k` fields get — never from the raw // values. A dimension carrying a runtime-dependent length (a KV-cache // width, a dynamic batch) would otherwise re-split every anchored @@ -147,6 +147,11 @@ impl MatmulAutotuneKey { // bucket shares the kernel choice benchmarked on one representative, // and the launch still derives the legal line size from the real // tensors. + // + // The stride factors are the exception: [`stride_factor`] must split + // out real strides under the 16-byte async-copy/TMA threshold, or a + // winner tuned on an aligned bucket member is an invalid config for + // the under-aligned ones. See its doc comment. let m_anchored = anchor(m, None, None, None); let n_anchored = anchor(n, None, None, None); let k_anchored = anchor(k, None, None, None); @@ -173,21 +178,21 @@ impl MatmulAutotuneKey { // its transposed view. Batch strides are products of these, so their // alignment can only be higher. let lhs_stride_factor = match matrix_layout_lhs { - MatrixBatchLayout::Contiguous => stride_factor(k_anchored, elem_lhs), + MatrixBatchLayout::Contiguous => stride_factor(k, k_anchored, elem_lhs), // TMA can't handle discontiguous batches because they're all combined into one dim MatrixBatchLayout::MildlyPermuted { transposed: true, batch_swap: false, - } => stride_factor(m_anchored, elem_lhs), + } => stride_factor(m, m_anchored, elem_lhs), _ => 0, }; let rhs_stride_factor = match matrix_layout_rhs { - MatrixBatchLayout::Contiguous => stride_factor(n_anchored, elem_rhs), + MatrixBatchLayout::Contiguous => stride_factor(n, n_anchored, elem_rhs), // TMA can't handle discontiguous batches because they're all combined into one dim MatrixBatchLayout::MildlyPermuted { transposed: true, batch_swap: false, - } => stride_factor(k_anchored, elem_rhs), + } => stride_factor(k, k_anchored, elem_rhs), _ => 0, }; @@ -219,9 +224,38 @@ impl MatmulAutotuneKey { } } -/// Stride alignment (in powers of two of bytes) of the canonical row stride -/// `cols` — the tightest non-contiguous stride of the layout. -fn stride_factor(cols: usize, elem: ElemType) -> u8 { +/// Async-copy and TMA global readers require 16-byte (2^4) aligned strides; +/// selection-time validation checks the *real* strides against this threshold +/// (`validate_async_copy_with_problem` / `validate_tma_with_problem`). +const READER_ALIGN_FACTOR: u8 = 4; + +/// Stride alignment (in powers of two of bytes) of the canonical row stride — +/// the tightest non-contiguous stride of the layout (`cols` for a contiguous +/// `[.., rows, cols]`, `rows` for its transposed view). +/// +/// Aligned dims key from the *anchored* value so raw lengths inside one +/// anchored bucket share a key (see the bucket comment in [`from_parts`]). +/// A dim whose real stride is under the 16-byte reader threshold must NOT +/// share that bucket: async-copy/TMA algorithms validate the real strides at +/// setup, so a winner tuned on an aligned representative is an invalid config +/// for the under-aligned members of its bucket — the launch then fails after +/// selection instead of never being selected. Clamping the anchored factor to +/// the threshold keeps the invariant `factor < 4 ⟺ real stride is under +/// 16 bytes` even when the anchor base is not a power of two (autotune levels +/// 0 and 2). +/// +/// [`from_parts`]: MatmulAutotuneKey::from_parts +fn stride_factor(cols_raw: usize, cols_anchored: usize, elem: ElemType) -> u8 { + let raw = raw_stride_factor(cols_raw, elem); + if raw < READER_ALIGN_FACTOR { + raw + } else { + raw_stride_factor(cols_anchored, elem).max(READER_ALIGN_FACTOR) + } +} + +/// [`stride_factor`] of one dim, without bucketing. +fn raw_stride_factor(cols: usize, elem: ElemType) -> u8 { let bytes = (cols * elem.size_bits()) / 8; bytes.trailing_zeros().min(MAX_STRIDE_FACTOR) as u8 } @@ -264,14 +298,41 @@ mod tests { /// endlessly for problems the anchor treats as equal. #[test] fn raw_lengths_within_a_bucket_share_one_key() { - // 65..=128 all anchor to the 128 bucket, at every alignment class: - // odd, 2-, 4-, 8-, 16-aligned, and the exact power of two. - let reference = key(64, 69, 64); - for kv in [70, 76, 88, 96, 112, 128] { + // 65..=128 all anchor to the 128 bucket. Every 16-byte-aligned raw + // length shares the bucket key: 76/88/96/112 (f32 row strides of + // 304/352/384/448 bytes) and the exact power of two. + let reference = key(64, 76, 64); + for kv in [88, 96, 112, 128] { assert_eq!(reference, key(64, kv, 64), "kv {kv} split the bucket"); } } + /// Raw lengths whose real row stride is under the 16-byte async-copy/TMA + /// threshold must NOT share the bucket of the aligned lengths: the reader + /// algorithms validate the real strides at setup, so an async-copy winner + /// tuned on an aligned representative is an invalid config for the + /// under-aligned members and the launch fails after selection. On an async + /// runner that failure is only warn-logged, which surfaced as silent NaN + /// weights in training (see the PR description for the full chain). + #[test] + fn under_aligned_lengths_split_from_the_aligned_bucket() { + let aligned = key(64, 128, 64); + // 126: f32 row stride 504 bytes = 8-byte aligned; anchors to 128. + assert_ne!( + aligned, + key(64, 126, 64), + "under-aligned kv shared the bucket" + ); + // 69 (odd) and 70 (2-aligned) split too, and by alignment class. + assert_ne!(aligned, key(64, 69, 64)); + assert_ne!(aligned, key(64, 70, 64)); + assert_ne!( + key(64, 69, 64), + key(64, 70, 64), + "distinct under-aligned classes merged" + ); + } + /// The bucket maximum shares its bucket's key. The scale thresholds sit /// exactly on the pow2 bucket ceilings (512, 2048), so deriving the scale /// from the raw dims used to split the (256, 512] bucket into @@ -315,10 +376,19 @@ mod tests { None, ) }; - // 65..=128 all anchor to the 128 bucket, at every alignment class. - let reference = key_t(69); - for m in [70, 76, 88, 96, 112, 128] { + // 65..=128 all anchor to the 128 bucket; the 16-byte-aligned raw + // lengths share its key. + let reference = key_t(76); + for m in [88, 96, 112, 128] { assert_eq!(reference, key_t(m), "m {m} split the transposed bucket"); } + // Under-aligned rows split out, exactly as in the contiguous case: + // this is the `dW = x^T @ dy` shape of a Linear(126, _) backward, + // whose real canonical stride is 126 * 4 = 504 bytes. + assert_ne!( + reference, + key_t(126), + "under-aligned transposed m shared the bucket" + ); } }