[Bugfix] Fix mHC block-M prenorm GEMM cross-row reduction carry-over (#49429)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Nick Hill committed
Jul 26, 2026 at 21:42 UTC
74d3b799e193756a8d42b59d223e38a39fc6be55
1 file changed
+8
-6
vllm/model_executor/kernels/mhc/tilelang_kernels.py
+8
-6
@@ -856,21 +856,23 @@ def hc_prenorm_gemm_block_m_tilelang(
856
T.sync_threads()
857
858
if warp_id == 0:
859
+ reduced_acc = T.alloc_local((block_m,), T.float32)
860
+ reduced_sqr = T.alloc_local((block_m,), T.float32)
861
+ T.clear(reduced_acc)
862
+ T.clear(reduced_sqr)
863
for i_m in T.unroll(block_m):
864
token_idx = i_mt * block_m + i_m
865
if token_idx < num_tokens:
866
if lane < tile_n:
863
- reduced_acc = T.alloc_var(T.float32, init=0.0)
867
for i_w in T.unroll(num_warps):
865
- reduced_acc += warp_acc[i_w, i_m, lane]
868
+ reduced_acc[i_m] += warp_acc[i_w, i_m, lane]
869
out_idx = i_t * tile_n + lane
870
if out_idx < n_out:
868
- out[0, token_idx, out_idx] = reduced_acc
871
+ out[0, token_idx, out_idx] = reduced_acc[i_m]
872
if lane == 0 and i_t == 0:
870
- reduced_sqr = T.alloc_var(T.float32, init=0.0)
873
for i_w in T.unroll(num_warps):
872
- reduced_sqr += warp_sqr[i_w, i_m]
873
- sqrsum[0, token_idx] = reduced_sqr
874
+ reduced_sqr[i_m] += warp_sqr[i_w, i_m]
875
+ sqrsum[0, token_idx] = reduced_sqr[i_m]
876
877
if ENABLE_PDL:
878
T.pdl_trigger()