A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 2, 23:15 UTC
mhc_post - remove sts & add vectorized copies (#43437)
Signed-off-by: george <george@inferact.ai> Co-authored-by: george <george@inferact.ai>
gnovack committed
May 22, 2026 at 13:44 UTC
08cb46789d31e07f515243d646ca7f5df2d8d3e7
1 file changed
+3
-5
vllm/_tilelang_ops.py
+3
-5
@@ -331,7 +331,6 @@ def mhc_post_tilelang(
331
d: T.Tensor((n, h), T.bfloat16) # type: ignore[no-redef, valid-type]
332
x: T.Tensor((n, hc, h), T.bfloat16) # type: ignore[no-redef, valid-type]
333
with T.Kernel(n, threads=n_thr) as i_n:
334
- x_shared = T.alloc_shared((hc, h_blk), T.bfloat16)
334
b_shared = T.alloc_shared((hc, h_blk), T.bfloat16)
335
d_shared = T.alloc_shared(h_blk, T.bfloat16)
336
@@ -345,7 +344,7 @@ def mhc_post_tilelang(
344
T.copy(a[i_n, 0, 0], a_local)
345
T.copy(c[i_n, 0], c_local)
346
348
- for i0_h in T.Pipelined(T.ceildiv(h, h_blk), num_stages=2):
347
+ for i0_h in T.Serial(T.ceildiv(h, h_blk)):
348
T.copy(b[i_n, 0, i0_h * h_blk], b_shared)
349
T.copy(d[i_n, i0_h * h_blk], d_shared)
350
@@ -353,11 +352,10 @@ def mhc_post_tilelang(
352
T.copy(d_shared, d_local)
353
for i_hco, i1_h in T.Parallel(hc, h_blk):
354
x_local[i_hco, i1_h] = c_local[i_hco] * d_local[i1_h]
356
- for i_hci in T.serial(hc):
355
+ for i_hci in T.vectorized(hc):
356
x_local[i_hco, i1_h] += a_local[i_hci, i_hco] * b_local[i_hci, i1_h]
358
- T.copy(x_local, x_shared)
357
360
- T.copy(x_shared, x[i_n, 0, i0_h * h_blk])
358
+ T.copy(x_local, x[i_n, 0, i0_h * h_blk])
359
T.pdl_trigger()
360
361