Align causal GQA threadgroup memory length to 16 bytes - #44
Open
jordanswilson wants to merge 1 commit into
Open
Conversation
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 <noreply@anthropic.com>
LachlanStuart
added a commit
to LachlanStuart/h3.c
that referenced
this pull request
Sep 6, 2026
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
1 participant
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.
setThreadgroupMemoryLengthrequires the length to be a multiple of 16 bytes, buth3_gpu_gqa_causal_bf16sizes the causal score buffer assequence * sizeof(float)— out of spec whenever the prompt's token count is not a multiple of 4.The driver tolerates this silently, so ordinary runs work. But with the Metal API validation layer active (
MTL_DEBUG_LAYER=1— which any child process inherits from an Xcode-launched app with default scheme diagnostics), the first text encoder dispatch aborts:(140 bytes = a 35-token prompt.)
Fix
Round the threadgroup memory length up to the next 16-byte multiple. The kernel never reads the padding, and the rounding happens before the
maxThreadgroupMemoryLengthcapacity check so that check stays accurate. Behavior for already-aligned sequence lengths is unchanged.Verification
The bundled
h3_text_testscan't catch this — its fixture sequence length of 32 is already a multiple of 4. I verified with a minimal fixture-free repro callingh3_gpu_gqa_causal_bf16at sequence 35 (synthetic bf16 tensors, same head geometry as the Qwen text encoder):MTL_DEBUG_LAYER=1: aborts with the exact assertion above (SIGABRT)MTL_DEBUG_LAYER=1: passes, all outputs finiteHappy to add the repro as a test if you'd like it.
🤖 Generated with Claude Code