[Fix][Relax][LLM] Synchronize merge_state_inplace before overwriting the LSE

[akaashrp]

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.