Skip to content

Metal: one pointer-alignment rule, and tests that hold ops to it (#10) - #11

Merged
yhirose merged 6 commits into
masterfrom
metal-align
Nov 1, 2026
Merged

yhirose merged 6 commits into
masterfrom
metal-align

Conversation

@yhirose

@yhirose yhirose commented Oct 31, 2026 •

Copy link
Copy Markdown
Owner

Closes #10 on Metal.

What Apple GPUs do with a misaligned pointer (measured on the M1 Pro, probe kernels):

  • A load or store of 4 bytes or more is right from any 4-byte-aligned address, including float4 and float2 off 16 or 8 bytes.
  • From a 2-byte-aligned address it can silently move the wrong bytes (no fault). The compiler may also merge two narrower loads into such a one (two half4 loads merged into 16 bytes misread; one alone did not).
  • Apple's documentation asks only that a buffer's offset be a multiple of the bound argument's element.

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_ in metal.h states the rule once. bf16_rows_vec_ uses it in place of W.off % 16 (a W on 4 bytes now takes the vectorized loader), and quant_ok_ for Q4_K's 32-bit reads. Q6_K / Q8_0 keep their 2-byte check (16-bit reads), and gemv_quant's 16-byte activation stays as CUDA's shared contract.
  • The f16 prefill attention stages its cache four halves at a time without the check. It can't refuse (no composition takes an f16 cache), and it reads a cache 2 bytes off right. A test holds it to that.

Tests:

  • The check_misaligned tests from 229abb7 pass on Metal (cpu / gpu / auto).
  • gemm_bf16_nt and gemm_bf16_nt_f16acc must run (must_run) with A, C, or W 2 or 4 bytes off. Both backends with bf16_gemm compute any alignment, and the f16-sum op falls back to the f32 GEMM. Refusing in Metal's gemm_bf16_nt fails the A and C cases. Making word_aligned_ always true fails both W-off-4-bytes cases.
  • The f16 prefill attention must run and match with the cache one half off. Staging zeros whenever the address is 2 bytes off fails it.
  • Suite time is unchanged (about 2.3-2.5 s per mode).

CUDA (run on the RTX 3090, where the new f16 attention test first failed):

  • b11a820: the tiled prefill attention picks its vector loads by each operand's base and widens one element at a time otherwise, so an f16 (or f32) cache off its vectors runs instead of being refused. TL_ALIGN declares only the element.
  • 606fa74: the bf16 NT GEMMs declare their elements' alignment too, and check_cuda's self_guarded list is gone.
  • 6634e45: DOCTEST_CONFIG_TREAT_CHAR_STAR_AS_STRING, so a failing check_misaligned logs its case name instead of an address.
  • ctest 9/9 and check_cuda ALL OK on the RTX 3090.

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.
@yhirose yhirose changed the title Metal: pointer alignment follow-ups (#10) Metal: one pointer-alignment rule, and tests that hold ops to it (#10) Oct 31, 2026
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".
@yhirose
yhirose marked this pull request as ready for review November 1, 2026 01:38
@yhirose
yhirose merged commit f865312 into master Nov 1, 2026
9 checks passed
@yhirose
yhirose deleted the metal-align branch November 1, 2026 01:39
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.

Metal: declare kernels' pointer alignment as CUDA now does

1 participant