Repository navigation
Metal: one pointer-alignment rule, and tests that hold ops to it (#10) - #11
Merged
Merged
Conversation
Both backends with gemm_bf16_nt (CUDA since 7fdb056, Metal through its reader) compute any K and alignment, and qwen.h's prefill ignores the result, so a refusal would leave its chunk unwritten. The misaligned A and C checks now require the op to run, and a third hands it W one bf16 off, off the 4 bytes Metal's wider loads need (its vectorized loader misreads there; bf16_rows_vec_ sends such a W to the reader). On the M1 Pro, refusing a misaligned A or C, or sending that W to the vectorized loader, each fails one check.
Measured on the M1 Pro, a load or store of 4 bytes or more is right from any 4-byte-aligned address, float4 off 16 bytes included, and silently moves the wrong bytes from a 2-byte-aligned one; the compiler may also merge two narrower loads into such a one. Apple asks only that an offset be a multiple of the bound argument's element, so f32 buffers always pass and only 1- and 2-byte ones read wider need checking. word_aligned_ in metal.h states that once. bf16_rows_vec_ uses it in place of W.off % 16 (so a W on 4 bytes takes the vectorized loader too) and quant_ok_ for Q4_K's 32-bit reads. The bf16 prefill GEMM's misaligned cases move to a table and add W off 16 bytes; the f16-sum form gets both W cases, which it must also run (its shared op falls back to the f32 GEMM) within f16 rounding. With word_aligned_ always true, both W-off-4-bytes cases fail.
Its staging loads four halves at a time with no alignment check, which Metal's rule (word_aligned_) would ask of a 2-byte buffer. It must not refuse, since no composition takes an f16 cache, and on the M1 Pro it reads a cache 2 bytes off right. The test holds it to that; staging zeros whenever the address is 2 bytes off fails it.
An f16 cache has no composition to fall back to, so the tiled prefill must run on whatever offset it is handed. It declared K and V on four elements (float4 / uint2 staging), so on the RTX 3090 an f16 cache one half off was refused and 0f42e43's test failed (ctest cpu, gpu, auto). The core now picks the vector load per operand by its base (rows are AD-contiguous, so the base decides every row) and widens one element at a time otherwise; TL_ALIGN declares only the element. An f32 cache one float off now runs here too instead of taking the GEMM form. On the RTX 3090: ctest 9/9, check_cuda ALL OK. Zeroing the scalar path fails both misaligned prefill checks (f16 and f32). Registers 138/108/56 -> 134/110/64 (f16 64/128/256), 146/114 -> 132/110 (f32), no spills; batched attention at T=2048 stays 0.160 ms; Qwen2.5-0.5B chat output unchanged.
…ment Two conventions said "this kernel guards its own alignment": the bf16 NT GEMMs had no TL_ALIGN and sat in check_cuda's self_guarded list, while the tiled prefill attention declares its elements' sizes. The GEMMs now declare theirs too (A f32, B bf16, C f32), and the list is gone: every kernel that moves memory by vectors has a TL_ALIGN, and one that picks its width by the address declares the element. A pointer off its own element (an f32 operand 2 bytes off) is now refused at launch instead of launched. On the RTX 3090: ctest 9/9, check_cuda ALL OK; dropping one of the three declarations fails check_cuda (vector access undeclared).
doctest prints a const char* as its address by default, so a failing check_misaligned logged "0x5713d072eea8" instead of the case it ran (and test_array's CAPTURE(f.name) did the same). DOCTEST_CONFIG_TREAT_CHAR_STAR_AS_STRING must agree across every translation unit, so it is set where the suite is built: the tensorlib_test target and test/wasm/build.sh. No check compares char pointers, which the option would turn into string comparisons. On the RTX 3090: ctest 9/9; zeroing the tiled prefill's scalar path now logs "prefill attention, K and V off their vectors" and "f16 prefill attention, K and V off 4 bytes".
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.
Closes #10 on Metal.
What Apple GPUs do with a misaligned pointer (measured on the M1 Pro, probe kernels):
So f32 buffers are always safe, and only 1- and 2-byte buffers read wider than their element need a check. Metal's rule is therefore "4 bytes", not CUDA's "the vector's width", and only a handful of launchers apply it. Per-kernel declarations like CUDA's would restate those few checks, so instead:
word_aligned_inmetal.hstates the rule once.bf16_rows_vec_uses it in place ofW.off % 16(a W on 4 bytes now takes the vectorized loader), andquant_ok_for Q4_K's 32-bit reads. Q6_K / Q8_0 keep their 2-byte check (16-bit reads), andgemv_quant's 16-byte activation stays as CUDA's shared contract.Tests:
check_misalignedtests from 229abb7 pass on Metal (cpu / gpu / auto).gemm_bf16_ntandgemm_bf16_nt_f16accmust run (must_run) with A, C, or W 2 or 4 bytes off. Both backends withbf16_gemmcompute any alignment, and the f16-sum op falls back to the f32 GEMM. Refusing in Metal'sgemm_bf16_ntfails the A and C cases. Makingword_aligned_always true fails both W-off-4-bytes cases.CUDA (run on the RTX 3090, where the new f16 attention test first failed):
TL_ALIGNdeclares only the element.self_guardedlist is gone.DOCTEST_CONFIG_TREAT_CHAR_STAR_AS_STRING, so a failingcheck_misalignedlogs its case name instead of an address.