Skip to content

[None][perf] DSV4 o_proj: fuse fp8/UE8M0 quantize into o_a proj and introduce splitk cutedsl o_b proj #15257

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
mingyangHao wants to merge 7 commits into from
Closed

[None][perf] DSV4 o_proj: fuse fp8/UE8M0 quantize into o_a proj and introduce splitk cutedsl o_b proj #15257

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
6c0eaa6
[None][perf] fuse DSV4 FP8 o-projection split-K into mHC
mingyangHao Jun 29, 2026
41ec8fb
[None][perf] optimize DSV4 FP8 o_a decode epilogue
mingyangHao Jun 30, 2026
7716c86
[None][perf] Replace DSV4 O_b split-K with CuTe DSL
mingyangHao Jul 4, 2026
b03f0b9
[None][cleanup] Tidy DSV4 O-projection fusion
mingyangHao Jul 7, 2026
8348419
[None][cleanup] Derive DSV4 O-projection tactics
mingyangHao Jul 7, 2026
5f70bd2
[None][fix] Use legal MXF8 scale-vector encoding
mingyangHao Jul 7, 2026
c797edf
[None][fix] Enable DSV4 fused O-projection by default
mingyangHao Jul 8, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
110 changes: 83 additions & 27 deletions cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -108,11 +108,11 @@ __device__ __forceinline__ void stsm_x4_b16_rout(void* smem_dst, uint32_t a, uin

template <uint32_t SHAPE_N, uint32_t HIDDEN, uint32_t HC_MULT, uint32_t BLOCK_M, uint32_t BLOCK_N, uint32_t BLOCK_K,
uint32_t kSwizzleCDMode, uint32_t N_B_STAGES, uint32_t N_INPUT_STAGES, uint32_t kNumMMAThreads,
uint32_t kNumPmapThreads, uint32_t kNumSplits = 1, bool kEarlyRelease = false>
uint32_t kNumPmapThreads, uint32_t kNumSplits = 1, bool kEarlyRelease = false, uint32_t kXSplit = 1>
__global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) fused_tf32_pmap_gemm_rout_atomic_impl(
const uint32_t shape_m, const __grid_constant__ cute::TmaDescriptor tensor_map_residual,
const __grid_constant__ cute::TmaDescriptor tensor_map_x, const __grid_constant__ cute::TmaDescriptor tensor_map_b,
const __grid_constant__ cute::TmaDescriptor tensor_map_residual_out,
uint32_t const shape_m, __grid_constant__ const cute::TmaDescriptor tensor_map_residual,
__grid_constant__ const cute::TmaDescriptor tensor_map_x, __grid_constant__ const cute::TmaDescriptor tensor_map_b,
__grid_constant__ const cute::TmaDescriptor tensor_map_residual_out,
float* __restrict__ D, // [M, SHAPE_N] (caller memsets to 0)
float const* __restrict__ post_mix, float const* __restrict__ comb_mix, float* __restrict__ sqr_sum)
{ // [M] (caller memsets to 0)
Expand Down Expand Up @@ -175,6 +175,7 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) fused_tf3
auto smem_res = PatternVisitor(
[&, base = cursor](uint32_t const& i) { return reinterpret_cast<nv_bfloat16*>(base + i * SMEM_RES_PER_ISTG); });
cursor += N_INPUT_STAGES * SMEM_RES_PER_ISTG;
// Stream split-x tiles through the existing x buffers.
auto smem_x_stg = PatternVisitor(
[&, base = cursor](uint32_t const& i) { return reinterpret_cast<nv_bfloat16*>(base + i * SMEM_X_PER_ISTG); });
cursor += N_INPUT_STAGES * SMEM_X_PER_ISTG;
Expand All @@ -192,13 +193,18 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) fused_tf3
auto full_input = PatternVisitor([=](uint32_t const& i) { return barrier_start_ptr + 2 * N_B_STAGES + i; });
auto empty_input
= PatternVisitor([=](uint32_t const& i) { return barrier_start_ptr + 2 * N_B_STAGES + N_INPUT_STAGES + i; });
auto full_cast = PatternVisitor(
constexpr uint32_t kXBarriers = (kXSplit > 1) ? 2 * N_INPUT_STAGES : 0;
auto full_x = PatternVisitor(
[=](uint32_t const& i) { return barrier_start_ptr + 2 * N_B_STAGES + 2 * N_INPUT_STAGES + i; });
auto empty_x = PatternVisitor(
[=](uint32_t const& i) { return barrier_start_ptr + 2 * N_B_STAGES + 3 * N_INPUT_STAGES + i; });
auto full_cast = PatternVisitor(
[=](uint32_t const& i) { return barrier_start_ptr + 2 * N_B_STAGES + 2 * N_INPUT_STAGES + kXBarriers + i; });
auto empty_cast = PatternVisitor([=](uint32_t const& i)
{ return barrier_start_ptr + 2 * N_B_STAGES + 2 * N_INPUT_STAGES + kNumCastStages + i; });
auto tmem_full_barrier = barrier_start_ptr + 2 * N_B_STAGES + 2 * N_INPUT_STAGES + 2 * kNumCastStages;
{ return barrier_start_ptr + 2 * N_B_STAGES + 2 * N_INPUT_STAGES + kXBarriers + kNumCastStages + i; });
auto tmem_full_barrier = barrier_start_ptr + 2 * N_B_STAGES + 2 * N_INPUT_STAGES + kXBarriers + 2 * kNumCastStages;

cursor += (2 * N_B_STAGES + 2 * N_INPUT_STAGES + 2 * kNumCastStages + 1) * sizeof(Barrier);
cursor += (2 * N_B_STAGES + 2 * N_INPUT_STAGES + kXBarriers + 2 * kNumCastStages + 1) * sizeof(Barrier);
auto tmem_ptr_in_smem = reinterpret_cast<uint32_t*>(cursor);

if (warp_idx == 1 and cute::elect_one_sync())
Expand All @@ -215,6 +221,15 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) fused_tf3
full_input[i]->init(1);
empty_input[i]->init(kNumPmapThreads);
}
if constexpr (kXSplit > 1)
{
#pragma unroll
for (uint32_t i = 0; i < N_INPUT_STAGES; ++i)
{
full_x[i]->init(1);
empty_x[i]->init(kNumPmapThreads);
}
}
#pragma unroll
for (uint32_t i = 0; i < kNumCastStages; ++i)
{
Expand Down Expand Up @@ -278,6 +293,7 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) fused_tf3
uint32_t b_stage = 0;
uint32_t i_stage = 0;
uint32_t s = 0;
uint32_t xs = 0;
for (uint32_t ht = 0; ht < H_TILES_PER_SPLIT; ++ht)
{
const uint32_t h_tile = h_tile_start + ht;
Expand All @@ -290,10 +306,27 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) fused_tf3
tma_copy<BLOCK_K, BLOCK_M, kSwizzleResMode>(&tensor_map_residual, full_input[i_stage],
smem_res[i_stage] + j * BLOCK_M * BLOCK_K, j * HIDDEN + h_idx, m_idx);
}
tma_copy<BLOCK_K, BLOCK_M, kSwizzleXMode>(
&tensor_map_x, full_input[i_stage], smem_x_stg[i_stage], h_idx, m_idx);
constexpr uint32_t kInputBytes = SMEM_RES_PER_ISTG + SMEM_X_PER_ISTG;
full_input[i_stage]->arrive_and_expect_tx(kInputBytes);
if constexpr (kXSplit == 1)
{
tma_copy<BLOCK_K, BLOCK_M, kSwizzleXMode>(
&tensor_map_x, full_input[i_stage], smem_x_stg[i_stage], h_idx, m_idx);
constexpr uint32_t kInputBytes = SMEM_RES_PER_ISTG + SMEM_X_PER_ISTG;
full_input[i_stage]->arrive_and_expect_tx(kInputBytes);
}
else
{
full_input[i_stage]->arrive_and_expect_tx(SMEM_RES_PER_ISTG);
#pragma unroll
for (uint32_t xp = 0; xp < kXSplit; ++xp)
{
uint32_t const xslot = xs % N_INPUT_STAGES;
empty_x[xslot]->wait(((xs / N_INPUT_STAGES) & 1) ^ 1);
tma_copy<BLOCK_K, BLOCK_M, kSwizzleXMode>(
&tensor_map_x, full_x[xslot], smem_x_stg[xslot], h_idx, xp * shape_m + m_idx);
full_x[xslot]->arrive_and_expect_tx(SMEM_X_PER_ISTG);
++xs;
}
}

#pragma unroll
for (uint32_t hc = 0; hc < HC_MULT; ++hc)
Expand Down Expand Up @@ -448,19 +481,6 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) fused_tf3
const uint32_t i_stage = ht % N_INPUT_STAGES;
full_input[i_stage]->wait((ht / N_INPUT_STAGES) & 1);

uint32_t x_vals[2][kNumLoads];
{
uint8_t const* x_base
= reinterpret_cast<uint8_t*>(smem_x_stg[i_stage]) + sub_warp_idx * BLOCK_M_PER_WARP * kSwizzleXMode;
#pragma unroll
for (uint32_t i = 0; i < kNumLoads; i += 2)
{
auto smem_ptr = x_base + get_swizzled_smem_offset<kSwizzleXMode>(i + lane_idx / 16, lane_idx % 16);
deep_gemm::sm90::SM90_U32x4_LDSM_N::copy(x_vals[0][i + 0], x_vals[1][i + 0], x_vals[0][i + 1],
x_vals[1][i + 1], const_cast<uint8_t*>(smem_ptr));
}
}

uint32_t r_vals[HC_MULT][2][kNumLoads];
#pragma unroll
for (uint32_t j = 0; j < HC_MULT; ++j)
Expand All @@ -480,12 +500,48 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) fused_tf3
float2 xf[2][kNumLoads];
#pragma unroll
for (uint32_t u = 0; u < 2; ++u)
{
#pragma unroll
for (uint32_t i = 0; i < kNumLoads; ++i)
xf[u][i] = make_float2(0.f, 0.f);
#pragma unroll
for (uint32_t xp = 0; xp < kXSplit; ++xp)
{
nv_bfloat16* x_buf;
if constexpr (kXSplit == 1)
{
xf[u][i] = __bfloat1622float2(*reinterpret_cast<nv_bfloat162*>(&x_vals[u][i]));
x_buf = smem_x_stg[i_stage];
}
else
{
uint32_t const xsg = ht * kXSplit + xp;
uint32_t const xslot = xsg % N_INPUT_STAGES;
full_x[xslot]->wait((xsg / N_INPUT_STAGES) & 1);
x_buf = smem_x_stg[xslot];
}

uint32_t x_vals[2][kNumLoads];
uint8_t const* x_base
= reinterpret_cast<uint8_t*>(x_buf) + sub_warp_idx * BLOCK_M_PER_WARP * kSwizzleXMode;
#pragma unroll
for (uint32_t i = 0; i < kNumLoads; i += 2)
{
auto smem_ptr = x_base + get_swizzled_smem_offset<kSwizzleXMode>(i + lane_idx / 16, lane_idx % 16);
deep_gemm::sm90::SM90_U32x4_LDSM_N::copy(x_vals[0][i + 0], x_vals[1][i + 0], x_vals[0][i + 1],
x_vals[1][i + 1], const_cast<uint8_t*>(smem_ptr));
}
if constexpr (kXSplit > 1)
{
empty_x[(ht * kXSplit + xp) % N_INPUT_STAGES]->arrive();
}
#pragma unroll
for (uint32_t u = 0; u < 2; ++u)
#pragma unroll
for (uint32_t i = 0; i < kNumLoads; ++i)
{
float2 value = __bfloat1622float2(*reinterpret_cast<nv_bfloat162*>(&x_vals[u][i]));
xf[u][i].x += value.x;
xf[u][i].y += value.y;
}
}

if constexpr (kEarlyRelease)
Expand Down
90 changes: 74 additions & 16 deletions cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -293,14 +293,53 @@ static constexpr uint32_t fhcSmemSize()
using FusedRoutFn = void (*)(
uint32_t, CUtensorMap, CUtensorMap, CUtensorMap, CUtensorMap, float*, float const*, float const*, float*);

template <uint32_t Hidden, uint32_t KS>
template <uint32_t Hidden, uint32_t KS, uint32_t XS = 1>
static FusedRoutFn fhcInstance()
{
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>;
/*kEarlyRelease=*/false, XS>;
}

template <uint32_t Hidden, uint32_t KS, uint32_t XS>
static FusedRoutFn fhcXSplitInstanceIfSupported()
{
if constexpr (isSupportedFhcMmaKS<Hidden, KS>())
{
return fhcInstance<Hidden, KS, XS>();
}
else
{
TLLM_CHECK_WITH_INFO(false, "mhcFusedHcLaunch: unsupported (kNumSplits=%u, hidden=%u)", KS, Hidden);
return nullptr;
}
}

template <uint32_t Hidden, uint32_t XS>
static FusedRoutFn pickFhcXSplitKs(uint32_t ks)
{
switch (ks)
{
case 1: return fhcInstance<Hidden, 1, XS>();
case 2: return fhcXSplitInstanceIfSupported<Hidden, 2, XS>();
case 4: return fhcXSplitInstanceIfSupported<Hidden, 4, XS>();
case 8: return fhcXSplitInstanceIfSupported<Hidden, 8, XS>();
case 16: return fhcXSplitInstanceIfSupported<Hidden, 16, XS>();
default: TLLM_CHECK_WITH_INFO(false, "mhcFusedHcLaunch: unsupported kNumSplits=%u for split x", ks); return nullptr;
}
}

template <uint32_t Hidden>
static FusedRoutFn pickFhcXSplit(uint32_t ks, uint32_t xs)
{
switch (xs)
{
case 2: return pickFhcXSplitKs<Hidden, 2>(ks);
case 4: return pickFhcXSplitKs<Hidden, 4>(ks);
default: TLLM_CHECK_WITH_INFO(false, "mhcFusedHcLaunch: unsupported x split=%u", xs); return nullptr;
}
}

template <uint32_t Hidden, uint32_t KS>
Expand Down Expand Up @@ -375,7 +414,7 @@ static void mhcFusedHcLaunchImpl(__nv_bfloat16 const* x_prev, __nv_bfloat16 cons
__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, __nv_bfloat16 const* norm_weight, float norm_eps,
cudaStream_t stream)
cudaStream_t stream, int x_num_splits = 1)
{
if (M <= 0)
return;
Expand Down Expand Up @@ -407,7 +446,8 @@ static void mhcFusedHcLaunchImpl(__nv_bfloat16 const* x_prev, __nv_bfloat16 cons
/*swizzleBytes=*/128, sizeof(__nv_bfloat16));

CUtensorMap desc_x = getCachedTma2D(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),
static_cast<uint32_t>(x_num_splits) * m_u, FHC_BLOCK_K, FHC_BLOCK_M,
static_cast<uint64_t>(Hidden) * sizeof(__nv_bfloat16),
/*swizzleBytes=*/128, sizeof(__nv_bfloat16));

CUtensorMap desc_b = getCachedTma2D(const_cast<float*>(w_t), CU_TENSOR_MAP_DATA_TYPE_TFLOAT32, SHAPE_K, FHC_SHAPE_N,
Expand All @@ -419,8 +459,11 @@ static void mhcFusedHcLaunchImpl(__nv_bfloat16 const* x_prev, __nv_bfloat16 cons
/*swizzleBytes=*/128, sizeof(__nv_bfloat16));

// ---- Step 1: fused post-mapping + TF32 GEMM + sqrsum + residual_out ----
constexpr uint32_t fused_smem = fhcSmemSize();
FusedRoutFn fa = pickFhc<Hidden>(ks);
// Split-x adds barriers but reuses the x data buffers.
uint32_t const extra_x_smem = (x_num_splits > 1) ? (2u * FHC_N_INPUT_STG * sizeof(uint64_t)) : 0u;
uint32_t const fused_smem = fhcSmemSize() + extra_x_smem;
FusedRoutFn fa
= (x_num_splits > 1) ? pickFhcXSplit<Hidden>(ks, static_cast<uint32_t>(x_num_splits)) : pickFhc<Hidden>(ks);
TLLM_CUDA_CHECK(cudaFuncSetAttribute(
reinterpret_cast<void const*>(fa), cudaFuncAttributeMaxDynamicSharedMemorySize, fused_smem));

Expand All @@ -443,7 +486,7 @@ void mhcFusedHcLaunch(__nv_bfloat16 const* x_prev, __nv_bfloat16 const* residual
__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, __nv_bfloat16 const* norm_weight, float norm_eps, cudaStream_t stream)
int sinkhorn_repeat, __nv_bfloat16 const* norm_weight, float norm_eps, cudaStream_t stream, int x_num_splits)
{
if (M <= 0)
return;
Expand All @@ -457,13 +500,13 @@ void mhcFusedHcLaunch(__nv_bfloat16 const* x_prev, __nv_bfloat16 const* residual
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, norm_weight, norm_eps, stream);
hc_post_mult_value, sinkhorn_repeat, norm_weight, norm_eps, stream, x_num_splits);
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, norm_weight, norm_eps, stream);
hc_post_mult_value, sinkhorn_repeat, norm_weight, norm_eps, stream, x_num_splits);
return;
default: return;
}
Expand All @@ -481,19 +524,21 @@ void mhcFusedHcLaunch(__nv_bfloat16 const* x_prev, __nv_bfloat16 const* residual
using FmaKsplitFn = void (*)(__nv_bfloat16 const*, __nv_bfloat16 const*, float const*, float const*, float const*,
float*, float*, int, int, int, __nv_bfloat16*);

template <int TN, int KS>
template <int TN, int KS, int XS = 1>
static FmaKsplitFn fhcFmaInstance()
{
return &fused_fma_kernels::fused_pmap_gemm_fma_ksplit<TN, KS, /*BF16_VEC_OVERRIDE=*/0, /*WRITE_RESIDUAL=*/true>;
return &fused_fma_kernels::fused_pmap_gemm_fma_ksplit<TN, KS, /*BF16_VEC_OVERRIDE=*/0,
/*WRITE_RESIDUAL=*/true, XS>;
}

// Valid (tile_n, num_k_splits) combinations the fused_hc FMA path supports.
// Keep this limited to the small/mid-M sweet spots from profile_fair_report v4.
static FmaKsplitFn pickFhcFma(int tile_n, int ks)
template <int XS>
static FmaKsplitFn pickFhcFmaXSplit(int tile_n, int ks)
{
#define FHCFMA_CASE(TN, KS) \
if (tile_n == (TN) && ks == (KS)) \
return fhcFmaInstance<TN, KS>()
return fhcFmaInstance<TN, KS, XS>()

FHCFMA_CASE(1, 1);
FHCFMA_CASE(1, 2);
Expand All @@ -513,16 +558,29 @@ static FmaKsplitFn pickFhcFma(int tile_n, int ks)
FHCFMA_CASE(12, 1);
FHCFMA_CASE(24, 1);
#undef FHCFMA_CASE
TLLM_CHECK_WITH_INFO(false, "mhcFusedHcFmaLaunch: unsupported (tile_n=%d, ks=%d)", tile_n, ks);
TLLM_CHECK_WITH_INFO(false, "mhcFusedHcFmaLaunch: unsupported (tile_n=%d, ks=%d, x_splits=%d)", tile_n, ks, XS);
return nullptr;
}

static FmaKsplitFn pickFhcFma(int tile_n, int ks, int x_num_splits)
{
switch (x_num_splits)
{
case 1: return pickFhcFmaXSplit<1>(tile_n, ks);
case 2: return pickFhcFmaXSplit<2>(tile_n, ks);
case 4: return pickFhcFmaXSplit<4>(tile_n, ks);
default:
TLLM_CHECK_WITH_INFO(false, "mhcFusedHcFmaLaunch: unsupported x_num_splits=%d", x_num_splits);
return nullptr;
}
}

void mhcFusedHcFmaLaunch(__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 tile_n, 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, __nv_bfloat16 const* norm_weight, float norm_eps, cudaStream_t stream)
int sinkhorn_repeat, __nv_bfloat16 const* norm_weight, float norm_eps, cudaStream_t stream, int x_num_splits)
{
if (M <= 0)
return;
Expand All @@ -538,7 +596,7 @@ void mhcFusedHcFmaLaunch(__nv_bfloat16 const* x_prev, __nv_bfloat16 const* resid
int const N = static_cast<int>(FHC_SHAPE_N);

// ---- Step 1: fused pmap + GEMM + sqrsum + residual_cur (FMA ksplit) ----
FmaKsplitFn fn = pickFhcFma(tile_n, num_k_splits);
FmaKsplitFn fn = pickFhcFma(tile_n, num_k_splits, x_num_splits);
dim3 const grid(static_cast<unsigned>(M), static_cast<unsigned>(N / tile_n), static_cast<unsigned>(num_k_splits));
dim3 const block(256);
tensorrt_llm::common::launchWithPdlWhenEnabled("fused_pmap_gemm_fma_ksplit", fn, grid, block, 0, stream,
Expand Down
Loading
Loading