From c43326ba42a00d148ace1a62e39a893e1af6ab73 Mon Sep 17 00:00:00 2001
From: Oseltamivir <bryansg2013@gmail.com>
Date: Sun, 3 May 2026 10:25:38 -0700
Subject: [PATCH 1/3] fix: support 7168 fused MHC hidden size

Signed-off-by: Oseltamivir <bryansg2013@gmail.com>
---
 .../kernels/mhcKernels/mhcFusedHcKernel.cu    | 215 +++++++++++++++---
 .../kernels/mhcKernels/mhcKernels.h           |  10 +-
 tensorrt_llm/_torch/modules/mhc/mhc_cuda.py   |  62 +++--
 3 files changed, 234 insertions(+), 53 deletions(-)

diff --git a/cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu b/cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu
index 334961bdfced..23c15c72a9bc 100644
--- a/cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu
+++ b/cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu
@@ -91,9 +91,11 @@ inline void fhcZeroWorkspaces(float* y_acc, uint32_t y_elems, float* r_acc, uint
 } // namespace
 
 // ---- mHC fused kernel shape constants (mirrors the Python module) ----
-static constexpr uint32_t FHC_SHAPE_N = 24;  // HC_MULT * (2 + HC_MULT) = 4 * 6 = 24
-static constexpr uint32_t FHC_HIDDEN = 4096; // only this hidden size is currently wired up
+// HC_MULT * (2 + HC_MULT) = 4 * 6 = 24.
+static constexpr uint32_t FHC_SHAPE_N = 24;
 static constexpr uint32_t FHC_HC_MULT = 4;
+static constexpr uint32_t FHC_HIDDEN_FLASH = 4096;
+static constexpr uint32_t FHC_HIDDEN_PRO = 7168;
 static constexpr uint32_t FHC_BLOCK_M = 64;
 static constexpr uint32_t FHC_BLOCK_N = 32;
 static constexpr uint32_t FHC_BLOCK_K = 64;
@@ -103,6 +105,38 @@ static constexpr uint32_t FHC_N_INPUT_STG = 2;
 static constexpr uint32_t FHC_NUM_MMA_TH = 128;
 static constexpr uint32_t FHC_NUM_PMAP_TH = 128;
 
+template <uint32_t Hidden>
+static constexpr bool isSupportedFhcHidden()
+{
+    return Hidden == FHC_HIDDEN_FLASH || Hidden == FHC_HIDDEN_PRO;
+}
+
+static bool isSupportedFhcHiddenRuntime(int hidden_size)
+{
+    return hidden_size == static_cast<int>(FHC_HIDDEN_FLASH) || hidden_size == static_cast<int>(FHC_HIDDEN_PRO);
+}
+
+// Validate the tcgen05 MMA fused-HC compile-time shape contract. Hidden must
+// be divisible into BLOCK_K tiles, kNumSplits must evenly divide those tiles,
+// and the hidden dimension must align with the per-token post-map vectorized
+// write granularity. Keep this in sync with the Python tactic filter.
+template <uint32_t Hidden, uint32_t KS>
+static constexpr bool isSupportedFhcMmaKS()
+{
+    static_assert(isSupportedFhcHidden<Hidden>(), "Unsupported fused-HC hidden size");
+    static_assert(KS > 0, "kNumSplits must be positive");
+
+    constexpr uint32_t hTilesPerHc = Hidden / FHC_BLOCK_K;
+    constexpr uint32_t blockSizeBf = FHC_NUM_MMA_TH + FHC_NUM_PMAP_TH;
+    constexpr uint32_t warpSizeBf = 32;
+    constexpr uint32_t numWarpsBf = blockSizeBf / warpSizeBf;
+    constexpr uint32_t toksPerCta = (FHC_BLOCK_M + KS - 1) / KS;
+    constexpr uint32_t warpsPerTok = (numWarpsBf > toksPerCta) ? (numWarpsBf / toksPerCta) : 1u;
+    constexpr uint32_t bf16VecLi = 8;
+
+    return Hidden % FHC_BLOCK_K == 0 && hTilesPerHc % KS == 0 && Hidden % (warpsPerTok * warpSizeBf * bf16VecLi) == 0;
+}
+
 static CUtensorMap makeTma2D(void* base, CUtensorMapDataType dtype, uint64_t gmemInner, uint64_t gmemOuter,
     uint32_t smemInner, uint32_t smemOuter, uint64_t gmemOuterStrideBytes, uint32_t swizzleBytes, uint32_t elemBytes)
 {
@@ -144,25 +178,42 @@ static constexpr uint32_t fhcSmemSize()
 using FusedRoutFn = void (*)(
     uint32_t, CUtensorMap, CUtensorMap, CUtensorMap, CUtensorMap, float*, float const*, float const*, float*);
 
-template <uint32_t KS>
+template <uint32_t Hidden, uint32_t KS>
 static FusedRoutFn fhcInstance()
 {
-    return &fused_mhc::fused_tf32_pmap_gemm_rout_atomic_impl<FHC_SHAPE_N, FHC_HIDDEN, FHC_HC_MULT, FHC_BLOCK_M,
-        FHC_BLOCK_N, FHC_BLOCK_K, FHC_SWIZZLE_CD, FHC_N_B_STAGES, FHC_N_INPUT_STG, FHC_NUM_MMA_TH, FHC_NUM_PMAP_TH, KS,
+    static_assert(isSupportedFhcHidden<Hidden>(), "Unsupported fused-HC hidden size");
+    static_assert(isSupportedFhcMmaKS<Hidden, KS>(), "Unsupported fused-HC MMA kNumSplits for hidden size");
+    return &fused_mhc::fused_tf32_pmap_gemm_rout_atomic_impl<FHC_SHAPE_N, Hidden, FHC_HC_MULT, FHC_BLOCK_M, FHC_BLOCK_N,
+        FHC_BLOCK_K, FHC_SWIZZLE_CD, FHC_N_B_STAGES, FHC_N_INPUT_STG, FHC_NUM_MMA_TH, FHC_NUM_PMAP_TH, KS,
         /*kEarlyRelease=*/false>;
 }
 
+template <uint32_t Hidden, uint32_t KS>
+static FusedRoutFn fhcInstanceIfSupported()
+{
+    if constexpr (isSupportedFhcMmaKS<Hidden, KS>())
+    {
+        return fhcInstance<Hidden, KS>();
+    }
+    else
+    {
+        TLLM_CHECK_WITH_INFO(false, "mhcFusedHcLaunch: unsupported kNumSplits=%u for hidden_size=%u", KS, Hidden);
+        return nullptr;
+    }
+}
+
+template <uint32_t Hidden>
 static FusedRoutFn pickFhc(uint32_t ks)
 {
     switch (ks)
     {
-    case 1: return fhcInstance<1>();
-    case 2: return fhcInstance<2>();
-    case 4: return fhcInstance<4>();
-    case 8: return fhcInstance<8>();
-    case 16: return fhcInstance<16>();
-    case 32: return fhcInstance<32>();
-    case 64: return fhcInstance<64>();
+    case 1: return fhcInstanceIfSupported<Hidden, 1>();
+    case 2: return fhcInstanceIfSupported<Hidden, 2>();
+    case 4: return fhcInstanceIfSupported<Hidden, 4>();
+    case 8: return fhcInstanceIfSupported<Hidden, 8>();
+    case 16: return fhcInstanceIfSupported<Hidden, 16>();
+    case 32: return fhcInstanceIfSupported<Hidden, 32>();
+    case 64: return fhcInstanceIfSupported<Hidden, 64>();
     default: TLLM_CHECK_WITH_INFO(false, "mhcFusedHcLaunch: unsupported kNumSplits=%u", ks); return nullptr;
     }
 }
@@ -197,20 +248,24 @@ static int selectBigFuseBS(int M)
     return 512;
 }
 
-void mhcFusedHcLaunch(__nv_bfloat16 const* x_prev, __nv_bfloat16 const* residual_prev, float const* post_mix_prev,
-    float const* comb_mix_prev, float const* w_t, float const* hc_scale, float const* hc_base,
-    __nv_bfloat16* residual_cur, float* post_mix_cur, float* comb_mix_cur, __nv_bfloat16* layer_input_cur,
-    float* y_acc_workspace, float* r_acc_workspace, int M, int hidden_size, int hc_mult, int num_k_splits,
-    int bigfuse_block_size, float rms_eps, float hc_pre_eps, float hc_sinkhorn_eps, float hc_post_mult_value,
-    int sinkhorn_repeat, cudaStream_t stream)
+template <uint32_t Hidden>
+static void mhcFusedHcLaunchImpl(__nv_bfloat16 const* x_prev, __nv_bfloat16 const* residual_prev,
+    float const* post_mix_prev, float const* comb_mix_prev, float const* w_t, float const* hc_scale,
+    float const* hc_base, __nv_bfloat16* residual_cur, float* post_mix_cur, float* comb_mix_cur,
+    __nv_bfloat16* layer_input_cur, float* y_acc_workspace, float* r_acc_workspace, int M, int hidden_size, int hc_mult,
+    int num_k_splits, int bigfuse_block_size, float rms_eps, float hc_pre_eps, float hc_sinkhorn_eps,
+    float hc_post_mult_value, int sinkhorn_repeat, cudaStream_t stream)
 {
     if (M <= 0)
         return;
 
+    static_assert(isSupportedFhcHidden<Hidden>(), "Unsupported fused-HC hidden size");
+    TLLM_CHECK_WITH_INFO(hidden_size == static_cast<int>(Hidden),
+        "mhcFusedHcLaunch: dispatched Hidden=%u but got hidden_size=%d", Hidden, hidden_size);
     TLLM_CHECK_WITH_INFO(hc_mult == static_cast<int>(FHC_HC_MULT),
         "mhcFusedHcLaunch: hc_mult=%d not supported (only %u)", hc_mult, FHC_HC_MULT);
 
-    constexpr uint32_t SHAPE_K = FHC_HC_MULT * FHC_HIDDEN;
+    constexpr uint32_t SHAPE_K = FHC_HC_MULT * Hidden;
 
     uint32_t const m_u = static_cast<uint32_t>(M);
     uint32_t const ks = (num_k_splits > 0) ? static_cast<uint32_t>(num_k_splits) : pickKSplits(M);
@@ -225,8 +280,8 @@ void mhcFusedHcLaunch(__nv_bfloat16 const* x_prev, __nv_bfloat16 const* residual
         SHAPE_K, m_u, FHC_BLOCK_K, FHC_BLOCK_M, static_cast<uint64_t>(SHAPE_K) * sizeof(__nv_bfloat16),
         /*swizzleBytes=*/128, sizeof(__nv_bfloat16));
 
-    CUtensorMap desc_x = makeTma2D(const_cast<__nv_bfloat16*>(x_prev), CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, FHC_HIDDEN,
-        m_u, FHC_BLOCK_K, FHC_BLOCK_M, static_cast<uint64_t>(FHC_HIDDEN) * sizeof(__nv_bfloat16),
+    CUtensorMap desc_x = makeTma2D(const_cast<__nv_bfloat16*>(x_prev), CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, Hidden, m_u,
+        FHC_BLOCK_K, FHC_BLOCK_M, static_cast<uint64_t>(Hidden) * sizeof(__nv_bfloat16),
         /*swizzleBytes=*/128, sizeof(__nv_bfloat16));
 
     CUtensorMap desc_b = makeTma2D(const_cast<float*>(w_t), CU_TENSOR_MAP_DATA_TYPE_TFLOAT32, SHAPE_K, FHC_SHAPE_N,
@@ -239,7 +294,7 @@ void mhcFusedHcLaunch(__nv_bfloat16 const* x_prev, __nv_bfloat16 const* residual
 
     // ---- Step 1: fused post-mapping + TF32 GEMM + sqrsum + residual_out ----
     constexpr uint32_t fused_smem = fhcSmemSize();
-    FusedRoutFn fa = pickFhc(ks);
+    FusedRoutFn fa = pickFhc<Hidden>(ks);
     TLLM_CUDA_CHECK(cudaFuncSetAttribute(
         reinterpret_cast<void const*>(fa), cudaFuncAttributeMaxDynamicSharedMemorySize, fused_smem));
 
@@ -257,6 +312,37 @@ void mhcFusedHcLaunch(__nv_bfloat16 const* x_prev, __nv_bfloat16 const* residual
         hc_post_mult_value, sinkhorn_repeat, /*num_splits=*/1, /*block_size=*/bs, stream);
 }
 
+void mhcFusedHcLaunch(__nv_bfloat16 const* x_prev, __nv_bfloat16 const* residual_prev, float const* post_mix_prev,
+    float const* comb_mix_prev, float const* w_t, float const* hc_scale, float const* hc_base,
+    __nv_bfloat16* residual_cur, float* post_mix_cur, float* comb_mix_cur, __nv_bfloat16* layer_input_cur,
+    float* y_acc_workspace, float* r_acc_workspace, int M, int hidden_size, int hc_mult, int num_k_splits,
+    int bigfuse_block_size, float rms_eps, float hc_pre_eps, float hc_sinkhorn_eps, float hc_post_mult_value,
+    int sinkhorn_repeat, cudaStream_t stream)
+{
+    if (M <= 0)
+        return;
+
+    TLLM_CHECK_WITH_INFO(isSupportedFhcHiddenRuntime(hidden_size),
+        "mhcFusedHcLaunch: unsupported hidden_size=%d; supported hidden sizes are 4096 and 7168", hidden_size);
+
+    switch (hidden_size)
+    {
+    case static_cast<int>(FHC_HIDDEN_FLASH):
+        mhcFusedHcLaunchImpl<FHC_HIDDEN_FLASH>(x_prev, residual_prev, post_mix_prev, comb_mix_prev, w_t, hc_scale,
+            hc_base, residual_cur, post_mix_cur, comb_mix_cur, layer_input_cur, y_acc_workspace, r_acc_workspace, M,
+            hidden_size, hc_mult, num_k_splits, bigfuse_block_size, rms_eps, hc_pre_eps, hc_sinkhorn_eps,
+            hc_post_mult_value, sinkhorn_repeat, stream);
+        return;
+    case static_cast<int>(FHC_HIDDEN_PRO):
+        mhcFusedHcLaunchImpl<FHC_HIDDEN_PRO>(x_prev, residual_prev, post_mix_prev, comb_mix_prev, w_t, hc_scale,
+            hc_base, residual_cur, post_mix_cur, comb_mix_cur, layer_input_cur, y_acc_workspace, r_acc_workspace, M,
+            hidden_size, hc_mult, num_k_splits, bigfuse_block_size, rms_eps, hc_pre_eps, hc_sinkhorn_eps,
+            hc_post_mult_value, sinkhorn_repeat, stream);
+        return;
+    default: return;
+    }
+}
+
 // ===================================================================
 // FMA-path fused hyper-connection boundary launcher.
 //
@@ -317,6 +403,8 @@ void mhcFusedHcFmaLaunch(__nv_bfloat16 const* x_prev, __nv_bfloat16 const* resid
 
     TLLM_CHECK_WITH_INFO(hc_mult == static_cast<int>(FHC_HC_MULT),
         "mhcFusedHcFmaLaunch: hc_mult=%d not supported (only %u)", hc_mult, FHC_HC_MULT);
+    TLLM_CHECK_WITH_INFO(hidden_size % static_cast<int>(FHC_BLOCK_K) == 0,
+        "mhc fused-HC FMA path requires hidden_size to be divisible by %u, got %d", FHC_BLOCK_K, hidden_size);
     TLLM_CHECK_WITH_INFO(
         FHC_SHAPE_N % tile_n == 0, "mhcFusedHcFmaLaunch: SHAPE_N=%u not divisible by tile_n=%d", FHC_SHAPE_N, tile_n);
 
@@ -357,29 +445,48 @@ using FusedAllInOneFn = void (*)(uint32_t, CUtensorMap, CUtensorMap, CUtensorMap
     __nv_bfloat16*, float*, float*, int*, float const*, float const*, float const*, float const*, float*, float*, float,
     float, float, float, uint32_t);
 
-template <uint32_t KS>
+template <uint32_t Hidden, uint32_t KS>
 static FusedAllInOneFn fhcAllInOneInstance()
 {
-    return &fused_mhc::fused_allinone_tf32_pmap_gemm_atomic_impl<FHC_SHAPE_N, FHC_HIDDEN, FHC_HC_MULT, FHC_BLOCK_M,
+    static_assert(isSupportedFhcHidden<Hidden>(), "Unsupported fused-HC hidden size");
+    static_assert(isSupportedFhcMmaKS<Hidden, KS>(), "Unsupported fused-HC MMA kNumSplits for hidden size");
+    return &fused_mhc::fused_allinone_tf32_pmap_gemm_atomic_impl<FHC_SHAPE_N, Hidden, FHC_HC_MULT, FHC_BLOCK_M,
         FHC_BLOCK_N, FHC_BLOCK_K, FHC_SWIZZLE_CD, FHC_N_B_STAGES, FHC_N_INPUT_STG, FHC_NUM_MMA_TH, FHC_NUM_PMAP_TH, KS>;
 }
 
+template <uint32_t Hidden, uint32_t KS>
+static FusedAllInOneFn fhcAllInOneInstanceIfSupported()
+{
+    if constexpr (isSupportedFhcMmaKS<Hidden, KS>())
+    {
+        return fhcAllInOneInstance<Hidden, KS>();
+    }
+    else
+    {
+        TLLM_CHECK_WITH_INFO(
+            false, "mhcFusedHcAllInOneLaunch: unsupported kNumSplits=%u for hidden_size=%u", KS, Hidden);
+        return nullptr;
+    }
+}
+
+template <uint32_t Hidden>
 static FusedAllInOneFn pickFhcAllInOne(uint32_t ks)
 {
     switch (ks)
     {
-    case 1: return fhcAllInOneInstance<1>();
-    case 2: return fhcAllInOneInstance<2>();
-    case 4: return fhcAllInOneInstance<4>();
-    case 8: return fhcAllInOneInstance<8>();
-    case 16: return fhcAllInOneInstance<16>();
-    case 32: return fhcAllInOneInstance<32>();
-    case 64: return fhcAllInOneInstance<64>();
+    case 1: return fhcAllInOneInstanceIfSupported<Hidden, 1>();
+    case 2: return fhcAllInOneInstanceIfSupported<Hidden, 2>();
+    case 4: return fhcAllInOneInstanceIfSupported<Hidden, 4>();
+    case 8: return fhcAllInOneInstanceIfSupported<Hidden, 8>();
+    case 16: return fhcAllInOneInstanceIfSupported<Hidden, 16>();
+    case 32: return fhcAllInOneInstanceIfSupported<Hidden, 32>();
+    case 64: return fhcAllInOneInstanceIfSupported<Hidden, 64>();
     default: TLLM_CHECK_WITH_INFO(false, "mhcFusedHcAllInOneLaunch: unsupported kNumSplits=%u", ks); return nullptr;
     }
 }
 
-void mhcFusedHcAllInOneLaunch(__nv_bfloat16 const* x_prev, __nv_bfloat16 const* residual_prev,
+template <uint32_t Hidden>
+static void mhcFusedHcAllInOneLaunchImpl(__nv_bfloat16 const* x_prev, __nv_bfloat16 const* residual_prev,
     float const* post_mix_prev, float const* comb_mix_prev, float const* w_t, float const* hc_scale,
     float const* hc_base, __nv_bfloat16* residual_cur, float* post_mix_cur, float* comb_mix_cur,
     __nv_bfloat16* layer_input_cur, float* y_acc_workspace, float* r_acc_workspace, int* done_counter_workspace, int M,
@@ -389,10 +496,13 @@ void mhcFusedHcAllInOneLaunch(__nv_bfloat16 const* x_prev, __nv_bfloat16 const*
     if (M <= 0)
         return;
 
+    static_assert(isSupportedFhcHidden<Hidden>(), "Unsupported fused-HC hidden size");
+    TLLM_CHECK_WITH_INFO(hidden_size == static_cast<int>(Hidden),
+        "mhcFusedHcAllInOneLaunch: dispatched Hidden=%u but got hidden_size=%d", Hidden, hidden_size);
     TLLM_CHECK_WITH_INFO(hc_mult == static_cast<int>(FHC_HC_MULT),
         "mhcFusedHcAllInOneLaunch: hc_mult=%d not supported (only %u)", hc_mult, FHC_HC_MULT);
 
-    constexpr uint32_t SHAPE_K = FHC_HC_MULT * FHC_HIDDEN;
+    constexpr uint32_t SHAPE_K = FHC_HC_MULT * Hidden;
 
     uint32_t const m_u = static_cast<uint32_t>(M);
     uint32_t const ks = (num_k_splits > 0) ? static_cast<uint32_t>(num_k_splits) : 1u;
@@ -407,8 +517,8 @@ void mhcFusedHcAllInOneLaunch(__nv_bfloat16 const* x_prev, __nv_bfloat16 const*
         SHAPE_K, m_u, FHC_BLOCK_K, FHC_BLOCK_M, static_cast<uint64_t>(SHAPE_K) * sizeof(__nv_bfloat16),
         /*swizzleBytes=*/128, sizeof(__nv_bfloat16));
 
-    CUtensorMap desc_x = makeTma2D(const_cast<__nv_bfloat16*>(x_prev), CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, FHC_HIDDEN,
-        m_u, FHC_BLOCK_K, FHC_BLOCK_M, static_cast<uint64_t>(FHC_HIDDEN) * sizeof(__nv_bfloat16),
+    CUtensorMap desc_x = makeTma2D(const_cast<__nv_bfloat16*>(x_prev), CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, Hidden, m_u,
+        FHC_BLOCK_K, FHC_BLOCK_M, static_cast<uint64_t>(Hidden) * sizeof(__nv_bfloat16),
         /*swizzleBytes=*/128, sizeof(__nv_bfloat16));
 
     CUtensorMap desc_b = makeTma2D(const_cast<float*>(w_t), CU_TENSOR_MAP_DATA_TYPE_TFLOAT32, SHAPE_K, FHC_SHAPE_N,
@@ -421,7 +531,7 @@ void mhcFusedHcAllInOneLaunch(__nv_bfloat16 const* x_prev, __nv_bfloat16 const*
 
     // ---- Launch the single all-in-one kernel ----
     constexpr uint32_t fused_smem = fhcAllInOneSmemSize();
-    FusedAllInOneFn fa = pickFhcAllInOne(ks);
+    FusedAllInOneFn fa = pickFhcAllInOne<Hidden>(ks);
     TLLM_CUDA_CHECK(cudaFuncSetAttribute(
         reinterpret_cast<void const*>(fa), cudaFuncAttributeMaxDynamicSharedMemorySize, fused_smem));
 
@@ -433,6 +543,37 @@ void mhcFusedHcAllInOneLaunch(__nv_bfloat16 const* x_prev, __nv_bfloat16 const*
         static_cast<uint32_t>(sinkhorn_repeat));
 }
 
+void mhcFusedHcAllInOneLaunch(__nv_bfloat16 const* x_prev, __nv_bfloat16 const* residual_prev,
+    float const* post_mix_prev, float const* comb_mix_prev, float const* w_t, float const* hc_scale,
+    float const* hc_base, __nv_bfloat16* residual_cur, float* post_mix_cur, float* comb_mix_cur,
+    __nv_bfloat16* layer_input_cur, float* y_acc_workspace, float* r_acc_workspace, int* done_counter_workspace, int M,
+    int hidden_size, int hc_mult, int num_k_splits, float rms_eps, float hc_pre_eps, float hc_sinkhorn_eps,
+    float hc_post_mult_value, int sinkhorn_repeat, cudaStream_t stream)
+{
+    if (M <= 0)
+        return;
+
+    TLLM_CHECK_WITH_INFO(isSupportedFhcHiddenRuntime(hidden_size),
+        "mhcFusedHcAllInOneLaunch: unsupported hidden_size=%d; supported hidden sizes are 4096 and 7168", hidden_size);
+
+    switch (hidden_size)
+    {
+    case static_cast<int>(FHC_HIDDEN_FLASH):
+        mhcFusedHcAllInOneLaunchImpl<FHC_HIDDEN_FLASH>(x_prev, residual_prev, post_mix_prev, comb_mix_prev, w_t,
+            hc_scale, hc_base, residual_cur, post_mix_cur, comb_mix_cur, layer_input_cur, y_acc_workspace,
+            r_acc_workspace, done_counter_workspace, M, hidden_size, hc_mult, num_k_splits, rms_eps, hc_pre_eps,
+            hc_sinkhorn_eps, hc_post_mult_value, sinkhorn_repeat, stream);
+        return;
+    case static_cast<int>(FHC_HIDDEN_PRO):
+        mhcFusedHcAllInOneLaunchImpl<FHC_HIDDEN_PRO>(x_prev, residual_prev, post_mix_prev, comb_mix_prev, w_t,
+            hc_scale, hc_base, residual_cur, post_mix_cur, comb_mix_cur, layer_input_cur, y_acc_workspace,
+            r_acc_workspace, done_counter_workspace, M, hidden_size, hc_mult, num_k_splits, rms_eps, hc_pre_eps,
+            hc_sinkhorn_eps, hc_post_mult_value, sinkhorn_repeat, stream);
+        return;
+    default: return;
+    }
+}
+
 // ===================================================================
 // All-in-one single-kernel fused hyper-connection launcher (FMA).
 //
@@ -505,6 +646,8 @@ void mhcFusedHcFmaAllInOneLaunch(__nv_bfloat16 const* x_prev, __nv_bfloat16 cons
 
     TLLM_CHECK_WITH_INFO(hc_mult == static_cast<int>(FHC_HC_MULT),
         "mhcFusedHcFmaAllInOneLaunch: hc_mult=%d not supported (only %u)", hc_mult, FHC_HC_MULT);
+    TLLM_CHECK_WITH_INFO(hidden_size % static_cast<int>(FHC_BLOCK_K) == 0,
+        "mhc fused-HC FMA path requires hidden_size to be divisible by %u, got %d", FHC_BLOCK_K, hidden_size);
     TLLM_CHECK_WITH_INFO(FHC_SHAPE_N % tile_n == 0,
         "mhcFusedHcFmaAllInOneLaunch: SHAPE_N=%u not divisible by tile_n=%d", FHC_SHAPE_N, tile_n);
 
diff --git a/cpp/tensorrt_llm/kernels/mhcKernels/mhcKernels.h b/cpp/tensorrt_llm/kernels/mhcKernels/mhcKernels.h
index c64e494b50a1..f1bf3bd74f50 100644
--- a/cpp/tensorrt_llm/kernels/mhcKernels/mhcKernels.h
+++ b/cpp/tensorrt_llm/kernels/mhcKernels/mhcKernels.h
@@ -51,7 +51,10 @@ void mhcPostMappingLaunch(__nv_bfloat16 const* residual, __nv_bfloat16 const* x,
 //   y_acc_workspace: fp32 [M, 24]
 //   r_acc_workspace: fp32 [M]
 //
-// Shape constraints (B200 / sm_100a): hidden_size == 4096, hc_mult == 4.
+// Shape constraints (B200/B300 / sm_100a):
+//   hc_mult == 4
+//   hidden_size in {4096, 7168} for SM100/tcgen05 MMA fused-HC paths.
+// FMA fused-HC paths use runtime hidden_size but still require hidden_size % 64 == 0.
 // Passing num_k_splits == 0 or bigfuse_block_size == 0 falls back to the
 // internal heuristics.
 void mhcFusedHcLaunch(__nv_bfloat16 const* x_prev, __nv_bfloat16 const* residual_prev, float const* post_mix_prev,
@@ -101,7 +104,10 @@ void mhcFusedHcFmaLaunch(__nv_bfloat16 const* x_prev, __nv_bfloat16 const* resid
 //   r_acc_workspace: fp32 [M]
 //   done_counter_workspace: int32 [ceil(M / 64)]
 //
-// Shape constraints (B200 / sm_100a): hidden_size == 4096, hc_mult == 4.
+// Shape constraints (B200/B300 / sm_100a):
+//   hc_mult == 4
+//   hidden_size in {4096, 7168} for SM100/tcgen05 MMA fused-HC paths.
+// FMA fused-HC paths use runtime hidden_size but still require hidden_size % 64 == 0.
 // Passing num_k_splits == 0 falls back to internal heuristics.
 void mhcFusedHcAllInOneLaunch(__nv_bfloat16 const* x_prev, __nv_bfloat16 const* residual_prev,
     float const* post_mix_prev, float const* comb_mix_prev, float const* w_t, float const* hc_scale,
diff --git a/tensorrt_llm/_torch/modules/mhc/mhc_cuda.py b/tensorrt_llm/_torch/modules/mhc/mhc_cuda.py
index 881d672d0993..560abfc5122c 100644
--- a/tensorrt_llm/_torch/modules/mhc/mhc_cuda.py
+++ b/tensorrt_llm/_torch/modules/mhc/mhc_cuda.py
@@ -32,14 +32,17 @@
 
 @lru_cache(maxsize=1)
 def _fused_hc_mma_supported() -> bool:
-    """tcgen05 TF32 MMA paths (Path B / Path D, "fused_*_mma") require SM100+.
+    """tcgen05 TF32 MMA paths (Path B / Path D, "fused_*_mma") require SM100.
 
-    On pre-SM100 GPUs only the FMA paths (Path E / Path F, "fused_*_fma") are
-    safe to run. Called lazily so the module can still be imported on a host
-    without a CUDA device (e.g. CPU-only lint / typecheck).
+    Match the C++ compile guard for the tcgen05 kernels:
+    __CUDA_ARCH__ >= 1000 && __CUDA_ARCH__ < 1100. On other GPU generations,
+    only the FMA paths (Path E / Path F, "fused_*_fma") are safe to run.
+    Called lazily so the module can still be imported on a host without a CUDA
+    device (e.g. CPU-only lint / typecheck).
     """
     try:
-        return get_sm_version() >= 100
+        sm_version = get_sm_version()
+        return 100 <= sm_version < 110
     except Exception:
         return False
 
@@ -516,6 +519,33 @@ def mhc_post_mapping_cuda(
     "fused_all_fma": 3,  # 1-kernel fp32 FMA all-in-one    (Path F)
 }
 
+# The SM100/tcgen05 MMA fused-HC C++ kernels are statically instantiated.
+# FMA fused-HC paths use runtime hidden_size, but MMA paths must be explicitly
+# compiled for each supported hidden size.
+_FUSED_HC_MMA_SUPPORTED_HIDDEN_SIZES = {4096, 7168}
+
+
+def _fused_hc_mma_ks_supported(hidden_size: int, ks: int) -> bool:
+    if hidden_size not in _FUSED_HC_MMA_SUPPORTED_HIDDEN_SIZES:
+        return False
+
+    block_k = 64
+    block_m = 64
+    num_warps = 8
+    warp_size = 32
+    bf16_vec = 8
+
+    if hidden_size % block_k != 0:
+        return False
+    h_tiles = hidden_size // block_k
+    if h_tiles % ks != 0:
+        return False
+
+    toks_per_cta = (block_m + ks - 1) // ks
+    warps_per_tok = num_warps // toks_per_cta if num_warps > toks_per_cta else 1
+    return hidden_size % (warps_per_tok * warp_size * bf16_vec) == 0
+
+
 # Tactics supported by the half-fused FMA path in `mhcFusedHcFmaLaunch`
 # (must stay in sync with the C++ pickFhcFma() table).
 _FUSED_HC_HALF_FMA_TN_KS = (
@@ -733,12 +763,11 @@ def _alloc_fused_hc_outputs(
 _FUSED_HC_FALLBACK_TACTIC_FMA = ("fused_half_fma", 2, 1, 256, 1)
 
 
-def _get_fused_hc_fallback_tactic():
-    return (
-        _FUSED_HC_FALLBACK_TACTIC_MMA
-        if _fused_hc_mma_supported()
-        else _FUSED_HC_FALLBACK_TACTIC_FMA
-    )
+def _get_fused_hc_fallback_tactic(hidden_size: int | None = None):
+    mma_ok = _fused_hc_mma_supported()
+    if hidden_size is not None:
+        mma_ok = mma_ok and hidden_size in _FUSED_HC_MMA_SUPPORTED_HIDDEN_SIZES
+    return _FUSED_HC_FALLBACK_TACTIC_MMA if mma_ok else _FUSED_HC_FALLBACK_TACTIC_FMA
 
 
 class MhcFusedHcRunner(TunableRunner):
@@ -810,7 +839,10 @@ def get_valid_tactics(self, inputs, profile: OptimizationProfile, **kwargs):
         tactics = []
         # The MMA (tcgen05) paths require SM100+. On older archs only the FMA
         # paths are compilable/runnable — we simply never emit MMA tactics.
-        mma_ok = _fused_hc_mma_supported()
+        mma_ks = tuple(
+            ks for ks in _FUSED_HC_HALF_MMA_KS if _fused_hc_mma_ks_supported(self.hidden_size, ks)
+        )
+        mma_ok = _fused_hc_mma_supported() and bool(mma_ks)
         # Path F (fused_all_fma, 1-kernel FMA) — preferred at small M (<=32)
         # where MMA can't fill BLOCK_M=64. Include for M <= 64 as the
         # crossover is measured per-M.
@@ -826,7 +858,7 @@ def get_valid_tactics(self, inputs, profile: OptimizationProfile, **kwargs):
         # M (>=64). Include when M >= 48 to overlap with Path F at the
         # crossover boundary.
         if mma_ok and M >= 48:
-            for ks in _FUSED_HC_ALL_MMA_KS:
+            for ks in mma_ks:
                 m_tiles = (M + 63) // 64
                 if m_tiles * ks > 148 * 4:
                     continue
@@ -842,7 +874,7 @@ def get_valid_tactics(self, inputs, profile: OptimizationProfile, **kwargs):
         # Half-fused MMA path (2-kernel) — always an option when tcgen05 is
         # available.
         if mma_ok:
-            for ks in _FUSED_HC_HALF_MMA_KS:
+            for ks in mma_ks:
                 m_tiles = (M + 63) // 64
                 if m_tiles * ks > 148 * 4:
                     continue
@@ -870,7 +902,7 @@ def forward(self, inputs, *, tactic=-1, **kwargs):
         hc_base_cur = hc_base_cur.to(torch.float32).contiguous()
 
         if tactic == -1:
-            tactic = _get_fused_hc_fallback_tactic()
+            tactic = _get_fused_hc_fallback_tactic(self.hidden_size)
         backend, tile_n, num_k_splits, bigfuse_bs, tile_m = tactic
         backend_code = _FUSED_HC_BACKEND_CODE[backend]
 

From c7cbaa45b5556fca5f23a2ca2fb6b6048df5f1bd Mon Sep 17 00:00:00 2001
From: Mingyang Hao <mingyangHao@users.noreply.github.com>
Date: Tue, 5 May 2026 00:51:06 -0700
Subject: [PATCH 2/3] test: add DeepSeek-V4 Pro MHC coverage

Signed-off-by: Mingyang Hao <mingyangHao@users.noreply.github.com>
---
 tests/unittest/_torch/modules/test_mhc.py | 108 +++++++++++++++++-----
 1 file changed, 86 insertions(+), 22 deletions(-)

diff --git a/tests/unittest/_torch/modules/test_mhc.py b/tests/unittest/_torch/modules/test_mhc.py
index 7cd05ae0554a..48eaadd259da 100644
--- a/tests/unittest/_torch/modules/test_mhc.py
+++ b/tests/unittest/_torch/modules/test_mhc.py
@@ -329,6 +329,49 @@ def test_mhc_pre_mapping(n: int, hidden_size: int, hc_mult: int):
     torch.testing.assert_close(layer_input_ref, layer_input_cuda, rtol=1e-4, atol=1e-3)
 
 
+@pytest.mark.parametrize("n", [64])
+@pytest.mark.parametrize("hidden_size", [7168])
+@pytest.mark.parametrize("hc_mult", [4])
+def test_mhc_pre_mapping_pro_hidden_size(n: int, hidden_size: int, hc_mult: int):
+    test_data = generate_pre_data(
+        n=n,
+        hc_mult=hc_mult,
+        hidden_size=hidden_size,
+    )
+
+    test_module = mHC(
+        mult=hc_mult,
+        hidden_size=hidden_size,
+        sinkhorn_iters=test_data["sinkhorn_repeat"],
+        dtype=None,
+        eps=test_data["hc_pre_eps"],
+        norm_eps=test_data["rms_eps"],
+        post_mult_value=test_data["hc_post_mult_value"],
+    ).cuda()
+    test_module.fn.copy_(test_data["fn"])
+    test_module.scale.copy_(test_data["hc_scale"])
+    test_module.base.copy_(test_data["hc_base"])
+
+    residual = test_data["residual"]
+
+    post_mix_cuda, comb_mix_cuda, layer_input_cuda = test_module.pre_mapping(residual)
+    post_mix_ref, comb_mix_ref, layer_input_ref = vanilla_pre_mapping(
+        residual,
+        test_data["fn"],
+        test_data["hc_scale"],
+        test_data["hc_base"],
+        hc_mult,
+        test_data["rms_eps"],
+        test_data["hc_pre_eps"],
+        test_data["hc_sinkhorn_eps"],
+        test_data["hc_post_mult_value"],
+        test_data["sinkhorn_repeat"],
+    )
+    torch.testing.assert_close(post_mix_ref, post_mix_cuda, rtol=1e-4, atol=1e-3)
+    torch.testing.assert_close(comb_mix_ref, comb_mix_cuda, rtol=1e-3, atol=5e-3)
+    torch.testing.assert_close(layer_input_ref, layer_input_cuda, rtol=1e-4, atol=1e-3)
+
+
 @pytest.mark.parametrize("n", [64, 128, 4096, 8192])
 @pytest.mark.parametrize("hidden_size", [7168])
 @pytest.mark.parametrize("hc_mult", [4])
@@ -476,8 +519,26 @@ def _unfused():
 }
 
 
+def test_mhc_fused_hc_mma_tactic_filter_hidden_sizes():
+    from tensorrt_llm._torch.modules.mhc.mhc_cuda import (
+        _FUSED_HC_HALF_MMA_KS,
+        _fused_hc_mma_ks_supported,
+    )
+
+    supported_by_hidden_size = {
+        hidden_size: {
+            ks for ks in _FUSED_HC_HALF_MMA_KS if _fused_hc_mma_ks_supported(hidden_size, ks)
+        }
+        for hidden_size in (4096, 7168, 8192)
+    }
+
+    assert supported_by_hidden_size[4096] == {1, 2, 4, 8, 16, 32, 64}
+    assert supported_by_hidden_size[7168] == {1, 2, 4, 8, 16}
+    assert supported_by_hidden_size[8192] == set()
+
+
 @pytest.mark.parametrize("n", list(_BACKEND_TACTICS_BY_M.keys()))
-@pytest.mark.parametrize("hidden_size", [4096])
+@pytest.mark.parametrize("hidden_size", [4096, 7168])
 @pytest.mark.parametrize("hc_mult", [4])
 def test_mhc_fused_hc_backends(n: int, hidden_size: int, hc_mult: int):
     """Every wired fused_hc backend sees bit-identical input and is checked
@@ -556,13 +617,13 @@ def make_runner_inputs():
             cur_module.base.detach().clone(),
         ]
 
-    backend_outputs = {}
+    tactic_outputs = {}
     for tactic in _BACKEND_TACTICS_BY_M[n]:
         backend = tactic[0]
         residual_cur, post_mix_cur, comb_mix_cur, layer_input_cur = runner(
             inputs=make_runner_inputs(), tactic=tactic
         )
-        backend_outputs[backend] = (
+        tactic_outputs[tactic] = (
             residual_cur,
             post_mix_cur.view(n, hc_mult, 1),
             comb_mix_cur.view(n, hc_mult, hc_mult),
@@ -577,74 +638,77 @@ def make_runner_inputs():
     fp32_tol = dict(rtol=1e-3, atol=5e-3)
 
     # (1) Every backend must match the torch reference.
-    for backend, (
+    for tactic, (
         residual_cur,
         post_mix_cur,
         comb_mix_cur,
         layer_input_cur,
-    ) in backend_outputs.items():
+    ) in tactic_outputs.items():
         torch.testing.assert_close(
             residual_cur_ref,
             residual_cur,
             **bf16_tol,
-            msg=f"[vs torch-ref] backend={backend} n={n} residual mismatch",
+            msg=f"[vs torch-ref] tactic={tactic} n={n} hidden={hidden_size} residual mismatch",
         )
         torch.testing.assert_close(
             post_mix_ref,
             post_mix_cur,
             **fp32_tol,
-            msg=f"[vs torch-ref] backend={backend} n={n} post_mix mismatch",
+            msg=f"[vs torch-ref] tactic={tactic} n={n} hidden={hidden_size} post_mix mismatch",
         )
         torch.testing.assert_close(
             comb_mix_ref,
             comb_mix_cur,
             **fp32_tol,
-            msg=f"[vs torch-ref] backend={backend} n={n} comb_mix mismatch",
+            msg=f"[vs torch-ref] tactic={tactic} n={n} hidden={hidden_size} comb_mix mismatch",
         )
         torch.testing.assert_close(
             layer_input_ref,
             layer_input_cur,
             **bf16_tol,
-            msg=f"[vs torch-ref] backend={backend} n={n} layer_input mismatch",
+            msg=f"[vs torch-ref] tactic={tactic} n={n} hidden={hidden_size} layer_input mismatch",
         )
 
     # (2) All backends must agree with one golden backend at the same tolerance
     # as vs the torch ref. Different backends vary only in tile shape and
     # reduction order, so cross-backend divergence would indicate a kernel
     # correctness bug rather than expected rounding drift.
-    gold = "fused_half_mma" if "fused_half_mma" in backend_outputs else next(iter(backend_outputs))
-    gr, gpm, gcm, gli = backend_outputs[gold]
-    for backend, (
+    gold = next(
+        (tactic for tactic in tactic_outputs if tactic[0] == "fused_half_mma"),
+        next(iter(tactic_outputs)),
+    )
+    gr, gpm, gcm, gli = tactic_outputs[gold]
+    for tactic, (
         residual_cur,
         post_mix_cur,
         comb_mix_cur,
         layer_input_cur,
-    ) in backend_outputs.items():
-        if backend == gold:
+    ) in tactic_outputs.items():
+        if tactic == gold:
             continue
         torch.testing.assert_close(
             gr,
             residual_cur,
             **bf16_tol,
-            msg=f"[vs {gold}] backend={backend} n={n} residual mismatch",
+            msg=f"[vs {gold}] tactic={tactic} n={n} hidden={hidden_size} residual mismatch",
         )
         torch.testing.assert_close(
             gpm,
             post_mix_cur,
             **fp32_tol,
-            msg=f"[vs {gold}] backend={backend} n={n} post_mix mismatch",
+            msg=f"[vs {gold}] tactic={tactic} n={n} hidden={hidden_size} post_mix mismatch",
         )
         torch.testing.assert_close(
             gcm,
             comb_mix_cur,
             **fp32_tol,
-            msg=f"[vs {gold}] backend={backend} n={n} comb_mix mismatch",
+            msg=f"[vs {gold}] tactic={tactic} n={n} hidden={hidden_size} comb_mix mismatch",
         )
         torch.testing.assert_close(
             gli,
             layer_input_cur,
             **bf16_tol,
-            msg=f"[vs {gold}] backend={backend} n={n} layer_input mismatch",
+            msg=f"[vs {gold}] tactic={tactic} n={n} hidden={hidden_size} layer_input mismatch",
         )
 
 
@@ -657,13 +721,13 @@ def make_runner_inputs():
         ("fused_all_fma", 2, 1, 0, 1),
     ],
 )
-def test_mhc_fused_hc_realistic_scale_regression(tactic):
+@pytest.mark.parametrize("hidden_size", [4096, 7168])
+def test_mhc_fused_hc_realistic_scale_regression(tactic, hidden_size: int):
     """Real-scale mHC data catches fused_hc RMS normalization regressions."""
     from tensorrt_llm._torch.modules.mhc.mhc_cuda import MhcFusedHcRunner
 
     n = 16
     hc_mult = 4
-    hidden_size = 4096
     pre_data = generate_realistic_pre_data(n=n, hc_mult=hc_mult, hidden_size=hidden_size)
 
     torch.random.manual_seed(17)
@@ -723,7 +787,7 @@ def test_mhc_fused_hc_realistic_scale_regression(tactic):
 
 
 @pytest.mark.parametrize("n", [128, 2048])
-@pytest.mark.parametrize("hidden_size", [4096])
+@pytest.mark.parametrize("hidden_size", [4096, 7168])
 @pytest.mark.parametrize("hc_mult", [4])
 def test_mhc_fused_hc_cuda_graph(n: int, hidden_size: int, hc_mult: int):
     """CUDA-graph capture/replay of mHC.fused_hc.
@@ -845,7 +909,7 @@ def _inputs():
 
 
 @pytest.mark.parametrize("m", [64, 128, 4096, 8192])
-@pytest.mark.parametrize("hidden_size", [4096])
+@pytest.mark.parametrize("hidden_size", [4096, 7168])
 @pytest.mark.parametrize("hc_mult", [4])
 def test_hc_head(m: int, hidden_size: int, hc_mult: int):
     test_data = generate_head_data(

From 70927b2bb4c346908736b120cc39c0ad2bae4a59 Mon Sep 17 00:00:00 2001
From: Patrice Castonguay <55748270+pcastonguay@users.noreply.github.com>
Date: Tue, 5 May 2026 10:32:36 -0700
Subject: [PATCH 3/3] Pre-commit fix

Signed-off-by: Patrice Castonguay <55748270+pcastonguay@users.noreply.github.com>
---
 cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu | 8 ++++----
 tests/unittest/_torch/modules/test_mhc.py               | 1 -
 2 files changed, 4 insertions(+), 5 deletions(-)

diff --git a/cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu b/cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu
index 23c15c72a9bc..7d5d2b6daf45 100644
--- a/cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu
+++ b/cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu
@@ -565,10 +565,10 @@ void mhcFusedHcAllInOneLaunch(__nv_bfloat16 const* x_prev, __nv_bfloat16 const*
             hc_sinkhorn_eps, hc_post_mult_value, sinkhorn_repeat, stream);
         return;
     case static_cast<int>(FHC_HIDDEN_PRO):
-        mhcFusedHcAllInOneLaunchImpl<FHC_HIDDEN_PRO>(x_prev, residual_prev, post_mix_prev, comb_mix_prev, w_t,
-            hc_scale, hc_base, residual_cur, post_mix_cur, comb_mix_cur, layer_input_cur, y_acc_workspace,
-            r_acc_workspace, done_counter_workspace, M, hidden_size, hc_mult, num_k_splits, rms_eps, hc_pre_eps,
-            hc_sinkhorn_eps, hc_post_mult_value, sinkhorn_repeat, stream);
+        mhcFusedHcAllInOneLaunchImpl<FHC_HIDDEN_PRO>(x_prev, residual_prev, post_mix_prev, comb_mix_prev, w_t, hc_scale,
+            hc_base, residual_cur, post_mix_cur, comb_mix_cur, layer_input_cur, y_acc_workspace, r_acc_workspace,
+            done_counter_workspace, M, hidden_size, hc_mult, num_k_splits, rms_eps, hc_pre_eps, hc_sinkhorn_eps,
+            hc_post_mult_value, sinkhorn_repeat, stream);
         return;
     default: return;
     }
diff --git a/tests/unittest/_torch/modules/test_mhc.py b/tests/unittest/_torch/modules/test_mhc.py
index 48eaadd259da..9882603cc39c 100644
--- a/tests/unittest/_torch/modules/test_mhc.py
+++ b/tests/unittest/_torch/modules/test_mhc.py
@@ -619,7 +619,6 @@ def make_runner_inputs():
 
     tactic_outputs = {}
     for tactic in _BACKEND_TACTICS_BY_M[n]:
-        backend = tactic[0]
         residual_cur, post_mix_cur, comb_mix_cur, layer_input_cur = runner(
             inputs=make_runner_inputs(), tactic=tactic
         )