[ROCm][Bugfix] Fix TileLang mHC fused RMSNorm on 64-wide wavefronts (#58419)
Signed-off-by: Djordje Ramic <djoramic@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
djramic committed
Sep 24, 2026 at 01:03 UTC
5fcc6e7c7380c67ab1e1305ed2197b127f0535e5
1 file changed
+20
-3
vllm/model_executor/kernels/mhc/tilelang_kernels.py
+20
-3
@@ -23,6 +23,23 @@ from vllm.utils.math_utils import cdiv
23
24
ENABLE_PDL = current_platform.is_arch_support_pdl() and current_platform.is_cuda()
25
26
+# TileLang's HIP AllReduce shuffle phase ignores a nonzero reduction-group
27
+# offset, so the 64-thread RMSNorm reduction must be wave64-aligned and use
28
+# threads 0..63.
29
+_MIX_ON_TRAILING_THREADS = current_platform.is_rocm()
30
+
31
+
32
+def _is_mix_thread():
33
+ if _MIX_ON_TRAILING_THREADS:
34
+ return T.get_thread_binding() >= 64
35
+ return T.get_thread_binding() < 32
36
+
37
+
38
+def _is_collapse_thread():
39
+ if _MIX_ON_TRAILING_THREADS:
40
+ return T.get_thread_binding() < 64
41
+ return T.get_thread_binding() >= 32
42
+
43
44
@cache
45
def compute_num_split(block_k: int, k: int | None, grid_size: int) -> int:
@@ -353,7 +370,7 @@ def mhc_pre_big_fuse_with_norm_tilelang(
370
mixes[j] *= rms[0]
371
T.copy(mixes, mixes_shared)
372
356
- if split_mode != "input" and T.get_thread_binding() < 32:
373
+ if split_mode != "input" and _is_mix_thread():
374
cm = T.alloc_fragment((hc_mult, hc_mult), T.float32)
375
for j in T.Parallel(hc_mult):
376
if save_pre_mix:
@@ -399,7 +416,7 @@ def mhc_pre_big_fuse_with_norm_tilelang(
416
417
for j, k in T.Parallel(hc_mult, hc_mult):
418
comb_mix[i, j * hc_mult + k] = cm[j, k]
402
- elif split_mode != "stats" and T.get_thread_binding() >= 32:
419
+ elif split_mode != "stats" and _is_collapse_thread():
420
pre_mix_shared = T.alloc_shared(hc_mult, T.float32)
421
for j in T.Parallel(hc_mult):
422
if use_pre_mix_in:
@@ -546,7 +563,7 @@ def mhc_pre_big_fuse_broadcast_with_norm_tilelang(
563
mixes_shared = T.alloc_shared(hc_mult3, T.float32)
564
T.copy(mixes, mixes_shared)
565
549
- if T.get_thread_binding() < 32:
566
+ if _is_mix_thread():
567
cm = T.alloc_fragment((hc_mult, hc_mult), T.float32)
568
for j in T.Parallel(hc_mult):
569
post_mix[i, j] = (