Mhc large m

aiter/ops/mhc.py

@@ -269,6 +269,7 @@ def mhc_pre(
     sinkhorn_repeat: int = 20,  # if 0, only do pre for hc_head
     norm_weight: Optional[torch.Tensor] = None,
     norm_eps: float = 1e-6,
+    large_m_splitk: bool = False,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
     m = residual.size(0)
     hc_mult = residual.size(1)
@@ -278,7 +279,10 @@ def mhc_pre(
         hc_mult3 == hc_mult and sinkhorn_repeat == 0
     )
     hc_hidden_size = hc_mult * hidden_size
-    selected_splitk, selected_tile_k = get_mhc_pre_splitk(m, hc_hidden_size)
+    if large_m_splitk:
+        selected_splitk, selected_tile_k = get_mhc_pre_splitk_large_m(m, hc_hidden_size)
+    else:
+        selected_splitk, selected_tile_k = get_mhc_pre_splitk(m, hc_hidden_size)
     device = residual.device
     out_pad = torch.empty(
         selected_splitk, m, (hc_mult3 + 31) // 32 * 32, dtype=dtypes.fp32, device=device
@@ -337,9 +341,17 @@ def mhc_post(
     residual: Tensor,
     post_layer_mix: Tensor,
     comb_res_mix: Tensor,
+    store_nt: int = -1,
 ) -> None: ...
 
 
+def get_mhc_pre_splitk_large_m(m: int, hc_hidden_size: int) -> tuple[int, int]:
+    """Split-K policy for gfx950 large-M post_pre kernel (M > 1024)."""
+    if get_gfx_runtime() == "gfx950" and m >= 8192 and hc_hidden_size % (8 * 64) == 0:
+        return 8, 64
+    return get_mhc_pre_splitk(m, hc_hidden_size)
+
+
 @compile_ops("module_mhc")
 def mhc_fused_post_pre_gemm_sqrsum(
     gemm_out_mul: Tensor,
@@ -384,6 +396,68 @@ def mhc_fused_post_pre_fake(
     return post_mix, comb_mix, layer_input_out, next_residual
 
 
+@torch_compile_guard(gen_fake=mhc_fused_post_pre_fake)
+def mhc_fused_post_pre_large_m(
+    layer_input: torch.Tensor,
+    residual_in: torch.Tensor,
+    post_layer_mix: torch.Tensor,
+    comb_res_mix: torch.Tensor,
+    fn: torch.Tensor,
+    hc_scale: torch.Tensor,
+    hc_base: torch.Tensor,
+    rms_eps: float = 1e-6,
+    hc_pre_eps: float = 1e-6,
+    hc_sinkhorn_eps: float = 1e-6,
+    hc_post_mult_value: float = 1.0,
+    sinkhorn_repeat: int = 20,
+    norm_weight: Optional[torch.Tensor] = None,
+    norm_eps: float = 1e-6,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """gfx950 large-M post+pre (M > 1024): upstream ``mhc_post`` + ``mhc_pre``."""
+    m = residual_in.size(0)
+
+    if post_layer_mix.ndim == 3:
+        post_layer_mix = post_layer_mix.contiguous()
+    elif not post_layer_mix.is_contiguous():
+        post_layer_mix = post_layer_mix.contiguous()
+    if not comb_res_mix.is_contiguous():
+        comb_res_mix = comb_res_mix.contiguous()
+    if not residual_in.is_contiguous():
+        residual_in = residual_in.contiguous()
+    if not layer_input.is_contiguous():
+        layer_input = layer_input.contiguous()
+    if not fn.is_contiguous():
+        fn = fn.contiguous()
+    if norm_weight is not None and not norm_weight.is_contiguous():
+        norm_weight = norm_weight.contiguous()
+
+    next_residual = torch.empty_like(residual_in)
+    post_store_nt = 0 if m > 8 * get_cu_num() else -1
+    mhc_post(
+        next_residual,
+        layer_input,
+        residual_in,
+        post_layer_mix,
+        comb_res_mix,
+        post_store_nt,
+    )
+    post_mix, comb_mix, layer_input_out = mhc_pre(
+        next_residual,
+        fn,
+        hc_scale,
+        hc_base,
+        rms_eps,
+        hc_pre_eps,
+        hc_sinkhorn_eps,
+        hc_post_mult_value,
+        sinkhorn_repeat,
+        norm_weight,
+        norm_eps,
+        large_m_splitk=True,
+    )
+    return post_mix, comb_mix, layer_input_out, next_residual
+
+
 @torch_compile_guard(gen_fake=mhc_fused_post_pre_fake)
 def mhc_fused_post_pre(
     layer_input: torch.Tensor,
@@ -449,6 +523,24 @@ def mhc_fused_post_pre(
         )
         return post_mix, comb_mix, layer_input_out, next_residual
 
+    if force_fused and arch == "gfx950" and m > fused_m_upper_bound:
+        return mhc_fused_post_pre_large_m(
+            layer_input,
+            residual_in,
+            post_layer_mix,
+            comb_res_mix,
+            fn,
+            hc_scale,
+            hc_base,
+            rms_eps,
+            hc_pre_eps,
+            hc_sinkhorn_eps,
+            hc_post_mult_value,
+            sinkhorn_repeat,
+            norm_weight,
+            norm_eps,
+        )
+
     assert layer_input.shape == (
         m,
         hidden_size,

csrc/include/mhc.h

@@ -43,8 +43,8 @@ void mhc_post(torch::Tensor& out,            // (m, hc_mult, hidden_size)
               torch::Tensor& x,              // (m, hidden_size)
               torch::Tensor& residual,       // (m, hc_mult, hidden_size)
               torch::Tensor& post_layer_mix, // (m, hc_mult)
-              torch::Tensor& comb_res_mix    // (m, hc_mult, hc_mult)
-);
+              torch::Tensor& comb_res_mix,   // (m, hc_mult, hc_mult)
+              int store_nt                   = -1);
 void mhc_fused_post_pre_gemm_sqrsum(
     torch::Tensor& gemm_out_mul,    // (split_k * hc_mult, m, hc_mult3)
     torch::Tensor& gemm_out_sqrsum, // (split_k * hc_mult, m)

csrc/include/rocm_ops.hpp

@@ -2100,7 +2100,8 @@ namespace py = pybind11;
           py::arg("x"),                         \
           py::arg("residual"),                  \
           py::arg("post_layer_mix"),            \
-          py::arg("comb_res_mix"));             \
+          py::arg("comb_res_mix"),              \
+          py::arg("store_nt")       = -1);      \
     m.def("mhc_fused_post_pre_gemm_sqrsum",     \
           &aiter::mhc_fused_post_pre_gemm_sqrsum, \
           "mhc_fused_post_pre_gemm_sqrsum",     \

csrc/kernels/mhc_kernels.cu

@@ -759,7 +759,7 @@ namespace aiter {
         const at::hip::OptionalHIPGuardMasqueradingAsCUDA device_guard(device_of(layer_input));
         const hipStream_t stream = at::hip::getCurrentHIPStream();
         const int cu_num = get_num_cu_func();
-        
+
         MHC_PRE_BIG_FUSE_KERNEL_DISPATCH(m);
     }
 
@@ -1066,24 +1066,33 @@ namespace aiter {
         MHC_POST_KERNEL_IMPL_(kernel_name, hidden_size, residual_block, false); \
     }
 
-#define MHC_POST_KERNEL_DISPATCH(hidden_size) \
+#define MHC_POST_KERNEL_DISPATCH_NT(hidden_size, store_nt_val) \
     if (arch_id != "gfx942" && hidden_size % 1024 == 0) { \
-        MHC_POST_KERNEL_IMPL(mhc_post_kernel, hidden_size, 1024); \   
+        MHC_POST_KERNEL_IMPL_(mhc_post_kernel, hidden_size, 1024, store_nt_val); \
     } else if (hidden_size % 512 == 0) { \
-        MHC_POST_KERNEL_IMPL(mhc_post_kernel, hidden_size, 512); \   
+        MHC_POST_KERNEL_IMPL_(mhc_post_kernel, hidden_size, 512, store_nt_val); \
     } else if (hidden_size % 256 == 0) { \
-        MHC_POST_KERNEL_IMPL(mhc_post_kernel_x2vgpr, hidden_size, 256); \
+        MHC_POST_KERNEL_IMPL_(mhc_post_kernel_x2vgpr, hidden_size, 256, store_nt_val); \
     } else { \
         AITER_CHECK(false, "hidden_size must be divisible by 256"); \
     }
-
+
+#define MHC_POST_KERNEL_DISPATCH(hidden_size) \
+    do { \
+        if (m > 8 * cu_num) { \
+            MHC_POST_KERNEL_DISPATCH_NT(hidden_size, true); \
+        } else { \
+            MHC_POST_KERNEL_DISPATCH_NT(hidden_size, false); \
+        } \
+    } while (0)
+
     void mhc_post(
         torch::Tensor& out,
         torch::Tensor& x, // (m, hc_mult, h)
         torch::Tensor& residual, // (m, hc_mult, hidden_size)
         torch::Tensor& post_layer_mix, // (m, hc_mult)
-        torch::Tensor& comb_res_mix // (m, hc_mult, hc_mult)
-    )
+        torch::Tensor& comb_res_mix, // (m, hc_mult, hc_mult)
+        int store_nt = -1)
     {
         int m = residual.size(0);
         int hc_mult = residual.size(1);
@@ -1096,8 +1105,13 @@ namespace aiter {
         const hipStream_t stream = at::hip::getCurrentHIPStream();
         const int cu_num = get_num_cu_func();
         const std::string arch_id = get_gpu_arch();
-
-        MHC_POST_KERNEL_DISPATCH(hidden_size);
+        if (store_nt < 0) {
+            MHC_POST_KERNEL_DISPATCH(hidden_size);
+        } else if (store_nt != 0) {
+            MHC_POST_KERNEL_DISPATCH_NT(hidden_size, true);
+        } else {
+            MHC_POST_KERNEL_DISPATCH_NT(hidden_size, false);
+        }
     }
 
 
@@ -2077,4 +2091,5 @@ namespace aiter {
         MHC_FUSED_POST_PRE_GEMM_SQRSUM_KERNEL_DISPATCH(tile_k);
     }
 
+
 } // namespace aiter

op_tests/test_mhc.py

@@ -10,10 +10,19 @@
 import torch
 import aiter
 from aiter import dtypes
+from aiter.jit.utils.chip_info import get_gfx_runtime
 import argparse
 import pandas as pd
 from typing import Optional
 
+try:
+    from aiter.ops.mhc import mhc_fused_post_pre_large_m
+except ImportError:
+    mhc_fused_post_pre_large_m = None
+
+# gfx950 large-M path (mhc_fused_post_pre_large_m) applies when M > 1024.
+LARGE_M_MIN = 1025
+
 try:
     from aiter.ops.triton.fusions.mhc import mhc_post_pre as triton_mhc_post_pre
 
@@ -745,7 +754,7 @@ def mhc_post_pre_unfused_hip(
 
 
 @benchmark()
-def test_mhc_post_pre(m, hidden_size, hc_mult, fuse_rmsnorm=False):
+def test_mhc_post_pre(m, hidden_size, hc_mult, fuse_rmsnorm=False, large_m=False):
     """Fused mhc_post + mhc_pre: HIP ``mhc_fused_post_pre`` vs ref / unfused HIP / Triton."""
     if hidden_size < 512:
         aiter.logger.info(
@@ -906,6 +915,32 @@ def test_mhc_post_pre(m, hidden_size, hc_mult, fuse_rmsnorm=False):
     elif not _HAS_TRITON_MHC_POST_PRE:
         aiter.logger.info("skip Triton mhc_post_pre: import unavailable")
 
+    if large_m:
+        if m < LARGE_M_MIN:
+            aiter.logger.info("skip large_m_us: m=%s < %s", m, LARGE_M_MIN)
+        elif get_gfx_runtime() != "gfx950":
+            aiter.logger.info(
+                "skip large_m_us: gfx=%s (gfx950 only)", get_gfx_runtime()
+            )
+        elif mhc_fused_post_pre_large_m is None:
+            aiter.logger.info("skip large_m_us: mhc_fused_post_pre_large_m unavailable")
+        else:
+            (_, _, layer_input_large_m, _), large_m_us = run_perftest(
+                mhc_fused_post_pre_large_m,
+                layer_input,
+                residual_in,
+                post_layer_mix,
+                comb_res_mix,
+                fn,
+                hc_scale,
+                hc_base,
+                **hip_kwargs,
+            )
+            ret["large_m_us"] = large_m_us
+            ret["hip_large_m_err"] = checkAllclose(
+                layer_input_ref, layer_input_large_m, msg="large_m/layer_input"
+            )
+
     return ret
 
 
@@ -953,6 +988,12 @@ def test_mhc_post_pre(m, hidden_size, hc_mult, fuse_rmsnorm=False):
     action="store_true",
     help="Fuse RMSNorm into mhc_pre / mhc_post_pre HIP paths (mutually exclusive with --hc_head).",
 )
+parser.add_argument(
+    "--largeM",
+    action="store_true",
+    help="In mhc_post_pre summary, add large_m_us / hip_large_m_err columns "
+    "(gfx950, M>1024, mhc_fused_post_pre_large_m).",
+)
 
 args = parser.parse_args()
 
@@ -995,6 +1036,7 @@ def test_mhc_post_pre(m, hidden_size, hc_mult, fuse_rmsnorm=False):
                         hidden_size=hidden_size,
                         hc_mult=hc_mult,
                         fuse_rmsnorm=args.fuse_rmsnorm,
+                        large_m=args.largeM,
                     )
                     if ret.get("skipped"):
                         continue