akaashrp opened a new pull request, #20449: URL: https://github.com/apache/tvm/pull/20449
Add a barrier between reading and overwriting the shared log-sum-exp in `merge_state_inplace`, and let one thread store it. Every `threadIdx.x` thread of a block reads `S[bx, ty + by * bdy]`, derives both rescaling factors from it, rescales its own `VEC_SIZE` columns of `V`, then writes the merged value back to that same element. Nothing separates the read from the write, so this holds only while the threads advance in lockstep. `bdx` is `head_dim // VEC_SIZE` with `VEC_SIZE` capped at 4, giving one warp per block at head_dim 128, two at 256 and four at 512. Past one warp a thread that finishes early overwrites `S` while another has yet to read it, and the late thread then rescales its slice of `V` with an already merged LSE and stores an LSE merged twice. The result is a large error rather than a rounding difference. On Metal at head_dim 512 with 8 heads and identical inputs on every iteration, 22 of 25 runs at N=512 and 24 of 25 at N=4096 disagreed with the first run, while head_dim 128 never did. Checked against a NumPy reference, one affected row had 365 of its 512 columns wrong, confined to warp spans 0, 2 and 3, and the returned LSE was off by 0.74. Gemma 4 E2B, whose head_dim is 512, hit this in roughly 3% of chunked prefills, moving logits by up to 1.7 without changing the sampled tokens; a device sync after every VM call did not remove it, which is what placed it inside a single kernel. `merge_state_inplace_cpu` walks rows serially and is unaffected, and the measurements above are from Metal, though the same argument applies to any backend whose blocks span more than one warp or subgroup. The new case in `tests/python/relax/test_frontend_nn_llm_kernel_config.py` builds the kernel for Metal and checks the generated source for a barrier between the last read of the LSE and the first store to it, plus the single-thread guard; it fails at head_dim 128, 256 and 512 without the fix. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
