From 1e0bdcbf10c4ee981ffee96d929c2e2160c8e3a4 Mon Sep 17 00:00:00 2001 From: jordanswilson Date: Thu, 20 Aug 2026 22:18:24 -0400 Subject: [PATCH] Align causal GQA threadgroup memory length to 16 bytes setThreadgroupMemoryLength requires a multiple of 16 bytes, but the causal GQA score buffer was sized as sequence * sizeof(float), which violates that whenever the token count is not a multiple of 4. The driver tolerates it silently, so plain runs work, but under the Metal API validation layer (MTL_DEBUG_LAYER=1, e.g. any process spawned from an Xcode-launched app with default scheme diagnostics) the first text encoder dispatch aborts: -[MTLDebugComputeCommandEncoder setThreadgroupMemoryLength:atIndex:]: failed assertion `length(140) must be a multiple of 16 bytes.' (140 bytes = a 35-token prompt.) Round the length up to 16; the kernel never reads the padding. The rounding happens before the maxThreadgroupMemoryLength check so the check stays accurate. The bundled text encoder test cannot catch this because its fixture sequence length of 32 is already a multiple of 4. Verified with a sequence-35 repro: pre-fix it aborts under MTL_DEBUG_LAYER=1 with the assertion above, post-fix it passes with and without validation. Co-Authored-By: Claude Fable 5 --- h3_gpu.m | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/h3_gpu.m b/h3_gpu.m index c61d04c6..288554fe 100644 --- a/h3_gpu.m +++ b/h3_gpu.m @@ -4286,7 +4286,9 @@ int h3_gpu_gqa_causal_bf16(h3_gpu *opaque, h3_gpu_tensor *output, if (getenv("H3_MPS_GQA") && h3_gpu_gqa_mps( gpu, output, query, key, value, sequence, query_heads, kv_heads, head_dim, scale)) return 1; - size_t score_bytes = (size_t)sequence * sizeof(float); + /* Metal requires threadgroup memory lengths to be 16-byte multiples; + * the API validation layer aborts on unaligned lengths. */ + size_t score_bytes = ((size_t)sequence * sizeof(float) + 15) & ~(size_t)15; id pipeline = h3_gpu_pipeline(gpu, @"h3_gqa_causal_bf16"); if (!pipeline) return 0;