Skip to content

Commit 5b16725

Browse files
mingyangHaomingyangHao
committed
[None][perf] fuse DSV4 FP8 o-projection split-K into mHC
Fuse o_a FP8 and UE8M0 quantization into the CuTe DSL epilogue. Add a TRT-LLM-owned SM100 O_b split-K kernel using DeepGEMM MMA/TMA building blocks, and fold split-partial reduction into the mHC half-MMA and half-FMA paths. Keep the optimization opt-in and cover dispatch, producer correctness, and fused reduction. Signed-off-by: Mingyang Hao <200044211+mingyangHao@users.noreply.github.com>
1 parent ec3da87 commit 5b16725

18 files changed

Lines changed: 2560 additions & 161 deletions

File tree

‎cpp/tensorrt_llm/kernels/mhcKernels/CMakeLists.txt‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -19,17 +19,16 @@
1919
# this directory (bench_*.cu, probe_*.cu, profile_*.cu, test_*.cu) are
2020
# standalone binaries with their own main() and must NOT be compiled into the
2121
# trtllm shared library.
22-
set(SRC_CU mhcKernels.cu mhcFusedHcKernel.cu)
22+
set(SRC_CU dsv4Fp8SplitKGemm.cu mhcKernels.cu mhcFusedHcKernel.cu)
2323
add_library(mhcKernels_src OBJECT ${SRC_CU})
2424

2525
set_property(TARGET mhcKernels_src PROPERTY POSITION_INDEPENDENT_CODE ON)
2626
set_property(TARGET mhcKernels_src PROPERTY CUDA_RESOLVE_DEVICE_SYMBOLS ON)
2727
target_compile_options(mhcKernels_src
2828
PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:--use_fast_math>)
2929

30-
# Expose DeepGEMM helper headers (sm100_utils.cuh, utils.cuh, tma_utils.cuh,
31-
# reduction.cuh) Used by the tcgen05.mma-based fused post-mapping + GEMM kernel
32-
# on SM100.
30+
# Expose the DeepGEMM SM100 MMA/TMA building blocks used by the fused
31+
# post-mapping kernels and DSV4 O_b split-K GEMM.
3332
target_include_directories(
3433
mhcKernels_src
3534
PRIVATE ${CMAKE_BINARY_DIR}/_deps/deepgemm-src/deep_gemm/include)

‎cpp/tensorrt_llm/kernels/mhcKernels/dsv4Fp8SplitKGemm.cu‎

Lines changed: 208 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,208 @@
1+
/*
2+
* Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
#include "dsv4Fp8SplitKGemm.cuh"
18+
#include "mhcKernels.h"
19+
20+
#include "tensorrt_llm/common/assert.h"
21+
#include "tensorrt_llm/common/cudaUtils.h"
22+
23+
#include <cuda.h>
24+
#include <cuda_runtime.h>
25+
26+
#include <algorithm>
27+
#include <cstdint>
28+
29+
TRTLLM_NAMESPACE_BEGIN
30+
31+
namespace kernels::mhc
32+
{
33+
namespace
34+
{
35+
36+
constexpr int kDsv4ProN = 7168;
37+
constexpr int kDsv4ProK = 16384;
38+
constexpr int kBlockN = 128;
39+
constexpr int kBlockK = 128;
40+
constexpr int kClusterSize = 2;
41+
constexpr int kNumSmsForSwizzle = 148;
42+
constexpr int kNumThreads = 256;
43+
constexpr int kSmemCapacity = 232448;
44+
constexpr int kMaxPipelineStages = 32;
45+
46+
CUtensorMap makeTma2D(void const* base, CUtensorMapDataType dtype, uint64_t gmemInner, uint64_t gmemOuter,
47+
uint32_t smemInner, uint32_t smemOuter, uint64_t gmemOuterStrideBytes, CUtensorMapSwizzle swizzle)
48+
{
49+
if (swizzle != CU_TENSOR_MAP_SWIZZLE_NONE)
50+
{
51+
uint32_t const elementBytes = dtype == CU_TENSOR_MAP_DATA_TYPE_BFLOAT16 ? 2u
52+
: dtype == CU_TENSOR_MAP_DATA_TYPE_INT32 ? 4u
53+
: 1u;
54+
smemInner = 128u / elementBytes;
55+
}
56+
57+
CUtensorMap tensorMap;
58+
uint64_t const gmemDims[2] = {gmemInner, gmemOuter};
59+
uint64_t const gmemStrides[1] = {gmemOuterStrideBytes};
60+
uint32_t const smemDims[2] = {smemInner, smemOuter};
61+
uint32_t const elementStrides[2] = {1, 1};
62+
CUresult const result = cuTensorMapEncodeTiled(&tensorMap, dtype, 2, const_cast<void*>(base), gmemDims, gmemStrides,
63+
smemDims, elementStrides, CU_TENSOR_MAP_INTERLEAVE_NONE, swizzle, CU_TENSOR_MAP_L2_PROMOTION_L2_256B,
64+
CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE);
65+
TLLM_CHECK_WITH_INFO(
66+
result == CUDA_SUCCESS, "cuTensorMapEncodeTiled failed for DSV4 O_b split-K (%d)", static_cast<int>(result));
67+
return tensorMap;
68+
}
69+
70+
template <int BlockM>
71+
constexpr int numPipelineStages()
72+
{
73+
constexpr int smemCd = 2 * 16 * kBlockN * static_cast<int>(sizeof(__nv_bfloat16));
74+
constexpr int smemBarrierReserve = (kMaxPipelineStages * 3 + 4) * 8;
75+
constexpr int smemTmemPointer = 4;
76+
constexpr int smemPerStage
77+
= (BlockM / kClusterSize) * kBlockK + kBlockN * kBlockK + 2 * 128 * static_cast<int>(sizeof(uint32_t));
78+
constexpr int stages = (kSmemCapacity - smemCd - smemBarrierReserve - smemTmemPointer) / smemPerStage;
79+
return stages < kMaxPipelineStages ? stages : kMaxPipelineStages;
80+
}
81+
82+
template <int BlockM>
83+
constexpr int dynamicSmemSize()
84+
{
85+
constexpr int smemCd = 2 * 16 * kBlockN * static_cast<int>(sizeof(__nv_bfloat16));
86+
constexpr int smemBarrierReserve = (kMaxPipelineStages * 3 + 4) * 8;
87+
constexpr int smemTmemPointer = 4;
88+
constexpr int smemPerStage
89+
= (BlockM / kClusterSize) * kBlockK + kBlockN * kBlockK + 2 * 128 * static_cast<int>(sizeof(uint32_t));
90+
return smemCd + numPipelineStages<BlockM>() * smemPerStage + smemBarrierReserve + smemTmemPointer;
91+
}
92+
93+
template <int BlockM, int SplitK>
94+
void launchInstance(CUtensorMap tensorMapA, CUtensorMap tensorMapB, CUtensorMap tensorMapSfa, CUtensorMap tensorMapSfb,
95+
CUtensorMap tensorMapD, __nv_bfloat16* partials, int M, int numSms, cudaStream_t stream)
96+
{
97+
using namespace deep_gemm;
98+
constexpr int kNumStages = numPipelineStages<BlockM>();
99+
constexpr int kDynamicSmem = dynamicSmemSize<BlockM>();
100+
static_assert(kNumStages > 0, "DSV4 O_b split-K pipeline has no shared-memory stages");
101+
static_assert(kDynamicSmem <= kSmemCapacity, "DSV4 O_b split-K shared-memory budget exceeded");
102+
103+
auto kernel = &dsv4_splitk::dsv4Fp8SplitKGemmKernel<cute::UMMA::Major::K, cute::UMMA::Major::K, 128, 128, 0,
104+
kDsv4ProN, kDsv4ProK, BlockM, kBlockN, kBlockK, 1, 128, 128, 128, kNumStages, 128, 128, kClusterSize, true,
105+
kNumSmsForSwizzle, true, GemmType::Normal, false, cutlass::float_e4m3_t, cutlass::float_e4m3_t,
106+
cutlass::bfloat16_t, epilogue::transform::EpilogueIdentity, SplitK>;
107+
108+
TLLM_CUDA_CHECK(cudaFuncSetAttribute(
109+
reinterpret_cast<void const*>(kernel), cudaFuncAttributeMaxDynamicSharedMemorySize, kDynamicSmem));
110+
111+
cudaLaunchConfig_t config{};
112+
config.gridDim = dim3(numSms, 1, 1);
113+
config.blockDim = dim3(kNumThreads, 1, 1);
114+
config.dynamicSmemBytes = kDynamicSmem;
115+
config.stream = stream;
116+
cudaLaunchAttribute attribute{};
117+
attribute.id = cudaLaunchAttributeClusterDimension;
118+
attribute.val.clusterDim.x = kClusterSize;
119+
attribute.val.clusterDim.y = 1;
120+
attribute.val.clusterDim.z = 1;
121+
config.attrs = &attribute;
122+
config.numAttrs = 1;
123+
124+
TLLM_CUDA_CHECK(cudaLaunchKernelEx(&config, kernel, static_cast<int*>(nullptr), static_cast<uint32_t>(M),
125+
static_cast<uint32_t>(kDsv4ProN), static_cast<uint32_t>(kDsv4ProK), tensorMapA, tensorMapB, tensorMapSfa,
126+
tensorMapSfb, tensorMapD, static_cast<void*>(partials)));
127+
}
128+
129+
template <int SplitK>
130+
void dispatchBlockM(int blockM, CUtensorMap tensorMapA, CUtensorMap tensorMapB, CUtensorMap tensorMapSfa,
131+
CUtensorMap tensorMapSfb, CUtensorMap tensorMapD, __nv_bfloat16* partials, int M, int numSms, cudaStream_t stream)
132+
{
133+
#define TRTLLM_DSV4_SPLITK_BLOCK_M(BLOCK_M) \
134+
case BLOCK_M: \
135+
launchInstance<BLOCK_M, SplitK>( \
136+
tensorMapA, tensorMapB, tensorMapSfa, tensorMapSfb, tensorMapD, partials, M, numSms, stream); \
137+
return
138+
switch (blockM)
139+
{
140+
TRTLLM_DSV4_SPLITK_BLOCK_M(16);
141+
TRTLLM_DSV4_SPLITK_BLOCK_M(32);
142+
TRTLLM_DSV4_SPLITK_BLOCK_M(48);
143+
TRTLLM_DSV4_SPLITK_BLOCK_M(64);
144+
TRTLLM_DSV4_SPLITK_BLOCK_M(80);
145+
TRTLLM_DSV4_SPLITK_BLOCK_M(96);
146+
TRTLLM_DSV4_SPLITK_BLOCK_M(112);
147+
TRTLLM_DSV4_SPLITK_BLOCK_M(128);
148+
default: TLLM_CHECK_WITH_INFO(false, "unsupported DSV4 O_b split-K block M: %d", blockM);
149+
}
150+
#undef TRTLLM_DSV4_SPLITK_BLOCK_M
151+
}
152+
153+
} // namespace
154+
155+
void dsv4Fp8SplitKGemmLaunch(void const* a, int32_t const* sfa, void const* b, int32_t const* sfb,
156+
__nv_bfloat16* partials, int M, int N, int K, int sfaKStride, int sfbKStride, int numSplits, cudaStream_t stream)
157+
{
158+
TLLM_CHECK_WITH_INFO(a != nullptr && sfa != nullptr && b != nullptr && sfb != nullptr && partials != nullptr,
159+
"DSV4 O_b split-K received a null pointer");
160+
TLLM_CHECK_WITH_INFO(M > 0, "DSV4 O_b split-K requires a positive M, got %d", M);
161+
TLLM_CHECK_WITH_INFO(N == kDsv4ProN && K == kDsv4ProK, "DSV4 O_b split-K supports N=%d and K=%d, got N=%d K=%d",
162+
kDsv4ProN, kDsv4ProK, N, K);
163+
TLLM_CHECK_WITH_INFO(
164+
numSplits == 2 || numSplits == 4, "DSV4 O_b split-K requires 2 or 4 splits, got %d", numSplits);
165+
TLLM_CHECK_WITH_INFO(sfaKStride >= M && sfaKStride % 4 == 0,
166+
"DSV4 O_b split-K activation scale stride must be a 4-aligned value >= M, got %d", sfaKStride);
167+
TLLM_CHECK_WITH_INFO(
168+
sfbKStride == N, "DSV4 O_b split-K weight scale stride must equal N=%d, got %d", N, sfbKStride);
169+
170+
int device = 0;
171+
TLLM_CUDA_CHECK(cudaGetDevice(&device));
172+
int major = 0;
173+
int minor = 0;
174+
int numSms = 0;
175+
TLLM_CUDA_CHECK(cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, device));
176+
TLLM_CUDA_CHECK(cudaDeviceGetAttribute(&minor, cudaDevAttrComputeCapabilityMinor, device));
177+
TLLM_CUDA_CHECK(cudaDeviceGetAttribute(&numSms, cudaDevAttrMultiProcessorCount, device));
178+
TLLM_CHECK_WITH_INFO(major == 10, "DSV4 O_b split-K requires the SM100 family, got SM%d%d", major, minor);
179+
TLLM_CHECK_WITH_INFO(
180+
numSms % kClusterSize == 0, "DSV4 O_b split-K requires an even SM count for 2-CTA clusters, got %d", numSms);
181+
182+
int const blockM = std::min(((M + 15) / 16) * 16, 128);
183+
CUtensorMap const tensorMapA = makeTma2D(a, CU_TENSOR_MAP_DATA_TYPE_UINT8, K, M, kBlockK, blockM / kClusterSize,
184+
static_cast<uint64_t>(K), CU_TENSOR_MAP_SWIZZLE_128B);
185+
CUtensorMap const tensorMapB = makeTma2D(
186+
b, CU_TENSOR_MAP_DATA_TYPE_UINT8, K, N, kBlockK, kBlockN, static_cast<uint64_t>(K), CU_TENSOR_MAP_SWIZZLE_128B);
187+
CUtensorMap const tensorMapSfa = makeTma2D(sfa, CU_TENSOR_MAP_DATA_TYPE_INT32, sfaKStride, K / 512, blockM, 1,
188+
static_cast<uint64_t>(sfaKStride) * sizeof(int32_t), CU_TENSOR_MAP_SWIZZLE_NONE);
189+
CUtensorMap const tensorMapSfb = makeTma2D(sfb, CU_TENSOR_MAP_DATA_TYPE_INT32, sfbKStride, K / 512, kBlockN, 1,
190+
static_cast<uint64_t>(sfbKStride) * sizeof(int32_t), CU_TENSOR_MAP_SWIZZLE_NONE);
191+
CUtensorMap const tensorMapD = makeTma2D(partials, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, N, M, kBlockN, 16,
192+
static_cast<uint64_t>(N) * sizeof(__nv_bfloat16), CU_TENSOR_MAP_SWIZZLE_128B);
193+
194+
if (numSplits == 2)
195+
{
196+
dispatchBlockM<2>(
197+
blockM, tensorMapA, tensorMapB, tensorMapSfa, tensorMapSfb, tensorMapD, partials, M, numSms, stream);
198+
}
199+
else
200+
{
201+
dispatchBlockM<4>(
202+
blockM, tensorMapA, tensorMapB, tensorMapSfa, tensorMapSfb, tensorMapD, partials, M, numSms, stream);
203+
}
204+
}
205+
206+
} // namespace kernels::mhc
207+
208+
TRTLLM_NAMESPACE_END

0 commit comments

Comments
 (0)