Optimize the mhc kernels

tile_kernels/mhc/expand_kernel.py

@@ -4,14 +4,16 @@
 
 
 @tilelang.jit
-def expand_to_mhc_fwd_tl(hidden: int, mhc_mult: int) -> tilelang.JITKernel:
+def expand_to_mhc_fwd_tl(
+    hidden: int,
+    mhc_mult: int,
+    blk_n: int = 32,
+    blk_h: int = 128,
+) -> tilelang.JITKernel:
     n = T.dynamic('num_tokens')
     h = hidden
     mhc = mhc_mult
 
-    blk_n = 32
-    blk_h = 128
-
     @T.prim_func
     def expand_to_mhc_fwd_kernel(
         x: T.Tensor[(n, h), T.bfloat16],

tile_kernels/mhc/multilayer_recompute_kernel.py

@@ -76,41 +76,18 @@ def kernel(
             post_mix_local = T.alloc_fragment(mhc, T.float32)
             comb_mix_local = T.alloc_fragment((mhc, mhc), T.float32)
 
-            layer_output_shared = T.alloc_shared((2, h_blk), T.bfloat16)
-            pre_mix_shared = T.alloc_shared((2, mhc), T.float32)
-            post_mix_shared = T.alloc_shared((2, mhc), T.float32)
-            comb_mix_shared = T.alloc_shared((2, mhc, mhc), T.float32)
-
             for i0_h in T.serial(h // h_blk):
                 T.copy(initial_residual[i_n, 0, i0_h * h_blk], res_local)
 
-                if L_post > 0:
-                    layer_output_tensor_0 = T.make_tensor(layer_output_ptrs[0], (n, h), T.bfloat16)
-                    pre_mix_tensor_0 = T.make_tensor(pre_mix_ptrs[0], (n, mhc), T.float32)
-                    post_mix_tensor_0 = T.make_tensor(post_mix_ptrs[0], (n, mhc), T.float32)
-                    comb_mix_tensor_0 = T.make_tensor(comb_mix_ptrs[0], (n, mhc, mhc), T.float32)
-                    T.copy(layer_output_tensor_0[i_n, i0_h * h_blk], layer_output_shared[0, :])
-                    T.copy(pre_mix_tensor_0[i_n, 0], pre_mix_shared[0, :])
-                    T.copy(post_mix_tensor_0[i_n, 0], post_mix_shared[0, :])
-                    T.copy(comb_mix_tensor_0[i_n, 0, 0], comb_mix_shared[0, :, :])
-
                 for i_layer in T.serial(L_post):
+                    layer_output_tensor = T.make_tensor(layer_output_ptrs[i_layer], (n, h), T.bfloat16)
+                    pre_mix_tensor = T.make_tensor(pre_mix_ptrs[i_layer], (n, mhc), T.float32)
+                    post_mix_tensor = T.make_tensor(post_mix_ptrs[i_layer], (n, mhc), T.float32)
+                    comb_mix_tensor = T.make_tensor(comb_mix_ptrs[i_layer], (n, mhc, mhc), T.float32)
                     layer_input_tensor = T.make_tensor(layer_input_ptrs[i_layer], (n, h), T.bfloat16)
                     output_residual_tensor = T.make_tensor(residual_ptrs[i_layer], (n, mhc, h), T.bfloat16)
 
-                    phase = i_layer % 2
-
-                    if i_layer + 1 < L_post:
-                        next_layer_output_tensor = T.make_tensor(layer_output_ptrs[i_layer + 1], (n, h), T.bfloat16)
-                        next_pre_mix_tensor = T.make_tensor(pre_mix_ptrs[i_layer + 1], (n, mhc), T.float32)
-                        next_post_mix_tensor = T.make_tensor(post_mix_ptrs[i_layer + 1], (n, mhc), T.float32)
-                        next_comb_mix_tensor = T.make_tensor(comb_mix_ptrs[i_layer + 1], (n, mhc, mhc), T.float32)
-                        T.copy(next_layer_output_tensor[i_n, i0_h * h_blk], layer_output_shared[1 - phase, :])
-                        T.copy(next_pre_mix_tensor[i_n, 0], pre_mix_shared[1 - phase, :])
-                        T.copy(next_post_mix_tensor[i_n, 0], post_mix_shared[1 - phase, :])
-                        T.copy(next_comb_mix_tensor[i_n, 0, 0], comb_mix_shared[1 - phase, :, :])
-
-                    T.copy(pre_mix_shared[phase, :], pre_mix_local)
+                    T.copy(pre_mix_tensor[i_n, 0], pre_mix_local)
 
                     T.clear(layer_input_local)
                     for i_mhc in T.serial(mhc):
@@ -119,9 +96,9 @@ def kernel(
 
                     T.copy(layer_input_local, layer_input_tensor[i_n, i0_h * h_blk])
 
-                    T.copy(post_mix_shared[phase, :], post_mix_local)
-                    T.copy(comb_mix_shared[phase, :, :], comb_mix_local)
-                    T.copy(layer_output_shared[phase, :], layer_output_local)
+                    T.copy(post_mix_tensor[i_n, 0], post_mix_local)
+                    T.copy(comb_mix_tensor[i_n, 0, 0], comb_mix_local)
+                    T.copy(layer_output_tensor[i_n, i0_h * h_blk], layer_output_local)
                     for i_mhco, i1_h in T.Parallel(mhc, h_blk):
                         new_res_local[i_mhco, i1_h] = post_mix_local[i_mhco] * layer_output_local[i1_h]
                         for i_mhci in T.serial(mhc):
@@ -181,7 +158,18 @@ def mhc_multilayer_recompute(
         device=initial_residual.device,
     )
 
-    kernel = _mhc_multilayer_recompute_kernel(mhc_mult, hidden, num_layers, num_post)
+    # Two Wave64 groups improve the 4096/7168 hidden shapes; 2560 and 8192 retain one.
+    n_thr = 128 if hidden in (4096, 7168) else 64
+    # 8192 needs smaller hidden tiles to avoid its large per-CTA fragment footprint.
+    h_blk = 512 if hidden == 8192 else 2048
+    kernel = _mhc_multilayer_recompute_kernel(
+        mhc_mult,
+        hidden,
+        num_layers,
+        num_post,
+        n_thr=n_thr,
+        h_blk=h_blk,
+    )
     kernel(
         initial_residual.view(-1, mhc_mult, hidden),
         pre_mix_ptrs,

tile_kernels/mhc/norm_fn_kernel.py

@@ -68,29 +68,35 @@ def _mhc_pre_norm_fn_fwd_mul(
     mhc_mult3: int,
     n_rms_group: int,
     rms_group_size: int,
-    token_block: int = 32,
-    hidden_block: int = 256,
+    token_block: int = 16,
+    hidden_block: int = 128,
+    use_bf16_mma: bool = False,
+    fn_is_bf16: bool = False,
 ) -> tilelang.JITKernel:
     assert mhc_mult3 <= 32
+    assert not fn_is_bf16 or use_bf16_mma
     num_tokens = T.dynamic('num_tokens')
     assert rms_group_size % hidden_block == 0
+    fn_dtype = T.bfloat16 if fn_is_bf16 else T.float32
 
     @T.prim_func
     def _mhc_pre_norm_fn_fwd_mul_kernel(
         x: T.Tensor[(num_tokens, n_rms_group * rms_group_size), T.bfloat16],
-        fn: T.Tensor[(mhc_mult3, n_rms_group * rms_group_size), T.float32],
+        fn: T.Tensor[(mhc_mult3, n_rms_group * rms_group_size), fn_dtype],
         out: T.Tensor[(num_tokens, n_rms_group, mhc_mult3), T.float32],
         sqrsum: T.Tensor[(num_tokens, n_rms_group), T.float32],
     ) -> None:
         _ = mhc_mult3
         with T.Kernel(T.ceildiv(num_tokens, token_block), n_rms_group) as (pid_x, pid_y):
             out_frag = T.alloc_fragment((token_block, 32), T.float32)
-            sqrsum_part = T.alloc_fragment((token_block, 4), T.float32)
+            sqrsum_part = T.alloc_fragment(token_block, T.float32)
             T.clear(out_frag)
             T.clear(sqrsum_part)
             for pz in T.Pipelined(rms_group_size // hidden_block, num_stages=1):
                 x_smem_16 = T.alloc_shared((token_block, hidden_block), T.bfloat16)
-                fn_smem = T.alloc_shared((32, hidden_block), T.float32)
+                fn_smem = T.alloc_shared(
+                    (32, hidden_block), T.bfloat16 if use_bf16_mma else T.float32
+                )
 
                 T.annotate_layout({x_smem_16: tilelang.layout.make_swizzled_layout(x_smem_16)})
 
@@ -102,29 +108,58 @@ def _mhc_pre_norm_fn_fwd_mul_kernel(
                 x_frag = T.alloc_fragment((token_block, hidden_block), T.float32)
                 T.copy(x_frag_16, x_frag)
 
-                for jj in T.serial(hidden_block // 4):
-                    for i, j in T.Parallel(token_block, 4):
-                        sqrsum_part[i, j] += x_frag[i, jj * 4 + j] * x_frag[i, jj * 4 + j]
+                # Compute sum of squares: first square each element, then reduce
+                x_sq = T.alloc_fragment((token_block, hidden_block), T.float32)
+                for i, j in T.Parallel(token_block, hidden_block):
+                    x_sq[i, j] = x_frag[i, j] * x_frag[i, j]
+                sqrsum_blk = T.alloc_fragment(token_block, T.float32)
+                T.reduce_sum(x_sq, sqrsum_blk)
+                for i in T.Parallel(token_block):
+                    sqrsum_part[i] += sqrsum_blk[i]
 
                 T.gemm(
-                    x_frag,
+                    x_frag_16 if use_bf16_mma else x_frag,
                     fn_smem,
                     out_frag,
                     transpose_A=False,
                     transpose_B=True,
                     clear_accum=False,
                 )
-            sqrsum_l = T.alloc_fragment(token_block, T.float32)
-            T.reduce_sum(sqrsum_part, sqrsum_l)
             for i in T.Parallel(token_block):
-                sqrsum[pid_x * token_block + i, pid_y] = sqrsum_l[i]
+                sqrsum[pid_x * token_block + i, pid_y] = sqrsum_part[i]
             for i, j in T.Parallel(token_block, 32):
                 if j < 24:
                     out[pid_x * token_block + i, pid_y, j] = out_frag[i, j]
 
     return _mhc_pre_norm_fn_fwd_mul_kernel
 
 
+@tilelang.jit(pass_configs=_PASS_CONFIGS)
+def _mhc_pre_norm_fn_fwd_sqsum(hidden: int, hidden_block: int = 128) -> tilelang.JITKernel:
+    num_tokens = T.dynamic('num_tokens')
+    assert hidden % hidden_block == 0
+
+    @T.prim_func
+    def _mhc_pre_norm_fn_fwd_sqsum_kernel(
+        x: T.Tensor[(num_tokens, hidden), T.bfloat16],
+        sqsum: T.Tensor[num_tokens, T.float32],
+    ) -> None:
+        with T.Kernel(num_tokens, threads=hidden_block) as pid:
+            sum_reducer = T.alloc_reducer(1, T.float32, replication='all')
+            T.clear(sum_reducer)
+            for i0_h in T.serial(hidden // hidden_block):
+                x_frag = T.alloc_fragment(hidden_block, T.bfloat16)
+                T.copy(x[pid, i0_h * hidden_block], x_frag)
+                for i1_h in T.Parallel(hidden_block):
+                    x_val = T.cast(x_frag[i1_h], T.float32)
+                    sum_reducer[0] += x_val * x_val
+            T.finalize_reducer(sum_reducer)
+            if T.get_thread_binding() == 0:
+                sqsum[pid] = sum_reducer[0]
+
+    return _mhc_pre_norm_fn_fwd_sqsum_kernel
+
+
 @tilelang.jit(pass_configs=_PASS_CONFIGS)
 def _mhc_pre_norm_fn_fwd_norm(
     mhc_mult3: int,
@@ -212,8 +247,9 @@ def _mhc_pre_norm_fn_bwd_mul(
     mhc_mult3: int,
     n_rms_group: int,
     rms_group_size: int,
-    token_block: int = 128,
-    hidden_block: int = 64,
+    token_block: int = 64,
+    hidden_block: int = 32,
+    x_grad_is_zero: bool = False,
 ) -> tilelang.JITKernel:
     assert mhc_mult3 <= 32
     num_tokens = T.dynamic('num_tokens')
@@ -256,7 +292,10 @@ def _mhc_pre_norm_fn_bwd_mul_kernel(
                         padded_grad[i, j] = 0
 
                 x_grad_frag = T.alloc_fragment((token_block, hidden_block), T.float32)
-                T.copy(x_grad[px * token_block, yz], x_grad_frag)
+                if x_grad_is_zero:
+                    T.fill(x_grad_frag, 0)
+                else:
+                    T.copy(x_grad[px * token_block, yz], x_grad_frag)
 
                 T.gemm(
                     padded_grad,

tile_kernels/mhc/post_kernel.py

@@ -14,7 +14,7 @@
         tilelang.PassConfigKey.TL_DISABLE_VECTORIZE_256: True,
     },
 )
-def _mhc_post_fwd(mhc: int, hidden: int, n_thr: int = 128, h_blk: int = 1024) -> tilelang.JITKernel:
+def _mhc_post_fwd(mhc: int, hidden: int, n_thr: int = 512, h_blk: int = 1024) -> tilelang.JITKernel:
     n = T.dynamic('num_tokens')
     h = hidden
 
@@ -79,12 +79,6 @@ def _mhc_post_bwd_kernel(
         dd: T.Tensor[(n, h), T.bfloat16],
     ) -> None:
         with T.Kernel(n, threads=n_thr) as pid_n:
-            dx_shared = T.alloc_shared((4, h_blk), T.bfloat16)
-            b_shared = T.alloc_shared((4, h_blk), T.bfloat16)
-            db_shared = T.alloc_shared((4, h_blk), T.bfloat16)
-            d_shared = T.alloc_shared(h_blk, T.bfloat16)
-            dd_shared = T.alloc_shared(h_blk, T.bfloat16)
-
             dx_local = T.alloc_fragment((4, h_blk), T.float32)
             b_local = T.alloc_fragment((4, h_blk), T.float32)
             db_local = T.alloc_fragment((4, h_blk), T.float32)
@@ -102,13 +96,9 @@ def _mhc_post_bwd_kernel(
             T.clear(dc_reducer)
 
             for i0_h in T.Pipelined(T.ceildiv(h, h_blk), num_stages=3):
-                T.copy(dx[pid_n, 0, i0_h * h_blk], dx_shared, disable_tma=True)
-                T.copy(b[pid_n, 0, i0_h * h_blk], b_shared, disable_tma=True)
-                T.copy(d[pid_n, i0_h * h_blk], d_shared, disable_tma=True)
-
-                T.copy(dx_shared, dx_local)
-                T.copy(b_shared, b_local)
-                T.copy(d_shared, d_local)
+                T.copy(dx[pid_n, 0, i0_h * h_blk], dx_local, disable_tma=True)
+                T.copy(b[pid_n, 0, i0_h * h_blk], b_local, disable_tma=True)
+                T.copy(d[pid_n, i0_h * h_blk], d_local, disable_tma=True)
 
                 # da and db
                 T.clear(db_local)
@@ -125,11 +115,8 @@ def _mhc_post_bwd_kernel(
                         dc_reducer[i_mhc] += d_local[i1_h] * dx_local[i_mhc, i1_h]
                         dd_local[i1_h] += c_local[i_mhc] * dx_local[i_mhc, i1_h]
 
-                T.copy(db_local, db_shared)
-                T.copy(dd_local, dd_shared)
-
-                T.copy(db_shared, db[pid_n, 0, i0_h * h_blk], disable_tma=True)
-                T.copy(dd_shared, dd[pid_n, i0_h * h_blk], disable_tma=True)
+                T.copy(db_local, db[pid_n, 0, i0_h * h_blk], disable_tma=True)
+                T.copy(dd_local, dd[pid_n, i0_h * h_blk], disable_tma=True)
 
             T.finalize_reducer(da_reducer)
             T.finalize_reducer(dc_reducer)
@@ -186,7 +173,8 @@ def mhc_post_bwd(
     mhc = d_o.shape[2]
     h = d_o.shape[3]
 
-    bwd_kernel = _mhc_post_bwd(mhc, h)
+    # A second Wave64 improves the larger hidden shapes, while 1280 remains launch-overhead bound.
+    bwd_kernel = _mhc_post_bwd(mhc, h, n_thr=256 if h > 1280 else 128)
     (
         d_comb_res_mix,
         d_residual,

tile_kernels/mhc/pre_apply_mix_kernel.py

@@ -35,21 +35,17 @@ def _mhc_pre_apply_mix_fwd_kernel(
             T.copy(mix[pid_n, 0], mixl)
 
             for i0_h in T.Pipelined(h // h_blk, num_stages=2):
-                xs = T.alloc_shared((mhc, h_blk), T.bfloat16)
                 xl = T.alloc_fragment((mhc, h_blk), T.float32)
-                T.copy(x[pid_n, 0, i0_h * h_blk], xs, disable_tma=True)
-                T.copy(xs, xl, disable_tma=True)
+                T.copy(x[pid_n, 0, i0_h * h_blk], xl, disable_tma=True)
 
-                os = T.alloc_shared(h_blk, T.bfloat16)
                 ol = T.alloc_fragment(h_blk, T.float32)
                 T.clear(ol)
 
                 for i_mhc in T.serial(mhc):
                     for i1_h in T.Parallel(h_blk):
                         ol[i1_h] += mixl[i_mhc] * xl[i_mhc, i1_h]
 
-                T.copy(ol, os, disable_tma=True)
-                T.copy(os, o[pid_n, i0_h * h_blk], disable_tma=True)
+                T.copy(ol, o[pid_n, i0_h * h_blk], disable_tma=True)
 
     return _mhc_pre_apply_mix_fwd_kernel
 
@@ -60,6 +56,7 @@ def _mhc_pre_apply_mix_bwd(
     hidden: int,
     n_thr: int = 128,
     h_blk: int = 1024,
+    x_grad_is_zero: bool = False,
 ) -> tilelang.JITKernel:
     n = T.dynamic('n')
     h = hidden
@@ -93,10 +90,13 @@ def _mhc_pre_apply_mix_bwd_kernel(
                 T.copy(x[pid_n, 0, i0_h * h_blk], xs, disable_tma=True)
                 T.copy(xs, xl, disable_tma=True)
 
-                xgs = T.alloc_shared((mhc, h_blk), T.bfloat16)
                 xgl = T.alloc_fragment((mhc, h_blk), T.float32)
-                T.copy(x_grad[pid_n, 0, i0_h * h_blk], xgs, disable_tma=True)
-                T.copy(xgs, xgl, disable_tma=True)
+                if x_grad_is_zero:
+                    T.fill(xgl, 0)
+                else:
+                    xgs = T.alloc_shared((mhc, h_blk), T.bfloat16)
+                    T.copy(x_grad[pid_n, 0, i0_h * h_blk], xgs, disable_tma=True)
+                    T.copy(xgs, xgl, disable_tma=True)
 
                 for i_mhc, i1_h in T.Parallel(mhc, h_blk):
                     mgl[i_mhc] += ogl[i1_h] * xl[i_mhc, i1_h]

tile_kernels/mhc/pre_split_mixes_kernel.py

@@ -72,7 +72,7 @@ def _mhc_pre_split_mixes_bwd(
     mhc_mult: int,
     mhc_post_mult_value: float,
     token_block_size: int,
-    num_sms: int = 148,
+    num_sms: int = 104,
     dtype: T.dtype = T.float32,
 ) -> tilelang.JITKernel:
     num_tokens = T.dynamic('num_tokens')

tile_kernels/modeling/mhc/ops/expand.py

@@ -14,8 +14,19 @@ def forward(
         if out is None:
             out = hidden.new_empty(*hidden.shape[:-1], mhc_mult, hidden.shape[-1])
         assert hidden.is_contiguous()
-        kernel = expand_to_mhc_fwd_tl(hidden.shape[-1], mhc_mult)
-        kernel(hidden.flatten(0, -2), out.flatten(0, -3))
+        hidden_flat = hidden.flatten(0, -2)
+        num_tokens = hidden_flat.shape[0]
+        use_blk_n64 = num_tokens % 64 == 0 and not (
+            num_tokens >= 8192 and hidden.shape[-1] == 1280
+        )
+        use_blk_h256 = use_blk_n64 and hidden.shape[-1] >= 4096
+        kernel = expand_to_mhc_fwd_tl(
+            hidden.shape[-1],
+            mhc_mult,
+            blk_n=64 if use_blk_n64 else 32,
+            blk_h=256 if use_blk_h256 else 128,
+        )
+        kernel(hidden_flat, out.flatten(0, -3))
         return out
 
     @staticmethod

tile_kernels/modeling/mhc/ops/head_compute_mix.py

@@ -33,6 +33,9 @@ def backward(
         input_mix, mhc_scale, mhc_base = ctx.saved_tensors
 
         num_sms = get_num_sms()
+        num_tokens = input_mix.numel() // input_mix.shape[-1]
+        if num_sms == 104 and num_tokens % 128 == 0:
+            num_sms = 128
         input_mix_grad = torch.empty_like(input_mix)
         mhc_scale_grad_partial = torch.empty(
             num_sms,