This is an automated email from the ASF dual-hosted git repository.
tlopex pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 9713d6ee95 [Fix][Relax] Synchronize merge_state_inplace before
overwriting the LSE (#20449)
9713d6ee95 is described below
commit 9713d6ee95e64c063328b6d7747521b0fdf401be
Author: Akaash Parthasarathy <[email protected]>
AuthorDate: Sat Sep 26 08:12:07 2026 -0700
[Fix][Relax] Synchronize merge_state_inplace before overwriting the LSE
(#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.
---
.../tvm/relax/frontend/nn/llm/_decode_kernels.py | 8 +++++-
.../relax/test_frontend_nn_llm_kernel_config.py | 31 ++++++++++++++++++++++
2 files changed, 38 insertions(+), 1 deletion(-)
diff --git a/python/tvm/relax/frontend/nn/llm/_decode_kernels.py
b/python/tvm/relax/frontend/nn/llm/_decode_kernels.py
index fc966e7e54..d5098a35c5 100644
--- a/python/tvm/relax/frontend/nn/llm/_decode_kernels.py
+++ b/python/tvm/relax/frontend/nn/llm/_decode_kernels.py
@@ -490,7 +490,13 @@ def _merge_state_inplace(num_heads, head_dim, v_dtype,
target: Target, global_sy
V[bx, ty + by * bdy, tx * VEC_SIZE + vec] =
v_vec[vec]
# store s
- S[bx, ty + by * bdy] = T.log2(s_val[0] +
s_other_val[0]) + s_max[0]
+ # Every threadIdx.x thread of this block reads the
same S element
+ # above, so none may overwrite it until all of
them have read it.
+ # Without this barrier the threads only stay in
step while bdx fits
+ # in one warp, and a head_dim above 128 makes bdx
wider than that.
+ T.tvm_storage_sync("shared")
+ if tx == 0:
+ S[bx, ty + by * bdy] = T.log2(s_val[0] +
s_other_val[0]) + s_max[0]
func = merge_state_inplace
if global_symbol:
diff --git a/tests/python/relax/test_frontend_nn_llm_kernel_config.py
b/tests/python/relax/test_frontend_nn_llm_kernel_config.py
index b0561dffc6..0183481693 100644
--- a/tests/python/relax/test_frontend_nn_llm_kernel_config.py
+++ b/tests/python/relax/test_frontend_nn_llm_kernel_config.py
@@ -19,6 +19,7 @@ import pytest
import tvm
import tvm.testing
+from tvm.relax.frontend.nn.llm._decode_kernels import _merge_state_inplace
from tvm.relax.frontend.nn.llm._kernel_common import (
_get_prefill_kernel_config,
_get_prefill_shared_memory_usage,
@@ -240,5 +241,35 @@ def
test_wide_head_tree_attention_has_legal_metal_schedule(kernel):
assert 24_928 <= int(target.attrs["max_shared_memory_per_block"])
[email protected]("head_dim", [128, 256, 512])
+def test_merge_state_waits_before_overwriting_the_shared_lse(head_dim):
+ """merge_state_inplace must synchronize between reading and overwriting S.
+
+ Every threadIdx.x thread of a block reads the same S element and later
writes it back. The
+ threads only advance in step while bdx fits in one warp, and bdx is
head_dim // VEC_SIZE, so a
+ head_dim above 128 spreads them over several warps. Without a barrier a
warp that finishes
+ early overwrites S while another warp has yet to read it, and that warp
then rescales its slice
+ of V with an already merged LSE, corrupting whole warp-sized spans of the
output.
+ """
+ target = tvm.target.Target("metal")
+ func = _merge_state_inplace(8, head_dim, "float16", target,
global_symbol="merge")
+ lines = [
+ line.strip()
+ for line in tvm.tirx.build(func,
target=target).imports[0].inspect_source().splitlines()
+ ]
+
+ loads = [i for i, line in enumerate(lines) if "= S_ptr[" in line]
+ stores = [i for i, line in enumerate(lines) if line.startswith("S_ptr[")]
+ assert loads and stores, "expected the kernel to both read and write the
shared LSE"
+
+ last_load, first_store = max(loads), min(stores)
+ assert any(
+ last_load < i < first_store for i, line in enumerate(lines) if
"threadgroup_barrier" in line
+ ), "no barrier between reading the shared LSE and overwriting it"
+ assert any("threadIdx.x) == 0" in line for line in
lines[last_load:first_store]), (
+ "the shared LSE should be stored by a single thread"
+ )
+
+
if __name__ == "__main__":
tvm.testing.main()
