From 08cb46789d31e07f515243d646ca7f5df2d8d3e7 Mon Sep 17 00:00:00 2001 From: gnovack Date: Fri, 22 May 2026 13:44:29 -0700 Subject: [PATCH] mhc_post - remove sts & add vectorized copies (#43437) Signed-off-by: george Co-authored-by: george --- vllm/_tilelang_ops.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/vllm/_tilelang_ops.py b/vllm/_tilelang_ops.py index 011f4518eab..1a073e58d14 100644 --- a/vllm/_tilelang_ops.py +++ b/vllm/_tilelang_ops.py @@ -331,7 +331,6 @@ def mhc_post_tilelang( d: T.Tensor((n, h), T.bfloat16) # type: ignore[no-redef, valid-type] x: T.Tensor((n, hc, h), T.bfloat16) # type: ignore[no-redef, valid-type] with T.Kernel(n, threads=n_thr) as i_n: - x_shared = T.alloc_shared((hc, h_blk), T.bfloat16) b_shared = T.alloc_shared((hc, h_blk), T.bfloat16) d_shared = T.alloc_shared(h_blk, T.bfloat16) @@ -345,7 +344,7 @@ def mhc_post_tilelang( T.copy(a[i_n, 0, 0], a_local) T.copy(c[i_n, 0], c_local) - for i0_h in T.Pipelined(T.ceildiv(h, h_blk), num_stages=2): + for i0_h in T.Serial(T.ceildiv(h, h_blk)): T.copy(b[i_n, 0, i0_h * h_blk], b_shared) T.copy(d[i_n, i0_h * h_blk], d_shared) @@ -353,11 +352,10 @@ def mhc_post_tilelang( T.copy(d_shared, d_local) for i_hco, i1_h in T.Parallel(hc, h_blk): x_local[i_hco, i1_h] = c_local[i_hco] * d_local[i1_h] - for i_hci in T.serial(hc): + for i_hci in T.vectorized(hc): x_local[i_hco, i1_h] += a_local[i_hci, i_hco] * b_local[i_hci, i1_h] - T.copy(x_local, x_shared) - T.copy(x_shared, x[i_n, 0, i0_h * h_blk]) + T.copy(x_local, x[i_n, 0, i0_h * h_blk]) T.pdl_trigger()