Skip to content

[None][perf] Optimize Blackwell fused MHC half-MMA kernel #16799

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

Merged
mingyangHao merged 6 commits into from
Jul 27, 2026
Merged

[None][perf] Optimize Blackwell fused MHC half-MMA kernel #16799

Changes from 3 commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
c878fac
perf: optimize Blackwell fused MHC half-MMA kernel
MengmSun Jul 23, 2026
342f8d1
Merge branch 'main' into perf/blackwell-mhc-mma
MengmSun Jul 23, 2026
bd9cf29
chore: minimize Blackwell fused MHC MMA diff
MengmSun Jul 24, 2026
6e214bd
perf: update
MengmSun Jul 24, 2026
4a52d3d
[None][chore] Apply clang-format to fused_tf32_pmap_gemm.cuh
MengmSun Jul 26, 2026
fca43ef
[None][chore] Drop mhc_ prefix from packed FP32x2 PTX wrappers
MengmSun Jul 27, 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
239 changes: 145 additions & 94 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 @@ -99,6 +99,65 @@ __device__ __forceinline__ uint32_t get_swizzled_smem_offset(uint32_t const& off
return row * 128 + col * kSwizzleBase;
}

__device__ __forceinline__ uint64_t float2ToBits(float2 const& value)
{
return *reinterpret_cast<uint64_t const*>(&value);
}

__device__ __forceinline__ float2 bitsToFloat2(uint64_t const value)
{
return *reinterpret_cast<float2 const*>(&value);
}

__device__ __forceinline__ float2 multiplyFloat2(float const scalar, float2 const value)
{
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) && (__CUDA_ARCH__ < 1100)
uint64_t result;
float2 const scalar2{scalar, scalar};
asm("mul.rn.ftz.f32x2 %0, %1, %2;" : "=l"(result) : "l"(float2ToBits(scalar2)), "l"(float2ToBits(value)));
return bitsToFloat2(result);
#else
return float2{scalar * value.x, scalar * value.y};
#endif
}

__device__ __forceinline__ float2 fmaFloat2(float const scalar, float2 const value, float2 const accumulator)
{
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) && (__CUDA_ARCH__ < 1100)
uint64_t result;
float2 const scalar2{scalar, scalar};
asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;"
: "=l"(result)
: "l"(float2ToBits(scalar2)), "l"(float2ToBits(value)), "l"(float2ToBits(accumulator)));
return bitsToFloat2(result);
#else
return float2{fmaf(scalar, value.x, accumulator.x), fmaf(scalar, value.y, accumulator.y)};
#endif
}

__device__ __forceinline__ float2 fmaFloat2(float2 const lhs, float2 const rhs, float2 const accumulator)
{
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) && (__CUDA_ARCH__ < 1100)
uint64_t result;
asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;"
: "=l"(result)
: "l"(float2ToBits(lhs)), "l"(float2ToBits(rhs)), "l"(float2ToBits(accumulator)));
return bitsToFloat2(result);
#else
return float2{fmaf(lhs.x, rhs.x, accumulator.x), fmaf(lhs.y, rhs.y, accumulator.y)};
#endif
}

__device__ __forceinline__ float4 loadCoefficientRow(
float const* base, uint32_t const row, uint32_t const rowStride, uint32_t const vectorIndex, bool const valid)
{
if (!valid)
{
return float4{0.0F, 0.0F, 0.0F, 0.0F};
}
return __ldg(reinterpret_cast<float4 const*>(base + row * rowStride) + vectorIndex);
}

__device__ __forceinline__ void stsm_x4_b16_rout(void* smem_dst, uint32_t a, uint32_t b, uint32_t c, uint32_t d)
{
asm volatile(
Expand Down Expand Up @@ -178,10 +237,7 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) fused_tf3
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;
auto smem_post = reinterpret_cast<float*>(cursor);
cursor += SMEM_POST_SIZE;
auto smem_comb = reinterpret_cast<float*>(cursor);
cursor += SMEM_COMB_SIZE;
cursor += SMEM_POST_SIZE + SMEM_COMB_SIZE;
auto smem_rc = reinterpret_cast<nv_bfloat16*>(cursor); // [HC_MULT][BLOCK_M][BLOCK_K] bf16, single-buffered
cursor += SMEM_RC_SIZE;

Expand Down Expand Up @@ -237,39 +293,6 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) fused_tf3
const uint32_t h_tile_start = k_split_idx * H_TILES_PER_SPLIT;
constexpr uint32_t num_total_stages = H_TILES_PER_SPLIT * HC_MULT;

// Prologue: pmap warp group loads post_mix, comb_mix into SMEM
if (warp_idx >= kNumMMAThreads / 32)
{
const uint32_t pmap_tid = threadIdx.x - kNumMMAThreads;
#pragma unroll
for (uint32_t t = 0; t < 2; ++t)
{
uint32_t idx = pmap_tid + t * kNumPmapThreads;
if (idx < BLOCK_M * HC_MULT)
{
uint32_t m = idx / HC_MULT;
uint32_t hc = idx % HC_MULT;
uint32_t gmem_m = m_offset + m;
float v = (gmem_m < shape_m) ? post_mix[gmem_m * HC_MULT + hc] : 0.f;
smem_post[idx] = v;
}
}
#pragma unroll
for (uint32_t t = 0; t < 8; ++t)
{
uint32_t idx = pmap_tid + t * kNumPmapThreads;
if (idx < BLOCK_M * HC_MULT * HC_MULT)
{
uint32_t m = idx / (HC_MULT * HC_MULT);
uint32_t jk = idx % (HC_MULT * HC_MULT);
uint32_t gmem_m = m_offset + m;
float v = (gmem_m < shape_m) ? comb_mix[gmem_m * HC_MULT * HC_MULT + jk] : 0.f;
smem_comb[idx] = v;
}
}
}
__syncthreads();

if (warp_idx < kNumMMAThreads / 32)
{
// ----- TMA warp (warp 0) -----
Expand Down Expand Up @@ -417,24 +440,26 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) fused_tf3

float pm_u[HC_MULT], pm_l[HC_MULT];
float cm_u[HC_MULT][HC_MULT], cm_l[HC_MULT][HC_MULT];
#pragma unroll
for (uint32_t hc = 0; hc < HC_MULT; ++hc)
{
pm_u[hc] = smem_post[upper_row * HC_MULT + hc];
pm_l[hc] = smem_post[lower_row * HC_MULT + hc];
}
#pragma unroll
for (uint32_t j = 0; j < HC_MULT; ++j)
{
static_assert(HC_MULT == 4, "float4 row loads require HC_MULT == 4");
uint32_t const globalUpperRow = m_offset + upper_row;
uint32_t const globalLowerRow = m_offset + lower_row;
bool const upperRowValid = globalUpperRow < shape_m;
bool const lowerRowValid = globalLowerRow < shape_m;
*reinterpret_cast<float4*>(pm_u) = loadCoefficientRow(post_mix, globalUpperRow, HC_MULT, 0, upperRowValid);
*reinterpret_cast<float4*>(pm_l) = loadCoefficientRow(post_mix, globalLowerRow, HC_MULT, 0, lowerRowValid);
#pragma unroll
for (uint32_t hc = 0; hc < HC_MULT; ++hc)
for (uint32_t j = 0; j < HC_MULT; ++j)
{
cm_u[j][hc] = smem_comb[upper_row * HC_MULT * HC_MULT + j * HC_MULT + hc];
cm_l[j][hc] = smem_comb[lower_row * HC_MULT * HC_MULT + j * HC_MULT + hc];
*reinterpret_cast<float4*>(cm_u[j])
= loadCoefficientRow(comb_mix, globalUpperRow, HC_MULT * HC_MULT, j, upperRowValid);
*reinterpret_cast<float4*>(cm_l[j])
= loadCoefficientRow(comb_mix, globalLowerRow, HC_MULT * HC_MULT, j, lowerRowValid);
}
}

float sqr_u = 0.f, sqr_l = 0.f;
float2 squareSumUpper{0.0F, 0.0F};
float2 squareSumLower{0.0F, 0.0F};
constexpr uint32_t kNumBankGroupBytes = 16;
constexpr uint32_t kNumElemsPerBankGroup = kNumBankGroupBytes / sizeof(nv_bfloat16);
constexpr uint32_t kNumLoads = BLOCK_K / kNumElemsPerBankGroup;
Expand Down Expand Up @@ -493,71 +518,91 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) fused_tf3
empty_input[i_stage]->arrive();
}

// Wait for previous ht's residual_out TMA_STOREs to drain before we
// overwrite single-buffered smem_rc with new hc values.
// All four cast stages share the same phase at a tile boundary.
static_assert(kNumCastStages == HC_MULT, "The tile-batched cast protocol requires one stage per HC");
#pragma unroll
for (uint32_t castStage = 0; castStage < kNumCastStages; ++castStage)
{
empty_cast[castStage]->wait(((s / kNumCastStages) & 1) ^ 1);
}

// smem_rc is single-buffered, so the previous tile's stores must
// complete before this tile overwrites it.
if (ht > 0)
{
cute::tma_store_wait<0>();
}

#pragma unroll
for (uint32_t hc = 0; hc < HC_MULT; ++hc)
for (uint32_t i = 0; i < kNumLoads; i += 2)
{
const uint32_t cast_stage_idx = s % kNumCastStages;
empty_cast[cast_stage_idx]->wait(((s / kNumCastStages) & 1) ^ 1);

uint32_t rc_u_buf[kNumLoads], rc_l_buf[kNumLoads];
float2 residualUpper[HC_MULT][2];
float2 residualLower[HC_MULT][2];
#pragma unroll
for (uint32_t i = 0; i < kNumLoads; ++i)
for (uint32_t j = 0; j < HC_MULT; ++j)
{
float2 nu{pm_u[hc] * xf[0][i].x, pm_u[hc] * xf[0][i].y};
float2 nl{pm_l[hc] * xf[1][i].x, pm_l[hc] * xf[1][i].y};
residualUpper[j][0] = __bfloat1622float2(*reinterpret_cast<nv_bfloat162*>(&r_vals[j][0][i]));
residualUpper[j][1] = __bfloat1622float2(*reinterpret_cast<nv_bfloat162*>(&r_vals[j][0][i + 1]));
residualLower[j][0] = __bfloat1622float2(*reinterpret_cast<nv_bfloat162*>(&r_vals[j][1][i]));
residualLower[j][1] = __bfloat1622float2(*reinterpret_cast<nv_bfloat162*>(&r_vals[j][1][i + 1]));
}
#pragma unroll
for (uint32_t hc = 0; hc < HC_MULT; ++hc)
{
float2 newUpper0 = multiplyFloat2(pm_u[hc], xf[0][i]);
float2 newUpper1 = multiplyFloat2(pm_u[hc], xf[0][i + 1]);
float2 newLower0 = multiplyFloat2(pm_l[hc], xf[1][i]);
float2 newLower1 = multiplyFloat2(pm_l[hc], xf[1][i + 1]);
#pragma unroll
for (uint32_t j = 0; j < HC_MULT; ++j)
{
float2 ruj = __bfloat1622float2(*reinterpret_cast<nv_bfloat162*>(&r_vals[j][0][i]));
float2 rlj = __bfloat1622float2(*reinterpret_cast<nv_bfloat162*>(&r_vals[j][1][i]));
nu.x = fmaf(cm_u[j][hc], ruj.x, nu.x);
nu.y = fmaf(cm_u[j][hc], ruj.y, nu.y);
nl.x = fmaf(cm_l[j][hc], rlj.x, nl.x);
nl.y = fmaf(cm_l[j][hc], rlj.y, nl.y);
newUpper0 = fmaFloat2(cm_u[j][hc], residualUpper[j][0], newUpper0);
newUpper1 = fmaFloat2(cm_u[j][hc], residualUpper[j][1], newUpper1);
newLower0 = fmaFloat2(cm_l[j][hc], residualLower[j][0], newLower0);
newLower1 = fmaFloat2(cm_l[j][hc], residualLower[j][1], newLower1);
}
nv_bfloat162 b_up = __float22bfloat162_rn(nu);
nv_bfloat162 b_lo = __float22bfloat162_rn(nl);
uint32_t b_up_bits = *reinterpret_cast<uint32_t*>(&b_up);
uint32_t b_lo_bits = *reinterpret_cast<uint32_t*>(&b_lo);
rc_u_buf[i] = b_up_bits;
rc_l_buf[i] = b_lo_bits;
float2 ru = __bfloat1622float2(b_up);
float2 rl = __bfloat1622float2(b_lo);
sqr_u = fmaf(ru.x, ru.x, sqr_u);
sqr_u = fmaf(ru.y, ru.y, sqr_u);
sqr_l = fmaf(rl.x, rl.x, sqr_l);
sqr_l = fmaf(rl.y, rl.y, sqr_l);
cute::SM100_TMEM_STORE_16dp256b1x::copy(*reinterpret_cast<uint32_t*>(&ru.x),
*reinterpret_cast<uint32_t*>(&ru.y), *reinterpret_cast<uint32_t*>(&rl.x),
*reinterpret_cast<uint32_t*>(&rl.y), cast_stage_idx * BLOCK_K + i * 8);
}
cutlass::arch::fence_view_async_tmem_store();
tcgen05_before_thread_sync();
full_cast[cast_stage_idx]->arrive();
++s;

// STSM bf16 new_r values into smem_rc[hc] sub-region for this warp.
uint8_t* rc_base = reinterpret_cast<uint8_t*>(smem_rc) + hc * SMEM_RC_PER_HC
+ sub_warp_idx * BLOCK_M_PER_WARP * kSwizzleRoutMode;
#pragma unroll
for (uint32_t i = 0; i < kNumLoads; i += 2)
{
nv_bfloat162 const bf16Upper0 = __float22bfloat162_rn(newUpper0);
nv_bfloat162 const bf16Upper1 = __float22bfloat162_rn(newUpper1);
nv_bfloat162 const bf16Lower0 = __float22bfloat162_rn(newLower0);
nv_bfloat162 const bf16Lower1 = __float22bfloat162_rn(newLower1);
float2 const roundedUpper0 = __bfloat1622float2(bf16Upper0);
float2 const roundedUpper1 = __bfloat1622float2(bf16Upper1);
float2 const roundedLower0 = __bfloat1622float2(bf16Lower0);
float2 const roundedLower1 = __bfloat1622float2(bf16Lower1);
squareSumUpper = fmaFloat2(roundedUpper0, roundedUpper0, squareSumUpper);
squareSumUpper = fmaFloat2(roundedUpper1, roundedUpper1, squareSumUpper);
squareSumLower = fmaFloat2(roundedLower0, roundedLower0, squareSumLower);
squareSumLower = fmaFloat2(roundedLower1, roundedLower1, squareSumLower);

cute::SM100_TMEM_STORE_16dp256b1x::copy(*reinterpret_cast<uint32_t const*>(&roundedUpper0.x),
*reinterpret_cast<uint32_t const*>(&roundedUpper0.y),
*reinterpret_cast<uint32_t const*>(&roundedLower0.x),
*reinterpret_cast<uint32_t const*>(&roundedLower0.y), hc * BLOCK_K + i * 8);
cute::SM100_TMEM_STORE_16dp256b1x::copy(*reinterpret_cast<uint32_t const*>(&roundedUpper1.x),
*reinterpret_cast<uint32_t const*>(&roundedUpper1.y),
*reinterpret_cast<uint32_t const*>(&roundedLower1.x),
*reinterpret_cast<uint32_t const*>(&roundedLower1.y), hc * BLOCK_K + (i + 1) * 8);

uint8_t* rc_base = reinterpret_cast<uint8_t*>(smem_rc) + hc * SMEM_RC_PER_HC
+ sub_warp_idx * BLOCK_M_PER_WARP * kSwizzleRoutMode;
auto smem_ptr
= rc_base + get_swizzled_smem_offset<kSwizzleRoutMode>(i + lane_idx / 16, lane_idx % 16);
stsm_x4_b16_rout(smem_ptr, rc_u_buf[i + 0], rc_l_buf[i + 0], rc_u_buf[i + 1], rc_l_buf[i + 1]);
stsm_x4_b16_rout(smem_ptr, *reinterpret_cast<uint32_t const*>(&bf16Upper0),
*reinterpret_cast<uint32_t const*>(&bf16Lower0),
*reinterpret_cast<uint32_t const*>(&bf16Upper1),
*reinterpret_cast<uint32_t const*>(&bf16Lower1));
}
}
if constexpr (!kEarlyRelease)

cutlass::arch::fence_view_async_tmem_store();
tcgen05_before_thread_sync();
#pragma unroll
for (uint32_t castStage = 0; castStage < kNumCastStages; ++castStage)
{
empty_input[i_stage]->arrive();
full_cast[castStage]->arrive();
}
s += kNumCastStages;

// Emit HC_MULT TMA_STOREs of residual_cur: one per hc slice, per-warp rows.
cute::tma_store_fence();
Expand All @@ -574,12 +619,18 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) fused_tf3
cute::tma_store_arrive();
}
}
if constexpr (!kEarlyRelease)
{
empty_input[i_stage]->arrive();
}
}

// Drain any in-flight residual_out TMA stores before exit.
cute::tma_store_wait<0>();

// Warp-reduce sqr across 4 col_lanes then atomicAdd to global.
float sqr_u = squareSumUpper.x + squareSumUpper.y;
float sqr_l = squareSumLower.x + squareSumLower.y;
sqr_u += __shfl_xor_sync(0xffffffff, sqr_u, 1);
sqr_u += __shfl_xor_sync(0xffffffff, sqr_u, 2);
sqr_l += __shfl_xor_sync(0xffffffff, sqr_l, 1);
Expand Down