From 6c0eaa6c20017fbdf21460896f559736e4d7a8d3 Mon Sep 17 00:00:00 2001 From: Mingyang Hao <200044211+mingyangHao@users.noreply.github.com> Date: Mon, 29 Jun 2026 02:33:52 -0700 Subject: [PATCH 1/7] [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> --- .../kernels/mhcKernels/CMakeLists.txt | 7 +- .../kernels/mhcKernels/dsv4Fp8SplitKGemm.cu | 208 ++++++ .../kernels/mhcKernels/dsv4Fp8SplitKGemm.cuh | 670 ++++++++++++++++++ .../mhcKernels/dsv4Fp8SplitKScheduler.cuh | 142 ++++ .../mhcKernels/fused_tf32_pmap_gemm.cuh | 111 ++- .../kernels/mhcKernels/mhcFusedHcKernel.cu | 150 +++- .../kernels/mhcKernels/mhcKernels.h | 20 +- .../kernels/mhcKernels/mhc_fused_fma.cuh | 70 +- cpp/tensorrt_llm/thop/CMakeLists.txt | 1 + cpp/tensorrt_llm/thop/dsv4Fp8SplitKGemmOp.cpp | 82 +++ cpp/tensorrt_llm/thop/mhcOp.cpp | 46 +- .../_torch/custom_ops/cute_dsl_custom_ops.py | 165 +++++ .../blockwise_gemm/blockwise_gemm.py | 554 ++++++++++++++- tensorrt_llm/_torch/modules/attention.py | 90 +++ .../_torch/modules/mhc/hyper_connection.py | 14 +- tensorrt_llm/_torch/modules/mhc/mhc_cuda.py | 87 ++- .../deepseek_v4/test_deepseek_v4_o_proj.py | 236 +++++- tests/unittest/_torch/modules/test_mhc.py | 57 ++ 18 files changed, 2549 insertions(+), 161 deletions(-) create mode 100644 cpp/tensorrt_llm/kernels/mhcKernels/dsv4Fp8SplitKGemm.cu create mode 100644 cpp/tensorrt_llm/kernels/mhcKernels/dsv4Fp8SplitKGemm.cuh create mode 100644 cpp/tensorrt_llm/kernels/mhcKernels/dsv4Fp8SplitKScheduler.cuh create mode 100644 cpp/tensorrt_llm/thop/dsv4Fp8SplitKGemmOp.cpp diff --git a/cpp/tensorrt_llm/kernels/mhcKernels/CMakeLists.txt b/cpp/tensorrt_llm/kernels/mhcKernels/CMakeLists.txt index 58dc760fcf9d..5ec2c5cb7a93 100644 --- a/cpp/tensorrt_llm/kernels/mhcKernels/CMakeLists.txt +++ b/cpp/tensorrt_llm/kernels/mhcKernels/CMakeLists.txt @@ -19,7 +19,7 @@ # this directory (bench_*.cu, probe_*.cu, profile_*.cu, test_*.cu) are # standalone binaries with their own main() and must NOT be compiled into the # trtllm shared library. -set(SRC_CU mhcKernels.cu mhcFusedHcKernel.cu) +set(SRC_CU dsv4Fp8SplitKGemm.cu mhcKernels.cu mhcFusedHcKernel.cu) add_library(mhcKernels_src OBJECT ${SRC_CU}) set_property(TARGET mhcKernels_src PROPERTY POSITION_INDEPENDENT_CODE ON) @@ -27,9 +27,8 @@ set_property(TARGET mhcKernels_src PROPERTY CUDA_RESOLVE_DEVICE_SYMBOLS ON) target_compile_options(mhcKernels_src PRIVATE $<$:--use_fast_math>) -# Expose DeepGEMM helper headers (sm100_utils.cuh, utils.cuh, tma_utils.cuh, -# reduction.cuh) Used by the tcgen05.mma-based fused post-mapping + GEMM kernel -# on SM100. +# Expose the DeepGEMM SM100 MMA/TMA building blocks used by the fused +# post-mapping kernels and DSV4 O_b split-K GEMM. target_include_directories( mhcKernels_src PRIVATE ${CMAKE_BINARY_DIR}/_deps/deepgemm-src/deep_gemm/include) diff --git a/cpp/tensorrt_llm/kernels/mhcKernels/dsv4Fp8SplitKGemm.cu b/cpp/tensorrt_llm/kernels/mhcKernels/dsv4Fp8SplitKGemm.cu new file mode 100644 index 000000000000..5597fce8d6b2 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/mhcKernels/dsv4Fp8SplitKGemm.cu @@ -0,0 +1,208 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "dsv4Fp8SplitKGemm.cuh" +#include "mhcKernels.h" + +#include "tensorrt_llm/common/assert.h" +#include "tensorrt_llm/common/cudaUtils.h" + +#include +#include + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::mhc +{ +namespace +{ + +constexpr int kDsv4ProN = 7168; +constexpr int kDsv4ProK = 16384; +constexpr int kBlockN = 128; +constexpr int kBlockK = 128; +constexpr int kClusterSize = 2; +constexpr int kNumSmsForSwizzle = 148; +constexpr int kNumThreads = 256; +constexpr int kSmemCapacity = 232448; +constexpr int kMaxPipelineStages = 32; + +CUtensorMap makeTma2D(void const* base, CUtensorMapDataType dtype, uint64_t gmemInner, uint64_t gmemOuter, + uint32_t smemInner, uint32_t smemOuter, uint64_t gmemOuterStrideBytes, CUtensorMapSwizzle swizzle) +{ + if (swizzle != CU_TENSOR_MAP_SWIZZLE_NONE) + { + uint32_t const elementBytes = dtype == CU_TENSOR_MAP_DATA_TYPE_BFLOAT16 ? 2u + : dtype == CU_TENSOR_MAP_DATA_TYPE_INT32 ? 4u + : 1u; + smemInner = 128u / elementBytes; + } + + CUtensorMap tensorMap; + uint64_t const gmemDims[2] = {gmemInner, gmemOuter}; + uint64_t const gmemStrides[1] = {gmemOuterStrideBytes}; + uint32_t const smemDims[2] = {smemInner, smemOuter}; + uint32_t const elementStrides[2] = {1, 1}; + CUresult const result = cuTensorMapEncodeTiled(&tensorMap, dtype, 2, const_cast(base), gmemDims, gmemStrides, + smemDims, elementStrides, CU_TENSOR_MAP_INTERLEAVE_NONE, swizzle, CU_TENSOR_MAP_L2_PROMOTION_L2_256B, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TLLM_CHECK_WITH_INFO( + result == CUDA_SUCCESS, "cuTensorMapEncodeTiled failed for DSV4 O_b split-K (%d)", static_cast(result)); + return tensorMap; +} + +template +constexpr int numPipelineStages() +{ + constexpr int smemCd = 2 * 16 * kBlockN * static_cast(sizeof(__nv_bfloat16)); + constexpr int smemBarrierReserve = (kMaxPipelineStages * 3 + 4) * 8; + constexpr int smemTmemPointer = 4; + constexpr int smemPerStage + = (BlockM / kClusterSize) * kBlockK + kBlockN * kBlockK + 2 * 128 * static_cast(sizeof(uint32_t)); + constexpr int stages = (kSmemCapacity - smemCd - smemBarrierReserve - smemTmemPointer) / smemPerStage; + return stages < kMaxPipelineStages ? stages : kMaxPipelineStages; +} + +template +constexpr int dynamicSmemSize() +{ + constexpr int smemCd = 2 * 16 * kBlockN * static_cast(sizeof(__nv_bfloat16)); + constexpr int smemBarrierReserve = (kMaxPipelineStages * 3 + 4) * 8; + constexpr int smemTmemPointer = 4; + constexpr int smemPerStage + = (BlockM / kClusterSize) * kBlockK + kBlockN * kBlockK + 2 * 128 * static_cast(sizeof(uint32_t)); + return smemCd + numPipelineStages() * smemPerStage + smemBarrierReserve + smemTmemPointer; +} + +template +void launchInstance(CUtensorMap tensorMapA, CUtensorMap tensorMapB, CUtensorMap tensorMapSfa, CUtensorMap tensorMapSfb, + CUtensorMap tensorMapD, __nv_bfloat16* partials, int M, int numSms, cudaStream_t stream) +{ + using namespace deep_gemm; + constexpr int kNumStages = numPipelineStages(); + constexpr int kDynamicSmem = dynamicSmemSize(); + static_assert(kNumStages > 0, "DSV4 O_b split-K pipeline has no shared-memory stages"); + static_assert(kDynamicSmem <= kSmemCapacity, "DSV4 O_b split-K shared-memory budget exceeded"); + + auto kernel = &dsv4_splitk::dsv4Fp8SplitKGemmKernel; + + TLLM_CUDA_CHECK(cudaFuncSetAttribute( + reinterpret_cast(kernel), cudaFuncAttributeMaxDynamicSharedMemorySize, kDynamicSmem)); + + cudaLaunchConfig_t config{}; + config.gridDim = dim3(numSms, 1, 1); + config.blockDim = dim3(kNumThreads, 1, 1); + config.dynamicSmemBytes = kDynamicSmem; + config.stream = stream; + cudaLaunchAttribute attribute{}; + attribute.id = cudaLaunchAttributeClusterDimension; + attribute.val.clusterDim.x = kClusterSize; + attribute.val.clusterDim.y = 1; + attribute.val.clusterDim.z = 1; + config.attrs = &attribute; + config.numAttrs = 1; + + TLLM_CUDA_CHECK(cudaLaunchKernelEx(&config, kernel, static_cast(nullptr), static_cast(M), + static_cast(kDsv4ProN), static_cast(kDsv4ProK), tensorMapA, tensorMapB, tensorMapSfa, + tensorMapSfb, tensorMapD, static_cast(partials))); +} + +template +void dispatchBlockM(int blockM, CUtensorMap tensorMapA, CUtensorMap tensorMapB, CUtensorMap tensorMapSfa, + CUtensorMap tensorMapSfb, CUtensorMap tensorMapD, __nv_bfloat16* partials, int M, int numSms, cudaStream_t stream) +{ +#define TRTLLM_DSV4_SPLITK_BLOCK_M(BLOCK_M) \ + case BLOCK_M: \ + launchInstance( \ + tensorMapA, tensorMapB, tensorMapSfa, tensorMapSfb, tensorMapD, partials, M, numSms, stream); \ + return + switch (blockM) + { + TRTLLM_DSV4_SPLITK_BLOCK_M(16); + TRTLLM_DSV4_SPLITK_BLOCK_M(32); + TRTLLM_DSV4_SPLITK_BLOCK_M(48); + TRTLLM_DSV4_SPLITK_BLOCK_M(64); + TRTLLM_DSV4_SPLITK_BLOCK_M(80); + TRTLLM_DSV4_SPLITK_BLOCK_M(96); + TRTLLM_DSV4_SPLITK_BLOCK_M(112); + TRTLLM_DSV4_SPLITK_BLOCK_M(128); + default: TLLM_CHECK_WITH_INFO(false, "unsupported DSV4 O_b split-K block M: %d", blockM); + } +#undef TRTLLM_DSV4_SPLITK_BLOCK_M +} + +} // namespace + +void dsv4Fp8SplitKGemmLaunch(void const* a, int32_t const* sfa, void const* b, int32_t const* sfb, + __nv_bfloat16* partials, int M, int N, int K, int sfaKStride, int sfbKStride, int numSplits, cudaStream_t stream) +{ + TLLM_CHECK_WITH_INFO(a != nullptr && sfa != nullptr && b != nullptr && sfb != nullptr && partials != nullptr, + "DSV4 O_b split-K received a null pointer"); + TLLM_CHECK_WITH_INFO(M > 0, "DSV4 O_b split-K requires a positive M, got %d", M); + TLLM_CHECK_WITH_INFO(N == kDsv4ProN && K == kDsv4ProK, "DSV4 O_b split-K supports N=%d and K=%d, got N=%d K=%d", + kDsv4ProN, kDsv4ProK, N, K); + TLLM_CHECK_WITH_INFO( + numSplits == 2 || numSplits == 4, "DSV4 O_b split-K requires 2 or 4 splits, got %d", numSplits); + TLLM_CHECK_WITH_INFO(sfaKStride >= M && sfaKStride % 4 == 0, + "DSV4 O_b split-K activation scale stride must be a 4-aligned value >= M, got %d", sfaKStride); + TLLM_CHECK_WITH_INFO( + sfbKStride == N, "DSV4 O_b split-K weight scale stride must equal N=%d, got %d", N, sfbKStride); + + int device = 0; + TLLM_CUDA_CHECK(cudaGetDevice(&device)); + int major = 0; + int minor = 0; + int numSms = 0; + TLLM_CUDA_CHECK(cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, device)); + TLLM_CUDA_CHECK(cudaDeviceGetAttribute(&minor, cudaDevAttrComputeCapabilityMinor, device)); + TLLM_CUDA_CHECK(cudaDeviceGetAttribute(&numSms, cudaDevAttrMultiProcessorCount, device)); + TLLM_CHECK_WITH_INFO(major == 10, "DSV4 O_b split-K requires the SM100 family, got SM%d%d", major, minor); + TLLM_CHECK_WITH_INFO( + numSms % kClusterSize == 0, "DSV4 O_b split-K requires an even SM count for 2-CTA clusters, got %d", numSms); + + int const blockM = std::min(((M + 15) / 16) * 16, 128); + CUtensorMap const tensorMapA = makeTma2D(a, CU_TENSOR_MAP_DATA_TYPE_UINT8, K, M, kBlockK, blockM / kClusterSize, + static_cast(K), CU_TENSOR_MAP_SWIZZLE_128B); + CUtensorMap const tensorMapB = makeTma2D( + b, CU_TENSOR_MAP_DATA_TYPE_UINT8, K, N, kBlockK, kBlockN, static_cast(K), CU_TENSOR_MAP_SWIZZLE_128B); + CUtensorMap const tensorMapSfa = makeTma2D(sfa, CU_TENSOR_MAP_DATA_TYPE_INT32, sfaKStride, K / 512, blockM, 1, + static_cast(sfaKStride) * sizeof(int32_t), CU_TENSOR_MAP_SWIZZLE_NONE); + CUtensorMap const tensorMapSfb = makeTma2D(sfb, CU_TENSOR_MAP_DATA_TYPE_INT32, sfbKStride, K / 512, kBlockN, 1, + static_cast(sfbKStride) * sizeof(int32_t), CU_TENSOR_MAP_SWIZZLE_NONE); + CUtensorMap const tensorMapD = makeTma2D(partials, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, N, M, kBlockN, 16, + static_cast(N) * sizeof(__nv_bfloat16), CU_TENSOR_MAP_SWIZZLE_128B); + + if (numSplits == 2) + { + dispatchBlockM<2>( + blockM, tensorMapA, tensorMapB, tensorMapSfa, tensorMapSfb, tensorMapD, partials, M, numSms, stream); + } + else + { + dispatchBlockM<4>( + blockM, tensorMapA, tensorMapB, tensorMapSfa, tensorMapSfb, tensorMapD, partials, M, numSms, stream); + } +} + +} // namespace kernels::mhc + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/mhcKernels/dsv4Fp8SplitKGemm.cuh b/cpp/tensorrt_llm/kernels/mhcKernels/dsv4Fp8SplitKGemm.cuh new file mode 100644 index 000000000000..def8418e0632 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/mhcKernels/dsv4Fp8SplitKGemm.cuh @@ -0,0 +1,670 @@ +/* + * MIT License + * + * Copyright (c) 2025 DeepSeek + * Copyright (c) 2026 NVIDIA CORPORATION. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in all + * copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + * + * Adapted from DeepGEMM's SM100 1D1D FP8/FP4 GEMM implementation. + */ + +#pragma once +#pragma clang diagnostic push +#pragma clang diagnostic ignored "-Wunknown-attributes" + +#include + +// DeepGEMM is normally JIT-compiled and exposes several host/device helpers +// containing PTX. TRT-LLM compiles this translation unit with clang as NVCC's +// host compiler, so keep those helpers out of the host pass. +#pragma push_macro("CUTLASS_HOST_DEVICE") +#undef CUTLASS_HOST_DEVICE +#define CUTLASS_HOST_DEVICE CUTLASS_DEVICE +#include +#pragma pop_macro("CUTLASS_HOST_DEVICE") + +#include +#include +#include +#include +#include +#include + +#include "dsv4Fp8SplitKScheduler.cuh" +#include "tensorrt_llm/common/config.h" + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::mhc::dsv4_splitk +{ + +using namespace deep_gemm; + +CUTLASS_DEVICE void clusterSyncWithRelaxedArrive() +{ + cute::cluster_arrive_relaxed(); + cute::cluster_wait(); +} + +template +CUTLASS_GLOBAL void __launch_bounds__(kNumNonEpilogueThreads + kNumEpilogueThreads, 1) dsv4Fp8SplitKGemmKernel( + int* grouped_layout, uint32_t shape_m, uint32_t shape_n, uint32_t shape_k, + __grid_constant__ const cute::TmaDescriptor tensor_map_a, __grid_constant__ const cute::TmaDescriptor tensor_map_b, + __grid_constant__ const cute::TmaDescriptor tensor_map_sfa, + __grid_constant__ const cute::TmaDescriptor tensor_map_sfb, + __grid_constant__ const cute::TmaDescriptor tensor_map_cd, void* gmem_split_partials) +{ +#if (defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 1000)) or defined(__CLION_IDE__) + using Barrier = cutlass::arch::ClusterTransactionBarrier; + using Allocator = cute::conditional_t; + + // C/D type: BF16 and FP32 are supported, with or without accumulation + DG_STATIC_ASSERT(cute::is_same_v or cute::is_same_v, + "Invalid C/D data dtype"); + DG_STATIC_ASSERT(kSplitKFactor == 1 or kGemmType == GemmType::Normal, "Split-K only supports normal GEMM"); + DG_STATIC_ASSERT(kGemmType == GemmType::Normal, "DSV4 split-K only supports normal GEMM"); + + // MMA Configs + constexpr uint32_t LAYOUT_AD_M = 128; + constexpr uint32_t UMMA_M = LAYOUT_AD_M * kNumMulticast; + constexpr uint32_t UMMA_N = kSwapAB ? BLOCK_M : BLOCK_N; + constexpr uint32_t UMMA_K = 32; + constexpr uint32_t LOAD_BLOCK_M = BLOCK_M / (kIsMulticastOnA ? kNumMulticast : 1); + constexpr uint32_t LOAD_BLOCK_N = BLOCK_N / (kIsMulticastOnA ? 1 : kNumMulticast); + DG_STATIC_ASSERT(BLOCK_K == 128, "Invalid block K"); + DG_STATIC_ASSERT(kNumMulticast == 1 or kNumMulticast == 2, "Only support 1/2 multicast"); + DG_STATIC_ASSERT((kSwapAB and BLOCK_N == LAYOUT_AD_M) + or (not kSwapAB and (BLOCK_M == 32 or BLOCK_M == 64 or BLOCK_M == LAYOUT_AD_M)), + "Invalid block size"); + + // SF configs + constexpr uint32_t kNumUTCCPAlignedElems = 128; + constexpr uint32_t SF_BLOCK_M = math::constexpr_align(BLOCK_M, kNumUTCCPAlignedElems); + constexpr uint32_t SF_BLOCK_N = math::constexpr_align(BLOCK_N, kNumUTCCPAlignedElems); + constexpr uint32_t kNumSFAStagesPerLoad = kGranKA == 32 ? 1 : 4; + constexpr uint32_t kNumSFBStagesPerLoad = kGranKB == 32 ? 1 : 4; + DG_STATIC_ASSERT(kGranKA == 32 or kGranKA == 128, "Invalid granularity K for A"); + DG_STATIC_ASSERT(kGranKB == 32 or kGranKB == 128, "Invalid granularity K for B"); + DG_STATIC_ASSERT( + (kGemmType != GemmType::KGroupedContiguous) or kGranKA == kGranKB, "K-grouped SF requires kGranKA == kGranKB"); + + // Epilogue configs + // Always enable pipeline for better performance + constexpr uint32_t kNumEpilogueStages = 2; + constexpr uint32_t kNumTMAStoreStages = 2; + // NOTES: To maximize epilogue threads utilization, process an entire BLOCK_N + // per store stage for swap-AB cases, and an entire BLOCK_M for non-swap cases + constexpr uint32_t STORE_BLOCK_M = kSwapAB ? 16 : cute::min(BLOCK_M, LAYOUT_AD_M); + constexpr uint32_t STORE_BLOCK_N = kSwapAB ? BLOCK_N : kSwizzleCDMode / sizeof(cd_dtype_t); + constexpr uint32_t kNumUMMAStoreThreads = kSwapAB ? kNumEpilogueThreads : STORE_BLOCK_M; + DG_STATIC_ASSERT(kNumUMMAStoreThreads % 32 == 0, "Invalid store block M"); + + // Share memory sizes + constexpr uint32_t SMEM_CD_SIZE_PER_STAGE = STORE_BLOCK_M * STORE_BLOCK_N * sizeof(cd_dtype_t); + constexpr uint32_t SMEM_CD_SIZE = SMEM_CD_SIZE_PER_STAGE * kNumTMAStoreStages; + constexpr uint32_t SMEM_A_SIZE_PER_STAGE = LOAD_BLOCK_M * BLOCK_K * sizeof(a_dtype_t); + constexpr uint32_t SMEM_B_SIZE_PER_STAGE = LOAD_BLOCK_N * BLOCK_K * sizeof(b_dtype_t); + constexpr uint32_t SMEM_SFA_SIZE_PER_STAGE = SF_BLOCK_M * sizeof(uint32_t); + constexpr uint32_t SMEM_SFB_SIZE_PER_STAGE = SF_BLOCK_N * sizeof(uint32_t); + DG_STATIC_ASSERT( + SMEM_CD_SIZE % 1024 == 0 and SMEM_A_SIZE_PER_STAGE % 1024 == 0 and SMEM_B_SIZE_PER_STAGE % 1024 == 0, + "Shared memory of A/B must be aligned to 1024 bytes"); + // NOTES: Make sure we have enough shared memory for UMMA padding + constexpr uint32_t UMMA_A_SIZE_PER_STAGE + = math::constexpr_align(LOAD_BLOCK_M, LAYOUT_AD_M) * BLOCK_K * sizeof(a_dtype_t); + DG_STATIC_ASSERT(UMMA_A_SIZE_PER_STAGE <= SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE * kNumStages, + "Memory Out of bound for UMMA"); + + // Tensor memory size and offsets + constexpr uint32_t kNumAccumTmemCols = UMMA_N * kNumEpilogueStages; + constexpr uint32_t kNumSFATmemCols = SF_BLOCK_M / 32; + constexpr uint32_t kNumSFBTmemCols = SF_BLOCK_N / 32; + constexpr uint32_t kNumTmemCols + = utils::get_num_aligned_tmem_cols(); + constexpr uint32_t kTmemStartColOfSFA = kNumAccumTmemCols; + constexpr uint32_t kTmemStartColOfSFB = kNumAccumTmemCols + kNumSFATmemCols; + DG_STATIC_ASSERT(32 <= kNumTmemCols and kNumTmemCols <= 512, "Invalid tensor memory columns"); + + // Synchronize the cluster before 2-CTA TMEM allocation + kNumMulticast > 1 ? clusterSyncWithRelaxedArrive() : void(); + + // Utils + bool const is_leader_cta = cute::block_rank_in_cluster() == 0; + auto const warp_idx = cutlass::canonical_warp_idx_sync(); + auto const lane_idx = ptx::get_lane_idx(); + + // Prefetch TMA descriptors at the very beginning + if (warp_idx == 0) + { + cute::prefetch_tma_descriptor(&tensor_map_a); + cute::prefetch_tma_descriptor(&tensor_map_b); + cute::prefetch_tma_descriptor(&tensor_map_sfa); + cute::prefetch_tma_descriptor(&tensor_map_sfb); + cute::prefetch_tma_descriptor(&tensor_map_cd); + } + + // Overwrite shape constants if the compiler gives + shape_m = SHAPE_M != 0 ? SHAPE_M : shape_m; + shape_n = SHAPE_N != 0 ? SHAPE_N : shape_n; + shape_k = SHAPE_K != 0 ? SHAPE_K : shape_k; + auto const shape_sfa_k = math::ceil_div(shape_k, kGranKA * 4); + auto const shape_sfb_k = math::ceil_div(shape_k, kGranKB * 4); + + // Align to 1024 bytes for swizzle-128B + extern __shared__ __align__(1024) uint8_t smem_buffer[]; + + // D/A/B shared memory + auto smem_cd = utils::PatternVisitor( + [&](uint32_t const& i) { return reinterpret_cast(smem_buffer + i * SMEM_CD_SIZE_PER_STAGE); }); + auto smem_a = utils::PatternVisitor([&](uint32_t const& i) + { return reinterpret_cast(smem_buffer + SMEM_CD_SIZE + i * SMEM_A_SIZE_PER_STAGE); }); + auto smem_b = utils::PatternVisitor( + [&](uint32_t const& i) + { + return reinterpret_cast( + smem_buffer + SMEM_CD_SIZE + kNumStages * SMEM_A_SIZE_PER_STAGE + i * SMEM_B_SIZE_PER_STAGE); + }); + + // SFA/SFB shared memory + auto sf_start_ptr = reinterpret_cast(smem_b[kNumStages]); + auto smem_sfa = utils::PatternVisitor( + [=](uint32_t const& i) { return reinterpret_cast(sf_start_ptr + i * SMEM_SFA_SIZE_PER_STAGE); }); + auto smem_sfb = utils::PatternVisitor( + [=](uint32_t const& i) + { + return reinterpret_cast( + sf_start_ptr + kNumStages * SMEM_SFA_SIZE_PER_STAGE + i * SMEM_SFB_SIZE_PER_STAGE); + }); + + // Barriers and tensor memory pointer + auto barrier_start_ptr = reinterpret_cast(smem_sfb[kNumStages]); + ; + auto full_barriers = utils::PatternVisitor([=](uint32_t const& i) { return barrier_start_ptr + (i); }); + auto empty_barriers + = utils::PatternVisitor([=](uint32_t const& i) { return barrier_start_ptr + (kNumStages + i); }); + auto with_sf_full_barriers + = utils::PatternVisitor([=](uint32_t const& i) { return barrier_start_ptr + (kNumStages * 2 + i); }); + auto tmem_full_barriers + = utils::PatternVisitor([=](uint32_t const& i) { return barrier_start_ptr + (kNumStages * 3 + i); }); + auto tmem_empty_barriers = utils::PatternVisitor( + [=](uint32_t const& i) { return barrier_start_ptr + (kNumStages * 3 + kNumEpilogueStages + i); }); + auto tmem_ptr_in_smem = reinterpret_cast(barrier_start_ptr + kNumStages * 3 + kNumEpilogueStages * 2); + + // Initialize barriers + if (warp_idx == 1 and cute::elect_one_sync()) + { +#pragma unroll + for (uint32_t i = 0; i < kNumStages; ++i) + { + // Arrive at all CTAs + full_barriers[i]->init(1); + empty_barriers[i]->init(1); + // Arrive only at the leader CTA + with_sf_full_barriers[i]->init(kNumMulticast * 32); + } +#pragma unroll + for (uint32_t i = 0; i < kNumEpilogueStages; ++i) + { + // Arrive at all CTAs + tmem_full_barriers[i]->init(1); + // Arrive only at the leader CTA + tmem_empty_barriers[i]->init(kNumMulticast * kNumUMMAStoreThreads); + } + + // Make initialized barrier visible in async proxy + cutlass::arch::fence_barrier_init(); + } + else if (warp_idx == 2) + { + // Allocate tensor memory + Allocator().allocate(kNumTmemCols, tmem_ptr_in_smem); + } + kNumMulticast > 1 ? clusterSyncWithRelaxedArrive() : __syncthreads(); + + // Wait for primary kernel completion + cudaGridDependencySynchronize(); + + // Block scheduler + uint32_t m_block_idx, n_block_idx; + auto scheduler = SplitKScheduler( + shape_m, shape_n, shape_k, grouped_layout); + + // Pipeline and TMA phases + uint32_t stage_idx = 0, phase = 0; + auto advance_pipeline = [&](uint32_t& k_block_idx) + { + ++k_block_idx; + + // Flip phases only if reach the next first stage + stage_idx = stage_idx == kNumStages - 1 ? 0 : stage_idx + 1; + phase ^= stage_idx == 0; + }; + + // Dispatch warps into different roles + if (warp_idx == 0 and cute::elect_one_sync()) + { + // TMA load warp + // Persistently schedule over blocks + while (scheduler.get_next_block(m_block_idx, n_block_idx)) + { + // Use dynamic load block M, when swap-AB is enabled + auto const load_block_m + = kSwapAB ? scheduler.get_aligned_effective_m_in_block(m_block_idx) / kNumMulticast : LOAD_BLOCK_M; + + // For k-grouped layout, the number of block K is variable + auto const num_all_k_blocks = math::ceil_div(scheduler.current_shape_k, BLOCK_K); + DG_TRAP_ONLY_DEVICE_ASSERT(num_all_k_blocks % kSplitKFactor == 0); + auto const num_total_k_blocks = num_all_k_blocks / kSplitKFactor; + auto const k_block_offset = scheduler.split_k_idx * num_total_k_blocks; + for (uint32_t k_block_idx = 0; k_block_idx < num_total_k_blocks; advance_pipeline(k_block_idx)) + { + // Wait consumer release + empty_barriers[stage_idx]->wait(phase ^ 1); + + // Compute offsets + // NOTES: the group is always concatenated with the outer dimension + uint32_t m_idx + = scheduler.template get_global_idx<(kGemmType == GemmType::MGroupedMasked), IndexType::MN>( + shape_m, BLOCK_M, m_block_idx); + uint32_t n_idx = scheduler.template get_global_idx<(kMajorB == cute::UMMA::Major::K), IndexType::MN>( + shape_n, BLOCK_N, n_block_idx, m_block_idx); + + // NOTES: `k_idx` is actually the k index default for K-major, while `k_b_idx` may be MN-major + // And for all m-grouped GEMMs, A must be K-majored + DG_STATIC_ASSERT(kGemmType == GemmType::Normal or kGemmType == GemmType::KGroupedContiguous + or kGemmType == GemmType::Batched or kMajorA == cute::UMMA::Major::K, + "Invalid major"); + uint32_t const global_k_block_idx = k_block_offset + k_block_idx; + uint32_t k_idx = global_k_block_idx * BLOCK_K; + uint32_t k_a_idx = scheduler.template get_global_idx<(kMajorA == cute::UMMA::Major::MN), IndexType::K>( + shape_k, BLOCK_K, global_k_block_idx, m_block_idx); + uint32_t k_b_idx = scheduler.template get_global_idx<(kMajorB == cute::UMMA::Major::MN), IndexType::K>( + shape_k, BLOCK_K, global_k_block_idx, m_block_idx); + + // Add 2 CTA offsets + if constexpr (kNumMulticast > 1) + { + m_idx += kIsMulticastOnA ? (cute::block_rank_in_cluster() * load_block_m) : 0; + n_idx += kIsMulticastOnA ? 0 : (cute::block_rank_in_cluster() * LOAD_BLOCK_N); + } + + // Issue TMAs + constexpr bool kIsBatchedMM = (kGemmType == GemmType::Batched); + uint32_t const batch_idx = (kIsBatchedMM ? scheduler.current_group_idx : 0); + if constexpr (kMajorA == cute::UMMA::Major::K) + tma::copy( + &tensor_map_a, full_barriers[stage_idx], smem_a[stage_idx], k_a_idx, m_idx, 1, batch_idx); + if constexpr (kMajorA == cute::UMMA::Major::MN) + tma::copy( + &tensor_map_a, full_barriers[stage_idx], smem_a[stage_idx], m_idx, k_a_idx, 1, batch_idx); + if constexpr (kMajorB == cute::UMMA::Major::K) + tma::copy( + &tensor_map_b, full_barriers[stage_idx], smem_b[stage_idx], k_b_idx, n_idx, 1, batch_idx); + if constexpr (kMajorB == cute::UMMA::Major::MN) + tma::copy( + &tensor_map_b, full_barriers[stage_idx], smem_b[stage_idx], n_idx, k_b_idx, 1, batch_idx); + auto num_arrival_bytes + = SMEM_A_SIZE_PER_STAGE / (std::is_same_v ? 1 : 2) + + SMEM_B_SIZE_PER_STAGE / (std::is_same_v ? 1 : 2); + + // Issue SFA and SFB TMAs at certain stages + // No swizzling, so one TMA for one SF is enough + if (k_block_idx % kNumSFAStagesPerLoad == 0) + { + uint32_t sfa_m_idx = m_block_idx * BLOCK_M; + uint32_t sfa_k_idx + = scheduler.template get_global_idx<(not is_m_grouped_contiguous(kGemmType)), IndexType::SF_K>( + shape_sfa_k, 1, math::ceil_div(k_idx, BLOCK_K * kNumSFAStagesPerLoad)); + tma::copy( + &tensor_map_sfa, full_barriers[stage_idx], smem_sfa[stage_idx], sfa_m_idx, sfa_k_idx); + num_arrival_bytes += BLOCK_M * sizeof(uint32_t); + } + if (k_block_idx % kNumSFBStagesPerLoad == 0) + { + uint32_t sfb_n_idx = n_block_idx * BLOCK_N; + uint32_t sfb_k_idx = scheduler.template get_global_idx( + shape_sfb_k, 1, math::ceil_div(k_idx, BLOCK_K * kNumSFBStagesPerLoad), m_block_idx); + tma::copy( + &tensor_map_sfb, full_barriers[stage_idx], smem_sfb[stage_idx], sfb_n_idx, sfb_k_idx); + num_arrival_bytes += BLOCK_N * sizeof(uint32_t); + } + + // Arrive at full barriers + full_barriers[stage_idx]->arrive_and_expect_tx(num_arrival_bytes); + } + } + } + else if (warp_idx == 1 and is_leader_cta) + { + // MMA issue warp + // NOTES: only the leader CTA will do this + // Make instruction descriptor + auto instr_desc = kSwapAB ? cute::UMMA::make_instr_desc_block_scaled() + : cute::UMMA::make_instr_desc_block_scaled(); + auto sf_desc = mma::sm100::make_sf_desc(nullptr); + + DG_STATIC_ASSERT(kNumStages <= 32, "Too many stages"); + auto a_desc = mma::sm100::make_umma_desc(smem_a[0], 0, 0); + auto b_desc = mma::sm100::make_umma_desc(smem_b[0], 0, 0); + uint32_t a_desc_lo = lane_idx < kNumStages ? a_desc.lo + lane_idx * SMEM_A_SIZE_PER_STAGE / 16 : 0u; + uint32_t b_desc_lo = lane_idx < kNumStages ? b_desc.lo + lane_idx * SMEM_B_SIZE_PER_STAGE / 16 : 0u; + + // Checks for MMA instructions + // NOTES: CUTLASS does not have such checks except the MMA traits, but we are not using these traits + DG_STATIC_ASSERT((UMMA_M == 64 and UMMA_N % 8 == 0 and 8 <= UMMA_N and UMMA_N <= 256) + or (UMMA_M == 128 and UMMA_N % 16 == 0 and 16 <= UMMA_N and UMMA_N <= 256) + or (UMMA_M == 256 and UMMA_N % 16 == 0 and 16 <= UMMA_N and UMMA_N <= 256), + "Invalid MMA instruction shape"); + + // Persistently schedule over blocks + while (scheduler.get_next_block(m_block_idx, n_block_idx)) + { + // Wait tensor memory empty barrier arrival + auto accum_stage_idx = scheduler.current_iter % kNumEpilogueStages; + auto accum_phase_idx = (scheduler.current_iter / kNumEpilogueStages) & 1; + tmem_empty_barriers[accum_stage_idx]->wait(accum_phase_idx ^ 1); + ptx::tcgen05_after_thread_sync(); + + // Empty barrier arrival + auto empty_barrier_arrive = [&](bool const& do_tmem_full_arrive) + { + auto umma_arrive = [](uint64_t const* barrier) + { + if constexpr (kNumMulticast == 1) + { + cutlass::arch::umma_arrive(barrier); + } + else + { + constexpr uint16_t kCTAMask = (1 << kNumMulticast) - 1; + cutlass::arch::umma_arrive_multicast_2x1SM(barrier, kCTAMask); + } + }; + umma_arrive(reinterpret_cast(empty_barriers[stage_idx])); + + // NOTES: the tensor memory accumulator pipeline has nothing to do with multicasting + if (do_tmem_full_arrive) + umma_arrive(reinterpret_cast(tmem_full_barriers[accum_stage_idx])); + __syncwarp(); + }; + + // Dynamic update of UMMA N based on effective M, when swap-AB is enabled + if constexpr (kSwapAB) + { + uint32_t umma_n = scheduler.get_aligned_effective_m_in_block(m_block_idx); + mma::sm100::update_instr_desc_with_umma_n(instr_desc, umma_n); + } + + // Launch MMAs + auto const num_all_k_blocks = math::ceil_div(scheduler.current_shape_k, BLOCK_K); + DG_TRAP_ONLY_DEVICE_ASSERT(num_all_k_blocks % kSplitKFactor == 0); + auto const num_total_k_blocks = num_all_k_blocks / kSplitKFactor; +#pragma unroll 4 + for (uint32_t k_block_idx = 0; k_block_idx < num_total_k_blocks; advance_pipeline(k_block_idx)) + { + // Wait TMA and SF-transpose arrival + with_sf_full_barriers[stage_idx]->wait(phase); + ptx::tcgen05_after_thread_sync(); + + auto const a_desc_base_lo = ptx::exchange(a_desc_lo, stage_idx); + auto const b_desc_base_lo = ptx::exchange(b_desc_lo, stage_idx); + if (cute::elect_one_sync()) + { + // Do SF copy at certain stages + // TODO: process shared memory descriptor by addition + using cute_utccp_t = cute::conditional_t; + uint32_t const sfa_stage_in_group_idx = k_block_idx % kNumSFAStagesPerLoad; + if (sfa_stage_in_group_idx == 0) + { +#pragma unroll + for (uint32_t i = 0; i < SF_BLOCK_M / kNumUTCCPAlignedElems; ++i) + { + auto smem_ptr = smem_sfa[stage_idx] + i * kNumUTCCPAlignedElems; + mma::sm100::replace_smem_desc_addr(sf_desc, smem_ptr); + cute_utccp_t::copy(sf_desc, kTmemStartColOfSFA + i * 4); + } + } + uint32_t const sfb_stage_in_group_idx = k_block_idx % kNumSFBStagesPerLoad; + if (sfb_stage_in_group_idx == 0) + { +#pragma unroll + for (uint32_t i = 0; i < SF_BLOCK_N / kNumUTCCPAlignedElems; ++i) + { + auto smem_ptr = smem_sfb[stage_idx] + i * kNumUTCCPAlignedElems; + mma::sm100::replace_smem_desc_addr(sf_desc, smem_ptr); + cute_utccp_t::copy(sf_desc, kTmemStartColOfSFB + i * 4); + } + } + + // Issue UMMA + using mma_t = cute::conditional_t; +#pragma unroll + for (uint32_t k = 0; k < BLOCK_K / UMMA_K; ++k) + { + uint32_t const sfa_id = (kGranKA == 32 ? k : sfa_stage_in_group_idx); + uint32_t const sfb_id = (kGranKB == 32 ? k : sfb_stage_in_group_idx); + auto const runtime_instr_desc = kSwapAB + ? mma::sm100::make_runtime_instr_desc_with_sf_id(instr_desc, sfb_id, sfa_id) + : mma::sm100::make_runtime_instr_desc_with_sf_id(instr_desc, sfa_id, sfb_id); + + a_desc.lo = mma::sm100::advance_umma_desc_lo( + a_desc_base_lo, 0, k * UMMA_K); + b_desc.lo = mma::sm100::advance_umma_desc_lo( + b_desc_base_lo, 0, k * UMMA_K); + if constexpr (kSwapAB) + { + mma_t::fma(b_desc, a_desc, accum_stage_idx * UMMA_N, k_block_idx > 0 or k > 0, + runtime_instr_desc, kTmemStartColOfSFB, kTmemStartColOfSFA); + } + else + { + mma_t::fma(a_desc, b_desc, accum_stage_idx * UMMA_N, k_block_idx > 0 or k > 0, + runtime_instr_desc, kTmemStartColOfSFA, kTmemStartColOfSFB); + } + } + } + __syncwarp(); + + // Commit to the mbarrier object + // No explicit `tcgen05.fence::before_thread_sync` is needed, as this is implicitly performed by + // `tcgen05.commit` + empty_barrier_arrive(k_block_idx == num_total_k_blocks - 1); + } + } + + // To safely deconstruct barriers, we need another round of waits + auto const iter_idx = scheduler.current_iter - 1; + if (kNumMulticast > 1 and iter_idx >= 0) + { + auto const accum_phase_idx = (iter_idx / kNumEpilogueStages) & 1; + tmem_empty_barriers[iter_idx % kNumEpilogueStages]->wait(accum_phase_idx); + } + } + else if (warp_idx == 2) + { + // UTCCP transposer + auto utccp_required_smem_warp_transpose = [&](uint32_t const* smem_ptr) + { + DG_STATIC_ASSERT(kNumUTCCPAlignedElems == 128, "Invalid aligned elements"); + uint32_t values[4]; +#pragma unroll + for (uint32_t i = 0; i < 4; ++i) + values[i] = ptx::ld_shared(smem_ptr + i * 32 + lane_idx); + __syncwarp(); + ptx::st_shared(smem_ptr + lane_idx * 4, values[0], values[1], values[2], values[3]); + }; + + while (scheduler.get_next_block(m_block_idx, n_block_idx)) + { + auto const num_all_k_blocks = math::ceil_div(scheduler.current_shape_k, BLOCK_K); + DG_TRAP_ONLY_DEVICE_ASSERT(num_all_k_blocks % kSplitKFactor == 0); + auto const num_total_k_blocks = num_all_k_blocks / kSplitKFactor; + for (uint32_t k_block_idx = 0; k_block_idx < num_total_k_blocks; advance_pipeline(k_block_idx)) + { + // Wait TMA arrival + full_barriers[stage_idx]->wait(phase); + + // Transpose for UTCCP at certain stages + if (k_block_idx % kNumSFAStagesPerLoad == 0) + { +#pragma unroll + for (uint32_t i = 0; i < SF_BLOCK_M / kNumUTCCPAlignedElems; ++i) + utccp_required_smem_warp_transpose(smem_sfa[stage_idx] + i * kNumUTCCPAlignedElems); + // TODO: figure out whether the proxy fence is valid for 2-CTA cases + cutlass::arch::fence_view_async_shared(); + } + if (k_block_idx % kNumSFBStagesPerLoad == 0) + { +#pragma unroll + for (uint32_t i = 0; i < SF_BLOCK_N / kNumUTCCPAlignedElems; ++i) + utccp_required_smem_warp_transpose(smem_sfb[stage_idx] + i * kNumUTCCPAlignedElems); + // TODO: figure out whether the proxy fence is valid for 2-CTA cases + cutlass::arch::fence_view_async_shared(); + } + + // Arrive + with_sf_full_barriers[stage_idx]->arrive(0u); + } + } + } + else if (warp_idx >= kNumNonEpilogueThreads / 32 + and warp_idx < (kNumNonEpilogueThreads + kNumUMMAStoreThreads) / 32) + { + // Epilogue warp groups + auto const epilogue_warp_idx = warp_idx - (kNumNonEpilogueThreads / 32); + + // NOTES: tensor memory addresses are simplified, as the hardware will ignore the warp index bits, + // i.e., no need for `tmem_ptr |= (epilogue_warp_idx * 32) << 16`. + // NOTES: we also forbid two CTAs to share the same SM and its tensor memory + DG_TRAP_ONLY_DEVICE_ASSERT(ptx::ld_shared(tmem_ptr_in_smem) == 0); + + // Share store pipeline between blocks + uint32_t tma_stage_idx = 0; + + // Persistently schedule over blocks + while (scheduler.get_next_block(m_block_idx, n_block_idx)) + { + auto accum_stage_idx = scheduler.current_iter % kNumEpilogueStages; + auto accum_phase_idx = (scheduler.current_iter / kNumEpilogueStages) & 1; + + // Wait UMMA arrival + tmem_full_barriers[accum_stage_idx]->wait(accum_phase_idx); + ptx::tcgen05_after_thread_sync(); + + auto const tmem_base_addr = accum_stage_idx * UMMA_N; + auto const base_m_idx + = scheduler.template get_global_idx<(not is_m_grouped_contiguous(kGemmType)), IndexType::MN>( + shape_m, BLOCK_M, m_block_idx); + auto const base_n_idx = n_block_idx * BLOCK_N; + if constexpr (kSplitKFactor > 1) + { + // The regular TMA epilogue is tied to its 2D CD descriptor. + // Like the SM120 split-K path, write accumulators directly to + // a contiguous workspace. BF16 matches the downstream mHC + // input and avoids a separate FP32-to-BF16 cast kernel. + DG_STATIC_ASSERT(kSwapAB, "Split-K workspace store requires swap-AB"); + DG_STATIC_ASSERT(cute::is_same_v or cute::is_same_v, + "Split-K workspace store requires FP32 or BF16 output"); + constexpr uint32_t kRowsPerTmemLoad = 8; + constexpr uint32_t kColsPerEpilogueWarp = 32; + auto const effective_m = scheduler.get_aligned_effective_m_in_block(m_block_idx); + auto const num_m_stores = effective_m / STORE_BLOCK_M; + auto const split_offset = static_cast(scheduler.split_k_idx) * shape_m * shape_n; + + for (uint32_t s = 0; s < num_m_stores; ++s) + { +#pragma unroll + for (uint32_t i = 0; i < STORE_BLOCK_M / kRowsPerTmemLoad; ++i) + { + uint32_t const tmem_addr = tmem_base_addr + s * STORE_BLOCK_M + i * kRowsPerTmemLoad; + uint32_t values[kRowsPerTmemLoad]; + cute::SM100_TMEM_LOAD_32dp32b8x::copy(tmem_addr, values[0], values[1], values[2], values[3], + values[4], values[5], values[6], values[7]); + cutlass::arch::fence_view_async_tmem_load(); + + uint32_t const col = base_n_idx + epilogue_warp_idx * kColsPerEpilogueWarp + lane_idx; +#pragma unroll + for (uint32_t row = 0; row < kRowsPerTmemLoad; ++row) + { + uint32_t const global_row = base_m_idx + s * STORE_BLOCK_M + i * kRowsPerTmemLoad + row; + if (global_row < shape_m and col < shape_n) + { + auto const idx = split_offset + static_cast(global_row) * shape_n + col; + auto const value = __uint_as_float(values[row]); + if constexpr (cute::is_same_v) + static_cast(gmem_split_partials)[idx] = value; + else + static_cast(gmem_split_partials)[idx] + = cutlass::bfloat16_t(value); + } + } + } + } + + ptx::tcgen05_before_thread_sync(); + tmem_empty_barriers[accum_stage_idx]->arrive(0u); + } + else if constexpr (kSwapAB) + { + auto const effective_m = scheduler.get_aligned_effective_m_in_block(m_block_idx); + epilogue::sm100_store_cd_swap_ab(smem_cd, tma_stage_idx, tmem_base_addr, base_m_idx, base_n_idx, + scheduler.current_group_idx, effective_m, epilogue_warp_idx, lane_idx, + tmem_empty_barriers[accum_stage_idx], tensor_map_cd); + } + else + { + epilogue::sm100_store_cd(smem_cd, tma_stage_idx, tmem_base_addr, base_m_idx, base_n_idx, + scheduler.current_group_idx, epilogue_warp_idx, lane_idx, tmem_empty_barriers[accum_stage_idx], + tensor_map_cd); + } + } + } + + // TODO: Remove redundant synchronization + kNumMulticast > 1 ? clusterSyncWithRelaxedArrive() : __syncthreads(); + + // Deallocate tensor memory + if (warp_idx == 0) + Allocator().free(0, kNumTmemCols); + +#else + if (blockIdx.x == 0 and threadIdx.x == 0) + DG_DEVICE_ASSERT(false and "This kernel only support sm_100f"); +#endif +} + +} // namespace kernels::mhc::dsv4_splitk + +TRTLLM_NAMESPACE_END + +#pragma clang diagnostic pop diff --git a/cpp/tensorrt_llm/kernels/mhcKernels/dsv4Fp8SplitKScheduler.cuh b/cpp/tensorrt_llm/kernels/mhcKernels/dsv4Fp8SplitKScheduler.cuh new file mode 100644 index 000000000000..62ff8c4771f5 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/mhcKernels/dsv4Fp8SplitKScheduler.cuh @@ -0,0 +1,142 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels::mhc::dsv4_splitk +{ + +enum class IndexType +{ + MN, + K, + SF_K, +}; + +template +static constexpr uint32_t getNum1DBlocksPerGroup() +{ + uint32_t best = 0; + uint32_t minUsage = cute::numeric_limits::max(); + for (uint32_t candidate : {8u, 16u}) + { + uint32_t const usage = IsMulticastOnA + ? candidate * BlockN + deep_gemm::math::constexpr_ceil_div(NumSms, candidate) * BlockM + : candidate * BlockM + deep_gemm::math::constexpr_ceil_div(NumSms, candidate) * BlockN; + if (usage < minUsage) + { + best = candidate; + minUsage = usage; + } + } + return best; +} + +// Normal-GEMM-only persistent scheduler owned by the DSV4 O_b path. Keeping +// the split index outside the M/N tile index preserves adjacent CTA pairing +// for the 2-CTA cluster-N multicast used by the kernel. +template ()> +struct SplitKScheduler +{ + int current_iter = -1; + uint32_t num_blocks; + uint32_t num_mn_blocks; + uint32_t num_m_blocks; + uint32_t num_n_blocks; + uint32_t num_blocks_in_group; + uint32_t split_k_idx = 0; + uint32_t current_group_idx = 0; + uint32_t current_shape_k; + + CUTLASS_DEVICE explicit SplitKScheduler( + uint32_t shapeM, uint32_t shapeN, uint32_t shapeK, int* groupedLayout = nullptr) + : num_m_blocks(deep_gemm::math::ceil_div(shapeM, BlockM)) + , num_n_blocks(deep_gemm::math::ceil_div(shapeN, BlockN)) + , current_shape_k(shapeK) + { + static_cast(groupedLayout); + num_mn_blocks = num_m_blocks * num_n_blocks; + num_blocks = num_mn_blocks * SplitKFactor; + } + + CUTLASS_DEVICE void getSwizzledBlockIdx(uint32_t blockIdx, uint32_t& mBlockIdx, uint32_t& nBlockIdx) + { + static_assert(Num1DBlocksPerGroup % NumMulticast == 0, "Invalid split-K scheduler group size"); + uint32_t const primaryBlocks = IsMulticastOnA ? num_n_blocks : num_m_blocks; + uint32_t const secondaryBlocks = IsMulticastOnA ? num_m_blocks : num_n_blocks; + uint32_t const blocksPerGroup = secondaryBlocks * Num1DBlocksPerGroup; + uint32_t const groupIdx = blockIdx / blocksPerGroup; + uint32_t const firstBlockIdx = groupIdx * Num1DBlocksPerGroup; + uint32_t const inGroupIdx = blockIdx % blocksPerGroup; + num_blocks_in_group = cute::min(Num1DBlocksPerGroup, primaryBlocks - firstBlockIdx); + + if constexpr (IsMulticastOnA) + { + mBlockIdx = inGroupIdx / num_blocks_in_group; + nBlockIdx = firstBlockIdx + inGroupIdx % num_blocks_in_group; + } + else + { + mBlockIdx = firstBlockIdx + inGroupIdx % num_blocks_in_group; + nBlockIdx = inGroupIdx / num_blocks_in_group; + } + } + + template + CUTLASS_DEVICE uint32_t get_global_idx( + uint32_t shapeDim, uint32_t blockSize, uint32_t blockIdx, uint32_t mBlockIdx = 0) const + { + static_cast(WithGroupOffset); + static_cast(Type); + static_cast(shapeDim); + static_cast(mBlockIdx); + return blockIdx * blockSize; + } + + CUTLASS_DEVICE uint32_t get_aligned_effective_m_in_block(uint32_t mBlockIdx) const + { + static_cast(mBlockIdx); + static_assert(BlockM % 16 == 0, "DSV4 split-K block M must be 16-aligned"); + return BlockM; + } + + CUTLASS_DEVICE bool get_next_block(uint32_t& mBlockIdx, uint32_t& nBlockIdx) + { + uint32_t const nextBlockIdx = static_cast(++current_iter) * gridDim.x + blockIdx.x; + if (nextBlockIdx >= num_blocks) + { + return false; + } + + uint32_t const mnBlockIdx = nextBlockIdx % num_mn_blocks; + split_k_idx = nextBlockIdx / num_mn_blocks; + getSwizzledBlockIdx(mnBlockIdx, mBlockIdx, nBlockIdx); + return true; + } +}; + +} // namespace kernels::mhc::dsv4_splitk + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh b/cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh index 9e74d4eb20f7..c47cf64c16f6 100644 --- a/cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh +++ b/cpp/tensorrt_llm/kernels/mhcKernels/fused_tf32_pmap_gemm.cuh @@ -108,11 +108,11 @@ __device__ __forceinline__ void stsm_x4_b16_rout(void* smem_dst, uint32_t a, uin template + 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) @@ -175,6 +175,8 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) fused_tf3 auto smem_res = PatternVisitor( [&, base = cursor](uint32_t const& i) { return reinterpret_cast(base + i * SMEM_RES_PER_ISTG); }); cursor += N_INPUT_STAGES * SMEM_RES_PER_ISTG; + // For kXSplit>1 these buffers form a ring of single-partial x tiles. The + // consumer drains and sums them in FP32, so no additional data SMEM is needed. auto smem_x_stg = PatternVisitor( [&, base = cursor](uint32_t const& i) { return reinterpret_cast(base + i * SMEM_X_PER_ISTG); }); cursor += N_INPUT_STAGES * SMEM_X_PER_ISTG; @@ -192,13 +194,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(cursor); if (warp_idx == 1 and cute::elect_one_sync()) @@ -215,6 +222,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) { @@ -278,6 +294,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; @@ -290,10 +307,27 @@ __global__ void __launch_bounds__(kNumMMAThreads + kNumPmapThreads, 1) fused_tf3 tma_copy(&tensor_map_residual, full_input[i_stage], smem_res[i_stage] + j * BLOCK_M * BLOCK_K, j * HIDDEN + h_idx, m_idx); } - tma_copy( - &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( + &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( + &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) @@ -448,19 +482,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(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(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(smem_ptr)); - } - } - uint32_t r_vals[HC_MULT][2][kNumLoads]; #pragma unroll for (uint32_t j = 0; j < HC_MULT; ++j) @@ -480,12 +501,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(&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(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(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(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(&x_vals[u][i])); + xf[u][i].x += value.x; + xf[u][i].y += value.y; + } } if constexpr (kEarlyRelease) diff --git a/cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu b/cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu index 86586edc099a..73fdc3e85d2e 100644 --- a/cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu +++ b/cpp/tensorrt_llm/kernels/mhcKernels/mhcFusedHcKernel.cu @@ -91,8 +91,67 @@ inline void fhcZeroWorkspaces(float* y_acc, uint32_t y_elems, float* r_acc, uint y_acc, y_elems, r_acc, r_elems, done_counter, done_elems); } +// O_b split-K emits BF16 partials in split-major layout. Each thread reduces +// one 16-byte vector, keeping all eight lanes in FP32 until the final store. +__global__ void fhcReduceSplitXKernel( + uint4 const* __restrict__ partials, uint4* __restrict__ reduced, uint32_t vectors_per_split, int num_splits) +{ + uint32_t const tid = blockIdx.x * blockDim.x + threadIdx.x; + uint32_t const stride = gridDim.x * blockDim.x; + for (uint32_t vector_idx = tid; vector_idx < vectors_per_split; vector_idx += stride) + { + float2 accum[4] = { + make_float2(0.0f, 0.0f), + make_float2(0.0f, 0.0f), + make_float2(0.0f, 0.0f), + make_float2(0.0f, 0.0f), + }; + for (int split = 0; split < num_splits; ++split) + { + uint4 const raw = __ldg(&partials[static_cast(split) * vectors_per_split + vector_idx]); + auto const* pairs = reinterpret_cast<__nv_bfloat162 const*>(&raw); +#pragma unroll + for (int pair = 0; pair < 4; ++pair) + { + float2 const value = __bfloat1622float2(pairs[pair]); + accum[pair].x += value.x; + accum[pair].y += value.y; + } + } + + uint4 packed; + auto* pairs = reinterpret_cast<__nv_bfloat162*>(&packed); +#pragma unroll + for (int pair = 0; pair < 4; ++pair) + { + pairs[pair] = __float22bfloat162_rn(accum[pair]); + } + reduced[vector_idx] = packed; + } +} + } // namespace +void mhcReduceSplitXLaunch( + __nv_bfloat16 const* partials, __nv_bfloat16* reduced, int M, int hidden_size, int num_splits, cudaStream_t stream) +{ + TLLM_CHECK_WITH_INFO(partials != nullptr && reduced != nullptr, "split-x reduction received a null pointer"); + TLLM_CHECK_WITH_INFO(M > 0 && hidden_size > 0, "split-x reduction requires positive M and hidden_size"); + TLLM_CHECK_WITH_INFO( + num_splits > 1 && num_splits <= 16, "split-x reduction supports num_splits in [2, 16], got %d", num_splits); + TLLM_CHECK_WITH_INFO( + hidden_size % 8 == 0, "split-x reduction requires hidden_size divisible by 8, got %d", hidden_size); + + uint64_t const elements = static_cast(M) * static_cast(hidden_size); + uint64_t const vectors64 = elements / 8; + TLLM_CHECK_WITH_INFO(vectors64 <= UINT32_MAX, "split-x reduction tensor is too large"); + uint32_t const vectors = static_cast(vectors64); + constexpr uint32_t kBlock = 256; + uint32_t const blocks = min(static_cast((vectors + kBlock - 1) / kBlock), 148u * 8u); + fhcReduceSplitXKernel<<>>( + reinterpret_cast(partials), reinterpret_cast(reduced), vectors, num_splits); +} + // ---- mHC fused kernel shape constants (mirrors the Python module) ---- // HC_MULT * (2 + HC_MULT) = 4 * 6 = 24. static constexpr uint32_t FHC_SHAPE_N = 24; @@ -293,14 +352,53 @@ static constexpr uint32_t fhcSmemSize() using FusedRoutFn = void (*)( uint32_t, CUtensorMap, CUtensorMap, CUtensorMap, CUtensorMap, float*, float const*, float const*, float*); -template +template static FusedRoutFn fhcInstance() { static_assert(isSupportedFhcHidden(), "Unsupported fused-HC hidden size"); static_assert(isSupportedFhcMmaKS(), "Unsupported fused-HC MMA kNumSplits for hidden size"); return &fused_mhc::fused_tf32_pmap_gemm_rout_atomic_impl; + /*kEarlyRelease=*/false, XS>; +} + +template +static FusedRoutFn fhcXSplitInstanceIfSupported() +{ + if constexpr (isSupportedFhcMmaKS()) + { + return fhcInstance(); + } + else + { + TLLM_CHECK_WITH_INFO(false, "mhcFusedHcLaunch: unsupported (kNumSplits=%u, hidden=%u)", KS, Hidden); + return nullptr; + } +} + +template +static FusedRoutFn pickFhcXSplitKs(uint32_t ks) +{ + switch (ks) + { + case 1: return fhcInstance(); + case 2: return fhcXSplitInstanceIfSupported(); + case 4: return fhcXSplitInstanceIfSupported(); + case 8: return fhcXSplitInstanceIfSupported(); + case 16: return fhcXSplitInstanceIfSupported(); + default: TLLM_CHECK_WITH_INFO(false, "mhcFusedHcLaunch: unsupported kNumSplits=%u for split x", ks); return nullptr; + } +} + +template