diff --git a/lmdeploy/turbomind/builders/ffn.py b/lmdeploy/turbomind/builders/ffn.py index 5d88232c68..b62bb6383a 100644 --- a/lmdeploy/turbomind/builders/ffn.py +++ b/lmdeploy/turbomind/builders/ffn.py @@ -20,6 +20,9 @@ 'fuse_w1w3', ] +_SM90_FP8_FUSED_SILU_BLOCK = 128 +_SM90_BF16_FUSED_SILU_BLOCK = 64 + # --------------------------------------------------------------------------- # @transform_output_dim / @transform_input_dim helpers # --------------------------------------------------------------------------- @@ -31,6 +34,24 @@ def _interleave_w1w3(w1: torch.Tensor, w3: torch.Tensor) -> torch.Tensor: return torch.stack([w1, w3], dim=-1).reshape(w1.shape[:-1] + (-1,)).contiguous() +@transform_output_dim +def _block_pack_w1w3(w1: torch.Tensor, w3: torch.Tensor, *, + groups: int) -> torch.Tensor: + """Interleave gate/up by logical output groups. + + The elements per group self-adapt for each Linear tensor kind. For FP8, one logical 128-wide weight group + corresponds to one scale element. + """ + assert groups > 0, f'groups must be positive, got {groups}' + assert w1.shape[-1] % groups == 0, ( + f'output dim {w1.shape[-1]} not divisible by groups {groups}') + assert w3.shape[-1] == w1.shape[-1], ( + f'w1/w3 output dims differ: {w1.shape[-1]} != {w3.shape[-1]}') + w1g = w1.unflatten(-1, (groups, -1)) + w3g = w3.unflatten(-1, (groups, -1)) + return torch.stack([w1g, w3g], dim=-2).flatten(-3, -1).contiguous() + + @transform_output_dim def _chunk_w1w3(w1: torch.Tensor, w3: torch.Tensor, *, tp: int) -> torch.Tensor: @@ -49,61 +70,67 @@ def _chunk_w1w3(w1: torch.Tensor, w3: torch.Tensor, *, # --------------------------------------------------------------------------- +def _is_sm90() -> bool: + """Whether the current CUDA device is Hopper SM90.""" + return (torch.cuda.is_available() + and torch.cuda.get_device_capability() == (9, 0)) + + def _should_fuse_silu(w1_linear: Linear, act_type: str, is_moe: bool = False) -> bool: - """Determine if fused SiLU (interleave) should be used for w1+w3 fusion. + """Determine if fused SiLU should be used for w1+w3 fusion. Gold standard condition (from GEMM kernel constraints — trust it): - act_type == SiLU && (int4 || mxfp4 || fp8 || moe) && !(fp8 && SM90) + act_type == SiLU && (int4 || mxfp4 || fp8 || moe) + + Packing depends on format: FP8 uses [g128|u128|...] and SM90 BF16 uses + [g64|u64|...]; int4/mxfp4 use element interleave. Controllable via + config.fuse_silu after commit. """ if act_type not in ('', 'silu', 'SiLU'): return False - # Dense bf16/fp16 without MoE -> chunk, not interleave - weight = w1_linear.tensors.get('weight') - is_quantized = weight is not None and weight.element_size() < 2 + # SM90 BF16 dense uses the native fused-SiLU kernel. Other dense + # bf16/fp16 paths keep chunk layout and apply activation separately. + weight = w1_linear.tensors['weight'] + is_quantized = weight.element_size() < 2 if not is_quantized and not is_moe: - return False - - # FP8 on SM90 -> chunk - fmt = w1_linear.weight_format - if fmt is not None and fmt.name == 'fp8': - if torch.cuda.is_available(): - cap = torch.cuda.get_device_capability() - if cap == (9, 0): - return False + return weight.dtype == torch.bfloat16 and _is_sm90() # SM100+ grouped bf16 MoE: CublasGroupedKernel has no fused GatedSilu - if is_moe: - weight = w1_linear.tensors.get('weight') - if weight is not None and weight.dtype == torch.bfloat16: - if torch.cuda.is_available(): - cap = torch.cuda.get_device_capability() - if cap >= (10, 0): - return False + if (is_moe and weight.dtype == torch.bfloat16 + and torch.cuda.is_available() + and torch.cuda.get_device_capability() >= (10, 0)): + return False return True -def _can_fuse_w1w3(w1: Linear, tp: int) -> bool: +def _fused_silu_block(w1: Linear) -> int | None: + """Return the gate/up block width required by the fused SiLU kernel.""" + if w1.weight_format.name == 'fp8' and _is_sm90(): + return _SM90_FP8_FUSED_SILU_BLOCK + if _is_sm90() and w1.tensors['weight'].dtype == torch.bfloat16: + return _SM90_BF16_FUSED_SILU_BLOCK + return None + + +def _can_fuse_w1w3(w1: Linear, tp: int, *, + pack_block: int | None = None) -> bool: """Check whether w1+w3 fusion is safe for the given TP. - Fusion (interleave or chunk) concatenates w1 and w3 along the output dim. - For block-quantized formats (e.g. FP8 with block_out=128), the fused - scale count ``2 * cdiv(N/tp, block_out)`` must equal - ``cdiv(2*N/tp, block_out)``. This holds iff ``(N/tp) % block_out == 0``. - When it doesn't, the fused module's C++ allocation won't match the - concatenated scales and we must commit w1/w3 separately. + Fusion (interleave, block-pack, or chunk) concatenates w1 and w3 along the + output dim. For block-quantized formats, the local output must align to + ``block_out`` so concatenated scale counts match the fused allocation. + Native SM90 fused SiLU additionally requires format-specific gate/up groups. + Misaligned projections are committed separately. """ - if tp <= 1: - return True - fmt = w1.weight_format - if fmt is None or fmt.block_out is None: - return True - w = w1.tensors.get('weight') - if w is None: - return True - return (w.size(-1) // tp) % fmt.block_out == 0 + w = w1.tensors['weight'] + if w.size(-1) % tp != 0: + return False + block_out = w1.weight_format.block_out or 1 + required_block = math.lcm(block_out, pack_block or 1) + return (w.size(-1) // tp) % required_block == 0 def fuse_w1w3( @@ -124,10 +151,15 @@ def fuse_w1w3( block-scale alignment check in ``_can_fuse_w1w3``. """ fused_silu = _should_fuse_silu(w1, act_type, is_moe) - can_fuse = _can_fuse_w1w3(w1, tp) + pack_block = _fused_silu_block(w1) if fused_silu else None + can_fuse = _can_fuse_w1w3(w1, tp, pack_block=pack_block) if can_fuse: - if fused_silu: + if pack_block is not None: + inter_size = w1.tensors['weight'].shape[-1] + w1w3 = _block_pack_w1w3( + w1, w3, groups=inter_size // pack_block) + elif fused_silu: w1w3 = _interleave_w1w3(w1, w3) else: w1w3 = _chunk_w1w3(w1, w3, tp=tp) diff --git a/src/turbomind/core/module.h b/src/turbomind/core/module.h index 3e5e3257f6..0454d59d85 100644 --- a/src/turbomind/core/module.h +++ b/src/turbomind/core/module.h @@ -177,6 +177,13 @@ class Param { return slot_ ? *slot_ : Tensor{}; } + /// Replace the parameter slot (e.g. after quant alloc or MakeStridedPtrs). + void set(Tensor t) + { + TM_CHECK(slot_ != nullptr); + *slot_ = std::move(t); + } + explicit operator bool() const { return slot_ && static_cast(*slot_); diff --git a/src/turbomind/kernels/gemm/CMakeLists.txt b/src/turbomind/kernels/gemm/CMakeLists.txt index 7a1b210762..1be306b1f7 100644 --- a/src/turbomind/kernels/gemm/CMakeLists.txt +++ b/src/turbomind/kernels/gemm/CMakeLists.txt @@ -19,27 +19,102 @@ set(GEMM2_KERNELS_SM90 tma.cu kernel/sm90_16816_4.cu kernel/sm90_16816_8.cu - kernel/sm90_16816_16.cu kernel/sm90_64n32_8.cu + kernel/sm90_64n16_16.cu ) -set(GEMM2_ARCH_90_ENABLED FALSE) -set(_sm90_archs "${CMAKE_CUDA_ARCHITECTURES}") -list(FILTER _sm90_archs INCLUDE REGEX "^90") -if(_sm90_archs) - set(GEMM2_ARCH_90_ENABLED TRUE) -else() +# Bucket the selected archs by kernel family (ranges mirror arch.h). +set(GEMM2_SM70_ARCHS "") +set(GEMM2_SM75_ARCHS "") +set(GEMM2_SM80_ARCHS "") +set(GEMM2_SM90_ARCHS "") +foreach(arch IN LISTS CMAKE_CUDA_ARCHITECTURES) + if(arch MATCHES "^([0-9]+)(a)?(-real)?$") + set(arch_number "${CMAKE_MATCH_1}") + if(arch_number GREATER_EQUAL 70 AND arch_number LESS 75) + list(APPEND GEMM2_SM70_ARCHS "${arch}") + elseif(arch_number GREATER_EQUAL 75 AND arch_number LESS 80) + list(APPEND GEMM2_SM75_ARCHS "${arch}") + elseif(arch_number GREATER_EQUAL 80 AND arch_number LESS 90) + list(APPEND GEMM2_SM80_ARCHS "${arch}") + elseif(arch_number EQUAL 90) + list(APPEND GEMM2_SM90_ARCHS "${arch}") + elseif(arch_number GREATER_EQUAL 100) + # SM100/SM120 run the SM80 (s16816 mma.sync) kernels for now. + list(APPEND GEMM2_SM80_ARCHS "${arch}") + endif() + endif() +endforeach() + +if(NOT GEMM2_SM90_ARCHS) # When building for SM100+ without explicit SM90, still compile SM90 CUTLASS # kernels so the fat binary can run MoE models on H100 (CUTLASS fused path). set(_sm100_archs "${CMAKE_CUDA_ARCHITECTURES}") list(FILTER _sm100_archs INCLUDE REGEX "^100") if(_sm100_archs) - set(GEMM2_ARCH_90_ENABLED TRUE) - set(_sm90_archs "90") + set(GEMM2_SM90_ARCHS "90") message(STATUS "GEMM: auto-enabling SM90 CUTLASS kernels for H100 backward compatibility") endif() endif() +message(STATUS "GEMM SM70 CUDA archs: ${GEMM2_SM70_ARCHS}") +message(STATUS "GEMM SM75 CUDA archs: ${GEMM2_SM75_ARCHS}") +message(STATUS "GEMM SM80 CUDA archs: ${GEMM2_SM80_ARCHS}") +message(STATUS "GEMM SM90 CUDA archs: ${GEMM2_SM90_ARCHS}") + +function(configure_gemm2_kernel_target target) + set_property(TARGET ${target} PROPERTY POSITION_INDEPENDENT_CODE ON) + set_property(TARGET ${target} PROPERTY CUDA_RESOLVE_DEVICE_SYMBOLS ON) + target_compile_options(${target} PRIVATE + $<$: + -Xptxas=-v + --generate-line-info + --threads 16> + ) + target_link_libraries(${target} PRIVATE parser nvidia::cutlass::cutlass CUDA::cuda_driver) +endfunction() + +# cublas.cu registers for every arch. +add_library(gemm2_kernels STATIC cublas.cu) +configure_gemm2_kernel_target(gemm2_kernels) + +set(GEMM2_KERNEL_TARGETS gemm2_kernels) + +if(GEMM2_SM70_ARCHS) + add_library(gemm2_kernels_sm70 STATIC ${GEMM2_KERNELS_SM70}) + set_property(TARGET gemm2_kernels_sm70 PROPERTY + CUDA_ARCHITECTURES "${GEMM2_SM70_ARCHS}") + configure_gemm2_kernel_target(gemm2_kernels_sm70) + list(APPEND GEMM2_KERNEL_TARGETS gemm2_kernels_sm70) +endif() + +if(GEMM2_SM75_ARCHS) + add_library(gemm2_kernels_sm75 STATIC ${GEMM2_KERNELS_SM75}) + set_property(TARGET gemm2_kernels_sm75 PROPERTY + CUDA_ARCHITECTURES "${GEMM2_SM75_ARCHS}") + configure_gemm2_kernel_target(gemm2_kernels_sm75) + list(APPEND GEMM2_KERNEL_TARGETS gemm2_kernels_sm75) +endif() + +if(GEMM2_SM80_ARCHS) + add_library(gemm2_kernels_sm80 STATIC ${GEMM2_KERNELS_SM80}) + set_property(TARGET gemm2_kernels_sm80 PROPERTY + CUDA_ARCHITECTURES "${GEMM2_SM80_ARCHS}") + configure_gemm2_kernel_target(gemm2_kernels_sm80) + list(APPEND GEMM2_KERNEL_TARGETS gemm2_kernels_sm80) +endif() + +if(GEMM2_SM90_ARCHS) + # SM90 kernels only compile for 90/90a; avoid building them for sm_100. + add_library(gemm2_sm90 STATIC ${GEMM2_KERNELS_SM90}) + set_target_properties(gemm2_sm90 PROPERTIES + CUDA_ARCHITECTURES "${GEMM2_SM90_ARCHS}" + ) + target_compile_definitions(gemm2_sm90 PRIVATE -DCUTE_SM90_EXTENDED_MMA_SHAPES_ENABLED) + configure_gemm2_kernel_target(gemm2_sm90) + list(APPEND GEMM2_KERNEL_TARGETS gemm2_sm90) +endif() + add_library(gemm2 gemm.cu kernel.cu @@ -55,23 +130,22 @@ add_library(gemm2 tuner/sampler.cu tuner/stopping_criterion.cc tuner/params.cc - ${GEMM2_KERNELS_SM70} - ${GEMM2_KERNELS_SM75} - ${GEMM2_KERNELS_SM80} - cublas.cu moe_utils_v2.cu test/test_utils.cu ) target_link_libraries(gemm2 PRIVATE parser nvidia::cutlass::cutlass CUDA::cuda_driver) +# Retain file-scope Registrar objects (self-registration). Same pattern as attention_kernels. +foreach(target IN LISTS GEMM2_KERNEL_TARGETS) + target_link_libraries(gemm2 PUBLIC $) +endforeach() + # cublasGemmGroupedBatchedEx (CUDA 12.5+): grouped batched GEMM for MoE on SM100 -set(_has_sm100 FALSE) set(_archs_100 "${CMAKE_CUDA_ARCHITECTURES}") list(FILTER _archs_100 INCLUDE REGEX "^100") if(_archs_100 AND CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL "12.5") - set(_has_sm100 TRUE) - target_compile_definitions(gemm2 PRIVATE ENABLE_CUBLAS_GROUPED=1) + target_compile_definitions(gemm2_kernels PRIVATE ENABLE_CUBLAS_GROUPED=1) message(STATUS "GEMM: ENABLE_CUBLAS_GROUPED=1 (cublasGemmGroupedBatchedEx for MoE on SM100)") endif() @@ -83,56 +157,3 @@ target_compile_options(gemm2 PRIVATE ) set_property(TARGET gemm2 PROPERTY POSITION_INDEPENDENT_CODE ON) set_property(TARGET gemm2 PROPERTY CUDA_RESOLVE_DEVICE_SYMBOLS ON) - -if(GEMM2_ARCH_90_ENABLED) - # SM90 kernels only compile for 90/90a; avoid building them for sm_100. - add_library(gemm2_sm90 STATIC ${GEMM2_KERNELS_SM90}) - set_target_properties(gemm2_sm90 PROPERTIES - CUDA_ARCHITECTURES "${_sm90_archs}" - POSITION_INDEPENDENT_CODE ON - CUDA_RESOLVE_DEVICE_SYMBOLS ON - ) - target_compile_definitions(gemm2_sm90 PRIVATE -DCUTE_SM90_EXTENDED_MMA_SHAPES_ENABLED) - target_compile_options(gemm2_sm90 PRIVATE - $<$: - -Xptxas=-v - --generate-line-info - --threads 16> - ) - target_link_libraries(gemm2_sm90 PRIVATE parser nvidia::cutlass::cutlass CUDA::cuda_driver) - target_link_libraries(gemm2 PRIVATE gemm2_sm90) - - target_compile_definitions(gemm2 PRIVATE GEMM2_ARCH_90_ENABLED) -endif() - -if (BUILD_TEST) - # add_executable(test_gemm_v2 - # test/test_gemm_v2.cc - # ../../models/llama/LlamaLinear.cu - # test/reference.cu) - # target_link_libraries(test_gemm_v2 PRIVATE gemm2 core cublas quantization_kernels gpt_kernels) - - # add_executable(test_moe_utils test/test_moe_utils.cu test/test_utils.cu) - # target_link_libraries(test_moe_utils PRIVATE gemm2 core cublas) - - # if (NOT MSVC) - # FetchContent_Declare( - # repo-nvbench - # GIT_REPOSITORY https://github.com/NVIDIA/nvbench.git - # GIT_TAG d8dced8a64d9ce305add92fa6d274fd49b569b7e - # ) - - # set(NVBench_ENABLE_EXAMPLES OFF) - # set(NVBench_ENABLE_TESTING OFF) - # set(BUILD_SHARED_LIBS OFF) - - # FetchContent_MakeAvailable(repo-nvbench) - - # add_executable(gemm_bench - # test/gemm_bench.cu - # # test/test_utils.cu - # test/quantization.cu - # test/reference.cu) - # target_link_libraries(gemm_bench PRIVATE gemm2 core nvbench::nvbench cublas) - # endif () -endif () diff --git a/src/turbomind/kernels/gemm/arch.h b/src/turbomind/kernels/gemm/arch.h index ad021aa95b..dda71e60ee 100644 --- a/src/turbomind/kernels/gemm/arch.h +++ b/src/turbomind/kernels/gemm/arch.h @@ -35,7 +35,7 @@ struct Sm100: Arch<1000, 1200> { static constexpr int value = 1000; }; -// SM12.x (e.g. sm_120): use same CUTLASS SM90 kernel family as pre-PR Sm90+ range +// SM12.x (e.g. sm_120): no native kernels; falls back to the SM80 s16816 family struct Sm120: Arch<1200, 1300> { static constexpr int value = 1200; }; @@ -50,9 +50,9 @@ inline bool is_arch_compatible(int karch, int darch) case 750: return Sm75::is_compatible(darch); case 800: - return Sm80::is_compatible(darch); + return Sm80::is_compatible(darch) || Sm100::is_compatible(darch) || Sm120::is_compatible(darch); case 900: - return Sm90::is_compatible(darch) || Sm120::is_compatible(darch); + return Sm90::is_compatible(darch); case 1000: return Sm100::is_compatible(darch); case 1200: diff --git a/src/turbomind/kernels/gemm/arch/mma_sm70.h b/src/turbomind/kernels/gemm/arch/mma_sm70.h index f5a1d4f0cf..30070e43ae 100644 --- a/src/turbomind/kernels/gemm/arch/mma_sm70.h +++ b/src/turbomind/kernels/gemm/arch/mma_sm70.h @@ -18,7 +18,7 @@ struct SM70_MMA_884 { static constexpr int kThreadCount = 32; - static constexpr auto kOpClass = OpClass::kMMA_s884; + static constexpr auto kOpClass = OpClass::kMMA_h884; using FragA = Array; using FragB = Array; diff --git a/src/turbomind/kernels/gemm/arch/mma_sm80.h b/src/turbomind/kernels/gemm/arch/mma_sm80.h index 8a197a1617..2372cb8f8c 100644 --- a/src/turbomind/kernels/gemm/arch/mma_sm80.h +++ b/src/turbomind/kernels/gemm/arch/mma_sm80.h @@ -17,7 +17,7 @@ struct SM80_MMA_16x8x16_F32_F16_F16_F32_TN { static constexpr int kThreadCount = 32; - static constexpr auto kOpClass = OpClass::kMMA_s16816; + static constexpr auto kOpClass = OpClass::kMMA_h16816; using FragA = Array; using FragB = Array; diff --git a/src/turbomind/kernels/gemm/context.cu b/src/turbomind/kernels/gemm/context.cu index 3912a0c4d8..192d3b6da5 100644 --- a/src/turbomind/kernels/gemm/context.cu +++ b/src/turbomind/kernels/gemm/context.cu @@ -169,9 +169,23 @@ std::vector Context::Populate(const Kernel& kernel, const PopulatePa const auto& desc = kernel.desc(); const auto& info = kernel.info(); - const int64_t tiled_shape_m = cdiv(m, desc.cta_tile.x * (desc.group_axis == 0 ? num : 1)); - const int64_t tiled_shape_n = cdiv(n, desc.cta_tile.y * (desc.group_axis == 1 ? num : 1)); - const int chunk_cnt_k = cdiv(k, kernel.chunk_size_k()); + // Along a group axis, `m`/`n` is the packed sum across groups. Estimate tiles *per + // group* from the average group size so cost tracks real per-group tiling + // (`num * cdiv(m_e, TILE)`). Using `cdiv(packed, TILE * num)` jumps at every + // `k * TILE * num` boundary and incorrectly demotes large TILE_M for MoE. + const int64_t tiled_shape_m = [&]() -> int64_t { + if (desc.group_axis == 0 && num > 1) { + return cdiv(std::max(m / num, 1), (int64_t)desc.cta_tile.x); + } + return cdiv((int64_t)m, (int64_t)desc.cta_tile.x); + }(); + const int64_t tiled_shape_n = [&]() -> int64_t { + if (desc.group_axis == 1 && num > 1) { + return cdiv(std::max(n / num, 1), (int64_t)desc.cta_tile.y); + } + return cdiv((int64_t)n, (int64_t)desc.cta_tile.y); + }(); + const int chunk_cnt_k = cdiv(k, kernel.chunk_size_k()); // Despite we only have sm_count * constant tensor cores, this is the granularity for scheduling const int concurrency = sm_count_ * kernel.info().max_active_ctas; diff --git a/src/turbomind/kernels/gemm/convert.h b/src/turbomind/kernels/gemm/convert.h index fe5ec90f5c..6be71a5778 100644 --- a/src/turbomind/kernels/gemm/convert.h +++ b/src/turbomind/kernels/gemm/convert.h @@ -26,8 +26,10 @@ std::array GetConverters(DataType data_type, bool grouped, int sm); +// TM_GEMM_WEIGHT_PACK: unset/-1 = auto, 0 = force plain, 1 = force pack +int WeightPackEnv(); + // Free with `cudaFree` void* MakeStridedPtrs(const std::vector>& ptrs, cudaStream_t stream); -void* MakeBlockedPtrs(const std::vector>& ptrs, cudaStream_t stream); } // namespace turbomind::gemm diff --git a/src/turbomind/kernels/gemm/convert_v3.cu b/src/turbomind/kernels/gemm/convert_v3.cu index 6f9b347f2f..dffa11e482 100644 --- a/src/turbomind/kernels/gemm/convert_v3.cu +++ b/src/turbomind/kernels/gemm/convert_v3.cu @@ -1,5 +1,6 @@ #include +#include #include "src/turbomind/core/check.h" #include "src/turbomind/kernels/gemm/arch.h" @@ -82,6 +83,24 @@ constexpr auto operator|(constant, constant) return constant{}; } +int WeightPackEnv() +{ + static const int v = [] { + const char* p = std::getenv("TM_GEMM_WEIGHT_PACK"); + if (!p) { + return -1; + } + if (p[0] == '0' && p[1] == '\0') { + return 0; + } + if (p[0] == '1' && p[1] == '\0') { + return 1; + } + return -1; + }(); + return v; +} + std::array GetConverters(DataType data_type, DataType weight_type, // DataType input_type, @@ -103,12 +122,22 @@ std::array GetConverters(DataType data_type, constexpr Sm75 sm75{}; constexpr Sm70 sm70{}; + const int pack_env = WeightPackEnv(); + if (pack_env == 0) { + return {}; + } + if (weight_type == kHalf || weight_type == kBfloat16) { constexpr Cvt W; if (grouped) { - // SM10.x only: CublasGroupedKernel (cublasGemmGroupedBatchedEx) expects standard (K,N) - if (sm >= 100 && sm < 120) - return {}; + if (pack_env != 1) { + // SM10.x: CublasGroupedKernel expects standard (K,N) + if (sm >= 100 && sm < 120) + return {}; + // SM90: plain B for native GMMA (LinearWeight prepare stores physical (N,K)) + if (sm >= 90 && sm < 100) + return {}; + } // clang-format off if (sm >= 80) return {W(sm8_, kRow, s16816h | B | _1), {}}; if (sm == 75) return {W(sm75, kRow, s16816h | B | _1), {}}; @@ -207,37 +236,4 @@ void* MakeStridedPtrs(const std::vector>& ptrs, cudaStream return ptr; } -namespace { - -template -__global__ void fill_blocked_ptrs(Array src, void** dst, int n) -{ - const int idx = threadIdx.x + blockIdx.x * blockDim.x; - if (idx < n) { - dst[idx] = src[idx]; - } -} - -} // namespace - -void* MakeBlockedPtrs(const std::vector>& ptrs, cudaStream_t stream) -{ - constexpr int N = 64; - Array src{}; - static_assert(sizeof(src) <= 4096); // max parameter size for cuda11 - void** dst{}; - cudaMallocAsync(&dst, sizeof(void*) * ptrs.size(), stream); - for (int i = 0; i < (int)ptrs.size(); i += N) { - const int n = std::min(ptrs.size() - i, N); - for (int j = 0; j < n; ++j) { - auto& [p, s] = ptrs[i + j]; - src[j] = p; - } - fill_blocked_ptrs<<<1, N, 0, stream>>>(src, dst, n); - dst += n; - } - TM_CUDA_CHECK(cudaGetLastError()); - return dst - ptrs.size(); -} - } // namespace turbomind::gemm diff --git a/src/turbomind/kernels/gemm/cublas.cu b/src/turbomind/kernels/gemm/cublas.cu index 461c0d5036..c4424fc6c3 100644 --- a/src/turbomind/kernels/gemm/cublas.cu +++ b/src/turbomind/kernels/gemm/cublas.cu @@ -8,7 +8,7 @@ #include "src/turbomind/kernels/gemm/desc.h" #include "src/turbomind/kernels/gemm/kernel.h" #include "src/turbomind/kernels/gemm/matrix_ptr.h" -#include "src/turbomind/kernels/gemm/registry.h" +#include "src/turbomind/kernels/gemm/registrar.h" #include "src/turbomind/kernels/gemm/types.h" #include "src/turbomind/utils/cuda_utils.h" @@ -56,11 +56,15 @@ public: const MatrixLayout& Cdesc, void* D, const MatrixLayout& Ddesc, + void* W, + const MatrixLayout& Wdesc, int swizzle, int splits, Workspace& workspace, cudaStream_t stream) override { + (void)W; + (void)Wdesc; cublasOperation_t transa = Adesc.order == kColMajor ? CUBLAS_OP_N : CUBLAS_OP_T; cublasOperation_t transb = Bdesc.order == kColMajor ? CUBLAS_OP_N : CUBLAS_OP_T; @@ -162,11 +166,6 @@ private: size_t workspace_size_{}; }; -void Registry::cublas_float() -{ - Add(std::make_unique()); -} - #if defined(ENABLE_CUBLAS_GROUPED) // Grouped GEMM via cublasGemmGroupedBatchedEx (CUDA 12.5+, SM100). @@ -227,13 +226,16 @@ public: const MatrixLayout& Cdesc, void* D, const MatrixLayout& Ddesc, + void* W, + const MatrixLayout& Wdesc, int swizzle, int splits, Workspace& workspace, cudaStream_t stream) override { - if (!Adesc.offsets || !Ddesc.offsets || Adesc.offsets == reinterpret_cast(1) - || Ddesc.offsets == reinterpret_cast(1)) { + (void)W; + (void)Wdesc; + if (!Adesc.offsets || !Ddesc.offsets) { fprintf( stderr, "[TM][GEMM] CublasGrouped: missing or invalid offsets (Adesc.offsets=%p Ddesc.offsets=%p) num=%d rows=%d\n", @@ -471,11 +473,17 @@ private: size_t workspace_size_{}; }; -void Registry::sm100_cublas_grouped_float() -{ - Add(std::make_unique()); -} - #endif // ENABLE_CUBLAS_GROUPED +namespace { +Registrar reg([](Collector& c, int arch) { + c.add(std::make_unique()); +#if defined(ENABLE_CUBLAS_GROUPED) + if (Sm100::is_compatible(arch)) { + c.add(std::make_unique()); + } +#endif +}); +} // namespace + } // namespace turbomind::gemm diff --git a/src/turbomind/kernels/gemm/desc.h b/src/turbomind/kernels/gemm/desc.h index 8c52d18926..cae8ecf961 100644 --- a/src/turbomind/kernels/gemm/desc.h +++ b/src/turbomind/kernels/gemm/desc.h @@ -3,6 +3,7 @@ #pragma once #include +#include #include #include @@ -88,9 +89,10 @@ inline std::string to_string(const GemmDesc& d) enum class OpClass { kSIMT, - kMMA_s884, - kMMA_s16816, - kGMMA_s64n16 + kMMA_h884, + kMMA_h16816, + kGMMA_h64n16, + kGMMA_q64n32 }; inline const char* to_string(OpClass op) @@ -98,10 +100,14 @@ inline const char* to_string(OpClass op) switch (op) { case OpClass::kSIMT: return "simt"; - case OpClass::kMMA_s884: - return "s884"; - case OpClass::kMMA_s16816: - return "s16816"; + case OpClass::kMMA_h884: + return "mma_h884"; + case OpClass::kMMA_h16816: + return "mma_h16816"; + case OpClass::kGMMA_h64n16: + return "gmma_h64n16"; + case OpClass::kGMMA_q64n32: + return "gmma_q64n32"; default: return "unknown_op_cls"; } @@ -111,6 +117,8 @@ inline const char* to_string(OpClass op) struct KernelDesc { int arch; OpClass op_class; + uint32_t algo; // opaque; each KernelImpl defines its own bit-field layout (0 = default) + Order raster; // tile-scheduler raster order DataType type_a; DataType type_b; DataType type_c; @@ -130,11 +138,13 @@ struct KernelDesc { int policy_b; int3 cta_tile; int3 mma_tile; + int3 atom_layout; int2 cluster_shape; int3 align; int2 c_tile; int stages; bool split_k; + bool supports_fused_silu; int group_axis; int backend; bool transpose; @@ -179,11 +189,13 @@ inline KernelDesc transpose(const KernelDesc& d) k.policy_a = d.policy_b; k.policy_b = d.policy_a; + k.raster = ~d.raster; auto swap = [](auto& v) { std::swap(v.x, v.y); }; swap(k.cta_tile); swap(k.mma_tile); + swap(k.atom_layout); swap(k.cluster_shape); swap(k.align); swap(k.c_tile); diff --git a/src/turbomind/kernels/gemm/dispatch_cache.cu b/src/turbomind/kernels/gemm/dispatch_cache.cu index 19bd2ceb84..7f5bd5c1d7 100644 --- a/src/turbomind/kernels/gemm/dispatch_cache.cu +++ b/src/turbomind/kernels/gemm/dispatch_cache.cu @@ -5,6 +5,8 @@ #include "src/turbomind/kernels/gemm/kernel.h" #include "src/turbomind/kernels/gemm/types.h" #include +#include +#include #include #include #include @@ -28,6 +30,8 @@ static inline decltype(auto) as_tuple(const KernelDesc& d) { return std::tie(d.arch, d.op_class, + d.algo, + d.raster, d.type_a, d.type_b, d.type_c, @@ -47,11 +51,13 @@ static inline decltype(auto) as_tuple(const KernelDesc& d) d.policy_b, d.cta_tile, d.mma_tile, + d.atom_layout, d.cluster_shape, d.align, d.c_tile, d.stages, d.split_k, + d.supports_fused_silu, d.backend, d.transpose, d.group_axis); @@ -69,6 +75,16 @@ static inline bool operator==(const KernelDesc& a, const KernelDesc& b) namespace { +constexpr char kDispatchCacheMagic[8] = {'T', 'M', 'G', 'E', 'M', 'M', '2', '\0'}; +constexpr std::uint32_t kDispatchCacheVersion = 4; + +struct Header { + char magic[sizeof(kDispatchCacheMagic)]; + std::uint32_t version; + std::uint32_t record_size; + std::uint64_t record_count; +}; + struct Record { GemmDesc gemm; KernelDesc kernel; @@ -81,6 +97,12 @@ struct Record { void ExportDispatchCache(std::ostream& os, const std::vector>& entries) { + Header header{}; + std::memcpy(header.magic, kDispatchCacheMagic, sizeof(header.magic)); + header.version = kDispatchCacheVersion; + header.record_size = sizeof(Record); + header.record_count = entries.size(); + os.write((const char*)&header, sizeof(header)); for (const auto& [g, spec] : entries) { Record record{}; @@ -98,16 +120,30 @@ void ImportDispatchCache(std::istream& is, const std::vector& kernels) { is.seekg(0, is.end); - const auto size_in_bytes = is.tellg(); + const std::streamoff size_in_bytes = is.tellg(); is.seekg(0, is.beg); - if (size_in_bytes % sizeof(Record)) { - std::cerr << "File size is not a multiple of record size, faild to import records.\n"; + if (size_in_bytes < static_cast(sizeof(Header))) { + std::cerr << "Dispatch cache has no supported format header.\n"; + return; } - const int n = size_in_bytes / sizeof(Record); + Header header{}; + is.read((char*)&header, sizeof(header)); + if (std::memcmp(header.magic, kDispatchCacheMagic, sizeof(header.magic)) != 0 + || header.version != kDispatchCacheVersion) { + std::cerr << "Unsupported dispatch cache format.\n"; + return; + } + + const std::streamoff payload_size = size_in_bytes - sizeof(Header); + if (header.record_size != sizeof(Record) || payload_size % sizeof(Record) + || header.record_count != static_cast(payload_size / sizeof(Record))) { + std::cerr << "Dispatch cache size does not match its header.\n"; + return; + } - for (int i = 0; i < n; ++i) { + for (std::uint64_t i = 0; i < header.record_count; ++i) { Record record; is.read((char*)&record, sizeof(Record)); @@ -152,13 +188,13 @@ inline decltype(auto) as_tuple(const GemmDesc& d) d.quant_a.group_size, d.quant_b.type, d.quant_b.group_size, + d.epilogue, d.batch_dim, d.group_axis, d.m, d.n, d.k, d.num); - // Note: `d.epilogue` is not used yet } } // namespace diff --git a/src/turbomind/kernels/gemm/gemm.cu b/src/turbomind/kernels/gemm/gemm.cu index 7ae668ae59..88c0741db5 100644 --- a/src/turbomind/kernels/gemm/gemm.cu +++ b/src/turbomind/kernels/gemm/gemm.cu @@ -12,6 +12,7 @@ #include "src/turbomind/kernels/gemm/tuner/sampler.h" #include "src/turbomind/kernels/gemm/types.h" #include +#include #include #include #include @@ -194,6 +195,15 @@ struct Gemm::Impl { specs = Sampler{*measurer_, tuning_.clusters}.Run(specs, launch_func, st); + if (std::getenv("TM_GEMM_TUNE_VERBOSE")) { + for (const auto& s : specs) { + std::cout << "[tune] " << to_string(ctx.desc()) << " " << s.kernel->name() // + << " swizzle=" << s.swizzle // + << " splits=" << s.splits // + << " measured=" << s.measured << "\n"; + } + } + // for (const auto& s : specs) { // std::cout << s.kernel->name() // // << " swizzle=" << s.swizzle // @@ -259,6 +269,8 @@ int Gemm::Run(const Operation& operation, const MatrixLayout& Cdesc, void* D, const MatrixLayout& Ddesc, + void* W, + const MatrixLayout& Wdesc, const Workspace& workspace, cudaStream_t stream) { @@ -289,6 +301,8 @@ int Gemm::Run(const Operation& operation, Cdesc, D, Ddesc, + W, + Wdesc, spec.swizzle, spec.splits, _workspace, diff --git a/src/turbomind/kernels/gemm/gemm.h b/src/turbomind/kernels/gemm/gemm.h index 37a29324b4..e2a999cf31 100644 --- a/src/turbomind/kernels/gemm/gemm.h +++ b/src/turbomind/kernels/gemm/gemm.h @@ -35,6 +35,8 @@ class Gemm { const MatrixLayout& Cdesc, void* D, const MatrixLayout& Ddesc, + void* W, + const MatrixLayout& Wdesc, const Workspace& workspace, cudaStream_t stream); diff --git a/src/turbomind/kernels/gemm/gemm_universal_sm90.h b/src/turbomind/kernels/gemm/gemm_universal_sm90.h index 1638257984..d8aea36712 100644 --- a/src/turbomind/kernels/gemm/gemm_universal_sm90.h +++ b/src/turbomind/kernels/gemm/gemm_universal_sm90.h @@ -131,7 +131,7 @@ struct GemmUniversalSm90 { static constexpr int CTA_N = MMA_ATOM_N * kWorkGroupN; static constexpr int CTA_K = 128; - static constexpr int WARPGORUPS = kWorkGroupM * kWorkGroupN; + static constexpr int WARPGROUPS = kWorkGroupM * kWorkGroupN; static constexpr int MMA_M = MMA_ATOM_M * kWorkGroupM; static constexpr int MMA_N = MMA_ATOM_N * kWorkGroupN; @@ -153,7 +153,7 @@ struct GemmUniversalSm90 { static constexpr int WARPGROUP_SIZE = 128; - static constexpr int CTA_SIZE = WARPGROUP_SIZE * (WARPGORUPS + 1); + static constexpr int CTA_SIZE = WARPGROUP_SIZE * (WARPGROUPS + 1); using Ta = __nv_fp8_e4m3; using Tb = __nv_fp8_e4m3; @@ -185,7 +185,7 @@ struct GemmUniversalSm90 { }; Source source; __align__(128) Array C; - __align__(128) float UV[WARPGORUPS][round_up(CTA_M_U * CTA_N_V, 32)]; + __align__(128) float UV[WARPGROUPS][round_up(CTA_M_U * CTA_N_V, 32)]; __align__(128) uint64_t producer_bar[Stages]; __align__(128) uint64_t consumer_bar[Stages]; }; @@ -213,7 +213,7 @@ struct GemmUniversalSm90 { PRAGMA_UNROLL for (int s = 0; s < Stages; ++s) { ProducerBar::init(&producer_bar[s], 1); - ConsumerBar::init(&consumer_bar[s], kClusterSize * WARPGORUPS); + ConsumerBar::init(&consumer_bar[s], kClusterSize * WARPGROUPS); } cutlass::arch::fence_view_async_shared(); if constexpr (kClusterSize > 1) { @@ -225,7 +225,7 @@ struct GemmUniversalSm90 { const int warpgroup_id = cutlass::canonical_warp_group_idx(); - if (warpgroup_id == WARPGORUPS) { + if (warpgroup_id == WARPGROUPS) { cutlass::arch::warpgroup_reg_dealloc<32>(); static_assert(CTA_M % kMulticastA == 0); @@ -241,7 +241,7 @@ struct GemmUniversalSm90 { auto& smem_U = storage.source.U; auto& smem_V = storage.source.V; - if (threadIdx.x == WARPGORUPS * WARPGROUP_SIZE) { + if (threadIdx.x == WARPGROUPS * WARPGROUP_SIZE) { cutlass::PipelineState write_state{0, 1, 0}; while (sched.next()) { auto [valid_cta_tile_p, cluster_tile_p] = sched.is_valid_tile(); @@ -458,7 +458,7 @@ struct GemmUniversalSm90 { cute::tma_store_wait<0>(); } - cutlass::arch::NamedBarrier(WARPGORUPS * WARPGROUP_SIZE).sync(); + cutlass::arch::NamedBarrier(WARPGROUPS * WARPGROUP_SIZE).sync(); // epilogue const int warp_id = threadIdx.x / WARP_SIZE; @@ -489,7 +489,7 @@ struct GemmUniversalSm90 { } } cute::tma_store_fence(); // visibility: smem -> async proxy - cutlass::arch::NamedBarrier(WARPGORUPS * WARPGROUP_SIZE).sync(); + cutlass::arch::NamedBarrier(WARPGROUPS * WARPGROUP_SIZE).sync(); if (threadIdx.x == 0) { cute::SM90_TMA_STORE_2D::copy(&tm_c, &smem_C, offset_m, offset_n); diff --git a/src/turbomind/kernels/gemm/gemm_universal_sm90_bf16.h b/src/turbomind/kernels/gemm/gemm_universal_sm90_bf16.h new file mode 100644 index 0000000000..e186484a7f --- /dev/null +++ b/src/turbomind/kernels/gemm/gemm_universal_sm90_bf16.h @@ -0,0 +1,1187 @@ +#pragma once + +/* + * SM90 dense BF16 GEMM — CuTe cooperative mainloop (PipelineTmaAsync + TiledMma). + * + * Operand swap (API vs GMMA): + * LlamaLinear / KernelImpl: A = activations (M_batch, K), B = weights (K, N_out) + * → host TMA: act as (M,K) row-major; weight transposed to (N,K). + * GMMA TileShape = (OUT, BATCH, K) = (TILE_N, TILE_M, TILE_K) + * → GMMA-A SMEM = weight (OUT × K), K-major SoT (SmemLayoutAtomA) + * → GMMA-B SMEM = act (BATCH × K), K-major SoT (SmemLayoutAtomB) + * Epilogue maps GMMA C (OUT, BATCH) → problem C (M_batch, N_out) = (BATCH, OUT). + * + * Layout contract (Task 1 SoT): + * SmemLayout{A,B} = tile_to_shape(SmemLayoutAtom*, Shape, Step<_1,_2,_3>) + * Host CUtensorMap uses CU_TENSOR_MAP_SWIZZLE_128B (KernelImplSm90Bf16) matching + * Layout_K_SW128_Atom for TILE_K=64 BF16. Multicast boxes are multiples of the + * 8-row SW128 atom so linear mc_offset*TILE_K stitching matches the composed SoT. + * + * Layout gate — producer-store TV vs consumer GMMA TV (same ComposedLayout SoT from + * gmma_bf16_sm90.h / GmmaBF16Traits::{SmemLayoutAtomA,SmemLayoutAtomB}): + * Weight (GMMA-A): SmemLayoutA = tile_to_shape(SmemLayoutAtomA, (TILE_N,TILE_K,Stages), + * Step<_1,_2,_3>). Producer TMA (host CUtensorMap SW128) writes the multicast box + * into stage pipe of that ComposedLayout; consumer TiledMma DescriptorIterator + * (partition_A → make_fragment_A) reads the same SoT — no alternate swizzle. + * Act (GMMA-B): SmemLayoutB = tile_to_shape(SmemLayoutAtomB, (TILE_M,TILE_K,Stages), + * Step<_1,_2,_3>). Dense: TMA store TV ↔ GMMA load TV share SmemLayoutAtomB. + * Indexed-A gather (grouped): full producer WG (128 threads) cooperative cp.async + * through SmemLayoutAtomB / SmemLayoutB_2D (`&sB(m,k)`); GMEM addressing only + * (idxs[m0+offset_m+m] → token row). All 128 threads are Pipeline Producers so + * producer_acquire runs lockstep (no mid-stage gather_bar). Dynamic TileScheduler + * folded onto warp0 (next → acquire → release); producers_bar unused here. + * Indexed load as iterator_sm80 (idxs→src_data_vec_, +=TILE_K); TMA before gather. + * + * Barrier contract (PipelineTmaAsync full = ClusterTransactionBarrier): + * Dense / blocked-A grouped: expect_tx(weight+act bytes); both operands TMA-arrive + * via complete_tx (act + weight). Grouped/flat blocked: TMA maps and scheduler + * offsets prepared by prepare_tma_descs_sm90_bf16; GEMM indexes by group_idx + * (no fence_acquire — that is only for in-kernel tensormap replace). + * Indexed-A (v3 / PTX Example-2 .noinc contract): expect_tx(weight bytes only). + * Full-barrier init arrive count = 1 (leader arrive_and_expect_tx) + N_noinc + * (one cpasync_barrier_arrive_noinc per gather thread). PTX: with .noinc, init + * MUST account for those arrive-ons. Per stage: acquire → gather → arrive_noinc + * → weight TMA .with(*bar). No producer cp_async_wait / no software act + * complete_tx. Indexed A does not TMA-multicast; each CTA gathers full TILE_M. + * Weight TMA multicast OK. + * + * Epilogue (CUTLASS Sm90TmaWarpSpecialized-style): + * EpilogueTile = (_128,_32) for STSM R2S into SW128 SMEM. Host TMA box is + * (_64,_32) — bf16 SW128 make_2d_tma_desc requires innermost 64. Device issues + * 2× SM90_TMA_STORE per epi tile along OUT. Warp0 after fence + EpilogueBarrier. + */ + +#include +#include +#include + +#include + +#include "cute/algorithm/gemm.hpp" +#include "cute/arch/cluster_sm90.hpp" +#include "cute/arch/copy_sm80.hpp" +#include "cute/arch/copy_sm90.hpp" +#include "cute/arch/copy_sm90_tma.hpp" +#include "cute/tensor.hpp" + +#include "cutlass/arch/barrier.h" +#include "cutlass/arch/reg_reconfig.h" +#include "cutlass/array.h" +#include "cutlass/cutlass.h" +#include "cutlass/pipeline/sm90_pipeline.hpp" + +#include "src/turbomind/core/data_type.h" + +#include "src/turbomind/kernels/core/array.h" +#include "src/turbomind/kernels/core/common.h" +#include "src/turbomind/kernels/core/layout.h" +#include "src/turbomind/kernels/core/smem.h" + +#include "src/turbomind/kernels/gemm/arch.h" +#include "src/turbomind/kernels/gemm/cp_async.h" +#include "src/turbomind/kernels/gemm/gmma_bf16_sm90.h" +#include "src/turbomind/kernels/gemm/matrix_ptr.h" +#include "src/turbomind/kernels/gemm/scheduler.cuh" +#include "src/turbomind/kernels/gemm/sm90_bf16_traits.h" +#include "src/turbomind/kernels/gemm/sm90_utils.h" +#include "src/turbomind/kernels/gemm/types.h" +#include "src/turbomind/kernels/gemm/utils.h" + +namespace turbomind::gemm { + +namespace detail { + +// CuTe TMA issue from host CUtensorMap (KernelImpl): Copy_Atom + .with(*bar[, mcast]). +// Descriptor stays host-built; in-kernel path matches PIPELINING.md / TMA.md contract. +// (copy_traits_sm90_tma.hpp arrives via cute/tensor.hpp → copy_atom.hpp; do not +// include it before tensor.hpp — that breaks CuTe's include order.) +template +__device__ void tma_load_with_barrier(const cute::TmaDescriptor* desc, + uint64_t* bar, + Element* smem, + int crd0, + int crd1, + uint16_t mcast_mask, + uint64_t cache_hint = (uint64_t)cute::TMA::CacheHintSm90::EVICT_NORMAL) +{ + constexpr int kNumBits = BoxMN * BoxK * (int)cute::sizeof_bits_v; + constexpr int kNumVals = BoxMN * BoxK; + + // Dummy Aux — descriptor is swapped in via .with(desc, *bar[, mask]). + using Aux = cute:: + AuxTmaParams, cute::Layout>, cute::Swizzle<0, 4, 3>>; + + auto g = cute::make_tensor(cute::make_inttuple_iter(crd0, crd1), cute::Layout>{}); + auto s = cute::make_tensor(cute::make_smem_ptr(smem), cute::Layout>{}); + + if constexpr (Multicast > 1) { + using Traits = cute::Copy_Traits, Aux>; + using Atom = cute::Copy_Atom; + Atom tma{Traits{cute::TmaDescriptor{}, Aux{}}}; + cute::copy(tma.with(desc, *bar, mcast_mask, (cute::TMA::CacheHintSm90)cache_hint), g, s); + } + else { + using Traits = cute::Copy_Traits, Aux>; + using Atom = cute::Copy_Atom; + Atom tma{Traits{cute::TmaDescriptor{}, Aux{}}}; + cute::copy(tma.with(desc, *bar, 0, (cute::TMA::CacheHintSm90)cache_hint), g, s); + (void)mcast_mask; + } +} + +// STSM atom aliases, tiered by per-WG accumulator vals/thread: big path must match +// pre-templatize (U32x4 / U16x8) exactly; small (4 vals) covers the narrow per-WG +// GMMA-N of 1x2 tiles with TILE_M = 16. (TILE_M = 8 on 1x2 would need 2-val atoms, +// but its per-WG GMMA-N of 4 is already below the GMMA atom minimum of 8.) +template +struct EpiStsmAtoms { + static_assert(kStsmVals >= 8); + using CopyAtomC = cute::Copy_Atom; + using CopyOpR2S = cute::SM90_U16x8_STSM_T; +}; +template<> +struct EpiStsmAtoms<4> { + using CopyAtomC = cute::Copy_Atom; + using CopyOpR2S = cute::SM90_U16x4_STSM_T; +}; + +__device__ __forceinline__ float silu_mul(float g, float u) +{ + return fdividef(g, 1.f + expf(-g)) * u; +} + +// Device TMA map helpers for MoE prepare kernel (copy → replace addr/dim1 → publish). +__device__ __forceinline__ void copy_tma_desc(CUtensorMap* dst, const CUtensorMap* src, int lane) +{ + constexpr int kWords = (int)(sizeof(CUtensorMap) / sizeof(uint2)); + if (lane < kWords) { + ((uint2*)dst)[lane] = ((const uint2*)src)[lane]; + } +} + +__device__ __forceinline__ void replace_tma_addr_dim1(CUtensorMap* desc, void* global_addr, int dim1) +{ + uint32_t uint_ptr = cast_smem_ptr_to_uint(desc); + // clang-format off + asm volatile("tensormap.replace.tile.global_address.shared::cta.b1024.b64 [%0], %1;" ::"r"(uint_ptr), "l"(global_addr)); + asm volatile("tensormap.replace.tile.global_dim.shared::cta.b1024.b32 [%0], 1, %1;" ::"r"(uint_ptr), "r"(dim1)); + // clang-format on +} + +__device__ __forceinline__ void publish_tma_desc(CUtensorMap* gmem_desc, CUtensorMap* smem_desc) +{ + uint32_t uint_ptr = cast_smem_ptr_to_uint(smem_desc); + // clang-format off + asm volatile("tensormap.cp_fenceproxy.global.shared::cta.tensormap::generic.release.gpu.sync.aligned [%0], [%1], 128;" :: "l"(gmem_desc), "r"(uint_ptr)); + // clang-format on +} + +template +__device__ __forceinline__ void rebase_publish_tma_descs(CUtensorMap* gmem_out, + CUtensorMap* smem_desc, + Array templates, + Array global_addrs, + Array dims, + int stride_desc_idx, + uint64_t stride_bytes, + int lane) +{ + PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + copy_tma_desc(&smem_desc[i], templates[i], lane); + } + __syncwarp(); + if (lane == 0) { + PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + replace_tma_addr_dim1(&smem_desc[i], global_addrs[i], dims[i]); + } + replace_tma_global_stride(&smem_desc[stride_desc_idx], stride_bytes); + } + __syncwarp(); + PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + publish_tma_desc(&gmem_out[i], &smem_desc[i]); + } + __syncwarp(); +} + +} // namespace detail + +// Rebase grouped TMA templates and materialize scheduler offsets in workspace. +// Indexed: [B, C]. Blocked: [A, B, C]. Flat input is one blocked group. +template +__global__ void __launch_bounds__(32, 1) prepare_tma_descs_sm90_bf16(const __grid_constant__ CUtensorMap tm_a, + const __grid_constant__ CUtensorMap tm_b, + const __grid_constant__ CUtensorMap tm_c, + MatrixParam param_A, + MatrixParam param_B, + MatrixParam param_C, + CUtensorMap* out, + int* offsets, + int M_total, + int N) +{ + constexpr int kNumAB = (kStridingA == Striding::kBlocked) ? 2 : 1; + constexpr int kNum = kNumAB + 1; + + __shared__ __align__(128) CUtensorMap smem_desc[kNum]; + + const int g = (int)blockIdx.x; + const int lane = (int)threadIdx.x & 31; + + using Ta = nv_bfloat16; + using Tb = nv_bfloat16; + using Tc = nv_bfloat16; + + const int m0 = param_A.offsets ? __ldg(param_A.offsets + g) : 0; + const int m1 = param_A.offsets ? __ldg(param_A.offsets + g + 1) : M_total; + const int M = m1 - m0; + + CUtensorMap* gmem_out = out + g * kNum; + + if (lane == 0) { + offsets[g] = m0; + if (g + 1 == gridDim.x) { + offsets[g + 1] = m1; + } + } + + if constexpr (kStridingA == Striding::kBlocked) { + Array templates; + templates[0] = &tm_a; + templates[1] = &tm_b; + templates[2] = &tm_c; + Array addrs; + const auto b = resolve(param_B, g); + addrs[0] = resolve(param_A, g).ptr.ptr; + addrs[1] = b.ptr.ptr; + addrs[2] = resolve(param_C, g).ptr.ptr; + Array dims; + dims[0] = M; + dims[1] = N; + dims[2] = M; + detail::rebase_publish_tma_descs<3>( + gmem_out, smem_desc, templates, addrs, dims, 1, (uint64_t)b.ptr.stride * sizeof(Tb), lane); + } + else { + // Indexed-A: gather activations; prepare weight B + output C only. + Array templates; + templates[0] = &tm_b; + templates[1] = &tm_c; + Array addrs; + const auto b = resolve(param_B, g); + addrs[0] = b.ptr.ptr; + addrs[1] = resolve(param_C, g).ptr.ptr; + Array dims; + dims[0] = N; + dims[1] = M; + detail::rebase_publish_tma_descs<2>( + gmem_out, smem_desc, templates, addrs, dims, 0, (uint64_t)b.ptr.stride * sizeof(Tb), lane); + } +} + +template +struct GemmUniversalSm90_Bf16 { + + static constexpr bool kDebug = false; + + static constexpr Order kRasterOrder = raster_order; + + // L2 eviction policy for mainloop weight loads. Instantiation axis (desc policy_b): + // 0 = EVICT_NORMAL, 1 = EVICT_FIRST (weight panel is streamed once when the problem + // has a single M-tile). Variants co-exist in the catalog; the tuner picks. + static constexpr int kL2HintW = kL2HintW_; + static constexpr uint64_t kWeightL2Policy = + kL2HintW ? (uint64_t)cute::TMA::CacheHintSm90::EVICT_FIRST : (uint64_t)cute::TMA::CacheHintSm90::EVICT_NORMAL; + + using Arch = Sm90; + using Tile = Tile_; + + static constexpr bool kSupportsFusedSilu = kSupportsFusedSilu_; + + // Problem CTA tile: M = batch (act rows), N = out (weight cols), K + static constexpr int TILE_M = Tile::TILE_M; + static constexpr int TILE_N = Tile::TILE_N; + static constexpr int TILE_K = Tile::TILE_K; + static_assert(TILE_N % 128 == 0); + static_assert(TILE_M >= 8 && TILE_M % 8 == 0); + static_assert(TILE_K == 64); // host TMA still SW128 / K-atom for this step + + using AtomLayoutMNK = typename Tile::AtomLayoutMNK; + // Fused SiLU: kAtomM == 1 (1x2/1x1) covers the full OUT extent per WG, so the + // epilogue's C-ownership R2S/TMA mapping applies directly. kAtomM == 2 (2x1) + // splits the [g64|u64] blocks across WGs (WG0 gate, WG1 up); the epilogue pairs + // them through an f32 smem staging buffer (run_epilogue), staged per whole CTA + // tile: kEpiBatch == TILE_M and the silu output exactly fills one epi tile. + static_assert(!kSupportsFusedSilu || cute::size<0>(AtomLayoutMNK{}) == 1 + || (cute::size<0>(AtomLayoutMNK{}) * 128 == TILE_N && TILE_M <= 32)); + + // Traits: OUT=N_out=TILE_N, BATCH=M_batch=TILE_M; AtomLayout from Tile_ + using Traits = GmmaBF16Traits; + using TiledMma = typename Traits::TiledMma; + + static constexpr int WARPGROUPS = cute::size(AtomLayoutMNK{}); // math WGs (cooperative) + + static constexpr int kMulticastA = multicast_a; // act along TILE_M + static constexpr int kMulticastB = multicast_b; // weight along TILE_N + + static constexpr int kClusterSize = kMulticastA * kMulticastB; + + static constexpr int Stages = Tile::Stages; + + static constexpr int WARPGROUP_SIZE = 128; + static constexpr int kMathGroupSize = WARPGROUP_SIZE * WARPGROUPS; + static constexpr int CTA_SIZE = WARPGROUP_SIZE * (WARPGROUPS + 1); + + static constexpr int K_PIPE_MMAS = 1; + + using Ta = nv_bfloat16; // API A = activations + using Tb = nv_bfloat16; // API B = weights + using Tc = nv_bfloat16; + + using Cluster = arch::Cluster; + + static constexpr auto is_grouped_gemm = is_grouped_gemm_; + + static constexpr Striding kStridingA = kStridingA_; + static constexpr Striding kStridingB = is_grouped_gemm_ ? Striding::kBlocked : Striding::kFlat; + static constexpr Striding kStridingC = is_grouped_gemm_ ? Striding::kBlocked : Striding::kFlat; + + static constexpr bool kIndexedGather = (kStridingA_ == Striding::kIndexed); + + // setmaxnreg: each WG ≤ 256, multiples of 8. Budgets from Tile (TMA vs indexed). + static constexpr int kProducerRegs = kIndexedGather ? Tile::kProducerRegsIndexed : Tile::kProducerRegsTma; + static constexpr int kMathRegs = kIndexedGather ? Tile::kMathRegsIndexed : Tile::kMathRegsTma; + static_assert(kProducerRegs >= 24 && kProducerRegs % 8 == 0); + static_assert(kMathRegs >= 24 && kMathRegs % 8 == 0 && kMathRegs <= 256); + static_assert(WARPGROUPS == 1 || WARPGROUPS == 2); + static_assert(WARPGROUPS != 2 || kProducerRegs + 2 * kMathRegs <= 504); + static_assert(WARPGROUPS != 1 || kProducerRegs + kMathRegs <= 512); + + using Scheduler = TileScheduler; + + using MainloopPipeline = cutlass::PipelineTmaAsync; + using PipelineState = typename MainloopPipeline::PipelineState; + using PipelineStorage = typename MainloopPipeline::SharedStorage; + + // SMEM SoT: weight → GMMA-A (OUT,K,PIPE); act → GMMA-B (BATCH,K,PIPE). K-major Step<_1,_2,_3>. + using SmemLayoutA = + decltype(cute::tile_to_shape(typename Traits::SmemLayoutAtomA{}, + cute::make_shape(cute::Int{}, cute::Int{}, cute::Int{}), + cute::Step{})); + using SmemLayoutB = + decltype(cute::tile_to_shape(typename Traits::SmemLayoutAtomB{}, + cute::make_shape(cute::Int{}, cute::Int{}, cute::Int{}), + cute::Step{})); + // Per-stage 2D view for indexed gather store TV (same atom as SmemLayoutB / GMMA-B). + using SmemLayoutB_2D = decltype(cute::tile_to_shape(typename Traits::SmemLayoutAtomB{}, + cute::make_shape(cute::Int{}, cute::Int{}), + cute::Step{})); + + static constexpr int kTmaTxBytesWeight = (int)sizeof(Tb) * (TILE_N * TILE_K); + static constexpr int kTmaTxBytesAct = (int)sizeof(Ta) * (TILE_M * TILE_K); + // Dense / blocked-A: expect_tx(weight+act); both via TMA complete_tx. + // Indexed-A: expect_tx(weight only); gather gated by cpasync_barrier_arrive_noinc. + static constexpr int kTmaTxBytes = + (kStridingA_ == Striding::kIndexed) ? kTmaTxBytesWeight : (kTmaTxBytesWeight + kTmaTxBytesAct); + + // Grouped: per-expert maps in workspace (indexed [B,C]; blocked [A,B,C]). + static constexpr int kTmaDescNumAB = is_grouped_gemm_ ? (kStridingA_ == Striding::kBlocked ? 2 : 1) : 0; + static constexpr int kTmaDescNumC = is_grouped_gemm_ ? 1 : 0; + static constexpr int kTmaDescNum = (kTmaDescNumAB + kTmaDescNumC) > 0 ? (kTmaDescNumAB + kTmaDescNumC) : 1; + + // Epilogue tile must match TiledMma::tile_size ownership (MMA.md / CUTLASS + // Sm90 TMA epi): kEpiOut = 64 * (WGs along OUT). AtomLayout<_2,_1,_1> → (128, ≤32); + // <_1,_2,_1> / <_1,_1,_1> → (64, TILE_M). Epi SMEM must cover full tile_size N + // or ThrN=1 STSM writes OOB past kEpiBatch=32. + static constexpr int kAtomM = Traits::kAtomM; + static constexpr int kAtomN = Traits::kAtomN; + static constexpr int kEpiOut = 64 * kAtomM; + static constexpr int kEpiBatch = (kAtomN == 2) ? TILE_M : ((TILE_M <= 32) ? TILE_M : 32); + static constexpr int kTmaOut = 64; // TMA box along OUT (64 bf16 = SW128B) + static constexpr int kSwizzleC = 128; + static constexpr int kFusedSiluBlock = 64; + static constexpr int kFragmentSize = (kEpiOut * kEpiBatch) / kMathGroupSize; + static_assert(TILE_N % kEpiOut == 0); + static_assert(TILE_M % kEpiBatch == 0); + static_assert(kEpiOut % kTmaOut == 0); + static_assert(kFragmentSize >= 1); + using EpilogueTile = cute::Shape, cute::Int>; + + using SmemLayoutAtomD = decltype( + gmma_ss_smem_selector, cute::Int>()); + using SmemLayoutD = + decltype(cute::tile_to_shape(SmemLayoutAtomD{}, + cute::make_shape(cute::Int{}, cute::Int{}, cute::_1{}), + cute::Step{})); + + // Retile atom (CUTLASS builder always uses STSM_N for C_atom); R2S is STSM_T (M-major). + // Atom size = per-WG accumulator vals/thread: 2x1 tiles split GMMA-M, so per-WG + // GMMA-N = TILE_M (TILE_M/2 vals); 1x2 splits GMMA-N too (TILE_M/4 vals). Matches the + // old kCValsPerThread>=8 criterion on every pre-existing tile; TILE_M=8 2x1 and + // 16x256_1x2 get U32x2 / U16x4. + static constexpr int kCValsPerThread = (kAtomN == 2) ? TILE_M / 4 : TILE_M / 2; + static_assert(kCValsPerThread >= 4, "STSM needs at least U32x2 / 4 bf16 vals per thread"); + using CopyAtomC = typename detail::EpiStsmAtoms<(kCValsPerThread >= 8) ? 8 : 4>::CopyAtomC; + using CopyOpR2S = typename detail::EpiStsmAtoms<(kCValsPerThread >= 8) ? 8 : 4>::CopyOpR2S; + + struct LayoutC { + static constexpr int S0 = kEpiBatch; + static constexpr int C0 = kTmaOut; // TMA box, not full epi OUT + static constexpr int C1 = 1; + }; + + struct SharedStorage { + cute::array_aligned> A; + cute::array_aligned> B; + // Shared epilogue D (CUTLASS Sm90TmaWarpSpecialized — one buffer for all math WGs). + cute::array_aligned, 128> D; + // f32 gate/up staging for 2x1 fused-SiLU (the [g64|u64] blocks are split + // across the two math WGs, so silu pairs are exchanged through smem). + static constexpr int kSiluStageElems = (kSupportsFusedSilu && kAtomM == 2) ? TILE_N * TILE_M : 1; + cute::array_aligned silu_stage; + PipelineStorage pipeline; + typename Scheduler::Storage sched; + int gather_alive; + int gather_k_iters; + int gather_m0; + int gather_M_group; + int gather_offset_m; + }; + + static constexpr int kSmemSize = (int)sizeof(SharedStorage); + + using ClusterShape = cute::Shape, cute::_1, cute::_1>; + + // Host: launch TMA-map/offset preparation and return the prepared offset table. + static int* PrepareTmaDescs(const CUtensorMap& tm_a, + const CUtensorMap& tm_b, + const CUtensorMap& tm_c, + const MatrixParam& param_A, + const MatrixParam& param_B, + const MatrixParam& param_C, + CUtensorMap* out, + int num_groups, + int M, + int N, + cudaStream_t stream) + { + if constexpr (!is_grouped_gemm_) { + return nullptr; + } + int* offsets = reinterpret_cast(out + num_groups * kTmaDescNum); + prepare_tma_descs_sm90_bf16 + <<>>(tm_a, tm_b, tm_c, param_A, param_B, param_C, out, offsets, M, N); + return offsets; + } + + __device__ void operator()(const CUtensorMap& tm_a, + const CUtensorMap& tm_b, + const CUtensorMap& tm_c, + const CUtensorMap& tm_u, + const CUtensorMap& tm_v, + const MatrixParam& param_A, + const MatrixParam& param_B, + const MatrixParam& param_U, + const MatrixParam& param_V, + const MatrixParam& param_C, + bool fuse_silu, + Scheduler sched, + CUtensorMap* tensormap_buf, + char* smem_buf) + { + (void)tm_u; + (void)tm_v; + (void)param_U; + (void)param_V; + + SharedStorage& storage = *reinterpret_cast(smem_buf); + + const int wg_idx = cutlass::canonical_warp_group_idx(); + + if (threadIdx.x == 0) { + sched.init_dyanmic(storage.sched, kClusterSize * (WARPGROUPS * 4 + 1)); + } + + typename MainloopPipeline::Params pp; + pp.transaction_bytes = (uint32_t)kTmaTxBytes; + pp.num_consumers = (uint32_t)kMathGroupSize; + // Indexed: PTX Example 2 — init must include arrive-ons from each + // cpasync_barrier_arrive_noinc (1 per gather thread) plus leader + // arrive_and_expect_tx. Dense/blocked: TMA-only, count = 1. + pp.num_producers = (kStridingA == Striding::kIndexed) ? (1 + WARPGROUP_SIZE) : 1; + pp.initializing_warp = 0; + + if (wg_idx == WARPGROUPS) { + const int warp_id_in_wg = (threadIdx.x / WARP_SIZE) % 4; + if constexpr (kStridingA == Striding::kIndexed) { + // All 128 gather threads are Producers (lockstep acquire + noinc). + pp.role = MainloopPipeline::ThreadCategory::Producer; + pp.is_leader = (warp_id_in_wg == 0) && (threadIdx.x % WARP_SIZE == 0); + } + else { + pp.role = (warp_id_in_wg == 0) ? MainloopPipeline::ThreadCategory::Producer : + MainloopPipeline::ThreadCategory::NonParticipant; + pp.is_leader = (warp_id_in_wg == 0) && (threadIdx.x % WARP_SIZE == 0); + } + } + else { + pp.role = MainloopPipeline::ThreadCategory::Consumer; + pp.is_leader = 0; + } + + MainloopPipeline pipeline(storage.pipeline, pp, ClusterShape{}); + + if (threadIdx.x == 0) { + cutlass::arch::fence_view_async_shared(); + } + (kClusterSize > 1) ? cute::cluster_sync() : __syncthreads(); + + if (wg_idx == WARPGROUPS) { + cutlass::arch::warpgroup_reg_dealloc(); + + static_assert(TILE_M % kMulticastA == 0); + static_assert(TILE_N % kMulticastB == 0); + + cutlass::arch::NamedBarrier producers_bar(WARP_SIZE * 2, 7); + + const int warp_id = cutlass::canonical_warp_idx_sync(); + const int warp_in_wg = warp_id % 4; + const bool cta_0 = cute::block_id_in_cluster().x == 0; + + if constexpr (kStridingA == Striding::kIndexed) { + // Full producer WG gather. Scheduler folded onto warp0. + // Full-barrier init = 1 + 128 (expect_tx + one noinc per gather thr). + cutlass::arch::NamedBarrier gather_bar( + /*num_threads=*/WARPGROUP_SIZE, cutlass::arch::ReservedNamedBarriers::FirstUserBarrier); + + Cluster cluster(cute::block_id_in_cluster().x); + + const int mc_offset_n = cluster.cta_m() * (TILE_N / kMulticastB); + + auto* smem_act = storage.B.data(); + auto* smem_weight = storage.A.data() + mc_offset_n * TILE_K; + + PipelineState write_state = cutlass::make_producer_start_state(); + typename Scheduler::ConsumerState sched_state = sched.init_consumer(storage.sched); + typename Scheduler::ProducerState prod_state = sched.init_producer(storage.sched); + int lane_predicate = 0; + const int lane_id = threadIdx.x % WARP_SIZE; + const int prod_tid = threadIdx.x - WARPGROUPS * WARPGROUP_SIZE; + + if (warp_in_wg == 0) { + lane_predicate = cute::elect_one_sync(); + } + + const Ta* act_gmem = (const Ta*)param_A.ptr; + const int ldA = param_A.stride; + const int* idxs = param_A.idxs; + const int K = sched.gemm_shape().z; + + constexpr int kVec = 8; + constexpr int nvec = TILE_M * (TILE_K / kVec); + + typename Scheduler::Tile* tile; + + while (true) { + const CUtensorMap* Bdesc = &tm_b; + uint16_t mask_B = 0; + int coord_n = 0; + int k_iters = 0; + + if (warp_in_wg == 0 && cta_0) { + (void)prod_state.next(); + } + + if (warp_in_wg == 0) { + const bool alive = sched_state.acquire(tile); + int m0 = 0, M_group = 0, offset_m = 0; + + if (alive && tile->is_valid_cluster) { + if constexpr (is_grouped_gemm) { + const int g = tile->group_idx; + Bdesc = &tensormap_buf[g * kTmaDescNum]; + } + + mask_B = cluster.mask_n(); + coord_n = tile->offset_n + mc_offset_n; + offset_m = tile->offset_m; + m0 = [&] { + if constexpr (is_grouped_gemm) { + return tile->m0; + } + return 0; + }(); + M_group = [&] { + if constexpr (is_grouped_gemm) { + return tile->m1 - tile->m0; + } + return sched.gemm_shape().x; + }(); + k_iters = sched.k_iters_; + } + + if (lane_id == 0) { + storage.gather_alive = alive ? 1 : 0; + storage.gather_k_iters = k_iters; + storage.gather_m0 = m0; + storage.gather_M_group = M_group; + storage.gather_offset_m = offset_m; + } + __syncwarp(); + } + + // Tile header only. + gather_bar.arrive_and_wait(); + + if (storage.gather_alive == 0) { + break; + } + + k_iters = storage.gather_k_iters; + const int m0 = storage.gather_m0; + const int M_group = storage.gather_M_group; + const int offset_m = storage.gather_offset_m; + int coord_k = 0; + + // Indexed load pattern (iterator_sm80): per-tile idxs → src_data_vec_, + // then Prefetch advances each base by TILE_K (no idxs reload). + // TILE_M=8: nvec=64 < 128 → idle producers must skip ZFILL (pred=false + // still zeros dst). TILE_M>=16: exact pre-templatize tight loop. + if constexpr (TILE_M == 8) { + constexpr int kSlots = (nvec + WARPGROUP_SIZE - 1) / WARPGROUP_SIZE; + static_assert(kSlots >= 1); + const Ta* src_data_vec_[kSlots]; + int m_own[kSlots]; + int kk_own[kSlots]; + bool pred_row[kSlots]; + bool in_rng_slot[kSlots]; + PRAGMA_UNROLL + for (int t = 0; t < kSlots; ++t) { + const int i = prod_tid + t * WARPGROUP_SIZE; + const bool in_rng = i < nvec; + const int m = in_rng ? (i / (TILE_K / kVec)) : 0; + const int kk = in_rng ? ((i % (TILE_K / kVec)) * kVec) : 0; + const int packed = m0 + offset_m + m; + const bool row_ok = in_rng && (offset_m + m) < M_group; + const int token = (idxs && row_ok) ? __ldg(idxs + packed) : packed; + m_own[t] = m; + kk_own[t] = kk; + pred_row[t] = row_ok; + in_rng_slot[t] = in_rng; + src_data_vec_[t] = act_gmem + (int64_t)token * ldA + kk; + } + + for (; k_iters > 0; --k_iters) { + pipeline.producer_acquire(write_state); + auto* bar = pipeline.producer_get_barrier(write_state); + const int pipe = write_state.index(); + + if (warp_in_wg == 0 && lane_predicate) { + detail::tma_load_with_barrier( + Bdesc, + bar, + smem_weight + pipe * TILE_N * TILE_K, + coord_k, + coord_n, + mask_B, + kWeightL2Policy); + } + + { + cute::Tensor sB = cute::make_tensor( + cute::make_smem_ptr(smem_act + pipe * TILE_M * TILE_K), SmemLayoutB_2D{}); + + PRAGMA_UNROLL + for (int t = 0; t < kSlots; ++t) { + if constexpr ((nvec % WARPGROUP_SIZE) != 0) { + if (!in_rng_slot[t]) { + continue; + } + } + const bool pred = pred_row[t] && (coord_k + kk_own[t]) < K; + auto* dst = &sB(m_own[t], kk_own[t]); + cute::SM80_CP_ASYNC_CACHEGLOBAL_ZFILL::copy( + *reinterpret_cast(src_data_vec_[t]), + *reinterpret_cast(dst), + pred); + src_data_vec_[t] += TILE_K; + } + cutlass::arch::cpasync_barrier_arrive_noinc(bar); + } + + ++write_state; + coord_k += TILE_K; + } + } + else { + // Identical to 7f8e860 gather (TILE_M=128 → kSlots=8). + constexpr int kSlots = nvec / WARPGROUP_SIZE; + static_assert(nvec % WARPGROUP_SIZE == 0); + const Ta* src_data_vec_[kSlots]; + int m_own[kSlots]; + int kk_own[kSlots]; + bool pred_row[kSlots]; + PRAGMA_UNROLL + for (int t = 0; t < kSlots; ++t) { + const int i = prod_tid + t * WARPGROUP_SIZE; + const int m = i / (TILE_K / kVec); + const int kk = (i % (TILE_K / kVec)) * kVec; + const int packed = m0 + offset_m + m; + const bool row_ok = (offset_m + m) < M_group; + const int token = (idxs && row_ok) ? __ldg(idxs + packed) : packed; + m_own[t] = m; + kk_own[t] = kk; + pred_row[t] = row_ok; + src_data_vec_[t] = act_gmem + (int64_t)token * ldA + kk; + } + + // acquire → weight TMA → gather → arrive_noinc. No producer wait. + for (; k_iters > 0; --k_iters) { + pipeline.producer_acquire(write_state); + auto* bar = pipeline.producer_get_barrier(write_state); + const int pipe = write_state.index(); + + if (warp_in_wg == 0 && lane_predicate) { + detail::tma_load_with_barrier( + Bdesc, + bar, + smem_weight + pipe * TILE_N * TILE_K, + coord_k, + coord_n, + mask_B, + kWeightL2Policy); + } + + { + cute::Tensor sB = cute::make_tensor( + cute::make_smem_ptr(smem_act + pipe * TILE_M * TILE_K), SmemLayoutB_2D{}); + + PRAGMA_UNROLL + for (int t = 0; t < kSlots; ++t) { + const bool pred = pred_row[t] && (coord_k + kk_own[t]) < K; + auto* dst = &sB(m_own[t], kk_own[t]); + cute::SM80_CP_ASYNC_CACHEGLOBAL_ZFILL::copy( + *reinterpret_cast(src_data_vec_[t]), + *reinterpret_cast(dst), + pred); + src_data_vec_[t] += TILE_K; + } + cutlass::arch::cpasync_barrier_arrive_noinc(bar); + } + + ++write_state; + coord_k += TILE_K; + } + } + + if (warp_in_wg == 0) { + sched_state.release(); + } + } + + if (warp_in_wg == 0) { + sched_state.release(); + if (cta_0) { + sched.tail(prod_state); + } + if (lane_predicate) { + pipeline.producer_tail(write_state); + } + } + } + else if (warp_in_wg == 0) { + Cluster cluster(cute::block_id_in_cluster().x); + + // API A = act → GMMA-B SMEM; API B = weight → GMMA-A SMEM + const int mc_offset_m = cluster.cta_n() * (TILE_M / kMulticastA); + const int mc_offset_n = cluster.cta_m() * (TILE_N / kMulticastB); + + auto* smem_act = storage.B.data(); + auto* smem_weight = storage.A.data() + mc_offset_n * TILE_K; + + PipelineState write_state = cutlass::make_producer_start_state(); + + auto sched_state = sched.init_consumer(storage.sched); + + int lane_predicate = cute::elect_one_sync(); + + typename Scheduler::Tile* tile; + + while (sched_state.acquire(tile)) { + + if (tile->is_valid_cluster) { + + const CUtensorMap* Adesc = &tm_a; + const CUtensorMap* Bdesc = &tm_b; + + if constexpr (is_grouped_gemm) { + // Descs published by prepare_moe_tma_descs on this stream; + // fence_acquire only needed after in-kernel tensormap replace. + const int g = tile->group_idx; + CUtensorMap* const descs = tensormap_buf + g * kTmaDescNum; + if constexpr (kStridingA == Striding::kBlocked) { + Adesc = &descs[0]; + Bdesc = &descs[1]; + } + else { + Bdesc = &descs[0]; + } + } + + const uint16_t mask_B = cluster.mask_n(); // weight multicast + + const int offset_m = tile->offset_m; + const int offset_n = tile->offset_n; + + int k_iter = sched.k_iters_; + + int coord_k = 0; + const int coord_m = offset_m + mc_offset_m; + const int coord_n = offset_n + mc_offset_n; + + if (lane_predicate) { + // Dense Flat-A or grouped Blocked-A: both operands TMA into SoT. + const uint16_t mask_A = cluster.mask_m(); + for (; k_iter > 0; --k_iter) { + pipeline.producer_acquire(write_state); + auto* bar = pipeline.producer_get_barrier(write_state); + const int pipe = write_state.index(); + + detail::tma_load_with_barrier( + Adesc, + bar, + smem_act + mc_offset_m * TILE_K + pipe * TILE_M * TILE_K, + coord_k, + coord_m, + mask_A); + detail::tma_load_with_barrier( + Bdesc, + bar, + smem_weight + pipe * TILE_N * TILE_K, + coord_k, + coord_n, + mask_B, + kWeightL2Policy); + + coord_k += TILE_K; + ++write_state; + } + } + } + + if constexpr (Scheduler::is_dynamic) { + if (cta_0) { + producers_bar.arrive_unaligned(); + } + } + + sched_state.release(); + } + + sched_state.release(); + + if (lane_predicate) { + pipeline.producer_tail(write_state); + } + } + else if (warp_in_wg == 1 && cta_0) { + auto state = sched.init_producer(storage.sched); + while (state.next()) { + if constexpr (Scheduler::is_dynamic) { + producers_bar.arrive_and_wait_unaligned(); + } + } + sched.tail(state); + } + } + else { + cutlass::arch::warpgroup_reg_alloc(); + + // mma_tid in [0, kMathGroupSize) — math WGs share one TiledMma TV (cooperative). + const int mma_tid = threadIdx.x; + + cute::Tensor sA = cute::make_tensor(cute::make_smem_ptr(storage.A.data()), SmemLayoutA{}); + cute::Tensor sB = cute::make_tensor(cute::make_smem_ptr(storage.B.data()), SmemLayoutB{}); + + TiledMma tiled_mma; + auto thr_mma = tiled_mma.get_thread_slice(mma_tid); + + cute::Tensor tCsA = thr_mma.partition_A(sA); + cute::Tensor tCsB = thr_mma.partition_B(sB); + cute::Tensor tCrA = thr_mma.make_fragment_A(tCsA); + cute::Tensor tCrB = thr_mma.make_fragment_B(tCsB); + + PipelineState pipe_state{}; + PipelineState pipe_release = pipe_state; + + auto sched_state = sched.init_consumer(storage.sched); + + typename Scheduler::Tile* tile; + sched_state.acquire(tile); + + // CUTLASS Sm90TmaWarpSpecialized R2S: as_position_independent + + // make_tiled_copy_C_atom → STSM_T; warp0 TMA after fence + EpilogueBarrier. + cute::Tensor sD = cute::as_position_independent_swizzle_tensor( + cute::make_tensor(cute::make_smem_ptr(storage.D.data()), SmemLayoutD{})); + + CopyAtomC tiled_copy_atom_c{}; + auto tiled_copy_C_atom = cute::make_tiled_copy_C_atom(tiled_copy_atom_c, tiled_mma); + auto tiled_r2s = + cute::make_tiled_copy_S(cute::Copy_Atom{}, tiled_copy_C_atom); + auto thr_r2s = tiled_r2s.get_slice(mma_tid); + + cute::Tensor tRS_sD = thr_r2s.partition_D(sD); // (R2S,R2S_M,R2S_N,PIPE) + auto tRS_rD_layout = cute::make_layout(cute::take<0, 3>(cute::shape(thr_r2s.partition_S(sD)))); + + const bool issue_tma_store = (mma_tid / WARP_SIZE) == 0; + + auto epi_synchronize = [&] { + cutlass::arch::NamedBarrier::sync(kMathGroupSize, + cutlass::arch::ReservedNamedBarriers::EpilogueBarrier); + }; + + while (tile->alive) { + + if (tile->is_valid_cta) { + cute::Tensor accum = + cute::partition_fragment_C(tiled_mma, cute::take<0, 2>(typename Traits::TileShape{})); + cute::clear(accum); + + int k_iter = sched.k_iters_; + + { + auto token = pipeline.consumer_try_wait(pipe_state); + pipeline.consumer_wait(pipe_state, token); + const int read = pipe_state.index(); + cute::warpgroup_fence_operand(accum); + cute::warpgroup_arrive(); + tiled_mma.accumulate_ = cute::GMMA::ScaleOut::Zero; + CUTE_UNROLL + for (int k_block = 0; k_block < cute::size<2>(tCrA); ++k_block) { + cute::gemm(tiled_mma, + tCrA(cute::_, cute::_, k_block, read), + tCrB(cute::_, cute::_, k_block, read), + accum); + tiled_mma.accumulate_ = cute::GMMA::ScaleOut::One; + } + cute::warpgroup_commit_batch(); + ++pipe_state; + --k_iter; + } + + tiled_mma.accumulate_ = cute::GMMA::ScaleOut::One; + + PRAGMA_NO_UNROLL + for (; k_iter > 0; --k_iter) { + auto token = pipeline.consumer_try_wait(pipe_state); + pipeline.consumer_wait(pipe_state, token); + const int read = pipe_state.index(); + cute::warpgroup_fence_operand(accum); + cute::warpgroup_arrive(); + cute::gemm(tiled_mma, + tCrA(cute::_, cute::_, cute::_, read), + tCrB(cute::_, cute::_, cute::_, read), + accum); + cute::warpgroup_commit_batch(); + cute::warpgroup_wait(); + cute::warpgroup_fence_operand(accum); + pipeline.consumer_release(pipe_release); + ++pipe_state; + ++pipe_release; + } + + cute::warpgroup_wait<0>(); + pipeline.consumer_release(pipe_release); + ++pipe_release; + + if (issue_tma_store) { + cute::tma_store_wait<0>(); + } + epi_synchronize(); + + const int offset_m = tile->offset_m; + const int offset_n = tile->offset_n; + + const void* Cdesc = &tm_c; + if constexpr (is_grouped_gemm) { + Cdesc = tensormap_buf + tile->group_idx * kTmaDescNum + kTmaDescNumAB; + } + + cute::Tensor tRS_rAcc = thr_r2s.retile_S(accum); // ((R2S,R2S_V),MMA_M,MMA_N) + cute::Tensor tRS_rD = cute::make_tensor(tRS_rD_layout); + + cute::Tensor tRS_rAcc_frg = cute::recast>(tRS_rAcc); + cute::Tensor tRS_rD_frg = cute::recast>(tRS_rD); + + const int mma_tile_m = cute::size<0>(typename Traits::TileShape{}) / cute::size<1>(tRS_rAcc); + const int mma_tile_n = cute::size<1>(typename Traits::TileShape{}) / cute::size<2>(tRS_rAcc); + constexpr int epi_tile_m = kEpiOut; + constexpr int epi_tile_n = kEpiBatch; + constexpr int epi_n_count = TILE_M / kEpiBatch; + constexpr int tma_m_count = kEpiOut / kTmaOut; + (void)epi_tile_m; + (void)mma_tile_m; + + auto run_epilogue = [&](auto fused_silu) { + constexpr bool kFuseSilu = decltype(fused_silu)::value; + static_assert(!kFuseSilu || kSupportsFusedSilu); + // 2x1 silu: the [g64|u64] blocks are split across the two math WGs + // (WG0 gate, WG1 up), so pairs are exchanged through the f32 + // silu_stage buffer and the silu output (TILE_N/2) fills exactly + // one epi tile. 1xN silu pairs gate/up in-register instead. + constexpr bool kStagedSilu = kFuseSilu && kAtomM == 2; + static_assert(!kStagedSilu || (TILE_M == kEpiBatch && TILE_N / 2 == kEpiOut)); + + float* stage = storage.silu_stage.data(); + (void)stage; + // R2S-element (OUT, BATCH) coords within the epi tile; the staged + // fill uses them to locate each output's gate/up in silu_stage. + // Source-side partition: fragment element j holds the value for + // coord tRS_cEpi(j) (same coords as retile_S(accum)). + cute::Tensor cEpi = cute::make_identity_tensor(EpilogueTile{}); + cute::Tensor tRS_cEpi = thr_r2s.partition_S(cEpi); + (void)tRS_cEpi; + + if constexpr (kStagedSilu) { + cute::Tensor cD = + cute::make_identity_tensor(cute::make_shape(cute::Int{}, cute::Int{})); + cute::Tensor tCcD = thr_mma.partition_C(cD); + CUTE_UNROLL + for (int i = 0; i < cute::size(accum); ++i) { + stage[cute::get<0>(tCcD(i)) + cute::get<1>(tCcD(i)) * TILE_N] = accum(i); + } + epi_synchronize(); + } + + constexpr int kStoreOut = kFuseSilu ? (TILE_N / 2) : TILE_N; + constexpr int epi_m_count = kStoreOut / kEpiOut; + static_assert(kStoreOut % kEpiOut == 0); + + CUTE_UNROLL + for (int epi_n = 0; epi_n < epi_n_count; ++epi_n) { + CUTE_UNROLL + for (int epi_m = 0; epi_m < epi_m_count; ++epi_m) { + // OUT strips of kEpiOut; in-register gate/up occupy adjacent + // strips (staged silu reads pairs from silu_stage instead). + const int mma_m = (kFuseSilu && !kStagedSilu) ? 2 * epi_m : epi_m; + const int mma_n = (epi_n * epi_tile_n) / mma_tile_n; + (void)mma_m; + (void)mma_n; + + const int epi_n_in_mma = epi_n % (mma_tile_n / epi_tile_n); + const int r2s_v = epi_n_in_mma * cute::size(tRS_rD_frg); + CUTE_UNROLL + for (int epi_v = 0; epi_v < cute::size(tRS_rD_frg); ++epi_v) { + cutlass::Array dst; + if constexpr (kStagedSilu) { + // tRS_cEpi has exactly kFragmentSize coords per thread + // (whole epi tile); staged tiles have epi_n_count == + // epi_m_count == 1 (assert above). + CUTE_UNROLL + for (int j = 0; j < kFragmentSize; ++j) { + const auto coord = tRS_cEpi(epi_v * kFragmentSize + j); + const int o = cute::get<0>(coord); + const int n = cute::get<1>(coord); + const int g_off = (o / kFusedSiluBlock) * (2 * kFusedSiluBlock) + + o % kFusedSiluBlock + n * TILE_N; + dst[j] = cutlass::bfloat16_t( + detail::silu_mul(stage[g_off], stage[g_off + kFusedSiluBlock])); + } + } + else if constexpr (kFuseSilu) { + // Block-pack [g64|u64]: up is the next MMA_M strip. + auto gate = tRS_rAcc_frg(cute::_, mma_m, mma_n)(r2s_v + epi_v); + auto up = tRS_rAcc_frg(cute::_, mma_m + 1, mma_n)(r2s_v + epi_v); + CUTE_UNROLL + for (int j = 0; j < kFragmentSize; ++j) { + dst[j] = cutlass::bfloat16_t(detail::silu_mul(gate[j], up[j])); + } + } + else { + auto src = tRS_rAcc_frg(cute::_, mma_m, mma_n)(r2s_v + epi_v); + CUTE_UNROLL + for (int j = 0; j < kFragmentSize; ++j) { + dst[j] = cutlass::bfloat16_t(src[j]); + } + } + tRS_rD_frg(epi_v) = dst; + } + + // Single D buffer: drain prior TMA before R2S (convert above + // overlaps the in-flight store). Do not wait right after arrive. + if (!(epi_n == 0 && epi_m == 0)) { + if (issue_tma_store) { + cute::tma_store_wait<0>(); + } + epi_synchronize(); + } + + cute::copy(tiled_r2s, tRS_rD, tRS_sD(cute::_, cute::_, cute::_, 0)); + + cutlass::arch::fence_view_async_shared(); + epi_synchronize(); + + if (issue_tma_store) { + constexpr int kAtomElems = kTmaOut * kEpiBatch; + CUTE_UNROLL + for (int tma_m = 0; tma_m < tma_m_count; ++tma_m) { + const int out_off = epi_m * kEpiOut + tma_m * kTmaOut; + const int store_n = kFuseSilu ? (offset_n / 2 + out_off) : (offset_n + out_off); + cute::SM90_TMA_STORE::copy(Cdesc, + storage.D.data() + tma_m * kAtomElems, + store_n, + offset_m + epi_n * kEpiBatch); + cute::tma_store_arrive(); + } + } + epi_synchronize(); + } + } + }; + + if constexpr (kSupportsFusedSilu) { + if (fuse_silu) { + run_epilogue(std::true_type{}); + } + else { + run_epilogue(std::false_type{}); + } + } + else { + run_epilogue(std::false_type{}); + } + } + else if (tile->is_valid_cluster) { + int k_iter = sched.k_iters_; + for (; k_iter > 0; --k_iter) { + auto token = pipeline.consumer_try_wait(pipe_state); + pipeline.consumer_wait(pipe_state, token); + pipeline.consumer_release(pipe_state); + ++pipe_state; + } + pipe_release = pipe_state; + } + + sched_state.release(); + sched_state.acquire(tile); + } + + sched_state.release(); + + if (issue_tma_store) { + cute::tma_store_wait<0>(); + } + } + } +}; + +} // namespace turbomind::gemm diff --git a/src/turbomind/kernels/gemm/gemm_universal_sm90_fp8_wa.h b/src/turbomind/kernels/gemm/gemm_universal_sm90_fp8_wa.h new file mode 100644 index 0000000000..0d8099c0d4 --- /dev/null +++ b/src/turbomind/kernels/gemm/gemm_universal_sm90_fp8_wa.h @@ -0,0 +1,1104 @@ +#pragma once + +#include +#include +#include + +#include +#include + +#include "cute/arch/cluster_sm90.hpp" +#include "cute/arch/copy_sm80.hpp" +#include "cute/arch/copy_sm90.hpp" +#include "cute/arch/copy_sm90_desc.hpp" +#include "cute/arch/copy_sm90_tma.hpp" +#include "cute/arch/mma_sm90_desc.hpp" +#include "cute/tensor.hpp" + +#include "cutlass/arch/barrier.h" +#include "cutlass/arch/reg_reconfig.h" +#include "cutlass/cutlass.h" +#include "cutlass/pipeline/sm90_pipeline.hpp" + +#include "src/turbomind/core/data_type.h" + +#include "src/turbomind/kernels/core/array_ops.h" +#include "src/turbomind/kernels/core/common.h" +#include "src/turbomind/kernels/core/smem.h" + +#include "src/turbomind/kernels/gemm/arch.h" +#include "src/turbomind/kernels/gemm/cp_async.h" +#include "src/turbomind/kernels/gemm/iterator_sm90.h" +#include "src/turbomind/kernels/gemm/matrix_ptr.h" +#include "src/turbomind/kernels/gemm/scheduler.cuh" +#include "src/turbomind/kernels/gemm/types.h" +#include "src/turbomind/kernels/gemm/utils.h" + +/* + * SM90 blockscaled FP8 GEMM — weight-as-A (WA) GMMA binding. + * + * LlamaLinear API: A=act (M,K), B=weight (K,N), U/V scales + * GMMA: A_smem=weight (OUT×K), B_smem=act (BATCH×K) + * TileShape (OUT,BATCH,K)=(TILE_N,TILE_M,TILE_K) + * Scales: V sparse on GMMA-M (OUT); U dense on GMMA-N (BATCH) + * + * Host launch reuses KernelImplSm90 (API TMA boxes unchanged). + */ +#include "src/turbomind/kernels/gemm/gmma_bf16_sm90.h" +#include "src/turbomind/kernels/gemm/gmma_fp8_sm90.h" +#include "src/turbomind/kernels/gemm/prepare_moe_tma_descs_sm90_fp8.h" +#include "src/turbomind/kernels/gemm/sm90_fp8_wa_traits.h" +#include "src/turbomind/kernels/gemm/sm90_utils.h" + +namespace turbomind::gemm { + +template +struct GemmUniversalSm90_Fp8Wa { + + static constexpr bool kDebug = false; + + using Arch = Sm90; + using Tile = Tile_; + + static constexpr bool kSupportsFusedSilu = kSupportsFusedSilu_; + + static constexpr int TILE_M = Tile::TILE_M; + static constexpr int TILE_N = Tile::TILE_N; + static constexpr int TILE_K = Tile::TILE_K; + + static constexpr int WG_M = Tile::WG_M; + static constexpr int WG_N = Tile::WG_N; + + static constexpr int WG_TILE_M = TILE_M / WG_M; // BATCH per WG + static constexpr int WG_TILE_N = TILE_N / WG_N; // OUT per WG + static_assert(TILE_M % WG_M == 0); + static_assert(TILE_N % WG_N == 0); + static_assert(WG_TILE_N % 64 == 0); // WGMMA atom M along OUT + + static constexpr int kSchedWarpGroups = 1; + + static constexpr int WARPGROUPS = WG_M * WG_N; + + static constexpr Order kRasterOrder = raster_order; + static constexpr int kAlgoFamily = 2; + + // GMMA-M=OUT=WG_TILE_N, GMMA-N=BATCH=WG_TILE_M + using GMMA = ScaledGmmaFP8_WA; + using AccumC = typename GMMA::AccumC; + using FragC = typename GMMA::FragC; + + // Fused SiLU: [g128|u128] along OUT. WG_N == 1 pairs GMMA-M atoms i with i+2 in + // register (ITER_M=4); WG_N == 2 stages gate/up through smem (ITER_M=2 per WG). + static_assert(!kSupportsFusedSilu || (TILE_N == 256 && GMMA::OP_M == 64 && GMMA::ITER_M == 4 / WG_N)); + + static constexpr int kMulticastA = multicast_a; + static constexpr int kMulticastB = multicast_b; + + static constexpr int kClusterSize = kMulticastA * kMulticastB; + + static constexpr int Stages = Tile::Stages; + + static constexpr bool kSplitK = false; + static constexpr int kChunkSizeK = TILE_K; + + static constexpr int WARPGROUP_SIZE = 128; + + static constexpr int kMathGroupSize = WARPGROUP_SIZE * WARPGROUPS; + + static constexpr int CTA_SIZE = WARPGROUP_SIZE * (WARPGROUPS + 1); + + using Ta = __nv_fp8_e4m3; + using Tb = __nv_fp8_e4m3; + using Tc = nv_bfloat16; + + using Tu = float; + using Tv = float; + using Tw = float; // dynamic output group scales (fused path) + + using Cluster = arch::Cluster; + + static constexpr auto is_grouped_gemm = is_grouped_gemm_; + + static constexpr Striding kStridingA = kStridingA_; + static constexpr Striding kStridingB = is_grouped_gemm_ ? Striding::kBlocked : Striding::kFlat; + static constexpr Striding kStridingC = is_grouped_gemm_ ? Striding::kBlocked : Striding::kFlat; + + // Indexed gather: A/U via cp.async; TMA only B (+ C store descs). + static constexpr bool kIndexedGather = (kStridingA_ == Striding::kIndexed); + + using Scheduler = TileScheduler; + + static constexpr int kMulticastU = is_grouped_gemm ? 1 : kMulticastA; + + using ProducerBar = cutlass::arch::ClusterTransactionBarrier; + using ConsumerBar = cutlass::arch::ClusterBarrier; + + static constexpr int kAlignmentU = 16 / sizeof(Tu); + static constexpr int kBoxU = TILE_M + (is_grouped_gemm ? kAlignmentU : 0); + + // Alignment requirement for SMEM addr. This forbids multicast factor 8. + static_assert(kMulticastU == 1 || sizeof(Tu) * kBoxU / kMulticastU % 128 == 0); + + static constexpr int kTmaTxBytesWeight = (int)sizeof(Tb) * (TILE_N * TILE_K); + static constexpr int kTmaTxBytesAct = (int)sizeof(Ta) * (TILE_M * TILE_K); + static constexpr int kTmaTxBytesU = (int)sizeof(Tu) * kBoxU; + // Dense / blocked-A: expect_tx(A+B+U). Indexed-A: expect_tx(B only); A/U via cp.async noinc. + static constexpr int kTmaTxBytes = kTmaTxBytesWeight + (kIndexedGather ? 0 : (kTmaTxBytesAct + kTmaTxBytesU)); + + // Dense: unused. Grouped indexed: [B, C]. Grouped blocked: [A, B, U, C]. + static constexpr int kTmaDescNum = !is_grouped_gemm_ ? 1 : (kIndexedGather ? 2 : 4); + static constexpr int kCdescIdx = kIndexedGather ? 1 : 3; + + // Indexed gather stores act into GMMA-B (BATCH×K). + using SmemLayoutB_2D = decltype(cute::tile_to_shape(cute::SM90::GMMA::Layout_K_SW128_Atom{}, + cute::make_shape(cute::Int{}, cute::Int{}), + cute::Step{})); + + // setmaxnreg: each WG ≤ 256, multiples of 8. Budgets come from Tile + // (TMA vs indexed). 2 math WGs pack to 504; 1 math WG packs to ≤512. + static constexpr int kProducerRegs = kIndexedGather ? Tile::kProducerRegsIndexed : Tile::kProducerRegsTma; + static constexpr int kMathRegs = kIndexedGather ? Tile::kMathRegsIndexed : Tile::kMathRegsTma; + static_assert(kProducerRegs >= 24 && kProducerRegs % 8 == 0); + static_assert(kMathRegs >= 24 && kMathRegs % 8 == 0 && kMathRegs <= 256); + static_assert(WARPGROUPS == 1 || WARPGROUPS == 2); + static_assert(WARPGROUPS != 2 || kProducerRegs + 2 * kMathRegs == 504); + static_assert(WARPGROUPS != 1 || kProducerRegs + kMathRegs <= 512); + + // ! SMEM addr must be SBO aligned for TMA load/store + struct SharedStorage { + // GMMA-A = weight (OUT×K); GMMA-B = act (BATCH×K) + __align__(1024) Array A; + __align__(1024) Array B; + // Fused amax reduce: per math WG, [warp][t0][scale_i], scale_i < OP_N/4 ≤ 64 + __align__(128) float fused_amax_scratch[WARPGROUPS][4 * 4 * 64]; + // f32 gate/up staging for WG_N == 2 fused SiLU (gate/up split across math WGs). + static constexpr int kSiluStageElems = (kSupportsFusedSilu && WG_N == 2) ? TILE_N * TILE_M : 1; + __align__(128) float silu_stage[kSiluStageElems]; + __align__(128) Tu U[Stages][round_up(kBoxU, 128)]; // at least 128 byte alignment + __align__(128) Tv V[Stages][2]; + __align__(8) uint64_t producer_bar[Stages]; + __align__(8) uint64_t consumer_bar[Stages]; + typename Scheduler::Storage sched; + int gather_alive; + int gather_k_iters; + int gather_m0; + int gather_M_group; + int gather_offset_m; + }; + + template + struct Output { + static_assert(!kFuseSilu || kSupportsFusedSilu); + + using Tc = std::conditional_t; + + static constexpr int kStoreOut = kFuseSilu ? TILE_N / 2 : TILE_N; + + struct LayoutC { + // One box covering the whole CTA tile (all math WGs share smem C). + static constexpr int S0 = TILE_M; + static constexpr int C0 = kStoreOut; + static constexpr int C1 = 1; + static constexpr int C1_store = 1; + }; + + static constexpr int kSwizzleC = 0; + }; + + static constexpr int kOutputOffset = round_up(sizeof(SharedStorage), 1024); + + static constexpr int GetSmemSize(bool fuse_silu) + { + return kOutputOffset + + (fuse_silu ? TILE_M * (TILE_N / 2) * (int)sizeof(__nv_fp8_e4m3) : + TILE_M * TILE_N * (int)sizeof(nv_bfloat16)); + } + + static constexpr int kSmemSize = GetSmemSize(false); + + static constexpr int OUTER_M = GMMA::OUTER_M; + // Weight-scale predicates along OUT (problem N / GMMA-M). + static constexpr int MMA_SUBTILE_M = WG_TILE_N / OUTER_M; + + // Host: rebase per-expert TMA maps into workspace before GEMM launch. + static int* PrepareTmaDescs(const CUtensorMap& tm_a, + const CUtensorMap& tm_b, + const CUtensorMap& tm_u, + const CUtensorMap& tm_c, + const MatrixParam& param_A, + const MatrixParam& param_B, + const MatrixParam& param_U, + const MatrixParam& param_C, + bool fuse_silu, + CUtensorMap* out, + int num_groups, + int M, + int N, + cudaStream_t stream) + { + if constexpr (!is_grouped_gemm_) { + return nullptr; + } + int* offsets = reinterpret_cast(out + num_groups * kTmaDescNum); + prepare_moe_tma_descs_sm90_fp8<<>>( + tm_a, tm_b, tm_u, tm_c, param_A, param_B, param_U, param_C, fuse_silu, out, offsets, M, N); + return offsets; + } + + __device__ void operator()(const CUtensorMap& tm_a, + const CUtensorMap& tm_b, + const CUtensorMap& tm_c, + const CUtensorMap& tm_u, + const CUtensorMap& tm_v, + const MatrixParam& param_A, + const MatrixParam& param_B, + const MatrixParam& param_U, + const MatrixParam& param_V, + const MatrixParam& param_C, + const MatrixParam& param_W, + bool fuse_silu, + Scheduler sched, + CUtensorMap* tensormap_buf, + char* smem_buf) + { + SharedStorage& storage = *reinterpret_cast(smem_buf); + + uint64_t* producer_bar = storage.producer_bar; + uint64_t* consumer_bar = storage.consumer_bar; + + constexpr int kProducerBarInit = kIndexedGather ? (1 + WARPGROUP_SIZE) : (1 + 1); + + if (threadIdx.x == 0) { + PRAGMA_UNROLL + for (int s = 0; s < Stages; ++s) { + ProducerBar::init(&producer_bar[s], kProducerBarInit); + ConsumerBar::init(&consumer_bar[s], WARPGROUPS * kClusterSize * 4); + } + sched.init_dyanmic(storage.sched, kClusterSize * (WARPGROUPS * 4 + 1)); + cutlass::arch::fence_view_async_shared(); + if constexpr (kClusterSize > 1) { + cutlass::arch::fence_barrier_init(); + } + } + + (kClusterSize > 1) ? cute::cluster_sync() : __syncthreads(); + + const int wg_idx = cutlass::canonical_warp_group_idx(); + + if (wg_idx == WARPGROUPS) { + cutlass::arch::warpgroup_reg_dealloc(); + + static_assert(TILE_M % kMulticastA == 0); + static_assert(TILE_N % kMulticastB == 0); + + cutlass::arch::NamedBarrier producers_bar(WARP_SIZE * 2, 7); + + const int warp_id = cutlass::canonical_warp_idx_sync(); + const int warp_in_wg = warp_id % 4; + const bool cta_0 = cute::block_id_in_cluster().x == 0; + + if constexpr (kIndexedGather) { + // Full producer WG gather. Scheduler folded onto warp0. + cutlass::arch::NamedBarrier gather_bar( + /*num_threads=*/WARPGROUP_SIZE, cutlass::arch::ReservedNamedBarriers::FirstUserBarrier); + + Cluster cluster(cute::block_id_in_cluster().x); + + const int mc_offset_n = cluster.cta_m() * (TILE_N / kMulticastB); + + // WA: weight→GMMA-A (storage.A), act→GMMA-B (storage.B) + auto* smem_weight = storage.A.data() + mc_offset_n * TILE_K; + auto* smem_act = storage.B.data(); + auto& smem_U = storage.U; + auto& smem_V = storage.V; + + cutlass::PipelineState write_state{0, 1, 0}; + + typename Scheduler::ConsumerState sched_state = sched.init_consumer(storage.sched); + typename Scheduler::ProducerState prod_state = sched.init_producer(storage.sched); + int lane_predicate = 0; + const int lane_id = threadIdx.x % WARP_SIZE; + const int prod_tid = threadIdx.x - WARPGROUPS * WARPGROUP_SIZE; + + if (warp_in_wg == 0) { + lane_predicate = cute::elect_one_sync(); + } + + const Ta* act_gmem = (const Ta*)param_A.ptr; + const int ldA = param_A.stride; + const int* idxs = param_A.idxs; + const Tu* u_gmem = (const Tu*)param_U.ptr; + const int ldU = param_U.stride; + const int K = sched.gemm_shape().z; + + constexpr int kVec = 16; // uint4 / sizeof(fp8) + constexpr int nvec = TILE_M * (TILE_K / kVec); + // TILE_M=8 → nvec=64 < WARPGROUP_SIZE: one slot, idle threads predicated off. + constexpr int kSlots = (nvec + WARPGROUP_SIZE - 1) / WARPGROUP_SIZE; + static_assert(kSlots >= 1); + // U: one float per producer thread along TILE_M (thread m loads row m). + static_assert(TILE_M <= WARPGROUP_SIZE); + + typename Scheduler::Tile* tile; + + while (true) { + const CUtensorMap* Bdesc = &tm_b; + uint16_t mask_B = 0; + int coord_n = 0; + int k_iters = 0; + const Tv* gmem_V0 = (const Tv*)param_V.ptr; + const Tv* gmem_V1 = nullptr; + int ldV = param_V.stride; + + if (warp_in_wg == 0 && cta_0) { + (void)prod_state.next(); + } + + if (warp_in_wg == 0) { + const bool alive = sched_state.acquire(tile); + int m0 = 0, M_group = 0, offset_m = 0; + + if (alive && tile->is_valid_cluster) { + if constexpr (is_grouped_gemm) { + const int g = tile->group_idx; + Bdesc = &tensormap_buf[g * kTmaDescNum]; + const auto v = resolve(param_V, g); + gmem_V0 = (const Tv*)v.ptr.ptr; + ldV = v.ptr.stride; + } + + mask_B = cluster.mask_n(); + coord_n = tile->offset_n + mc_offset_n; + offset_m = tile->offset_m; + m0 = [&] { + if constexpr (is_grouped_gemm) { + return tile->m0; + } + return 0; + }(); + M_group = [&] { + if constexpr (is_grouped_gemm) { + return tile->m1 - tile->m0; + } + return sched.gemm_shape().x; + }(); + k_iters = sched.k_iters_; + + gmem_V0 += (tile->offset_n / 128) * ldV; + gmem_V1 = gmem_V0; + if (tile->offset_n / 128 + 1 < cdiv(sched.gemm_shape().y, 128)) { + gmem_V1 += ldV; + } + } + + if (lane_id == 0) { + storage.gather_alive = alive ? 1 : 0; + storage.gather_k_iters = k_iters; + storage.gather_m0 = m0; + storage.gather_M_group = M_group; + storage.gather_offset_m = offset_m; + } + __syncwarp(); + } + + // Tile header only. + gather_bar.arrive_and_wait(); + + if (storage.gather_alive == 0) { + break; + } + + k_iters = storage.gather_k_iters; + const int m0 = storage.gather_m0; + const int M_group = storage.gather_M_group; + const int offset_m = storage.gather_offset_m; + int coord_k = 0; + + // iterator_sm80 style: idxs → src bases, then += TILE_K each K tile. + const Ta* src_data_vec_[kSlots]; + int m_own[kSlots]; + int kk_own[kSlots]; + bool pred_row[kSlots]; + const Tu* src_u_base; + int m_u; + bool pred_u; + const int u_pad = m0 % kAlignmentU; + + PRAGMA_UNROLL + for (int t = 0; t < kSlots; ++t) { + const int i = prod_tid + t * WARPGROUP_SIZE; + const bool in_vec = i < nvec; + const int m = in_vec ? i / (TILE_K / kVec) : 0; + const int kk = in_vec ? (i % (TILE_K / kVec)) * kVec : 0; + const int packed = m0 + offset_m + m; + const bool row_ok = in_vec && (offset_m + m) < M_group; + const int token = (idxs && row_ok) ? __ldg(idxs + packed) : packed; + m_own[t] = m; + kk_own[t] = kk; + pred_row[t] = row_ok; + src_data_vec_[t] = act_gmem + (int64_t)token * ldA + kk; + } + + { + // TILE_M < WARPGROUP_SIZE: idle threads must skip U ZFILL entirely + // (pred=false still zeros dst — see sm90_bf16). + const int m = prod_tid; + const bool in_m = m < TILE_M; + const int packed = m0 + offset_m + m; + const bool row_ok = in_m && (offset_m + m) < M_group; + const int token = (idxs && row_ok) ? __ldg(idxs + packed) : packed; + m_u = m; + pred_u = row_ok; + src_u_base = u_gmem + token; + } + + // Warp0-only weight TMA + V; full WG gathers A/U. + GmemIteratorSm90 gmem_B{ + (warp_in_wg == 0) ? Bdesc : &tm_b, {0, (warp_in_wg == 0) ? coord_n : 0}, {TILE_K, 0}}; + + for (; k_iters > 0; --k_iters) { + const int pipe = write_state.index(); + ConsumerBar::wait(&consumer_bar[pipe], write_state.phase()); + + if (warp_in_wg == 0 && lane_predicate) { + ProducerBar::arrive_and_expect_tx(&producer_bar[pipe], kTmaTxBytes); + gmem_B.Step(&producer_bar[pipe], &smem_weight[pipe * TILE_N * TILE_K], mask_B); + uint32_t uint_ptr_V = cast_smem_ptr_to_uint(smem_V[pipe]); + CP_ASYNC::apply(uint_ptr_V, gmem_V0, true); + CP_ASYNC::apply(uint_ptr_V + sizeof(Tv), gmem_V1, true); + ++gmem_V0; + ++gmem_V1; + } + + { + cute::Tensor sB = cute::make_tensor(cute::make_smem_ptr(smem_act + pipe * TILE_M * TILE_K), + SmemLayoutB_2D{}); + + PRAGMA_UNROLL + for (int t = 0; t < kSlots; ++t) { + const bool pred = pred_row[t] && (coord_k + kk_own[t]) < K; + auto* dst = &sB(m_own[t], kk_own[t]); + cute::SM80_CP_ASYNC_CACHEGLOBAL_ZFILL::copy( + *reinterpret_cast(src_data_vec_[t]), + *reinterpret_cast(dst), + pred); + src_data_vec_[t] += TILE_K; + } + + // U: one float per thread along TILE_M; ColMajor (token + k_col*ldU). + // Idle (m >= TILE_M): do not issue ZFILL — pred=false still zeros dst. + if (m_u < TILE_M) { + const bool pred = pred_u && (coord_k < K); + auto* dst = &smem_U[pipe][u_pad + m_u]; + const Tu* src = src_u_base + (coord_k / 128) * ldU; + cute::SM80_CP_ASYNC_CACHEALWAYS_ZFILL::copy( + *reinterpret_cast(src), *reinterpret_cast(dst), pred); + } + + cutlass::arch::cpasync_barrier_arrive_noinc(&producer_bar[pipe]); + } + + ++write_state; + coord_k += TILE_K; + } + + if (warp_in_wg == 0) { + sched_state.release(); + } + } + + if (warp_in_wg == 0) { + sched_state.release(); + if (cta_0) { + sched.tail(prod_state); + } + if constexpr (kClusterSize > 1) { + if (lane_predicate) { + for (int i = 0; i < Stages; ++i) { + ConsumerBar::wait(&consumer_bar[write_state.index()], write_state.phase()); + ++write_state; + } + } + __syncwarp(); + } + } + } + else if (warp_in_wg == 0) { + Cluster cluster(cute::block_id_in_cluster().x); + + const int mc_offset_m = cluster.cta_n() * (TILE_M / kMulticastA); + const int mc_offset_n = cluster.cta_m() * (TILE_N / kMulticastB); + + auto smem_weight = storage.A.data() + mc_offset_n * TILE_K; + auto smem_act = storage.B.data() + mc_offset_m * TILE_K; + auto& smem_U = storage.U; + auto& smem_V = storage.V; + + cutlass::PipelineState write_state{0, 1, 0}; + + auto sched_state = sched.init_consumer(storage.sched); + + int lane_predicate = cute::elect_one_sync(); + + typename Scheduler::Tile* tile; + + while (sched_state.acquire(tile)) { + + if (tile->is_valid_cluster) { + + const CUtensorMap* Adesc = &tm_a; + const CUtensorMap* Bdesc = &tm_b; + const CUtensorMap* Udesc = &tm_u; + + const Tv* gmem_V0 = (const Tv*)param_V.ptr; + const Tv* gmem_V1; + int ldV = param_V.stride; + + if constexpr (is_grouped_gemm) { + // Descs published by prepare_moe_tma_descs on this stream; + // fence_acquire only needed after in-kernel tensormap replace. + const int g = tile->group_idx; + CUtensorMap* const descs = tensormap_buf + g * kTmaDescNum; + if constexpr (kStridingA == Striding::kBlocked) { + Adesc = &descs[0]; + Bdesc = &descs[1]; + Udesc = &descs[2]; + } + else { + Bdesc = &descs[0]; + } + const auto v = resolve(param_V, g); + gmem_V0 = (const Tv*)v.ptr.ptr; + ldV = v.ptr.stride; + } + + if (lane_predicate) { + const int offset_k = 0; + + const uint16_t mask_A = cluster.mask_m(); + const uint16_t mask_B = cluster.mask_n(); + + const int offset_m = tile->offset_m; + const int offset_n = tile->offset_n; + + int k_iter = sched.k_iters_; + + GmemIteratorSm90 gmem_A{ + Adesc, {offset_k, offset_m + mc_offset_m}, {TILE_K, 0}}; + GmemIteratorSm90 gmem_B{ + Bdesc, {offset_k, offset_n + mc_offset_n}, {TILE_K, 0}}; + + const int mc_offset_u = kMulticastU > 1 ? mc_offset_m : 0; + // column-major + GmemIteratorSm90 gmem_U{ + Udesc, {offset_m + mc_offset_u, offset_k / 128}, {0, 1}}; + + gmem_V0 += (offset_n / 128) * ldV + (offset_k / 128); + gmem_V1 = gmem_V0; + if (offset_n / 128 + 1 < cdiv(sched.gemm_shape().y, 128)) { + gmem_V1 += ldV; + } + + for (; k_iter > 0; --k_iter) { + int pipe = write_state.index(); + ConsumerBar::wait(&consumer_bar[pipe], write_state.phase()); + ProducerBar::arrive_and_expect_tx(&producer_bar[pipe], kTmaTxBytes); + // API A=act → GMMA-B; API B=weight → GMMA-A + gmem_A.Step(&producer_bar[pipe], &smem_act[pipe * TILE_M * TILE_K], mask_A); + gmem_B.Step(&producer_bar[pipe], &smem_weight[pipe * TILE_N * TILE_K], mask_B); + gmem_U.Step(&producer_bar[pipe], smem_U[pipe] + mc_offset_u, mask_A); + uint32_t uint_ptr_V = cast_smem_ptr_to_uint(smem_V[pipe]); + CP_ASYNC::apply(uint_ptr_V, gmem_V0, true); + CP_ASYNC::apply(uint_ptr_V + sizeof(Tv), gmem_V1, true); + ++gmem_V0; + ++gmem_V1; + cutlass::arch::cpasync_barrier_arrive_noinc(&producer_bar[pipe]); + ++write_state; + } + } + } + + if constexpr (Scheduler::is_dynamic) { + if (cta_0) { + producers_bar.arrive_unaligned(); + } + } + + sched_state.release(); + + } // scheduler loop + + // release last tile + sched_state.release(); + + if constexpr (kClusterSize > 1) { + if (lane_predicate) { + for (int i = 0; i < Stages; ++i) { + ConsumerBar::wait(&consumer_bar[write_state.index()], write_state.phase()); + ++write_state; + } + } + __syncwarp(); + } + } + else if (warp_in_wg == 1 && cta_0) { + if constexpr (!kIndexedGather) { + auto state = sched.init_producer(storage.sched); + while (state.next()) { + if constexpr (Scheduler::is_dynamic) { + producers_bar.arrive_and_wait_unaligned(); + } + } + sched.tail(state); + } + } + } + else { + cutlass::arch::warpgroup_reg_alloc(); + + auto& smem_weight = storage.A; // GMMA-A + auto& smem_act = storage.B; // GMMA-B + auto& smem_U = storage.U; + auto& smem_V = storage.V; + + const int wg_idx_m = WG_M > 1 ? wg_idx % WG_M : 0; // along BATCH + const int wg_idx_n = WG_N > 1 ? wg_idx / WG_M : 0; // along OUT + + // GMMA-A steps along OUT (TILE_N); GMMA-B along BATCH (TILE_M). + auto smem_desc_A = make_smem_desc(&smem_weight[wg_idx_n * WG_TILE_N * TILE_K], 1); + auto smem_desc_B = make_smem_desc(&smem_act[wg_idx_m * WG_TILE_M * TILE_K], 1); + + SmemDescIterV2> 4)> smem_iter_A{smem_desc_A}; + SmemDescIterV2> 4)> smem_iter_B{smem_desc_B}; + + cutlass::arch::NamedBarrier barrier(WARPGROUP_SIZE, 2 + wg_idx); // per math WG + + // All-math-WG sync (smem C is shared; one full-tile TMA store per tile). + auto epi_synchronize = [&] { + cutlass::arch::NamedBarrier::sync(kMathGroupSize, + cutlass::arch::ReservedNamedBarriers::EpilogueBarrier); + }; + + cutlass::PipelineState pipe_state{}; + + const int warp_id = cutlass::canonical_warp_idx_sync(); + const int lane_id = cutlass::canonical_lane_idx(); + + auto consumer_arrive = [&] { + auto bar = &consumer_bar[pipe_state.index()]; + __syncwarp(); + if constexpr (kClusterSize > 1) { + ConsumerBar::arrive(bar, lane_id, lane_id < kClusterSize); + } + else { + if (lane_id == 0) { + ConsumerBar::arrive(bar); + } + } + }; + + auto sched_state = sched.init_consumer(storage.sched); + + typename Scheduler::Tile* tile; + + sched_state.acquire(tile); + + while (tile->alive) { + + if (tile->is_valid_cta) { + AccumC accum_C{}; + FragC frag_C; + + auto pred_W = Fetch_W(tile, wg_idx_n); + + float scale_w[2]; + auto Load_W = [&] { + scale_w[0] = smem_V[pipe_state.index()][0]; + scale_w[1] = smem_V[pipe_state.index()][1]; + }; + + // Dense act scales along GMMA-N (BATCH): this thread's owned columns. + int u_pad = 0; + if constexpr (is_grouped_gemm) { + u_pad = tile->m0 % kAlignmentU; + } + const int batch0 = wg_idx_m * WG_TILE_M; + typename GMMA::FragA act_scales{}; + auto Load_A = [&] { + const int pipe = pipe_state.index(); + // n = 2*t0 + 8*v2; t0 = lane%4 + const int lane_n0 = 2 * (lane_id % 4); + PRAGMA_UNROLL + for (int s = 0; s < GMMA::OP_N / 8; ++s) { + const int n0 = batch0 + lane_n0 + 8 * s; + act_scales[2 * s + 0] = smem_U[pipe][u_pad + n0]; + act_scales[2 * s + 1] = smem_U[pipe][u_pad + n0 + 1]; + } + }; + + auto gmma = [&] { // + GMMA::apply(smem_iter_A, smem_iter_B, frag_C, accum_C, scale_w, act_scales, pred_W); + }; + + if constexpr (is_grouped_gemm) { + if (threadIdx.x == 0) { // single-store issuer + cute::tma_store_wait<0>(); + } + epi_synchronize(); + } + + int k_iter = sched.k_iters_; + + ProducerBar::wait(&producer_bar[pipe_state.index()], pipe_state.phase()); + Load_W(); + Load_A(); + smem_iter_A.Reset(pipe_state.index()); + smem_iter_B.Reset(pipe_state.index()); + gmma(); + consumer_arrive(); + ++pipe_state; + --k_iter; + + ProducerBar::wait(&producer_bar[pipe_state.index()], pipe_state.phase()); + Load_W(); + Load_A(); + smem_iter_A.Reset(pipe_state.index()); + smem_iter_B.Reset(pipe_state.index()); + + PRAGMA_NO_UNROLL + for (; k_iter > 1; --k_iter) { + gmma(); + consumer_arrive(); + ++pipe_state; + ProducerBar::wait(&producer_bar[pipe_state.index()], pipe_state.phase()); + Load_W(); + Load_A(); + smem_iter_A.Reset(pipe_state.index()); + smem_iter_B.Reset(pipe_state.index()); + } + + gmma(); + + const int thread_idx = threadIdx.x % WARPGROUP_SIZE; + if constexpr (!is_grouped_gemm) { + if (threadIdx.x == 0) { // single-store issuer + cute::tma_store_wait<0>(); + } + epi_synchronize(); + } + + consumer_arrive(); + ++pipe_state; + + auto run_epilogue = [&](auto fused_silu) { + constexpr bool kFuseSilu = decltype(fused_silu)::value; + using OutputTraits = Output; + using OutputT = typename OutputTraits::Tc; + using LayoutC = typename OutputTraits::LayoutC; + + constexpr int kStoreOut = OutputTraits::kStoreOut; + OutputT* smem_C = reinterpret_cast(smem_buf + kOutputOffset); + + // Fused SiLU: silu(gate)*up on GMMA-M atoms {0,1}×{2,3}, then + // per-BATCH-row amax over fused OUT=128 → FP8 + W scales. + // (WG_N == 1 only; WG_N == 2 stages gate/up through smem below.) + if constexpr (kFuseSilu && WG_N == 1) { + static_assert(GMMA::ITER_N == 1); + constexpr float kQmax = 448.f; + constexpr int kScales = GMMA::kActScalesPerThread; + + PRAGMA_UNROLL + for (int i_m = 0; i_m < 2; ++i_m) { + auto& gate = accum_C[i_m][0][0][0][0][0]; + auto& up = accum_C[i_m + 2][0][0][0][0][0]; + constexpr int kNumRegs = (int)(sizeof(gate) / sizeof(float)); + PRAGMA_UNROLL + for (int i = 0; i < kNumRegs; ++i) { + const float g = gate[i]; + const float u = up[i]; + gate[i] = fdividef(g, 1.f + expf(-g)) * u; + } + } + + // Partial amax over this thread's OUT rows in gate atoms {0,1}. + // CRegisters: [0]=(m,n0), [1]=(m,n1), [2]=(m+8,n0), [3]=(m+8,n1). + float amax[kScales]; + PRAGMA_UNROLL + for (int i = 0; i < kScales; ++i) { + amax[i] = 0.f; + } + PRAGMA_UNROLL + for (int i_m = 0; i_m < 2; ++i_m) { + auto& gate = accum_C[i_m][0][0][0][0][0]; + PRAGMA_UNROLL + for (int c = 0, s = 0; c < GMMA::OP_N; c += 8, ++s) { + amax[2 * s + 0] = fmaxf(amax[2 * s + 0], fabsf(gate[c / 2 + 0])); + amax[2 * s + 0] = fmaxf(amax[2 * s + 0], fabsf(gate[c / 2 + 2])); + amax[2 * s + 1] = fmaxf(amax[2 * s + 1], fabsf(gate[c / 2 + 1])); + amax[2 * s + 1] = fmaxf(amax[2 * s + 1], fabsf(gate[c / 2 + 3])); + } + } + // Reduce across OUT-owning lanes (t1) within the warp. + PRAGMA_UNROLL + for (int i = 0; i < kScales; ++i) { + amax[i] = fmaxf(amax[i], __shfl_xor_sync(0xffffffffu, amax[i], 4)); + amax[i] = fmaxf(amax[i], __shfl_xor_sync(0xffffffffu, amax[i], 8)); + amax[i] = fmaxf(amax[i], __shfl_xor_sync(0xffffffffu, amax[i], 16)); + } + // Reduce across warps (t2) for full fused OUT=128. + { + const int t0 = lane_id % 4; + const int t1 = (lane_id / 4) % 8; + const int warp = warp_id % 4; + float* scratch = storage.fused_amax_scratch[wg_idx]; + if (t1 == 0) { + PRAGMA_UNROLL + for (int i = 0; i < kScales; ++i) { + scratch[(warp * 4 + t0) * 64 + i] = amax[i]; + } + } + barrier.sync(); + PRAGMA_UNROLL + for (int i = 0; i < kScales; ++i) { + float m = scratch[(0 * 4 + t0) * 64 + i]; + m = fmaxf(m, scratch[(1 * 4 + t0) * 64 + i]); + m = fmaxf(m, scratch[(2 * 4 + t0) * 64 + i]); + m = fmaxf(m, scratch[(3 * 4 + t0) * 64 + i]); + amax[i] = fmaxf(m, 1e-8f); + } + barrier.sync(); + } + + float inv[kScales]; + PRAGMA_UNROLL + for (int i = 0; i < kScales; ++i) { + inv[i] = kQmax / amax[i]; + } + PRAGMA_UNROLL + for (int i_m = 0; i_m < 2; ++i_m) { + auto& gate = accum_C[i_m][0][0][0][0][0]; + PRAGMA_UNROLL + for (int c = 0, s = 0; c < GMMA::OP_N; c += 8, ++s) { + gate[c / 2 + 0] *= inv[2 * s + 0]; + gate[c / 2 + 1] *= inv[2 * s + 1]; + gate[c / 2 + 2] *= inv[2 * s + 0]; + gate[c / 2 + 3] *= inv[2 * s + 1]; + } + } + + // W: one scale per BATCH row (act); n_group along OUT tiles. + if (param_W.ptr && (lane_id / 4) % 8 == 0 && (warp_id % 4) == 0) { + const int n_group = tile->offset_n / TILE_N; + const int t0 = lane_id % 4; + const int batch0 = wg_idx_m * WG_TILE_M; + Tw* W = reinterpret_cast(param_W.ptr); + const int ldW = param_W.stride; + int row_end = 0x7fffffff; + int row_base = tile->offset_m + batch0; + if constexpr (is_grouped_gemm) { + row_base += tile->m0; + row_end = tile->m1; + } + PRAGMA_UNROLL + for (int c = 0, s = 0; c < GMMA::OP_N; c += 8, ++s) { + const int row0 = row_base + 2 * t0 + c; + const int row1 = row0 + 1; + if (row0 < row_end && (batch0 + 2 * t0 + c) < TILE_M) { + W[(int64_t)n_group * ldW + row0] = amax[2 * s + 0] / kQmax; + } + if (row1 < row_end && (batch0 + 2 * t0 + c + 1) < TILE_M) { + W[(int64_t)n_group * ldW + row1] = amax[2 * s + 1] / kQmax; + } + } + } + } + + if constexpr (kFuseSilu && WG_N == 2) { + // WG_N == 2 staged fused SiLU: WG0's accum holds the gate + // slab, WG1's the up slab. Exchange raw f32 through smem, then + // each WG computes silu(g)*u for its fused half, cross-WG amax + // per BATCH row, quantize, write its half of smem C. + constexpr float kQmax = 448.f; + constexpr int kWgFused = (TILE_N / 2) / WG_N; + float* stage = storage.silu_stage; // [TILE_N][TILE_M] + + GMMA::foreach_C(accum_C, [&](auto& C, int m_atom, int n_atom) { + const int out0 = wg_idx_n * WG_TILE_N + m_atom * GMMA::OP_M; + const int batch0 = wg_idx_m * WG_TILE_M + n_atom * GMMA::OP_N; + const int m0 = out0 + (warp_id % 4) * 16 + lane_id / 4; + const int lane_n0 = 2 * (lane_id % 4); + PRAGMA_UNROLL + for (int c = 0; c < GMMA::OP_N; c += 8) { + const int n0 = batch0 + lane_n0 + c; + stage[(int64_t)m0 * TILE_M + n0] = C[c / 2 + 0]; + stage[(int64_t)m0 * TILE_M + (n0 + 1)] = C[c / 2 + 1]; + stage[(int64_t)(m0 + 8) * TILE_M + n0] = C[c / 2 + 2]; + stage[(int64_t)(m0 + 8) * TILE_M + (n0 + 1)] = C[c / 2 + 3]; + } + }); + epi_synchronize(); + + // Linear repartition: kWgFused×TILE_M items, one row per thread pair. + constexpr int kItems = kWgFused * TILE_M / WARPGROUP_SIZE; + float vals[kItems]; + float amax = 1e-8f; + const int base = thread_idx * kItems; + PRAGMA_UNROLL + for (int i = 0; i < kItems; ++i) { + const int idx = base + i; + const int row = idx / kWgFused; + const int of = wg_idx_n * kWgFused + idx % kWgFused; + const float g = stage[(int64_t)of * TILE_M + row]; + const float u = stage[(int64_t)(of + TILE_N / 2) * TILE_M + row]; + vals[i] = fdividef(g, 1.f + expf(-g)) * u; + amax = fmaxf(amax, fabsf(vals[i])); + } + amax = fmaxf(amax, __shfl_xor_sync(0xffffffffu, amax, 1)); + + float* scratch = storage.fused_amax_scratch[wg_idx]; + const int row = base / kWgFused; + if ((thread_idx & 1) == 0) { + scratch[row] = amax; + } + epi_synchronize(); + const float row_amax = fmaxf( + fmaxf(storage.fused_amax_scratch[0][row], storage.fused_amax_scratch[1][row]), 1e-8f); + const float inv = kQmax / row_amax; + + PRAGMA_UNROLL + for (int i = 0; i < kItems; ++i) { + const int idx = base + i; + const int r = idx / kWgFused; + const int c = wg_idx_n * kWgFused + idx % kWgFused; + smem_C[(int64_t)r * kStoreOut + c] = OutputT(vals[i] * inv); + } + + // W: one scale per BATCH row (spans both WGs); WG0 pair leaders store. + if (param_W.ptr && wg_idx == 0 && (thread_idx & 1) == 0) { + const int n_group = tile->offset_n / TILE_N; + Tw* W = reinterpret_cast(param_W.ptr); + const int ldW = param_W.stride; + int row_end = 0x7fffffff; + int row_g = tile->offset_m + row; + if constexpr (is_grouped_gemm) { + row_g += tile->m0; + row_end = tile->m1; + } + if (row_g < row_end && row < TILE_M) { + W[(int64_t)n_group * ldW + row_g] = row_amax / kQmax; + } + } + } + else { + // WA CRegisters (OUT, BATCH) → problem smem C (BATCH, kStoreOut) row-major. + static_assert(GMMA::OP_N % 8 == 0); + GMMA::foreach_C(accum_C, [&](auto& C, int m_atom, int n_atom) { + if constexpr (kFuseSilu) { + if (m_atom >= 2) { + return; // up atoms consumed + } + } + // m_atom along OUT, n_atom along BATCH + const int out0 = wg_idx_n * (kStoreOut / WG_N) + m_atom * GMMA::OP_M; + const int batch0 = wg_idx_m * WG_TILE_M + n_atom * GMMA::OP_N; + const int m0 = out0 + (warp_id % 4) * 16 + lane_id / 4; + const int lane_n0 = 2 * (lane_id % 4); + PRAGMA_UNROLL + for (int c = 0, s = 0; c < GMMA::OP_N; c += 8, ++s) { + const int n0 = batch0 + lane_n0 + c; + smem_C[(int64_t)n0 * kStoreOut + m0] = OutputT(C[c / 2 + 0]); + smem_C[(int64_t)(n0 + 1) * kStoreOut + m0] = OutputT(C[c / 2 + 1]); + smem_C[(int64_t)n0 * kStoreOut + (m0 + 8)] = OutputT(C[c / 2 + 2]); + smem_C[(int64_t)(n0 + 1) * kStoreOut + (m0 + 8)] = OutputT(C[c / 2 + 3]); + } + }); + } + + cutlass::arch::fence_view_async_shared(); + cute::tma_store_fence(); + epi_synchronize(); + + const int offset_m = tile->offset_m; + const int offset_n = tile->offset_n; + + const void* Cdesc = &tm_c; + + // One full-tile store for all math WGs (smem C is shared). + if (threadIdx.x == 0) { + if constexpr (is_grouped_gemm) { + Cdesc = tensormap_buf + tile->group_idx * kTmaDescNum + kCdescIdx; + } + const int store_n = kFuseSilu ? (offset_n / 2) : offset_n; + cute::SM90_TMA_STORE::copy(Cdesc, smem_C, store_n, offset_m); + cute::tma_store_arrive(); + } + if constexpr (is_grouped_gemm) { + if (threadIdx.x == 0) { + cute::tma_store_wait<0>(); + } + epi_synchronize(); + } + }; + + if constexpr (kSupportsFusedSilu) { + if (fuse_silu) { + run_epilogue(std::true_type{}); + } + else { + run_epilogue(std::false_type{}); + } + } + else { + run_epilogue(std::false_type{}); + } + } + else if (tile->is_valid_cluster) { + int k_iter = sched.k_iters_; + for (; k_iter > 0; --k_iter) { + ProducerBar::wait(&producer_bar[pipe_state.index()], pipe_state.phase()); + consumer_arrive(); + ++pipe_state; + } + } + + sched_state.release(); + sched_state.acquire(tile); + + } // scheduler loop + + // release last tile + sched_state.release(); + + if (threadIdx.x % WARPGROUP_SIZE == 0) { + cute::tma_store_wait<0>(); + } + } + + } // operator() + + // Weight-scale 128-block pred along OUT (problem N / GMMA-M). + __device__ auto Fetch_W(typename Scheduler::Tile* tile, int wg_idx_n) + { + constexpr int BLK_SUBTILE_M = 128 / OUTER_M; + static_assert(MMA_SUBTILE_M - 1 < BLK_SUBTILE_M + 1); + + Array pred_W{}; + // i=0 is covered too: for 1 WG offset < 128 keeps pred_W[0] false (as before); + // for WG_N == 2 the wg_idx_n term selects the second V scale block. + const int offset = tile->offset_n % 128 + wg_idx_n * WG_TILE_N; + PRAGMA_UNROLL + for (int i = 0; i < MMA_SUBTILE_M; ++i) { + pred_W[i] = (i * OUTER_M + offset) >= 128; + } + + return pred_W; + } +}; + +} // namespace turbomind::gemm diff --git a/src/turbomind/kernels/gemm/gemm_universal_sm90_v2.h b/src/turbomind/kernels/gemm/gemm_universal_sm90_v2.h deleted file mode 100644 index 1290e624bb..0000000000 --- a/src/turbomind/kernels/gemm/gemm_universal_sm90_v2.h +++ /dev/null @@ -1,746 +0,0 @@ -#pragma once - -#include -#include - -#include - -#include "cute/arch/cluster_sm90.hpp" -#include "cute/arch/copy_sm90.hpp" -#include "cute/arch/copy_sm90_tma.hpp" -#include "cute/arch/mma_sm90_desc.hpp" -#include "cute/arch/mma_sm90_gmma.hpp" -#include "cute/atom/mma_traits.hpp" -#include "cute/atom/mma_traits_sm90_gmma.hpp" - -#include "cutlass/arch/barrier.h" -#include "cutlass/arch/reg_reconfig.h" -#include "cutlass/cutlass.h" -#include "cutlass/pipeline/sm90_pipeline.hpp" - -#include "src/turbomind/core/data_type.h" - -#include "src/turbomind/kernels/core/array_ops.h" -#include "src/turbomind/kernels/core/common.h" -#include "src/turbomind/kernels/core/layout.h" -#include "src/turbomind/kernels/core/meta.h" -#include "src/turbomind/kernels/core/smem.h" - -#include "src/turbomind/kernels/gemm/arch.h" -#include "src/turbomind/kernels/gemm/iterator_sm90.h" -#include "src/turbomind/kernels/gemm/scheduler.cuh" -#include "src/turbomind/kernels/gemm/types.h" -#include "src/turbomind/kernels/gemm/utils.h" - -namespace turbomind::gemm { - -namespace GMMA = cute::SM90::GMMA; - -inline __device__ cute::GmmaDescriptor make_smem_desc(void* smem_ptr, int layout_type) -{ - auto uint_ptr = cast_smem_ptr_to_uint(smem_ptr); - - cute::GmmaDescriptor desc{}; - desc.bitfield.start_address_ = uint_ptr >> 4; - desc.bitfield.layout_type_ = layout_type; - desc.bitfield.leading_byte_offset_ = 0; - desc.bitfield.stride_byte_offset_ = 1024 >> 4; - desc.bitfield.base_offset_ = 0; - - return desc; -} - -template -struct SmemDescIterV2 { - union { - uint32_t u32_[2]; - uint64_t u64_; - }; - - uint32_t base_; - - __device__ SmemDescIterV2(uint64_t desc): u64_{desc}, base_{u32_[0]} {} - - __device__ void Advance(int stage) - { - u32_[0] += Step; - if (stage == Stages - 1) { - u32_[0] = base_; - } - } - - __device__ void Reset(int stage) - { - u32_[0] = base_ + stage * Step; - } - - __device__ SmemDescIterV2& operator+=(int offset) - { - u32_[0] += offset; - return *this; - } - - __device__ SmemDescIterV2& operator-=(int offset) - { - u32_[0] -= offset; - return *this; - } - - __device__ operator uint64_t() - { - return u64_; - } -}; - -template -inline __device__ void -wgmma_impl(uint64_t desc_a, uint64_t desc_b, float* frag_C, bool clear, std::index_sequence) -{ - return MMA_Atom::fma(desc_a, desc_b, frag_C[Is]..., clear ? GMMA::ScaleOut::Zero : GMMA::ScaleOut::One); -} - -template -inline __device__ void wgmma(uint64_t desc_a, uint64_t desc_b, float (&frag_C)[N], bool clear) -{ - return wgmma_impl(desc_a, desc_b, frag_C, clear, std::make_index_sequence{}); -} - -inline __device__ void warpgroup_fence_operand(float& reg) -{ - asm volatile("" : "+f"(reg)::"memory"); -} - -template -inline __device__ void warpgroup_fence_operand(float (&x)[M][N][K]) -{ - PRAGMA_UNROLL - for (int m = 0; m < M; ++m) { - PRAGMA_UNROLL - for (int n = 0; n < N; ++n) { - PRAGMA_UNROLL - for (int k = 0; k < K; ++k) { - warpgroup_fence_operand(x[m][n][k]); - } - } - } -} - -template -inline __device__ void warpgroup_fence_operand(float (&x)[N][K]) -{ - PRAGMA_UNROLL - for (int n = 0; n < N; ++n) { - PRAGMA_UNROLL - for (int k = 0; k < K; ++k) { - warpgroup_fence_operand(x[n][k]); - } - } -} - -template -__device__ void for_(std::index_sequence, Func func) -{ - return (func(constant{}), ...); -} - -namespace arch { - -template -struct Cluster { - static constexpr int M = M_; - static constexpr int N = N_; - - static constexpr int C = mk2cs(M, N).x; - static constexpr int S = mk2cs(M, N).y; - - static constexpr int size = M * N; - - static constexpr uint16_t kMaskC = (1 << C) - 1; - static constexpr uint16_t kMaskS = ((1 << size) - 1) / kMaskC; - - __device__ static ushort2 mask_cs(int cta_id) - { - const auto [c, s] = cta_cs(cta_id); - return make_ushort2(kMaskS << c, kMaskC << s * C); - } - - __device__ static ushort2 mask_mn(int cta_id) - { - auto [c, s] = mask_cs(cta_id); - return order == kColMajor ? ushort2{c, s} : ushort2{s, c}; - } - - __device__ static int2 cta_cs(int cta_id) - { - return {C > 1 ? cta_id % C : 0, S > 1 ? cta_id / C : 0}; - } - - __device__ static int2 cta_mn(int cta_id) - { - return cs2mk(cta_cs(cta_id)); - } - - int2 cta_mn_; - ushort2 mask_mn_; - - __device__ explicit Cluster(int cta_id): cta_mn_(cta_mn(cta_id)), mask_mn_(mask_mn(cta_id)) {} - - __device__ int cta_m() - { - return cta_mn_.x; - } - - __device__ int cta_n() - { - return cta_mn_.y; - } - - __device__ uint16_t mask_m() - { - return mask_mn_.x; - } - - __device__ uint16_t mask_n() - { - return mask_mn_.y; - } -}; - -} // namespace arch - -struct GemmUniversalSm90_v2 { - - static constexpr bool kDebug = false; - - using Arch = Sm90; - - // using MMA_Atom = GMMA::MMA_64x128x16_F32BF16BF16_SS; - using MMA_Atom = GMMA::MMA_64x96x32_F32E4M3E4M3_SS_TN<>; - static constexpr typename cute::MMA_Traits::Shape_MNK MMA_Shape{}; - - static constexpr int MMA_ATOM_M = cute::get<0>(MMA_Shape); - static constexpr int MMA_ATOM_N = cute::get<1>(MMA_Shape); - static constexpr int MMA_ATOM_K = cute::get<2>(MMA_Shape); - - static constexpr int WARPGORUPS = 2; - - static constexpr int TILE_M = 128; - static constexpr int TILE_N = MMA_ATOM_N; - static constexpr int TILE_K = 128; - - static constexpr int MMA_ITER_M = TILE_M / MMA_ATOM_M; - static constexpr int MMA_ITER_N = TILE_N / MMA_ATOM_N; - static constexpr int MMA_ITER_K = TILE_K / MMA_ATOM_K; - - static constexpr int kMulticastA = 1; - static constexpr int kMulticastB = 2; - - static constexpr int kClusterSize = kMulticastA * kMulticastB; - - static constexpr int Stages = 4; - - static constexpr bool kSplitK = false; - static constexpr int kChunkSizeK = TILE_K; - - static constexpr int WARPGROUP_SIZE = 128; - - static constexpr int CTA_SIZE = WARPGROUP_SIZE * (WARPGORUPS + 1); - - using Ta = __nv_fp8_e4m3; - using Tb = __nv_fp8_e4m3; - using Tc = nv_bfloat16; - - using Tu = float; - using Tv = float; - - using Cluster = arch::Cluster; - - using Scheduler = TileScheduler; - - using ProducerBar = cutlass::arch::ClusterTransactionBarrier; - using ConsumerBar = cutlass::arch::ClusterBarrier; - - static constexpr int MAX_K = 32768; - - static constexpr int TILE_M_U = cdiv(TILE_M, 1); - static constexpr int CTA_K_U = cdiv(TILE_K, 128); - - static constexpr int kTmaTxBytes = - sizeof(Ta) * (TILE_M * TILE_K) + sizeof(Tb) * (TILE_N * TILE_K) + sizeof(Tu) * TILE_M_U * CTA_K_U; - - // ! Smem addr must be SBO aligned for TMA load/store - struct SharedStorage { - struct Source { - __align__(1024) Array A; - __align__(1024) Array B; - __align__(1024) Tu U[Stages][round_up(TILE_M_U * CTA_K_U, 32)]; - __align__(1024) Tv V[2][WARPGORUPS][cdiv(MAX_K, 128)]; - }; - Source source; - __align__(1024) Array C; - __align__(128) uint64_t producer_bar[Stages]; - __align__(128) uint64_t consumer_bar[Stages]; - int pipe_count[WARPGORUPS]; - }; - - static constexpr int kSmemSize = sizeof(SharedStorage); - - static constexpr int kSwizzleC = 2 * std::gcd(TILE_N, 128 / sizeof(Tc)); - - using LayoutC = std::conditional_t= 32, - SmemLayoutV2, - SmemLayoutV2>; - - __device__ void operator()(const CUtensorMap& tm_a, - const CUtensorMap& tm_b, - const CUtensorMap& tm_c, - const CUtensorMap& tm_u, - const CUtensorMap& tm_v, - const void* U_, - int ldU, - const void* V_, - int ldV, - Scheduler sched, - char* smem_buf) - { - SharedStorage& storage = *reinterpret_cast(smem_buf); - - uint64_t* producer_bar = storage.producer_bar; - uint64_t* consumer_bar = storage.consumer_bar; - - if (threadIdx.x == 0) { - PRAGMA_UNROLL - for (int s = 0; s < Stages; ++s) { - ProducerBar::init(&producer_bar[s], 1); - ConsumerBar::init(&consumer_bar[s], kClusterSize * 4); - } - cutlass::arch::fence_view_async_shared(); - if constexpr (kClusterSize > 1) { - cutlass::arch::fence_barrier_init(); - } - PRAGMA_UNROLL - for (int i = 0; i < WARPGORUPS; ++i) { - storage.pipe_count[i] = 0; - } - } - - (kClusterSize > 1) ? cute::cluster_sync() : __syncthreads(); - - const int warpgroup_id = cutlass::canonical_warp_group_idx(); - - if (warpgroup_id == WARPGORUPS) { - cutlass::arch::warpgroup_reg_dealloc<40>(); - - static_assert(TILE_M % kMulticastA == 0); - static_assert(TILE_N % kMulticastB == 0); - - if (threadIdx.x == WARPGORUPS * WARPGROUP_SIZE) { - - Cluster cluster(cute::block_id_in_cluster().x); - - const int mc_offset_m = cluster.cta_n() * (TILE_M / kMulticastA); - const int mc_offset_n = cluster.cta_m() * (TILE_N / kMulticastB); - - auto smem_A = storage.source.A.data() + mc_offset_m * TILE_K; - auto smem_B = storage.source.B.data() + mc_offset_n * TILE_K; - auto& smem_U = storage.source.U; - - sched.grid_init(); - - cutlass::PipelineState write_state{0, 1, 0}; - - while (sched.next()) { - auto [valid_cta_tile_p, cluster_tile_p] = sched.is_valid_tile(); - - if (!cluster_tile_p) { - // OOB tile caused by swizzle pattern - continue; - } - - const auto tile_offset = sched.tile_offset(); - const auto [iter_k_beg, iter_k_end] = sched.iter_k_range(); - - const int offset_k = iter_k_beg * TILE_K; - - const uint16_t mask_A = cluster.mask_m(); - const uint16_t mask_B = cluster.mask_n(); - - const int offset_m = tile_offset.x * TILE_M; - const int offset_n = tile_offset.y * TILE_N; - - int k_iter = iter_k_end - iter_k_beg; - - GmemIteratorSm90 gmem_A{&tm_a, {offset_k, offset_m + mc_offset_m}, {TILE_K, 0}}; - GmemIteratorSm90 gmem_B{&tm_b, {offset_k, offset_n + mc_offset_n}, {TILE_K, 0}}; - - // column-major - GmemIteratorSm90 gmem_U{&tm_u, {offset_m + mc_offset_m, offset_k / 128}, {0, 1}}; - - for (; k_iter > 0; --k_iter) { - int pipe = write_state.index(); - ConsumerBar::wait(&consumer_bar[pipe], write_state.phase()); - ProducerBar::arrive_and_expect_tx(&producer_bar[pipe], kTmaTxBytes); - gmem_A.Load(&producer_bar[pipe], &smem_A[pipe * TILE_M * TILE_K], mask_A); - gmem_B.Load(&producer_bar[pipe], &smem_B[pipe * TILE_N * TILE_K], mask_B); - gmem_U.Load(&producer_bar[pipe], &smem_U[pipe][0] + mc_offset_m, mask_A); - ++write_state; - } - } - } - } - else { - cutlass::arch::warpgroup_reg_alloc<232>(); - - sched.grid_init(WARPGORUPS); - - auto& smem_A = storage.source.A; - auto& smem_B = storage.source.B; - auto& smem_U = storage.source.U; - - auto smem_desc_A = make_smem_desc(&smem_A, 1); - auto smem_desc_B = make_smem_desc(&smem_B, 1); - - SmemDescIterV2> 4)> smem_iter_A{smem_desc_A}; - SmemDescIterV2> 4)> smem_iter_B{smem_desc_B}; - - constexpr int kStepMA = (sizeof(Ta) * MMA_ATOM_M * TILE_K) >> 4; - constexpr int kStepNB = (sizeof(Tb) * MMA_ATOM_N * TILE_K) >> 4; - constexpr int kStepKA = (sizeof(Ta) * MMA_ATOM_K) >> 4; - constexpr int kStepKB = (sizeof(Tb) * MMA_ATOM_K) >> 4; - - auto math_barrier_sync = [&](int phase, int alive = 1) { - constexpr int base = (int)cutlass::arch::ReservedNamedBarriers::FirstUserBarrier; - constexpr int threads = WARPGORUPS * WARPGROUP_SIZE; - int res; - asm volatile("{\n" - " .reg.pred p;\n" - " setp.ne.b32 p, %3, 0;\n" - " barrier.cta.red.or.pred p, %1, %2, p;\n" - " selp.s32 %0, 1, 0, p;\n" - "}\n" - : "=r"(res) - : "r"(base + warpgroup_id ^ phase), "r"(threads), "r"(alive)); - return res; - }; - - cutlass::arch::NamedBarrier wg_barrier(WARPGROUP_SIZE, warpgroup_id + 2); // 2,3 - - sched.next(warpgroup_id); - - if (warpgroup_id == 1) { - math_barrier_sync(1); - } - - while (sched.next(WARPGORUPS)) { - auto [cta_tile_p, cluster_tile_p] = sched.is_valid_tile(); - - if (!cluster_tile_p) { - // OOB tile caused by swizzle pattern - continue; - } - - MMA_Atom::CRegisters frag_C[MMA_ITER_M][MMA_ITER_N]; - MMA_Atom::CRegisters accum_C[MMA_ITER_M][MMA_ITER_N]{}; - - const auto tile_offset = sched.tile_offset(); - const auto [iter_k_beg, iter_k_end] = sched.iter_k_range(); - - const auto [M, N, K, L] = sched.gemm_shape(); - - const int offset_m = tile_offset.x * TILE_M; - const int offset_n = tile_offset.y * TILE_N; - const int offset_k = 0; - - int k_iter = iter_k_end - iter_k_beg; - - const int warp_id = threadIdx.x / WARP_SIZE; - const int lane_id = threadIdx.x % WARP_SIZE; - - const int wg_lane = threadIdx.x % WARPGROUP_SIZE; - - cutlass::PipelineState pipe_state{}; - - auto consumer_arrive = [&] { - __syncwarp(); - if constexpr (kClusterSize > 1) { - ConsumerBar::arrive(&consumer_bar[pipe_state.index()], lane_id, lane_id < kClusterSize); - } - else { - if (lane_id == 0) { - ConsumerBar::arrive(&consumer_bar[pipe_state.index()]); - } - } - }; - - if constexpr (kClusterSize > 1) { - if (!cta_tile_p) { // other CTAs in the cluster are still alive - math_barrier_sync(0); - pipe_state.advance(storage.pipe_count[warpgroup_id ^ 1]); - for (; k_iter > 0; --k_iter) { - ProducerBar::wait(&producer_bar[pipe_state.index()], pipe_state.phase()); - consumer_arrive(); - ++pipe_state; - } - if (wg_lane == 0) { - storage.pipe_count[warpgroup_id] = pipe_state.count(); - } - math_barrier_sync(1); - continue; - } - } - - auto Copy = [k = cdiv(K, 128)](Tv* dst, const Tv* src) { - for (int i = threadIdx.x % WARPGROUP_SIZE; i < k; i += WARPGROUP_SIZE) { - dst[i] = __ldg(&src[i]); - } - }; - auto gmem_V = (const Tv*)V_ + (offset_n / 128) * ldV + (offset_k / 128); - Copy(storage.source.V[0][warpgroup_id], gmem_V); - - uint32_t pred_V{}; - int iter_V{}; - - constexpr int OUTER_N = std::gcd(MMA_ATOM_N, 128); - if constexpr (OUTER_N != 128) { - - static_assert(MMA_ATOM_N <= 128 + OUTER_N, "MMA inst is crossing more than 2 scale blocks"); - - constexpr uint32_t mask = (1UL << (TILE_M / OUTER_N)) - 1; - - int phase = 128 - offset_n % 128; - pred_V = (mask << (phase / OUTER_N)) & mask; - - if (pred_V && offset_n / 128 + 1 < cdiv(N, 128)) { - Copy(storage.source.V[1][warpgroup_id], gmem_V + ldV); - } - - // if constexpr (kWorkGroupN > 1) { - // constexpr int tiles = MMA_ATOM_N / OUTER_N; - // pred_V = (pred_V >> (warp_group_id_n * tiles)) & ((1 << tiles) - 1); - // } - } - - float scale_V[2]; - auto Load_V = [&] { - scale_V[0] = storage.source.V[0][warpgroup_id][iter_V]; - if (pred_V) { - scale_V[1] = storage.source.V[1][warpgroup_id][iter_V]; - } - ++iter_V; - }; - - float scale_U[MMA_ITER_M][2]; - const int offset_U = warp_id % 4 * 16 + lane_id / 4; - auto Load_U = [&] { - for (int m = 0; m < MMA_ITER_M; ++m) { - scale_U[m][0] = smem_U[pipe_state.index()][offset_U + m * MMA_ATOM_M]; - scale_U[m][1] = smem_U[pipe_state.index()][offset_U + m * MMA_ATOM_M + 8]; - } - }; - - auto scale_accum = [&](int m) { // cta_n = mma_iter_n * wg_n * mma_atom_n - float scales[2][2]; - scales[0][0] = scale_U[m][0] * scale_V[0]; - scales[1][0] = scale_U[m][1] * scale_V[0]; - scales[0][1] = scale_U[m][0] * scale_V[1]; - scales[1][1] = scale_U[m][1] * scale_V[1]; - PRAGMA_UNROLL - for (int n = 0; n < MMA_ITER_N; ++n) { - PRAGMA_UNROLL - for (int c0 = 0; c0 < MMA_ATOM_N; c0 += OUTER_N) { - bool pred = (pred_V & (1U << (c0 / OUTER_N))); - PRAGMA_UNROLL - for (int cc = 0; cc < OUTER_N; cc += 8) { - int c = c0 + cc; - // clang-format off - accum_C[m][n][c / 2 + 0] += (pred ? scales[0][1] : scales[0][0]) * frag_C[m][n][c / 2 + 0]; - accum_C[m][n][c / 2 + 1] += (pred ? scales[0][1] : scales[0][0]) * frag_C[m][n][c / 2 + 1]; - accum_C[m][n][c / 2 + 2] += (pred ? scales[1][1] : scales[1][0]) * frag_C[m][n][c / 2 + 2]; - accum_C[m][n][c / 2 + 3] += (pred ? scales[1][1] : scales[1][0]) * frag_C[m][n][c / 2 + 3]; - // clang-format on - } - } - } - - }; - - auto gmma = [&](int m) { - PRAGMA_UNROLL - for (int k = 0; k < MMA_ITER_K; ++k) { - PRAGMA_UNROLL - for (int n = 0; n < MMA_ITER_N; ++n) { - wgmma(smem_iter_A, smem_iter_B, frag_C[m][n], k == 0); - smem_iter_B += kStepNB; - } - smem_iter_B -= MMA_ITER_N * kStepNB; - smem_iter_A += kStepKA; - smem_iter_B += kStepKB; - } - smem_iter_A -= MMA_ITER_K * kStepKA; - smem_iter_B -= MMA_ITER_K * kStepKB; - smem_iter_A += kStepMA; - cute::warpgroup_commit_batch(); - }; - - static_assert(MMA_ITER_N == 1); - - math_barrier_sync(0); - - pipe_state.advance(storage.pipe_count[warpgroup_id ^ 1]); - - smem_iter_A.Reset(pipe_state.index()); - smem_iter_B.Reset(pipe_state.index()); - Load_V(); - ProducerBar::wait(&producer_bar[pipe_state.index()], pipe_state.phase()); - Load_U(); - cute::warpgroup_arrive(); - gmma(0); - gmma(1); - cute::warpgroup_wait<1>(); - scale_accum(0); - cute::warpgroup_wait<0>(); - scale_accum(1); - consumer_arrive(); - ++pipe_state; - --k_iter; - - Load_V(); - ProducerBar::wait(&producer_bar[pipe_state.index()], pipe_state.phase()); - Load_U(); - smem_iter_A.Reset(pipe_state.index()); - smem_iter_B.Reset(pipe_state.index()); - - for (; k_iter > 1; --k_iter) { - cute::warpgroup_arrive(); - gmma(0); - gmma(1); - cute::warpgroup_wait<1>(); - scale_accum(0); - cute::warpgroup_wait<0>(); - scale_accum(1); - consumer_arrive(); - ++pipe_state; - Load_V(); - ProducerBar::wait(&producer_bar[pipe_state.index()], pipe_state.phase()); - Load_U(); - smem_iter_A.Reset(pipe_state.index()); - smem_iter_B.Reset(pipe_state.index()); - } - - cute::warpgroup_arrive(); - gmma(0); - gmma(1); - cute::warpgroup_wait<1>(); - scale_accum(0); - cute::warpgroup_wait<0>(); - scale_accum(1); - consumer_arrive(); - ++pipe_state; - - if (wg_lane == 0) { - storage.pipe_count[warpgroup_id] = pipe_state.count(); - } - - math_barrier_sync(1); - - // epilogue - PRAGMA_UNROLL - for (int m = 0; m < MMA_ITER_M; ++m) { - PRAGMA_UNROLL - for (int n = 0; n < MMA_ITER_N; ++n) { - - constexpr int N = LayoutC::C0; - constexpr int SW_bits = log2(kSwizzleC / 16); - - static_assert(!SW_bits || MMA_ATOM_N % LayoutC::C0 == 0); - - const int m0 = m * MMA_ATOM_M; - const int n0 = n * MMA_ATOM_N; - - PRAGMA_UNROLL - for (int i = 0; i < MMA_ATOM_N; i += 16) { - __align__(16) Array tvec = cast(*(Array*)&accum_C[m][n][i / 2]); - - int mm = m0 + warp_id % 4 * 16 + (lane_id & 8); - int nn = n0 + i / N * N; - - int addr = ((nn / N) * TILE_M * N) + (mm * N) + (nn % N); - - int s = lane_id % 8; - int c = (lane_id & 16) / 2 + i % N; - - addr += Swizzle::apply(s * N + c); - - auto& uvec = (Array&)tvec; - cute::SM90_U32x4_STSM_N::copy( - uvec[0], uvec[1], uvec[2], uvec[3], (cutlass::uint128_t&)storage.C[addr]); - } - } - } - - cute::tma_store_fence(); // visibility: smem -> async proxy - - wg_barrier.sync(); - - const int wg_thread_id = threadIdx.x % WARPGROUP_SIZE; - - if (wg_thread_id < LayoutC::C1) { - const int tma_n = wg_thread_id * LayoutC::C0; - cute::SM90_TMA_STORE::copy( - &tm_c, &storage.C[wg_thread_id * TILE_M * LayoutC::C0], offset_n + tma_n, offset_m); - cute::tma_store_arrive(); - cute::tma_store_wait<0>(); - } - - wg_barrier.sync(); - - } // scheduler loop - - if (warpgroup_id == 0) { - math_barrier_sync(0, 0); - while (math_barrier_sync(1, 0)) { - math_barrier_sync(0, 0); - } - } - else { - while (math_barrier_sync(0, 0)) { - math_barrier_sync(1, 0); - } - } - - if (threadIdx.x % WARPGROUP_SIZE < LayoutC::C1) { - cute::tma_store_wait<0>(); - } - } - - if constexpr (kClusterSize > 1) { - cute::cluster_arrive(); - cute::cluster_wait(); - } - - } // operator() -}; - -extern __shared__ char smem_buf[]; - -template -__global__ void __launch_bounds__(Kernel::CTA_SIZE, 1) gemm_kernel_sm90(const __grid_constant__ CUtensorMap tm_a, - const __grid_constant__ CUtensorMap tm_b, - const __grid_constant__ CUtensorMap tm_c, - const __grid_constant__ CUtensorMap tm_u, - const __grid_constant__ CUtensorMap tm_v, - const void* U_, - int ldU, - const void* V_, - int ldV, - typename Kernel::Scheduler sched) -{ -#if __CUDA_ARCH__ - if constexpr (Kernel::Arch::is_compatible(__CUDA_ARCH__)) { - Kernel kernel; - kernel(tm_a, tm_b, tm_c, tm_u, tm_v, U_, ldU, V_, ldV, sched, smem_buf); - } -#endif -} - -} // namespace turbomind::gemm diff --git a/src/turbomind/kernels/gemm/gemm_universal_sm90_v3.h b/src/turbomind/kernels/gemm/gemm_universal_sm90_v3.h index ac5144512b..48f08be075 100644 --- a/src/turbomind/kernels/gemm/gemm_universal_sm90_v3.h +++ b/src/turbomind/kernels/gemm/gemm_universal_sm90_v3.h @@ -1,16 +1,19 @@ #pragma once #include +#include #include #include #include #include "cute/arch/cluster_sm90.hpp" +#include "cute/arch/copy_sm80.hpp" #include "cute/arch/copy_sm90.hpp" #include "cute/arch/copy_sm90_desc.hpp" #include "cute/arch/copy_sm90_tma.hpp" #include "cute/arch/mma_sm90_desc.hpp" +#include "cute/tensor.hpp" #include "cutlass/arch/barrier.h" #include "cutlass/arch/reg_reconfig.h" @@ -21,8 +24,6 @@ #include "src/turbomind/kernels/core/array_ops.h" #include "src/turbomind/kernels/core/common.h" -#include "src/turbomind/kernels/core/layout.h" - #include "src/turbomind/kernels/core/smem.h" #include "src/turbomind/kernels/gemm/arch.h" @@ -33,49 +34,73 @@ #include "src/turbomind/kernels/gemm/types.h" #include "src/turbomind/kernels/gemm/utils.h" +#include "src/turbomind/kernels/gemm/gmma_bf16_sm90.h" +#include "src/turbomind/kernels/gemm/prepare_moe_tma_descs_sm90_fp8.h" #include "src/turbomind/kernels/gemm/scaled_gmma_fp8_sm90.h" #include "src/turbomind/kernels/gemm/sm90_utils.h" +#include "src/turbomind/kernels/gemm/sm90_v3_traits.h" namespace turbomind::gemm { -template +template struct GemmUniversalSm90_v3 { static constexpr bool kDebug = false; using Arch = Sm90; + using Tile = Tile_; + + static constexpr bool kSupportsFusedSilu = kSupportsFusedSilu_; + + static constexpr int TILE_M = Tile::TILE_M; + static constexpr int TILE_N = Tile::TILE_N; + static constexpr int TILE_K = Tile::TILE_K; - static constexpr int TILE_M = 128; - static constexpr int TILE_N = 192; - static constexpr int TILE_K = 128; + // Fused SiLU pairs OP_N=128 gate/up atoms inside TILE_N=256. + static_assert(!kSupportsFusedSilu || (TILE_N == 256 && Tile::kMaxOpN == 128)); - static constexpr int WG_M = 2; - static constexpr int WG_N = 1; + static constexpr int WG_M = Tile::WG_M; + static constexpr int WG_N = Tile::WG_N; static constexpr int WG_TILE_M = TILE_M / WG_M; static constexpr int WG_TILE_N = TILE_N / WG_N; + static_assert(TILE_M % WG_M == 0); + static_assert(TILE_N % WG_N == 0); + static_assert(WG_TILE_M % 64 == 0); // WGMMA atom M static constexpr int kSchedWarpGroups = 1; - static constexpr int WARPGORUPS = WG_M * WG_N; + static constexpr int WARPGROUPS = WG_M * WG_N; - using GMMA = ScaledGmmaFP8_TN; + static constexpr Order kRasterOrder = raster_order; + static constexpr int kAlgoFamily = 1; + + using GMMA = ScaledGmmaFP8_TN; + using AccumC = typename GMMA::AccumC; + using FragC = typename GMMA::FragC; + using FragU = typename GMMA::FragU; static constexpr int kMulticastA = multicast_a; static constexpr int kMulticastB = multicast_b; static constexpr int kClusterSize = kMulticastA * kMulticastB; - static constexpr int Stages = 4; + static constexpr int Stages = Tile::Stages; static constexpr bool kSplitK = false; static constexpr int kChunkSizeK = TILE_K; static constexpr int WARPGROUP_SIZE = 128; - static constexpr int kMathGroupSize = WARPGROUP_SIZE * WARPGORUPS; + static constexpr int kMathGroupSize = WARPGROUP_SIZE * WARPGROUPS; - static constexpr int CTA_SIZE = WARPGROUP_SIZE * (WARPGORUPS + 1); + static constexpr int CTA_SIZE = WARPGROUP_SIZE * (WARPGROUPS + 1); using Ta = __nv_fp8_e4m3; using Tb = __nv_fp8_e4m3; @@ -83,11 +108,19 @@ struct GemmUniversalSm90_v3 { using Tu = float; using Tv = float; + using Tw = float; // dynamic output group scales (fused path) using Cluster = arch::Cluster; static constexpr auto is_grouped_gemm = is_grouped_gemm_; + static constexpr Striding kStridingA = kStridingA_; + static constexpr Striding kStridingB = is_grouped_gemm_ ? Striding::kBlocked : Striding::kFlat; + static constexpr Striding kStridingC = is_grouped_gemm_ ? Striding::kBlocked : Striding::kFlat; + + // Indexed gather: A/U via cp.async; TMA only B (+ C store descs). + static constexpr bool kIndexedGather = (kStridingA_ == Striding::kIndexed); + using Scheduler = TileScheduler; static constexpr int kMulticastU = is_grouped_gemm ? 1 : kMulticastA; @@ -101,32 +134,122 @@ struct GemmUniversalSm90_v3 { // Alignment requirement for SMEM addr. This forbids multicast factor 8. static_assert(kMulticastU == 1 || sizeof(Tu) * kBoxU / kMulticastU % 128 == 0); - static constexpr int kTmaTxBytes = - sizeof(Ta) * (TILE_M * TILE_K) + sizeof(Tb) * (TILE_N * TILE_K) + sizeof(Tu) * kBoxU; + static constexpr int kTmaTxBytesWeight = (int)sizeof(Tb) * (TILE_N * TILE_K); + static constexpr int kTmaTxBytesAct = (int)sizeof(Ta) * (TILE_M * TILE_K); + static constexpr int kTmaTxBytesU = (int)sizeof(Tu) * kBoxU; + // Dense / blocked-A: expect_tx(A+B+U). Indexed-A: expect_tx(B only); A/U via cp.async noinc. + static constexpr int kTmaTxBytes = kTmaTxBytesWeight + (kIndexedGather ? 0 : (kTmaTxBytesAct + kTmaTxBytesU)); + + // Dense: unused. Grouped indexed: [B, C]. Grouped blocked: [A, B, U, C]. + static constexpr int kTmaDescNum = !is_grouped_gemm_ ? 1 : (kIndexedGather ? 2 : 4); + static constexpr int kCdescIdx = kIndexedGather ? 1 : 3; + + // SW128 K-major SoT for indexed gather store TV (matches TMA/GMMA layout_type=1). + using SmemLayoutA_2D = decltype(cute::tile_to_shape(cute::SM90::GMMA::Layout_K_SW128_Atom{}, + cute::make_shape(cute::Int{}, cute::Int{}), + cute::Step{})); + + // setmaxnreg: each WG ≤ 256, multiples of 8. Budgets come from Tile + // (TMA vs indexed). 2 math WGs pack to 504; 1 math WG packs to ≤512. + static constexpr int kProducerRegs = kIndexedGather ? Tile::kProducerRegsIndexed : Tile::kProducerRegsTma; + static constexpr int kMathRegs = kIndexedGather ? Tile::kMathRegsIndexed : Tile::kMathRegsTma; + static_assert(kProducerRegs >= 24 && kProducerRegs % 8 == 0); + static_assert(kMathRegs >= 24 && kMathRegs % 8 == 0 && kMathRegs <= 256); + static_assert(WARPGROUPS == 1 || WARPGROUPS == 2); + static_assert(WARPGROUPS != 2 || kProducerRegs + 2 * kMathRegs == 504); + static_assert(WARPGROUPS != 1 || kProducerRegs + kMathRegs <= 512); // ! SMEM addr must be SBO aligned for TMA load/store struct SharedStorage { __align__(1024) Array A; __align__(1024) Array B; - __align__(1024) Array C; __align__(128) Tu U[Stages][round_up(kBoxU, 128)]; // at least 128 byte alignment __align__(128) Tv V[Stages][2]; - __align__(128) CUtensorMap tensor_map[5]; __align__(8) uint64_t producer_bar[Stages]; __align__(8) uint64_t consumer_bar[Stages]; typename Scheduler::Storage sched; + int gather_alive; + int gather_k_iters; + int gather_m0; + int gather_M_group; + int gather_offset_m; }; - static constexpr int kSmemSize = sizeof(SharedStorage); + template + struct Output { + static_assert(!kFuseSilu || kSupportsFusedSilu); + + using Tc = std::conditional_t; + using ElementC = std::conditional_t; + + static constexpr int kStoreN = kFuseSilu ? TILE_N / 2 : TILE_N; + static constexpr int kWgStoreN = kFuseSilu ? WG_TILE_N / 2 : WG_TILE_N; + static constexpr int kEpiN = kFuseSilu ? 64 : 32; + + using SmemLayoutAtomC = + decltype(gmma_ss_smem_selector, cute::Int>()); + using SmemLayoutC = + decltype(cute::tile_to_shape(SmemLayoutAtomC{}, + cute::make_shape(cute::Int{}, cute::Int{}), + cute::Step{})); - static constexpr int kSwizzleC = 2 * std::gcd(WG_TILE_N, 128 / sizeof(Tc)); + struct LayoutC { + static constexpr int S0 = WG_TILE_M; + static constexpr int C0 = kEpiN; + static constexpr int C1 = WG_TILE_N / C0; + static constexpr int C1_store = kWgStoreN / C0; + }; - using LayoutC = std::conditional_t= 32, - SmemLayoutV2, - SmemLayoutV2>; + static_assert(WG_TILE_N % LayoutC::C0 == 0); + static_assert(kWgStoreN % LayoutC::C0 == 0); + static_assert(decltype(cute::size<1>(SmemLayoutAtomC{}))::value == LayoutC::C0); - static constexpr int OUTER_N = GMMA::OUTER_N; - static constexpr int MMA_SUBTILE_N = GMMA::OP_N / OUTER_N; + static constexpr int kSwizzleC = LayoutC::C0 * (int)sizeof(Tc); + }; + + static constexpr int kOutputOffset = round_up(sizeof(SharedStorage), 1024); + + static constexpr int GetSmemSize(bool fuse_silu) + { + return kOutputOffset + + (fuse_silu ? TILE_M * (TILE_N / 2) * (int)sizeof(__nv_fp8_e4m3) : + TILE_M * TILE_N * (int)sizeof(nv_bfloat16)); + } + + static constexpr int kSmemSize = GetSmemSize(false); + + // Epi-only single-atom TiledMma: TV map for make_tiled_copy_C / STSM_N (mainloop unchanged). + using EpiTiledMma = decltype( + cute::make_tiled_mma(typename GMMA::Operation{}, cute::Layout>{})); + + static constexpr int OUTER_N = GMMA::OUTER_N; + // V-scale predicates span the full WG tile N (not just one MMA atom). + static constexpr int MMA_SUBTILE_N = WG_TILE_N / OUTER_N; + + // Host: rebase per-expert TMA maps into workspace before GEMM launch. + static int* PrepareTmaDescs(const CUtensorMap& tm_a, + const CUtensorMap& tm_b, + const CUtensorMap& tm_u, + const CUtensorMap& tm_c, + const MatrixParam& param_A, + const MatrixParam& param_B, + const MatrixParam& param_U, + const MatrixParam& param_C, + bool fuse_silu, + CUtensorMap* out, + int num_groups, + int M, + int N, + cudaStream_t stream) + { + if constexpr (!is_grouped_gemm_) { + return nullptr; + } + int* offsets = reinterpret_cast(out + num_groups * kTmaDescNum); + prepare_moe_tma_descs_sm90_fp8<<>>( + tm_a, tm_b, tm_u, tm_c, param_A, param_B, param_U, param_C, fuse_silu, out, offsets, M, N); + return offsets; + } __device__ void operator()(const CUtensorMap& tm_a, const CUtensorMap& tm_b, @@ -138,6 +261,8 @@ struct GemmUniversalSm90_v3 { const MatrixParam& param_U, const MatrixParam& param_V, const MatrixParam& param_C, + const MatrixParam& param_W, + bool fuse_silu, Scheduler sched, CUtensorMap* tensormap_buf, char* smem_buf) @@ -147,13 +272,15 @@ struct GemmUniversalSm90_v3 { uint64_t* producer_bar = storage.producer_bar; uint64_t* consumer_bar = storage.consumer_bar; + constexpr int kProducerBarInit = kIndexedGather ? (1 + WARPGROUP_SIZE) : (1 + 1); + if (threadIdx.x == 0) { PRAGMA_UNROLL for (int s = 0; s < Stages; ++s) { - ProducerBar::init(&producer_bar[s], 1 + 1); - ConsumerBar::init(&consumer_bar[s], WARPGORUPS * kClusterSize * 4); + ProducerBar::init(&producer_bar[s], kProducerBarInit); + ConsumerBar::init(&consumer_bar[s], WARPGROUPS * kClusterSize * 4); } - sched.init_dyanmic(storage.sched, kClusterSize * (WARPGORUPS * 4 + 1)); + sched.init_dyanmic(storage.sched, kClusterSize * (WARPGROUPS * 4 + 1)); cutlass::arch::fence_view_async_shared(); if constexpr (kClusterSize > 1) { cutlass::arch::fence_barrier_init(); @@ -164,19 +291,242 @@ struct GemmUniversalSm90_v3 { const int wg_idx = cutlass::canonical_warp_group_idx(); - if (wg_idx == WARPGORUPS) { - cutlass::arch::warpgroup_reg_dealloc<40>(); + if (wg_idx == WARPGROUPS) { + cutlass::arch::warpgroup_reg_dealloc(); static_assert(TILE_M % kMulticastA == 0); static_assert(TILE_N % kMulticastB == 0); cutlass::arch::NamedBarrier producers_bar(WARP_SIZE * 2, 7); - const int warp_id = cutlass::canonical_warp_idx_sync(); - const bool cta_0 = cute::block_id_in_cluster().x == 0; + const int warp_id = cutlass::canonical_warp_idx_sync(); + const int warp_in_wg = warp_id % 4; + const bool cta_0 = cute::block_id_in_cluster().x == 0; + + if constexpr (kIndexedGather) { + // Full producer WG gather. Scheduler folded onto warp0. + cutlass::arch::NamedBarrier gather_bar( + /*num_threads=*/WARPGROUP_SIZE, cutlass::arch::ReservedNamedBarriers::FirstUserBarrier); + + Cluster cluster(cute::block_id_in_cluster().x); + + const int mc_offset_n = cluster.cta_m() * (TILE_N / kMulticastB); + + auto* smem_A = storage.A.data(); + auto* smem_B = storage.B.data() + mc_offset_n * TILE_K; + auto& smem_U = storage.U; + auto& smem_V = storage.V; + + cutlass::PipelineState write_state{0, 1, 0}; + + typename Scheduler::ConsumerState sched_state = sched.init_consumer(storage.sched); + typename Scheduler::ProducerState prod_state = sched.init_producer(storage.sched); + int lane_predicate = 0; + const int lane_id = threadIdx.x % WARP_SIZE; + const int prod_tid = threadIdx.x - WARPGROUPS * WARPGROUP_SIZE; + + if (warp_in_wg == 0) { + lane_predicate = cute::elect_one_sync(); + } + + const Ta* act_gmem = (const Ta*)param_A.ptr; + const int ldA = param_A.stride; + const int* idxs = param_A.idxs; + const Tu* u_gmem = (const Tu*)param_U.ptr; + const int ldU = param_U.stride; + const int K = sched.gemm_shape().z; + + constexpr int kVec = 16; // uint4 / sizeof(fp8) + constexpr int nvec = TILE_M * (TILE_K / kVec); + constexpr int kSlots = nvec / WARPGROUP_SIZE; + static_assert(nvec % WARPGROUP_SIZE == 0); + // U: one float per producer thread along TILE_M (thread m loads row m). + static_assert(TILE_M <= WARPGROUP_SIZE); - if (warp_id % 4 == 0) { + typename Scheduler::Tile* tile; + + while (true) { + const CUtensorMap* Bdesc = &tm_b; + uint16_t mask_B = 0; + int coord_n = 0; + int k_iters = 0; + const Tv* gmem_V0 = (const Tv*)param_V.ptr; + const Tv* gmem_V1 = nullptr; + int ldV = param_V.stride; + + if (warp_in_wg == 0 && cta_0) { + (void)prod_state.next(); + } + + if (warp_in_wg == 0) { + const bool alive = sched_state.acquire(tile); + int m0 = 0, M_group = 0, offset_m = 0; + + if (alive && tile->is_valid_cluster) { + if constexpr (is_grouped_gemm) { + const int g = tile->group_idx; + Bdesc = &tensormap_buf[g * kTmaDescNum]; + const auto v = resolve(param_V, g); + gmem_V0 = (const Tv*)v.ptr.ptr; + ldV = v.ptr.stride; + } + + mask_B = cluster.mask_n(); + coord_n = tile->offset_n + mc_offset_n; + offset_m = tile->offset_m; + m0 = [&] { + if constexpr (is_grouped_gemm) { + return tile->m0; + } + return 0; + }(); + M_group = [&] { + if constexpr (is_grouped_gemm) { + return tile->m1 - tile->m0; + } + return sched.gemm_shape().x; + }(); + k_iters = sched.k_iters_; + + gmem_V0 += (tile->offset_n / 128) * ldV; + gmem_V1 = gmem_V0; + if (tile->offset_n / 128 + 1 < cdiv(sched.gemm_shape().y, 128)) { + gmem_V1 += ldV; + } + } + + if (lane_id == 0) { + storage.gather_alive = alive ? 1 : 0; + storage.gather_k_iters = k_iters; + storage.gather_m0 = m0; + storage.gather_M_group = M_group; + storage.gather_offset_m = offset_m; + } + __syncwarp(); + } + + // Tile header only. + gather_bar.arrive_and_wait(); + + if (storage.gather_alive == 0) { + break; + } + + k_iters = storage.gather_k_iters; + const int m0 = storage.gather_m0; + const int M_group = storage.gather_M_group; + const int offset_m = storage.gather_offset_m; + int coord_k = 0; + + // iterator_sm80 style: idxs → src bases, then += TILE_K each K tile. + const Ta* src_data_vec_[kSlots]; + int m_own[kSlots]; + int kk_own[kSlots]; + bool pred_row[kSlots]; + const Tu* src_u_base; + int m_u; + bool pred_u; + const int u_pad = m0 % kAlignmentU; + + PRAGMA_UNROLL + for (int t = 0; t < kSlots; ++t) { + const int i = prod_tid + t * WARPGROUP_SIZE; + const int m = i / (TILE_K / kVec); + const int kk = (i % (TILE_K / kVec)) * kVec; + const int packed = m0 + offset_m + m; + const bool row_ok = (offset_m + m) < M_group; + const int token = (idxs && row_ok) ? __ldg(idxs + packed) : packed; + m_own[t] = m; + kk_own[t] = kk; + pred_row[t] = row_ok; + src_data_vec_[t] = act_gmem + (int64_t)token * ldA + kk; + } + + { + // TILE_M < WARPGROUP_SIZE: idle threads must skip U ZFILL entirely + // (pred=false still zeros dst — see sm90_bf16). + const int m = prod_tid; + const bool in_m = m < TILE_M; + const int packed = m0 + offset_m + m; + const bool row_ok = in_m && (offset_m + m) < M_group; + const int token = (idxs && row_ok) ? __ldg(idxs + packed) : packed; + m_u = m; + pred_u = row_ok; + src_u_base = u_gmem + token; + } + + // Warp0-only weight TMA + V; full WG gathers A/U. + GmemIteratorSm90 gmem_B{ + (warp_in_wg == 0) ? Bdesc : &tm_b, {0, (warp_in_wg == 0) ? coord_n : 0}, {TILE_K, 0}}; + + for (; k_iters > 0; --k_iters) { + const int pipe = write_state.index(); + ConsumerBar::wait(&consumer_bar[pipe], write_state.phase()); + + if (warp_in_wg == 0 && lane_predicate) { + ProducerBar::arrive_and_expect_tx(&producer_bar[pipe], kTmaTxBytes); + gmem_B.Step(&producer_bar[pipe], &smem_B[pipe * TILE_N * TILE_K], mask_B); + uint32_t uint_ptr_V = cast_smem_ptr_to_uint(smem_V[pipe]); + CP_ASYNC::apply(uint_ptr_V, gmem_V0, true); + CP_ASYNC::apply(uint_ptr_V + sizeof(Tv), gmem_V1, true); + ++gmem_V0; + ++gmem_V1; + } + + { + cute::Tensor sA = cute::make_tensor(cute::make_smem_ptr(smem_A + pipe * TILE_M * TILE_K), + SmemLayoutA_2D{}); + + PRAGMA_UNROLL + for (int t = 0; t < kSlots; ++t) { + const bool pred = pred_row[t] && (coord_k + kk_own[t]) < K; + auto* dst = &sA(m_own[t], kk_own[t]); + cute::SM80_CP_ASYNC_CACHEGLOBAL_ZFILL::copy( + *reinterpret_cast(src_data_vec_[t]), + *reinterpret_cast(dst), + pred); + src_data_vec_[t] += TILE_K; + } + + // U: one float per thread along TILE_M; ColMajor (token + k_col*ldU). + // Idle (m >= TILE_M): do not issue ZFILL — pred=false still zeros dst. + if (m_u < TILE_M) { + const bool pred = pred_u && (coord_k < K); + auto* dst = &smem_U[pipe][u_pad + m_u]; + const Tu* src = src_u_base + (coord_k / 128) * ldU; + cute::SM80_CP_ASYNC_CACHEALWAYS_ZFILL::copy( + *reinterpret_cast(src), *reinterpret_cast(dst), pred); + } + + cutlass::arch::cpasync_barrier_arrive_noinc(&producer_bar[pipe]); + } + + ++write_state; + coord_k += TILE_K; + } + + if (warp_in_wg == 0) { + sched_state.release(); + } + } + if (warp_in_wg == 0) { + sched_state.release(); + if (cta_0) { + sched.tail(prod_state); + } + if constexpr (kClusterSize > 1) { + if (lane_predicate) { + for (int i = 0; i < Stages; ++i) { + ConsumerBar::wait(&consumer_bar[write_state.index()], write_state.phase()); + ++write_state; + } + } + __syncwarp(); + } + } + } + else if (warp_in_wg == 0) { Cluster cluster(cute::block_id_in_cluster().x); const int mc_offset_m = cluster.cta_n() * (TILE_M / kMulticastA); @@ -187,10 +537,6 @@ struct GemmUniversalSm90_v3 { auto& smem_U = storage.U; auto& smem_V = storage.V; - if constexpr (is_grouped_gemm) { - init_tma_descs<3>({&tm_a, &tm_b, &tm_u}, storage.tensor_map); - } - cutlass::PipelineState write_state{0, 1, 0}; auto sched_state = sched.init_consumer(storage.sched); @@ -209,37 +555,24 @@ struct GemmUniversalSm90_v3 { const Tv* gmem_V0 = (const Tv*)param_V.ptr; const Tv* gmem_V1; + int ldV = param_V.stride; if constexpr (is_grouped_gemm) { - const int g = tile->group_idx; - const int m0 = tile->m0; - const int m1 = tile->m1; - const int m = m1 - m0; - - Array global_addrs; - global_addrs[0] = (Ta*)param_A.ptr + m0 * (int64_t)param_A.stride; - global_addrs[1] = ((void**)param_B.ptr)[g]; - - const int beg_u = m0 / kAlignmentU * kAlignmentU; - const int end_u = round_up(m1, kAlignmentU); - global_addrs[2] = (Tu*)param_U.ptr + beg_u; - - Array dims; - dims[0] = m; - dims[1] = sched.gemm_shape().y; - dims[2] = end_u - beg_u; - - auto descs = update_tma_descs(tensormap_buf, storage.tensor_map, global_addrs, dims); - Adesc = &descs[0]; - Bdesc = &descs[1]; - Udesc = &descs[2]; - - gmem_V0 = ((Tv**)gmem_V0)[g]; - - PRAGMA_UNROLL - for (int i = 0; i < 3; ++i) { - cute::tma_descriptor_fence_acquire((cute::TmaDescriptor*)&descs[i]); + // Descs published by prepare_moe_tma_descs on this stream; + // fence_acquire only needed after in-kernel tensormap replace. + const int g = tile->group_idx; + CUtensorMap* const descs = tensormap_buf + g * kTmaDescNum; + if constexpr (kStridingA == Striding::kBlocked) { + Adesc = &descs[0]; + Bdesc = &descs[1]; + Udesc = &descs[2]; + } + else { + Bdesc = &descs[0]; } + const auto v = resolve(param_V, g); + gmem_V0 = (const Tv*)v.ptr.ptr; + ldV = v.ptr.stride; } if (lane_predicate) { @@ -263,10 +596,10 @@ struct GemmUniversalSm90_v3 { GmemIteratorSm90 gmem_U{ Udesc, {offset_m + mc_offset_u, offset_k / 128}, {0, 1}}; - gmem_V0 += (offset_n / 128) * param_V.stride + (offset_k / 128); + gmem_V0 += (offset_n / 128) * ldV + (offset_k / 128); gmem_V1 = gmem_V0; if (offset_n / 128 + 1 < cdiv(sched.gemm_shape().y, 128)) { - gmem_V1 += param_V.stride; + gmem_V1 += ldV; } for (; k_iter > 0; --k_iter) { @@ -310,24 +643,20 @@ struct GemmUniversalSm90_v3 { __syncwarp(); } } - else if (warp_id % 4 == 1 && cta_0) { - auto state = sched.init_producer(storage.sched); - while (state.next()) { - if constexpr (Scheduler::is_dynamic) { - producers_bar.arrive_and_wait_unaligned(); + else if (warp_in_wg == 1 && cta_0) { + if constexpr (!kIndexedGather) { + auto state = sched.init_producer(storage.sched); + while (state.next()) { + if constexpr (Scheduler::is_dynamic) { + producers_bar.arrive_and_wait_unaligned(); + } } + sched.tail(state); } - sched.tail(state); } } else { - cutlass::arch::warpgroup_reg_alloc<232>(); - - if constexpr (is_grouped_gemm) { - if (threadIdx.x % WARPGROUP_SIZE / WARP_SIZE == 0) { - init_tma_descs<1>({&tm_c}, storage.tensor_map + 3 + wg_idx); - } - } + cutlass::arch::warpgroup_reg_alloc(); auto& smem_A = storage.A; auto& smem_B = storage.B; @@ -372,8 +701,8 @@ struct GemmUniversalSm90_v3 { while (tile->alive) { if (tile->is_valid_cta) { - GMMA::AccumC accum_C{}; - GMMA::FragC frag_C; + AccumC accum_C{}; + FragC frag_C; auto pred_V = Fetch_V(tile, wg_idx_n); @@ -387,8 +716,8 @@ struct GemmUniversalSm90_v3 { if constexpr (is_grouped_gemm) { offset_U += tile->m0 % kAlignmentU; } - GMMA::FragU frag_U; - auto Load_U = [&] { + FragU frag_U; + auto Load_U = [&] { GMMA::foreach_m(frag_U, [&](auto& U, int m) { U[0] = smem_U[pipe_state.index()][offset_U + m * GMMA::OP_M]; U[1] = smem_U[pipe_state.index()][offset_U + m * GMMA::OP_M + 8]; @@ -400,16 +729,18 @@ struct GemmUniversalSm90_v3 { }; if constexpr (is_grouped_gemm) { - if (threadIdx.x % WARPGROUP_SIZE < LayoutC::C1) { - cute::tma_store_wait<0>(); + auto wait_tma_store = [&](auto fused_silu) { + constexpr bool kFuseSilu = decltype(fused_silu)::value; + using LayoutC = typename Output::LayoutC; + if (threadIdx.x % WARPGROUP_SIZE < LayoutC::C1) { + cute::tma_store_wait<0>(); + } + }; + if constexpr (kSupportsFusedSilu) { + fuse_silu ? wait_tma_store(std::true_type{}) : wait_tma_store(std::false_type{}); } - // No need to sync here as the update is warp synchronized - if (warp_id % 4 == 0) { - int m0 = tile->m0, m1 = tile->m1; - auto global_addr = (Tc*)param_C.ptr + m0 * (int64_t)param_C.stride; - int idx = 3 + wg_idx; - update_tma_descs<1>( - tensormap_buf + idx, storage.tensor_map + idx, {global_addr}, {m1 - m0}); + else { + wait_tma_store(std::false_type{}); } barrier.sync(); } @@ -448,8 +779,18 @@ struct GemmUniversalSm90_v3 { const int thread_idx = threadIdx.x % WARPGROUP_SIZE; if constexpr (!is_grouped_gemm) { - if (thread_idx < LayoutC::C1) { - cute::tma_store_wait<0>(); + auto wait_tma_store = [&](auto fused_silu) { + constexpr bool kFuseSilu = decltype(fused_silu)::value; + using LayoutC = typename Output::LayoutC; + if (thread_idx < LayoutC::C1_store) { + cute::tma_store_wait<0>(); + } + }; + if constexpr (kSupportsFusedSilu) { + fuse_silu ? wait_tma_store(std::true_type{}) : wait_tma_store(std::false_type{}); + } + else { + wait_tma_store(std::false_type{}); } barrier.sync(); } @@ -457,64 +798,178 @@ struct GemmUniversalSm90_v3 { consumer_arrive(); ++pipe_state; - Tc* smem_C = &storage.C[wg_idx_m * WG_TILE_M * TILE_N + wg_idx_n * WG_TILE_N]; + auto run_epilogue = [&](auto fused_silu) { + constexpr bool kFuseSilu = decltype(fused_silu)::value; + using OutputTraits = Output; + using OutputT = typename OutputTraits::Tc; + using ElementC = typename OutputTraits::ElementC; + using LayoutC = typename OutputTraits::LayoutC; + using SmemLayoutC = typename OutputTraits::SmemLayoutC; - // epilogue - GMMA::foreach_C(accum_C, [&](const auto& C, int m, int n) { - constexpr int N = LayoutC::C0; - constexpr int SW_bits = log2(kSwizzleC / 16); + OutputT* output_base = reinterpret_cast(smem_buf + kOutputOffset); + OutputT* smem_C = output_base + wg_idx_m * WG_TILE_M * OutputTraits::kStoreN + + wg_idx_n * OutputTraits::kWgStoreN; - static_assert(!SW_bits || GMMA::OP_N % LayoutC::C0 == 0); + // CUTLASS RowMajor epi: STSM_N into K-major swizzle panels. + // Fused SiLU: silu(gate)*up → per-row amax/quant (gs=128) → FP8 + W scales. static_assert(GMMA::OP_N % 16 == 0); + static_assert(GMMA::OP_N % LayoutC::C0 == 0); + + if constexpr (kFuseSilu) { + static_assert(!kFuseSilu || GMMA::ITER_N == 2); + static_assert(!kFuseSilu || GMMA::OP_N == 128); + constexpr float kQmax = 448.f; + // Gemm instantiates ScaledGmmaFP8 with PIPE/BATCH = 1. + // CRegisters: every 4 floats = [u0,u0,u1,u1] for 8 N-cols; + // each thread owns 2 M-rows (U[0]/U[1]); reduce amax across lane%4. + PRAGMA_UNROLL + for (int i_m = 0; i_m < GMMA::ITER_M; ++i_m) { + auto& gate = accum_C[i_m][0][0][0][0][0]; + auto& up = accum_C[i_m][1][0][0][0][0]; + constexpr int kNumRegs = (int)(sizeof(gate) / sizeof(float)); + static_assert(!kFuseSilu || kNumRegs == 64); + PRAGMA_UNROLL + for (int i = 0; i < kNumRegs; ++i) { + const float g = gate[i]; + const float u = up[i]; + gate[i] = fdividef(g, 1.f + expf(-g)) * u; + } + float amax0 = 0.f; + float amax1 = 0.f; + PRAGMA_UNROLL + for (int i = 0; i < kNumRegs; i += 4) { + amax0 = fmaxf(amax0, fabsf(gate[i + 0])); + amax0 = fmaxf(amax0, fabsf(gate[i + 1])); + amax1 = fmaxf(amax1, fabsf(gate[i + 2])); + amax1 = fmaxf(amax1, fabsf(gate[i + 3])); + } + amax0 = fmaxf(amax0, __shfl_xor_sync(0xffffffffu, amax0, 1)); + amax0 = fmaxf(amax0, __shfl_xor_sync(0xffffffffu, amax0, 2)); + amax1 = fmaxf(amax1, __shfl_xor_sync(0xffffffffu, amax1, 1)); + amax1 = fmaxf(amax1, __shfl_xor_sync(0xffffffffu, amax1, 2)); + amax0 = fmaxf(amax0, 1e-8f); + amax1 = fmaxf(amax1, 1e-8f); + const float scale0 = amax0 / kQmax; + const float scale1 = amax1 / kQmax; + const float inv0 = kQmax / amax0; + const float inv1 = kQmax / amax1; + PRAGMA_UNROLL + for (int i = 0; i < kNumRegs; i += 4) { + gate[i + 0] *= inv0; + gate[i + 1] *= inv0; + gate[i + 2] *= inv1; + gate[i + 3] *= inv1; + } + // W is a flat packed buffer (no per-expert TMA rebase). Global row + // is m0 + tile-local offset (C TMA is rebased; W is not). + if (param_W.ptr && (lane_id % 4) == 0) { + const int n_group = tile->offset_n / TILE_N; + int row0 = tile->offset_m + wg_idx_m * WG_TILE_M + (warp_id % 4) * 16 + lane_id / 4 + + i_m * GMMA::OP_M; + int row_end = 0x7fffffff; + if constexpr (is_grouped_gemm) { + row0 += tile->m0; + row_end = tile->m1; + } + Tw* W = reinterpret_cast(param_W.ptr); + const int ldW = param_W.stride; + if (row0 < row_end) { + W[(int64_t)n_group * ldW + row0] = scale0; + } + if (row0 + 8 < row_end) { + W[(int64_t)n_group * ldW + row0 + 8] = scale1; + } + } + } + } - const int m0 = m * GMMA::OP_M; - const int n0 = n * GMMA::OP_N; - - PRAGMA_UNROLL - for (int i = 0; i < GMMA::OP_N; i += 16) { - __align__(16) Array tvec = cast((Array&)C[i / 2]); - - // fill(tvec, Tc(255)); - - int mm = m0 + warp_id % 4 * 16 + (lane_id & 8); - int nn = n0 + i / N * N; + EpiTiledMma epi_tiled_mma{}; + // bf16: STSM_N. fused fp8: vectorizing R2S (STSM_N + e4m3 fails TV match). + using CopyAtomC = std::conditional_t< + kFuseSilu, + cute::Copy_Atom, ElementC>, + cute::Copy_Atom>; + auto tiled_copy_C = cute::make_tiled_copy_C(CopyAtomC{}, epi_tiled_mma); + auto thr_copy = tiled_copy_C.get_thread_slice(thread_idx); + + cute::Tensor sC = cute::as_position_independent_swizzle_tensor( + cute::make_tensor(cute::make_smem_ptr(reinterpret_cast(smem_C)), SmemLayoutC{})); + + auto tCr_layout = cute::layout(cute::partition_fragment_C( + epi_tiled_mma, cute::make_shape(cute::Int{}, cute::Int{}))); + + GMMA::foreach_C(accum_C, [&](auto& C, int m, int n) { + if constexpr (kFuseSilu) { + if (n != 0) { + return; // up atom consumed; store fused gate only + } + } + cute::Tensor tCr_f32 = + cute::make_tensor(reinterpret_cast(static_cast(&C)), tCr_layout); - int addr = ((nn / N) * WG_TILE_M * N) + (mm * N) + (nn % N); + cute::Tensor tCr_out = cute::make_tensor_like(tCr_f32); + CUTE_UNROLL + for (int i = 0; i < cute::size(tCr_f32); ++i) { + tCr_out(i) = ElementC(tCr_f32(i)); + } - int s = lane_id % 8; - int c = (lane_id & 16) / 2 + i % N; + cute::Tensor sC_atom = + cute::local_tile(sC, + cute::Shape, cute::Int>{}, + cute::make_coord(m, n)); - addr += Swizzle::apply(s * N + c); + cute::Tensor tCsC = thr_copy.partition_D(sC_atom); + cute::Tensor tCrS = thr_copy.retile_S(tCr_out); + cute::copy(tiled_copy_C, tCrS, tCsC); + }); - auto& uvec = (Array&)tvec; - cute::SM90_U32x4_STSM_N::copy(uvec[0], // - uvec[1], - uvec[2], - uvec[3], - (cutlass::uint128_t&)smem_C[addr]); + // AutoVectorizing R2S (fused fp8) needs an async-shared fence before TMA; + // STSM path is coherent with tma_store_fence alone. + if constexpr (kFuseSilu) { + cutlass::arch::fence_view_async_shared(); } - }); - - cute::tma_store_fence(); // visibility: smem -> async proxy + cute::tma_store_fence(); // visibility: smem -> async proxy - barrier.sync(); + barrier.sync(); - const int offset_m = tile->offset_m; - const int offset_n = tile->offset_n; + const int offset_m = tile->offset_m; + const int offset_n = tile->offset_n; - const void* Cdesc = &tm_c; + const void* Cdesc = &tm_c; - if (thread_idx < LayoutC::C1) { - const int tma_n = thread_idx * LayoutC::C0; + if (thread_idx < LayoutC::C1_store) { + const int tma_n = thread_idx * LayoutC::C0; + if constexpr (is_grouped_gemm) { + Cdesc = tensormap_buf + tile->group_idx * kTmaDescNum + kCdescIdx; + } + const int store_n = kFuseSilu ? + (offset_n / 2 + wg_idx_n * OutputTraits::kWgStoreN + tma_n) : + (offset_n + wg_idx_n * OutputTraits::kWgStoreN + tma_n); + cute::SM90_TMA_STORE::copy(Cdesc, + &smem_C[thread_idx * WG_TILE_M * LayoutC::C0], + store_n, + offset_m + wg_idx_m * WG_TILE_M); + cute::tma_store_arrive(); + } + // Grouped path skips the pre-epilogue tma_store_wait; drain before next tile. if constexpr (is_grouped_gemm) { - Cdesc = tensormap_buf + blockIdx.x * 5 + 3 + wg_idx; - cute::tma_descriptor_fence_acquire((cute::TmaDescriptor*)Cdesc); + if (thread_idx < LayoutC::C1_store) { + cute::tma_store_wait<0>(); + } + barrier.sync(); + } + }; + + if constexpr (kSupportsFusedSilu) { + if (fuse_silu) { + run_epilogue(std::true_type{}); + } + else { + run_epilogue(std::false_type{}); } - cute::SM90_TMA_STORE::copy(Cdesc, - &smem_C[thread_idx * WG_TILE_M * LayoutC::C0], - offset_n + wg_idx_n * WG_TILE_N + tma_n, - offset_m + wg_idx_m * WG_TILE_M); - cute::tma_store_arrive(); + } + else { + run_epilogue(std::false_type{}); } } else if (tile->is_valid_cluster) { @@ -534,62 +989,22 @@ struct GemmUniversalSm90_v3 { // release last tile sched_state.release(); - if (threadIdx.x % WARPGROUP_SIZE < LayoutC::C1) { - cute::tma_store_wait<0>(); - } - } - - } // operator() - - template - __device__ void init_tma_descs(Array param_desc, CUtensorMap* smem_desc) - { - const int lane_id = threadIdx.x % WARP_SIZE; - - if (lane_id < sizeof(CUtensorMap) / sizeof(uint2)) { - PRAGMA_UNROLL - for (int i = 0; i < N; ++i) { - ((uint2*)&smem_desc[i])[lane_id] = ((uint2*)param_desc[i])[lane_id]; - } - } - - __syncwarp(); - } - - template - __device__ CUtensorMap* - update_tma_descs(CUtensorMap* gmem_desc, CUtensorMap* smem_desc, Array global_addrs, Array dims) - { - const int lane_id = threadIdx.x % WARP_SIZE; - if (lane_id == 0) { - PRAGMA_UNROLL - for (int i = 0; i < N; ++i) { - uint32_t uint_ptr = cast_smem_ptr_to_uint(&smem_desc[i]); - // clang-format off - asm volatile("tensormap.replace.tile.global_address.shared::cta.b1024.b64 [%0], %1;" ::"r"(uint_ptr), "l"(global_addrs[i])); - if (i != 2) { - asm volatile("tensormap.replace.tile.global_dim.shared::cta.b1024.b32 [%0], 1, %1;" ::"r"(uint_ptr), "r"(dims[i])); - } else { // special case for U - asm volatile("tensormap.replace.tile.global_dim.shared::cta.b1024.b32 [%0], 0, %1;" ::"r"(uint_ptr), "r"(dims[i])); + auto wait_tma_store = [&](auto fused_silu) { + constexpr bool kFuseSilu = decltype(fused_silu)::value; + using LayoutC = typename Output::LayoutC; + if (threadIdx.x % WARPGROUP_SIZE < LayoutC::C1) { + cute::tma_store_wait<0>(); } - // clang-format on + }; + if constexpr (kSupportsFusedSilu) { + fuse_silu ? wait_tma_store(std::true_type{}) : wait_tma_store(std::false_type{}); + } + else { + wait_tma_store(std::false_type{}); } } - __syncwarp(); - - constexpr int kNumPerCta = 5; // a,b,u,c0,c1 - auto gmem_ptr = &gmem_desc[blockIdx.x * kNumPerCta]; - PRAGMA_UNROLL - for (int i = 0; i < N; ++i) { - uint32_t uint_ptr = cast_smem_ptr_to_uint(&smem_desc[i]); - // clang-format off - asm volatile("tensormap.cp_fenceproxy.global.shared::cta.tensormap::generic.release.gpu.sync.aligned [%0], [%1], 128;" :: "l"(gmem_ptr + i), "r"(uint_ptr)); - // clang-format on - } - - return gmem_ptr; - } + } // operator() __device__ auto Fetch_V(typename Scheduler::Tile* tile, int wg_idx_n) { diff --git a/src/turbomind/kernels/gemm/gemm_universal_sm90_v4.h b/src/turbomind/kernels/gemm/gemm_universal_sm90_v4.h deleted file mode 100644 index b69e768106..0000000000 --- a/src/turbomind/kernels/gemm/gemm_universal_sm90_v4.h +++ /dev/null @@ -1,788 +0,0 @@ -#pragma once - -#include -#include - -#include - -#include "cute/arch/cluster_sm90.hpp" -#include "cute/arch/copy_sm90.hpp" -#include "cute/arch/copy_sm90_desc.hpp" -#include "cute/arch/copy_sm90_tma.hpp" -#include "cute/arch/mma_sm90_desc.hpp" -#include "cute/arch/mma_sm90_gmma.hpp" -#include "cute/atom/mma_traits.hpp" -#include "cute/atom/mma_traits_sm90_gmma.hpp" - -#include "cutlass/arch/barrier.h" -#include "cutlass/arch/reg_reconfig.h" -#include "cutlass/cutlass.h" -#include "cutlass/pipeline/sm90_pipeline.hpp" - -#include "src/turbomind/core/data_type.h" - -#include "src/turbomind/kernels/core/array_ops.h" -#include "src/turbomind/kernels/core/common.h" -#include "src/turbomind/kernels/core/layout.h" -#include "src/turbomind/kernels/core/meta.h" -#include "src/turbomind/kernels/core/smem.h" - -#include "src/turbomind/kernels/gemm/arch.h" -#include "src/turbomind/kernels/gemm/iterator_sm90.h" -#include "src/turbomind/kernels/gemm/matrix_ptr.h" -#include "src/turbomind/kernels/gemm/scheduler.cuh" -#include "src/turbomind/kernels/gemm/types.h" -#include "src/turbomind/kernels/gemm/utils.h" - -namespace turbomind::gemm { - -namespace GMMA = cute::SM90::GMMA; - -inline __device__ cute::GmmaDescriptor make_smem_desc(void* smem_ptr, int layout_type) -{ - auto uint_ptr = cast_smem_ptr_to_uint(smem_ptr); - - cute::GmmaDescriptor desc{}; - desc.bitfield.start_address_ = uint_ptr >> 4; - desc.bitfield.layout_type_ = layout_type; - desc.bitfield.leading_byte_offset_ = 0; - desc.bitfield.stride_byte_offset_ = 1024 >> 4; - desc.bitfield.base_offset_ = 0; - - return desc; -} - -template -struct SmemDescIterV2 { - union { - uint32_t u32_[2]; - uint64_t u64_; - }; - - uint32_t base_; - - __device__ SmemDescIterV2(uint64_t desc): u64_{desc}, base_{u32_[0]} {} - - __device__ void Advance(int stage) - { - u32_[0] += Step; - if (stage == Stages - 1) { - u32_[0] = base_; - } - } - - __device__ void Reset(int stage) - { - u32_[0] = base_ + stage * Step; - } - - __device__ SmemDescIterV2& operator+=(int offset) - { - u32_[0] += offset; - return *this; - } - - __device__ SmemDescIterV2& operator-=(int offset) - { - u32_[0] -= offset; - return *this; - } - - __device__ operator uint64_t() - { - return u64_; - } -}; - -template -inline __device__ void -wgmma_impl(uint64_t desc_a, uint64_t desc_b, float* frag_C, bool clear, std::index_sequence) -{ - return MMA_Atom::fma(desc_a, desc_b, frag_C[Is]..., clear ? GMMA::ScaleOut::Zero : GMMA::ScaleOut::One); -} - -template -inline __device__ void wgmma(uint64_t desc_a, uint64_t desc_b, float (&frag_C)[N], bool clear) -{ - return wgmma_impl(desc_a, desc_b, frag_C, clear, std::make_index_sequence{}); -} - -inline __device__ void warpgroup_fence_operand(float& reg) -{ - asm volatile("" : "+f"(reg)::"memory"); -} - -template -inline __device__ void warpgroup_fence_operand(float (&x)[M][N][K]) -{ - PRAGMA_UNROLL - for (int m = 0; m < M; ++m) { - PRAGMA_UNROLL - for (int n = 0; n < N; ++n) { - PRAGMA_UNROLL - for (int k = 0; k < K; ++k) { - warpgroup_fence_operand(x[m][n][k]); - } - } - } -} - -template -inline __device__ void warpgroup_fence_operand(float (&x)[N][K]) -{ - PRAGMA_UNROLL - for (int n = 0; n < N; ++n) { - PRAGMA_UNROLL - for (int k = 0; k < K; ++k) { - warpgroup_fence_operand(x[n][k]); - } - } -} - -template -__device__ void for_(std::index_sequence, Func func) -{ - return (func(constant{}), ...); -} - -namespace arch { - -template -struct Cluster { - static constexpr int M = M_; - static constexpr int N = N_; - - static constexpr int C = mk2cs(M, N).x; - static constexpr int S = mk2cs(M, N).y; - - static constexpr int size = M * N; - - static constexpr uint16_t kMaskC = (1 << C) - 1; - static constexpr uint16_t kMaskS = ((1 << size) - 1) / kMaskC; - - __device__ static ushort2 mask_cs(int cta_id) - { - const auto [c, s] = cta_cs(cta_id); - return make_ushort2(kMaskS << c, kMaskC << s * C); - } - - __device__ static ushort2 mask_mn(int cta_id) - { - auto [c, s] = mask_cs(cta_id); - return order == kColMajor ? ushort2{c, s} : ushort2{s, c}; - } - - __device__ static int2 cta_cs(int cta_id) - { - return {C > 1 ? cta_id % C : 0, S > 1 ? cta_id / C : 0}; - } - - __device__ static int2 cta_mn(int cta_id) - { - return cs2mk(cta_cs(cta_id)); - } - - int2 cta_mn_; - ushort2 mask_mn_; - - __device__ explicit Cluster(int cta_id): cta_mn_(cta_mn(cta_id)), mask_mn_(mask_mn(cta_id)) {} - - __device__ int cta_m() - { - return cta_mn_.x; - } - - __device__ int cta_n() - { - return cta_mn_.y; - } - - __device__ uint16_t mask_m() - { - return mask_mn_.x; - } - - __device__ uint16_t mask_n() - { - return mask_mn_.y; - } -}; - -} // namespace arch - -struct GemmUniversalSm90_v3 { - - static constexpr bool kDebug = false; - - using Arch = Sm90; - - // using MMA_Atom = GMMA::MMA_64x128x16_F32BF16BF16_SS; - using MMA_Atom = GMMA::MMA_64x192x32_F32E4M3E4M3_SS_TN<>; - static constexpr typename cute::MMA_Traits::Shape_MNK MMA_Shape{}; - - static constexpr int MMA_ATOM_M = cute::get<0>(MMA_Shape); - static constexpr int MMA_ATOM_N = cute::get<1>(MMA_Shape); - static constexpr int MMA_ATOM_K = cute::get<2>(MMA_Shape); - - static constexpr int TILE_M = 128; - static constexpr int TILE_N = 192; - static constexpr int TILE_K = 128; - - static constexpr int WG_M = 2; - static constexpr int WG_N = 1; - - static constexpr int WG_TILE_M = TILE_M / WG_M; - static constexpr int WG_TILE_N = TILE_N / WG_N; - - static constexpr int kSchedWarpGroups = 1; - - static constexpr int WARPGORUPS = WG_M * WG_N; - - static constexpr int MMA_ITER_M = WG_TILE_M / MMA_ATOM_M; - static constexpr int MMA_ITER_N = WG_TILE_N / MMA_ATOM_N; - static constexpr int MMA_ITER_K = TILE_K / MMA_ATOM_K; - - static constexpr int kMulticastA = 1; - static constexpr int kMulticastB = 2; - - static constexpr int kClusterSize = kMulticastA * kMulticastB; - - static constexpr int Stages = 4; - - static constexpr bool kSplitK = false; - static constexpr int kChunkSizeK = TILE_K; - - static constexpr int WARPGROUP_SIZE = 128; - - static constexpr int CTA_SIZE = WARPGROUP_SIZE * (WARPGORUPS + 1); - - using Ta = __nv_fp8_e4m3; - using Tb = __nv_fp8_e4m3; - using Tc = nv_bfloat16; - - using Tu = float; - using Tv = float; - - using Cluster = arch::Cluster; - - using Scheduler = TileScheduler; - - using ProducerBar = cutlass::arch::ClusterTransactionBarrier; - using ConsumerBar = cutlass::arch::ClusterBarrier; - - static constexpr int MAX_K = 32768; - - static constexpr int TILE_M_U = cdiv(TILE_M, 1); - static constexpr int CTA_K_U = cdiv(TILE_K, 128); - - static constexpr int kTmaTxBytes = - sizeof(Ta) * (TILE_M * TILE_K) + sizeof(Tb) * (TILE_N * TILE_K) + sizeof(Tu) * TILE_M_U * CTA_K_U; - - // ! Smem addr must be SBO aligned for TMA load/store - struct SharedStorage { - struct Source { - __align__(1024) Array A; - __align__(1024) Array B; - __align__(1024) Tu U[Stages][TILE_M_U * CTA_K_U]; - // __align__(1024) Tv V[2][WARPGORUPS][cdiv(MAX_K, 128)]; - __align__(1024) Tv V[Stages][2 * cdiv(MAX_K, 128)]; - }; - Source source; - __align__(1024) Array C; - __align__(128) uint64_t producer_bar[Stages]; - __align__(128) uint64_t consumer_bar[Stages]; - __align__(128) CUtensorMap tma_desc_C[WARPGORUPS]; - }; - - static constexpr int kSmemSize = sizeof(SharedStorage); - - static constexpr int kSwizzleC = 2 * std::gcd(WG_TILE_N, 128 / sizeof(Tc)); - - using LayoutC = std::conditional_t= 32, - SmemLayoutV2, - SmemLayoutV2>; - - __device__ void operator()(const CUtensorMap& tm_a, - const CUtensorMap& tm_b, - const CUtensorMap& tm_c, - const CUtensorMap& tm_u, - const CUtensorMap& tm_v, - const MatrixParam& param_A, - const MatrixParam& param_B, - const MatrixParam& param_U, - const MatrixParam& param_V, - const MatrixParam& param_C, - uint2 box_V, - Scheduler sched, - CUtensorMap* tensormap_buf, - char* smem_buf) - { - SharedStorage& storage = *reinterpret_cast(smem_buf); - - uint64_t* producer_bar = storage.producer_bar; - uint64_t* consumer_bar = storage.consumer_bar; - - if (threadIdx.x == 0) { - PRAGMA_UNROLL - for (int s = 0; s < Stages; ++s) { - ProducerBar::init(&producer_bar[s], 1); - ConsumerBar::init(&consumer_bar[s], WARPGORUPS * kClusterSize * 4); - } - cutlass::arch::fence_view_async_shared(); - if constexpr (kClusterSize > 1) { - cutlass::arch::fence_barrier_init(); - } - } - - (kClusterSize > 1) ? cute::cluster_sync() : __syncthreads(); - - const int wg_idx = cutlass::canonical_warp_group_idx(); - - if (wg_idx == WARPGORUPS) { - cutlass::arch::warpgroup_reg_dealloc<40>(); - - static_assert(TILE_M % kMulticastA == 0); - static_assert(TILE_N % kMulticastB == 0); - - // if (threadIdx.x == WARPGORUPS * WARPGROUP_SIZE) { - if (threadIdx.x % WARPGROUP_SIZE / WARP_SIZE == 0) { - - Cluster cluster(cute::block_id_in_cluster().x); - - const int mc_offset_m = cluster.cta_n() * (TILE_M / kMulticastA); - const int mc_offset_n = cluster.cta_m() * (TILE_N / kMulticastB); - - auto smem_A = storage.source.A.data() + mc_offset_m * TILE_K; - auto smem_B = storage.source.B.data() + mc_offset_n * TILE_K; - auto& smem_U = storage.source.U; - auto& smem_V = storage.source.V; - - sched.grid_init(); - - cutlass::PipelineState write_state{0, 1, 0}; - cutlass::PipelineState v_state{0, 1, 0}; - - while (sched.next()) { - if (cute::elect_one_sync()) { - auto [valid_cta_tile_p, cluster_tile_p] = sched.is_valid_tile(); - - if (!cluster_tile_p) { - // OOB tile caused by swizzle pattern - continue; - } - - const auto tile_offset = sched.tile_offset(); - const auto [iter_k_beg, iter_k_end] = sched.iter_k_range(); - - const int offset_k = iter_k_beg * TILE_K; - - const uint16_t mask_A = cluster.mask_m(); - const uint16_t mask_B = cluster.mask_n(); - - const int offset_m = tile_offset.x * TILE_M; - const int offset_n = tile_offset.y * TILE_N; - - int k_iter = iter_k_end - iter_k_beg; - - GmemIteratorSm90 gmem_A{&tm_a, {offset_k, offset_m + mc_offset_m}, {TILE_K, 0}}; - GmemIteratorSm90 gmem_B{&tm_b, {offset_k, offset_n + mc_offset_n}, {TILE_K, 0}}; - GmemIteratorSm90 gmem_U{&tm_u, {offset_m + mc_offset_m, offset_k / 128}, {0, 1}}; - GmemIteratorSm90<1> gmem_V(&tm_v, {0, offset_n / 128}, {0, 0}); - - { - int pipe = write_state.index(); - ConsumerBar::wait(&consumer_bar[pipe], write_state.phase()); - const int v_bytes = sizeof(Tv) * box_V.x * box_V.y; - ProducerBar::arrive_and_expect_tx(&producer_bar[pipe], kTmaTxBytes + v_bytes); - gmem_A.Step(&producer_bar[pipe], &smem_A[pipe * TILE_M * TILE_K], mask_A); - gmem_B.Step(&producer_bar[pipe], &smem_B[pipe * TILE_N * TILE_K], mask_B); - gmem_U.Step(&producer_bar[pipe], &smem_U[pipe][0] + mc_offset_m, mask_A); - gmem_V.Step(&producer_bar[pipe], &smem_V[v_state.index()], 0); - ++write_state; - --k_iter; - } - - for (; k_iter > 0; --k_iter) { - int pipe = write_state.index(); - ConsumerBar::wait(&consumer_bar[pipe], write_state.phase()); - ProducerBar::arrive_and_expect_tx(&producer_bar[pipe], kTmaTxBytes); - gmem_A.Step(&producer_bar[pipe], &smem_A[pipe * TILE_M * TILE_K], mask_A); - gmem_B.Step(&producer_bar[pipe], &smem_B[pipe * TILE_N * TILE_K], mask_B); - gmem_U.Step(&producer_bar[pipe], &smem_U[pipe][0] + mc_offset_m, mask_A); - ++write_state; - } - - ++v_state; - } - } - } - } - else { - cutlass::arch::warpgroup_reg_alloc<232>(); - - sched.grid_init(kSchedWarpGroups); - - auto& smem_A = storage.source.A; - auto& smem_B = storage.source.B; - auto& smem_U = storage.source.U; - - const int wg_idx_m = WG_M > 1 ? wg_idx % WG_M : 0; - const int wg_idx_n = WG_N > 1 ? wg_idx / WG_M : 0; - - auto smem_desc_A = make_smem_desc(&smem_A[wg_idx_m * WG_TILE_M * TILE_K], 1); - auto smem_desc_B = make_smem_desc(&smem_B[wg_idx_n * WG_TILE_N * TILE_K], 1); - - SmemDescIterV2> 4)> smem_iter_A{smem_desc_A}; - SmemDescIterV2> 4)> smem_iter_B{smem_desc_B}; - - constexpr int kStepMA = (sizeof(Ta) * MMA_ATOM_M * TILE_K) >> 4; - constexpr int kStepNB = (sizeof(Tb) * MMA_ATOM_N * TILE_K) >> 4; - constexpr int kStepKA = (sizeof(Ta) * MMA_ATOM_K) >> 4; - constexpr int kStepKB = (sizeof(Tb) * MMA_ATOM_K) >> 4; - - cutlass::arch::NamedBarrier wg_barrier(WARPGROUP_SIZE, wg_idx + 2); // 2,3 - - auto epi_barrier = [&](int phase) { // 0, 1 - return EmptyBarrier{}; - // return cutlass::arch::NamedBarrier(WARPGORUPS * WARPGROUP_SIZE, wg_idx ^ phase); - }; - - if (wg_idx == 1) { - epi_barrier(1).arrive_unaligned(); - } - - cutlass::PipelineState pipe_state{}; - cutlass::PipelineState v_state{}; - - while (sched.next(kSchedWarpGroups)) { - auto [cta_tile_p, cluster_tile_p] = sched.is_valid_tile(); - - if (!cluster_tile_p) { - // OOB tile caused by swizzle pattern - continue; - } - - MMA_Atom::CRegisters frag_C[MMA_ITER_N]; - MMA_Atom::CRegisters accum_C[MMA_ITER_M][MMA_ITER_N]{}; - - const auto tile_offset = sched.tile_offset(); - const auto [iter_k_beg, iter_k_end] = sched.iter_k_range(); - - // const auto [M, N, K, L] = sched.gemm_shape(); - - const int offset_m = tile_offset.x * TILE_M; - const int offset_n = tile_offset.y * TILE_N; - - const int wg_offset_n = offset_n + wg_idx_n * WG_TILE_N; - - int k_iter = iter_k_end - iter_k_beg; - - const int warp_id = threadIdx.x / WARP_SIZE; - const int lane_id = threadIdx.x % WARP_SIZE; - - auto consumer_arrive = [&] { - __syncwarp(); - if constexpr (kClusterSize > 1) { - ConsumerBar::arrive(&consumer_bar[pipe_state.index()], lane_id, lane_id < kClusterSize); - } - else { - if (lane_id == 0) { - ConsumerBar::arrive(&consumer_bar[pipe_state.index()]); - } - } - __syncwarp(); - }; - - if constexpr (kClusterSize > 1) { - if (!cta_tile_p) { // other CTAs in the cluster are still alive - for (; k_iter > 0; --k_iter) { - ProducerBar::wait(&producer_bar[pipe_state.index()], pipe_state.phase()); - consumer_arrive(); - ++pipe_state; - } - epi_barrier(0).arrive_and_wait_unaligned(); - epi_barrier(1).arrive_unaligned(); - continue; - } - } - - // auto Copy = [k = cdiv(K, 128)](Tv* dst, const Tv* src) { - // PRAGMA_NO_UNROLL - // for (int i = threadIdx.x % WARPGROUP_SIZE; i < k; i += WARPGROUP_SIZE) { - // dst[i] = __ldg(&src[i]); - // } - // }; - // auto gmem_V = (const Tv*)param_V.ptr + (wg_offset_n / 128) * param_V.stride + (offset_k / 128); - // Copy(storage.source.V[0][wg_idx], gmem_V); - - uint32_t pred_V{}; - int iter_V{}; - - constexpr int OUTER_N = std::gcd(MMA_ATOM_N, 128); - if constexpr (OUTER_N != 128) { - - static_assert(MMA_ATOM_N <= 128 + OUTER_N, "MMA inst is crossing more than 2 scale blocks"); - - constexpr uint32_t mask = (1UL << (WG_TILE_N / OUTER_N)) - 1; - - int phase = 128 - wg_offset_n % 128; - pred_V = (mask << (phase / OUTER_N)) & mask; - - // if (pred_V && wg_offset_n / 128 + 1 < cdiv(N, 128)) { - // Copy(storage.source.V[1][wg_idx], gmem_V + param_V.stride); - // } - // if constexpr (WG_N > 1) { - // constexpr int tiles = MMA_ATOM_N / OUTER_N; - // pred_V = (pred_V >> (wg_idx_n * tiles)) & ((1 << tiles) - 1); - // } - } - - __syncwarp(); - - float scale_V[2]; - // auto Load_V = [&] { - // scale_V[0] = storage.source.V[0][wg_idx][iter_V]; - // if (pred_V) { - // scale_V[1] = storage.source.V[1][wg_idx][iter_V]; - // } - // ++iter_V; - // }; - auto Load_V = [&] { - // scale_V[0] = scale_V[1] = 1; - scale_V[0] = storage.source.V[v_state.index()][iter_V]; - if (pred_V) { - scale_V[1] = storage.source.V[v_state.index()][box_V.x + iter_V]; - } - ++iter_V; - }; - - float scale_U[MMA_ITER_M][2]; - const int offset_U = wg_idx_m * WG_TILE_M + warp_id % 4 * 16 + lane_id / 4; - auto Load_U = [&] { - for (int m = 0; m < MMA_ITER_M; ++m) { - scale_U[m][0] = smem_U[pipe_state.index()][offset_U + m * MMA_ATOM_M]; - scale_U[m][1] = smem_U[pipe_state.index()][offset_U + m * MMA_ATOM_M + 8]; - } - }; - - auto scale_accum = [&](int m) { // cta_n = mma_iter_n * wg_n * mma_atom_n - float scales[2][2]; - scales[0][0] = scale_U[m][0] * scale_V[0]; - scales[1][0] = scale_U[m][1] * scale_V[0]; - scales[0][1] = scale_U[m][0] * scale_V[1]; - scales[1][1] = scale_U[m][1] * scale_V[1]; - cute::warpgroup_wait<0>(); - PRAGMA_UNROLL - for (int n = 0; n < MMA_ITER_N; ++n) { - PRAGMA_UNROLL - for (int c0 = 0; c0 < MMA_ATOM_N; c0 += OUTER_N) { - bool pred = (pred_V & (1U << (c0 / OUTER_N))); - PRAGMA_UNROLL - for (int cc = 0; cc < OUTER_N; cc += 8) { - int c = c0 + cc; - // clang-format off - accum_C[m][n][c / 2 + 0] += (pred ? scales[0][1] : scales[0][0]) * frag_C[n][c / 2 + 0]; - accum_C[m][n][c / 2 + 1] += (pred ? scales[0][1] : scales[0][0]) * frag_C[n][c / 2 + 1]; - accum_C[m][n][c / 2 + 2] += (pred ? scales[1][1] : scales[1][0]) * frag_C[n][c / 2 + 2]; - accum_C[m][n][c / 2 + 3] += (pred ? scales[1][1] : scales[1][0]) * frag_C[n][c / 2 + 3]; - // clang-format on - } - } - } - }; - - auto gmma = [&](int m) { - PRAGMA_UNROLL - for (int k = 0; k < MMA_ITER_K; ++k) { - PRAGMA_UNROLL - for (int n = 0; n < MMA_ITER_N; ++n) { - wgmma(smem_iter_A, smem_iter_B, frag_C[n], k == 0); - smem_iter_B += kStepNB; - } - smem_iter_B -= MMA_ITER_N * kStepNB; - smem_iter_A += kStepKA; - smem_iter_B += kStepKB; - } - smem_iter_A -= MMA_ITER_K * kStepKA; - smem_iter_B -= MMA_ITER_K * kStepKB; - smem_iter_A += kStepMA; - cute::warpgroup_commit_batch(); - }; - - static_assert(MMA_ITER_N == 1); - - wg_barrier.sync(); - - ProducerBar::wait(&producer_bar[pipe_state.index()], pipe_state.phase()); - Load_V(); - Load_U(); - smem_iter_A.Reset(pipe_state.index()); - smem_iter_B.Reset(pipe_state.index()); - cute::warpgroup_arrive(); - gmma(0); - scale_accum(0); - consumer_arrive(); - ++pipe_state; - --k_iter; - - Load_V(); - ProducerBar::wait(&producer_bar[pipe_state.index()], pipe_state.phase()); - Load_U(); - smem_iter_A.Reset(pipe_state.index()); - smem_iter_B.Reset(pipe_state.index()); - - for (; k_iter > 1; --k_iter) { - cute::warpgroup_arrive(); - gmma(0); - scale_accum(0); - consumer_arrive(); - ++pipe_state; - Load_V(); - ProducerBar::wait(&producer_bar[pipe_state.index()], pipe_state.phase()); - Load_U(); - smem_iter_A.Reset(pipe_state.index()); - smem_iter_B.Reset(pipe_state.index()); - } - - cute::warpgroup_arrive(); - gmma(0); - scale_accum(0); - consumer_arrive(); - ++pipe_state; - ++v_state; - - const int wg_lane = threadIdx.x % WARPGROUP_SIZE; - - if (wg_lane < LayoutC::C1) { - cute::tma_store_wait<0>(); - } - - epi_barrier(0).arrive_and_wait_unaligned(); - - wg_barrier.sync(); - - // void* Cdesc{}; - // if (threadIdx.x % WARPGROUP_SIZE / WARP_SIZE == 0) { - // Cdesc = update_tma_desc(tm_c, tensormap_buf, &storage.tma_desc_C[wg_idx], wg_idx, param_C.ptr, - // M); - // } - - Tc* smem_C = &storage.C[wg_idx_m * WG_TILE_M * TILE_N + wg_idx_n * WG_TILE_N]; - - // epilogue - PRAGMA_UNROLL - for (int m = 0; m < MMA_ITER_M; ++m) { - PRAGMA_UNROLL - for (int n = 0; n < MMA_ITER_N; ++n) { - - constexpr int N = LayoutC::C0; - constexpr int SW_bits = log2(kSwizzleC / 16); - - static_assert(!SW_bits || MMA_ATOM_N % LayoutC::C0 == 0); - - const int m0 = m * MMA_ATOM_M; - const int n0 = n * MMA_ATOM_N; - - PRAGMA_UNROLL - for (int i = 0; i < MMA_ATOM_N; i += 16) { - __align__(16) Array tvec = cast(*(Array*)&accum_C[m][n][i / 2]); - - int mm = m0 + warp_id % 4 * 16 + (lane_id & 8); - int nn = n0 + i / N * N; - - int addr = ((nn / N) * WG_TILE_M * N) + (mm * N) + (nn % N); - - int s = lane_id % 8; - int c = (lane_id & 16) / 2 + i % N; - - addr += Swizzle::apply(s * N + c); - - auto& uvec = (Array&)tvec; - cute::SM90_U32x4_STSM_N::copy(uvec[0], // - uvec[1], - uvec[2], - uvec[3], - (cutlass::uint128_t&)smem_C[addr]); - } - } - } - - cute::tma_store_fence(); // visibility: smem -> async proxy - - wg_barrier.sync(); - - if (wg_lane < LayoutC::C1) { - const int tma_n = wg_lane * LayoutC::C0; - // cute::tma_descriptor_fence_acquire((cute::TmaDescriptor*)Cdesc); - cute::SM90_TMA_STORE::copy(&tm_c, - &smem_C[wg_lane * WG_TILE_M * LayoutC::C0], - offset_n + wg_idx_n * WG_TILE_N + tma_n, - offset_m + wg_idx_m * WG_TILE_M); - cute::tma_store_arrive(); - } - - epi_barrier(1).arrive_unaligned(); - - } // scheduler loop - - if (threadIdx.x % WARPGROUP_SIZE < LayoutC::C1) { - cute::tma_store_wait<0>(); - } - - if (wg_idx == 0) { - epi_barrier(0).arrive_and_wait_unaligned(); - } - } - - if constexpr (kClusterSize > 1) { - cute::cluster_arrive(); - cute::cluster_wait(); - } - - } // operator() - - struct EmptyBarrier { - __device__ EmptyBarrier(...) {} - __device__ void arrive_and_wait_unaligned() {} - __device__ void arrive_unaligned() {} - }; - - __device__ void* update_tma_desc(const CUtensorMap& param_desc, - CUtensorMap* gmem_desc, - CUtensorMap* smem_desc, - int index, - void* global_addr, - int dim) - { - uint32_t uint_ptr = cast_smem_ptr_to_uint(smem_desc); - - const int lane_id = threadIdx.x % WARP_SIZE; - - if (lane_id < sizeof(CUtensorMap) / sizeof(uint2)) { - ((uint2*)smem_desc)[lane_id] = ((uint2*)¶m_desc)[lane_id]; - } - - __syncwarp(); - - if (lane_id == 0) { - // clang-format off - asm volatile("tensormap.replace.tile.global_address.shared::cta.b1024.b64 [%0], %1;" ::"r"(uint_ptr), "l"(global_addr)); - asm volatile("tensormap.replace.tile.global_dim.shared::cta.b1024.b32 [%0], 1, %1;" ::"r"(uint_ptr), "r"(dim)); - // clang-format on - } - - __syncwarp(); - - constexpr int kNumPerCta = 4; - - auto gmem_ptr = (void*)&gmem_desc[blockIdx.x * kNumPerCta + index]; - - // clang-format off - asm volatile("tensormap.cp_fenceproxy.global.shared::cta.tensormap::generic.release.gpu.sync.aligned [%0], [%1], 128;" :: "l"(gmem_ptr), "r"(uint_ptr)); - // clang-format on - - return gmem_ptr; - } -}; - -} // namespace turbomind::gemm diff --git a/src/turbomind/kernels/gemm/gemm_universal_sm90_v5.h b/src/turbomind/kernels/gemm/gemm_universal_sm90_v5.h index 3a63bd487b..9018ddd5ea 100644 --- a/src/turbomind/kernels/gemm/gemm_universal_sm90_v5.h +++ b/src/turbomind/kernels/gemm/gemm_universal_sm90_v5.h @@ -35,6 +35,7 @@ #include "src/turbomind/kernels/gemm/types.h" #include "src/turbomind/kernels/gemm/utils.h" +#include "src/turbomind/kernels/gemm/prepare_moe_tma_descs_sm90_fp8.h" #include "src/turbomind/kernels/gemm/scaled_gmma_fp8_sm90.h" #include "src/turbomind/kernels/gemm/sm90_utils.h" @@ -47,7 +48,7 @@ struct GemmUniversalSm90_v5 { using Arch = Sm90; - static constexpr int WARPGORUPS = 4; + static constexpr int WARPGROUPS = 4; static constexpr int TILE_M = 128; static constexpr int TILE_N = 96; @@ -74,7 +75,7 @@ struct GemmUniversalSm90_v5 { static constexpr int WARPGROUP_SIZE = 128; static constexpr int kMathGroupSize = 256; - static constexpr int CTA_SIZE = WARPGROUP_SIZE * (WARPGORUPS + 1); + static constexpr int CTA_SIZE = WARPGROUP_SIZE * (WARPGROUPS + 1); using Ta = __nv_fp8_e4m3; using Tb = __nv_fp8_e4m3; @@ -106,7 +107,8 @@ struct GemmUniversalSm90_v5 { static constexpr int kTmaTxBytes = sizeof(Ta) * (TILE_M * TILE_K) + sizeof(Tb) * (TILE_N * TILE_K) + sizeof(Tu) * kBoxU; - static constexpr int kTmaDescNum = 7; + // Grouped: per-expert maps in workspace [A, B, U, C]. Dense unused. + static constexpr int kTmaDescNum = is_grouped_gemm_ ? 4 : 1; // ! Smem addr must be SBO aligned for TMA load/store struct SharedStorage { @@ -120,7 +122,6 @@ struct GemmUniversalSm90_v5 { __align__(1024) Array C; __align__(128) uint64_t producer_bar[Stages]; __align__(128) uint64_t consumer_bar[Stages]; - __align__(128) CUtensorMap tma_desc_buf[kTmaDescNum]; // int pipe_count[2]; typename Scheduler::Storage sched; }; @@ -136,6 +137,27 @@ struct GemmUniversalSm90_v5 { static constexpr int OUTER_N = GMMA::OUTER_N; static constexpr int MMA_SUBTILE_N = GMMA::OP_N / OUTER_N; + // Host: rebase per-expert A/B/U/C TMA maps into workspace before GEMM launch. + static void PrepareTmaDescs(const CUtensorMap& tm_a, + const CUtensorMap& tm_b, + const CUtensorMap& tm_u, + const CUtensorMap& tm_c, + const MatrixParam& param_A, + const MatrixParam& param_B, + const MatrixParam& param_U, + const MatrixParam& param_C, + CUtensorMap* out, + int num_experts, + int N, + cudaStream_t stream) + { + if constexpr (!is_grouped_gemm_) { + return; + } + prepare_moe_tma_descs_sm90_fp8 + <<>>(tm_a, tm_b, tm_u, tm_c, param_A, param_B, param_U, param_C, out, N); + } + __device__ void operator()(const CUtensorMap& tm_a, const CUtensorMap& tm_b, const CUtensorMap& tm_c, @@ -176,7 +198,7 @@ struct GemmUniversalSm90_v5 { const int wg_idx = cutlass::canonical_warp_group_idx(); - if (wg_idx == WARPGORUPS) { + if (wg_idx == WARPGROUPS) { cutlass::arch::warpgroup_reg_dealloc<32>(); static_assert(TILE_M % kMulticastA == 0); @@ -199,10 +221,6 @@ struct GemmUniversalSm90_v5 { auto& smem_U = storage.source.U; auto& smem_V = storage.source.V; - if constexpr (is_grouped_gemm) { - init_tma_descs<3>({&tm_a, &tm_b, &tm_u}, storage.tma_desc_buf); - } - cutlass::PipelineState write_state{0, 1, 0}; auto sched_state = sched.init_consumer(storage.sched); @@ -237,35 +255,14 @@ struct GemmUniversalSm90_v5 { const Tv* gmem_V1; if constexpr (is_grouped_gemm) { - const int g = tile->group_idx; - const int m0 = tile->m0; - const int m1 = tile->m1; - const int m = m1 - m0; - - Array global_addrs; - global_addrs[0] = (Ta*)param_A.ptr + m0 * (int64_t)param_A.stride; - global_addrs[1] = ((void**)param_B.ptr)[g]; - - const int beg_u = m0 / kAlignmentU * kAlignmentU; - const int end_u = round_up(m1, kAlignmentU); - global_addrs[2] = (Tu*)param_U.ptr + beg_u; - - Array dims; - dims[0] = m; - dims[1] = sched.gemm_shape().y; - dims[2] = end_u - beg_u; - - auto descs = update_tma_descs(tensormap_buf, storage.tma_desc_buf, global_addrs, dims); - Adesc = &descs[0]; - Bdesc = &descs[1]; - Udesc = &descs[2]; - - gmem_V0 = ((Tv**)gmem_V0)[g]; - - PRAGMA_UNROLL - for (int i = 0; i < 3; ++i) { - cute::tma_descriptor_fence_acquire((cute::TmaDescriptor*)&descs[i]); - } + // Descs published by prepare_moe_tma_descs on this stream; + // fence_acquire only needed after in-kernel tensormap replace. + const int g = tile->group_idx; + CUtensorMap* const descs = tensormap_buf + g * kTmaDescNum; + Adesc = &descs[0]; + Bdesc = &descs[1]; + Udesc = &descs[2]; + gmem_V0 = ((Tv**)gmem_V0)[g]; } if (lane_predicate) { @@ -351,12 +348,6 @@ struct GemmUniversalSm90_v5 { const int math_group_idx = wg_idx / 2; - if constexpr (is_grouped_gemm) { - if (threadIdx.x % WARPGROUP_SIZE / WARP_SIZE == 0) { - init_tma_descs<1>({&tm_c}, storage.tma_desc_buf + 3 + wg_idx); - } - } - auto& smem_A = storage.source.A; auto& smem_B = storage.source.B; auto& smem_U = storage.source.U; @@ -377,7 +368,7 @@ struct GemmUniversalSm90_v5 { auto math_barrier_sync = [&](int phase, int alive = 1) { constexpr int base = (int)cutlass::arch::ReservedNamedBarriers::FirstUserBarrier; const int barrier_id = base + math_group_idx ^ phase; - constexpr int threads = WARPGORUPS * WARPGROUP_SIZE; + constexpr int threads = WARPGROUPS * WARPGROUP_SIZE; int res = 0; asm volatile("{\n" " .reg.pred p;\n" @@ -459,15 +450,6 @@ struct GemmUniversalSm90_v5 { GMMA::apply(smem_iter_A, smem_iter_B, frag_C, accum_C, frag_U, scale_V, pred_V); }; - if constexpr (is_grouped_gemm) { - if (warp_id % 4 == 0) { - int m0 = tile->m0, m1 = tile->m1; - auto addr = (Tc*)param_C.ptr + m0 * (int64_t)param_C.stride; - int idx = 3 + wg_idx; - update_tma_descs<1>(tensormap_buf + idx, storage.tma_desc_buf + idx, {addr}, {m1 - m0}); - } - } - math_barrier_sync(0); pipe_state = {}; @@ -549,8 +531,7 @@ struct GemmUniversalSm90_v5 { const void* Cdesc = &tm_c; const int tma_n = thread_idx * LayoutC::C0; if constexpr (is_grouped_gemm) { - Cdesc = &tensormap_buf[blockIdx.x * kTmaDescNum + 3 + wg_idx]; - cute::tma_descriptor_fence_acquire((cute::TmaDescriptor*)Cdesc); + Cdesc = &tensormap_buf[tile->group_idx * kTmaDescNum + 3]; } cute::SM90_TMA_STORE::copy(Cdesc, &smem_C[thread_idx * WG_TILE_M * LayoutC::C0], @@ -613,55 +594,6 @@ struct GemmUniversalSm90_v5 { } // operator() - template - __device__ void init_tma_descs(Array param_desc, CUtensorMap* smem_desc) - { - const int lane_id = threadIdx.x % WARP_SIZE; - - if (lane_id < sizeof(CUtensorMap) / sizeof(uint2)) { - PRAGMA_UNROLL - for (int i = 0; i < N; ++i) { - ((uint2*)&smem_desc[i])[lane_id] = ((uint2*)param_desc[i])[lane_id]; - } - } - - __syncwarp(); - } - - template - __device__ CUtensorMap* - update_tma_descs(CUtensorMap* gmem_desc, CUtensorMap* smem_desc, Array global_addrs, Array dims) - { - const int lane_id = threadIdx.x % WARP_SIZE; - if (lane_id == 0) { - PRAGMA_UNROLL - for (int i = 0; i < N; ++i) { - uint32_t uint_ptr = cast_smem_ptr_to_uint(&smem_desc[i]); - // clang-format off - asm volatile("tensormap.replace.tile.global_address.shared::cta.b1024.b64 [%0], %1;" ::"r"(uint_ptr), "l"(global_addrs[i])); - if (i != 2) { - asm volatile("tensormap.replace.tile.global_dim.shared::cta.b1024.b32 [%0], 1, %1;" ::"r"(uint_ptr), "r"(dims[i])); - } else { // special case for U - asm volatile("tensormap.replace.tile.global_dim.shared::cta.b1024.b32 [%0], 0, %1;" ::"r"(uint_ptr), "r"(dims[i])); - } - // clang-format on - } - } - - __syncwarp(); - - auto gmem_ptr = &gmem_desc[blockIdx.x * kTmaDescNum]; - PRAGMA_UNROLL - for (int i = 0; i < N; ++i) { - uint32_t uint_ptr = cast_smem_ptr_to_uint(&smem_desc[i]); - // clang-format off - asm volatile("tensormap.cp_fenceproxy.global.shared::cta.tensormap::generic.release.gpu.sync.aligned [%0], [%1], 128;" :: "l"(gmem_ptr + i), "r"(uint_ptr)); - // clang-format on - } - - return gmem_ptr; - } - __device__ auto Fetch_V(const MatrixParam& param_V, int K, int N, diff --git a/src/turbomind/kernels/gemm/gmma_bf16_sm90.h b/src/turbomind/kernels/gemm/gmma_bf16_sm90.h new file mode 100644 index 0000000000..96e1992894 --- /dev/null +++ b/src/turbomind/kernels/gemm/gmma_bf16_sm90.h @@ -0,0 +1,104 @@ +#pragma once + +#include "cute/arch/mma_sm90.hpp" +#include "cute/atom/mma_atom.hpp" +#include "cute/atom/mma_traits_sm90_gmma.hpp" + +namespace turbomind::gemm { + +// Identical to cutlass::gemm::collective::detail::ss_smem_selector +// (sm90_common.inl:273-319). Including that .inl from TurboMind headers is +// awkward (it is meant to be pulled inside cutlass::gemm::collective::detail). +// Compose the same GMMA Layout_*_Atom SoT here — MMA.md §6 / SHARED_MEMORY.md. +template +CUTE_HOST_DEVICE constexpr auto gmma_ss_smem_selector() +{ + using namespace cute; + + auto BLK_MN0 = size<0>(BLK_MN{}); + auto BLK_K0 = size<0>(BLK_K{}); + + static_assert(BLK_MN0 % 8 == 0, "BLK_MN0 must be a multiple of 8."); + static_assert(BLK_K0 % 8 == 0, "BLK_K0 must be a multiple of 8."); + + if constexpr (major == GMMA::Major::MN) { + if constexpr (BLK_MN0 % size<0>(GMMA::Layout_MN_SW128_Atom{}) == 0) { + return GMMA::Layout_MN_SW128_Atom{}; + } + else if constexpr (BLK_MN0 % size<0>(GMMA::Layout_MN_SW64_Atom{}) == 0) { + return GMMA::Layout_MN_SW64_Atom{}; + } + else if constexpr (BLK_MN0 % size<0>(GMMA::Layout_MN_SW32_Atom{}) == 0) { + return GMMA::Layout_MN_SW32_Atom{}; + } + else if constexpr (BLK_MN0 % size<0>(GMMA::Layout_MN_INTER_Atom{}) == 0) { + return GMMA::Layout_MN_INTER_Atom{}; + } + else { + static_assert(BLK_MN0 % size<0>(GMMA::Layout_MN_INTER_Atom{}) == 0, + "BLK_MN0 must be a multiple of size<0>(GMMA::Layout_MN_INTER_Atom)"); + } + } + else if constexpr (major == GMMA::Major::K) { + // Prefer largest legal Swizzle: SW128 → SW64 → SW32 → INTER. + if constexpr (BLK_K0 % size<1>(GMMA::Layout_K_SW128_Atom{}) == 0) { + return GMMA::Layout_K_SW128_Atom{}; + } + else if constexpr (BLK_K0 % size<1>(GMMA::Layout_K_SW64_Atom{}) == 0) { + return GMMA::Layout_K_SW64_Atom{}; + } + else if constexpr (BLK_K0 % size<1>(GMMA::Layout_K_SW32_Atom{}) == 0) { + return GMMA::Layout_K_SW32_Atom{}; + } + else if constexpr (BLK_K0 % size<1>(GMMA::Layout_K_INTER_Atom{}) == 0) { + return GMMA::Layout_K_INTER_Atom{}; + } + else { + static_assert(BLK_K0 % size<1>(GMMA::Layout_K_INTER_Atom{}) == 0, + "BLK_K0 must be a multiple of size<1>(GMMA::Layout_K_INTER_Atom)"); + } + } +} + +// TILE_OUT = GMMA-M extent = problem N_out (weight output) → GMMA-A = weight +// TILE_BATCH = GMMA-N extent = problem M_batch (tokens) → GMMA-B = activation +// LlamaLinear API A/B stay act/weight; MMA operands are swapped relative to that API. +// AtomLayoutMNK comes from the Tile_ (CuTe Layout>). +template +struct GmmaBF16Traits { + using ElementA = cutlass::bfloat16_t; // GMMA-A = weight + using ElementB = cutlass::bfloat16_t; // GMMA-B = activation + using ElementC = float; + + using TileShape = cute::Shape, cute::Int, cute::Int>; + static constexpr auto MajorA = cute::GMMA::Major::K; + static constexpr auto MajorB = cute::GMMA::Major::K; + + using AtomLayoutMNK = AtomLayoutMNK_; + + static constexpr int kAtomM = cute::size<0>(AtomLayoutMNK{}); + static constexpr int kAtomN = cute::size<1>(AtomLayoutMNK{}); + static_assert(kAtomM * kAtomN >= 1); + // Layout<_2,_1,_1>: GMMA-M (OUT) must be multiple of 128 (2×64 atom). + // Layout<_1,_2,_1>: GMMA-M (OUT) multiple of 64; BATCH split across 2 WGs. + static_assert(TILE_OUT % (64 * kAtomM) == 0, "TILE_OUT vs AtomLayout M"); + static_assert(TILE_BATCH % kAtomN == 0, "TILE_BATCH vs AtomLayout N"); + static_assert(TILE_K % 8 == 0, "GMMA K divisibility (MMA.md)"); + + // ss_op_selector picks atom N from Tile_N only — it does NOT see AtomLayout. + // AtomLayout<_1,_2,_1> splits BATCH across WGs, so each WG's N is + // TILE_BATCH/kAtomN. Select the op from that per-WG tile; otherwise + // MMA_64x128 on TILE_BATCH=128 with 2 WGs along N OOBs SMEM (atom N=128 + // but each WG owns only 64). Same for M with AtomLayout<_2,_1,_1>. + using WgTileShape = cute::Shape, cute::Int, cute::Int>; + + using TiledMma = decltype(cute::make_tiled_mma( + cute::GMMA::ss_op_selector(), AtomLayoutMNK{})); + + // One SoT SMEM atom for TMA + GMMA (MMA.md §6). OUT×K and BATCH×K, K-major. + using SmemLayoutAtomA = decltype(gmma_ss_smem_selector, cute::Int>()); + using SmemLayoutAtomB = + decltype(gmma_ss_smem_selector, cute::Int>()); +}; + +} // namespace turbomind::gemm diff --git a/src/turbomind/kernels/gemm/gmma_fp8_sm90.h b/src/turbomind/kernels/gemm/gmma_fp8_sm90.h new file mode 100644 index 0000000000..056fc53dc3 --- /dev/null +++ b/src/turbomind/kernels/gemm/gmma_fp8_sm90.h @@ -0,0 +1,333 @@ +#pragma once + +#include + +#include "cute/arch/mma_sm90.hpp" +#include "cute/arch/mma_sm90_gmma.hpp" +#include "cute/atom/mma_atom.hpp" +#include "cute/atom/mma_traits.hpp" +#include "cute/atom/mma_traits_sm90_gmma.hpp" + +#include "src/turbomind/kernels/core/common.h" +#include "src/turbomind/kernels/core/meta.h" +#include "src/turbomind/kernels/gemm/gmma_bf16_sm90.h" +#include "src/turbomind/kernels/gemm/sm90_utils.h" + +namespace turbomind::gemm { + +/* + * Weight-as-A (WA) blockscaled FP8 GMMA — scale TV contract + * ========================================================= + * + * GMMA CLayout_64xN (MMA.md): thread t = t0 + 4*t1 + 32*t2, value (v0,v1,v2): + * m_gmma = t1 + 16*t2 + 8*v1 // OUT / weight axis after WA + * n_gmma = 2*t0 + v0 + 8*v2 // BATCH / act axis after WA + * + * Per 8-col stripe (v2), CRegisters pack as float[4]: + * [0]=(m0,n0), [1]=(m0,n1), [2]=(m0+8,n0), [3]=(m0+8,n1) + * + * Problem scales (SMEM shapes unchanged vs v3): + * U[row]: act kK scale, dense along problem M (= GMMA-N / BATCH) + * V[0|1]: weight kB scale, sparse along problem N/128 (= GMMA-M / OUT) + * + * WA apply (invert of ScaledGmmaFP8_TN::scale_batch_to_accum): + * sw0 = weight_scale_for(m0) // from V, 128-block along OUT + * sw1 = weight_scale_for(m0+8) // == sw0 when OUT atom is 64-aligned + * sa0 = act_scale[n0] // from U, dense along BATCH + * sa1 = act_scale[n1] + * accum[i] += sw * sa * frag[i] + * + * Load invert vs v3: + * v3 Load_U: dense along GMMA-M (act rows) + * v3 Load_V: sparse 2-wide along GMMA-N (weight 128-blocks) + * WA: sparse weight on GMMA-M; dense act on GMMA-N (per owned columns) + * + * RF: dense act scales grow with atom N (= BATCH). Prefer TILE_M 128→64→32. + */ + +// CuTe traits: weight=GMMA-A (OUT×K), act=GMMA-B (BATCH×K). SMEM atoms via gmma_ss_smem_selector. +template +struct ScaledGmmaFP8Traits { + using ElementA = cutlass::float_e4m3_t; // GMMA-A = weight + using ElementB = cutlass::float_e4m3_t; // GMMA-B = activation + using ElementC = float; + + using TileShape = cute::Shape, cute::Int, cute::Int>; + static constexpr auto MajorA = cute::GMMA::Major::K; + static constexpr auto MajorB = cute::GMMA::Major::K; + + using AtomLayoutMNK = AtomLayoutMNK_; + + static constexpr int kAtomM = cute::size<0>(AtomLayoutMNK{}); + static constexpr int kAtomN = cute::size<1>(AtomLayoutMNK{}); + static_assert(kAtomM * kAtomN >= 1); + static_assert(TILE_OUT % (64 * kAtomM) == 0, "TILE_OUT vs AtomLayout M"); + static_assert(TILE_BATCH % kAtomN == 0, "TILE_BATCH vs AtomLayout N"); + static_assert(TILE_K % 32 == 0, "FP8 GMMA K divisibility"); + + // Atom N from per-WG BATCH (not weight OUT) — same pitfall as GmmaBF16Traits. + using WgTileShape = cute::Shape, cute::Int, cute::Int>; + + using TiledMma = decltype(cute::make_tiled_mma( + cute::GMMA::ss_op_selector(), AtomLayoutMNK{})); + + using SmemLayoutAtomA = decltype(gmma_ss_smem_selector, cute::Int>()); + using SmemLayoutAtomB = + decltype(gmma_ss_smem_selector, cute::Int>()); +}; + +// Hand-rolled SS-TN WGMMA + WA scale-to-accum (CRegisters-aware). +// TILE_OUT → GMMA-M (weight); TILE_BATCH → GMMA-N (act). MAX_OP_N caps atom N (BATCH). +template +struct ScaledGmmaFP8_WA { + + static constexpr auto select_gmma_operation() + { + static_assert(TILE_OUT % (BATCH_M * PIPE_M) == 0); + static_assert(TILE_BATCH % (BATCH_N * PIPE_N) == 0); + + constexpr int M = TILE_OUT / (BATCH_M * PIPE_M); + constexpr int N = TILE_BATCH / (BATCH_N * PIPE_N); + + static_assert(M % 64 == 0); + + using namespace cute::SM90::GMMA; + + if constexpr (N % 256 == 0 && 256 <= MAX_OP_N) { + return type_c>; + } + else if constexpr (N % 224 == 0 && 224 <= MAX_OP_N) { + return type_c>; + } + else if constexpr (N % 192 == 0 && 192 <= MAX_OP_N) { + return type_c>; + } + else if constexpr (N % 160 == 0 && 160 <= MAX_OP_N) { + return type_c>; + } + else if constexpr (N % 128 == 0 && 128 <= MAX_OP_N) { + return type_c>; + } + else if constexpr (N % 96 == 0 && 96 <= MAX_OP_N) { + return type_c>; + } + else if constexpr (N % 64 == 0 && 64 <= MAX_OP_N) { + return type_c>; + } + else if constexpr (N % 32 == 0 && 32 <= MAX_OP_N) { + return type_c>; + } + else if constexpr (N % 16 == 0 && 16 <= MAX_OP_N) { + return type_c>; + } + else if constexpr (N % 8 == 0 && 8 <= MAX_OP_N) { + return type_c>; + } + else { + static_assert(N == 0, "unsupported WA BATCH configuration"); + } + } + + using Operation = typename decltype(select_gmma_operation())::type; + + static constexpr typename cute::MMA_Traits::Shape_MNK OP_Shape{}; + + static constexpr int OP_M = cute::get<0>(OP_Shape); + static constexpr int OP_N = cute::get<1>(OP_Shape); + static constexpr int OP_K = cute::get<2>(OP_Shape); + + static constexpr int ITER_M = TILE_OUT / OP_M / BATCH_M / PIPE_M; + static constexpr int ITER_N = TILE_BATCH / OP_N / BATCH_N / PIPE_N; + + // Weight scales along GMMA-M (OUT): 2 M-halves × sparse V select. + // Act scales along GMMA-N: OP_N/4 floats/thread (2 cols × OP_N/8 stripes). + using FragW = float[2]; + static constexpr int kActScalesPerThread = OP_N / 4; + using FragA = float[kActScalesPerThread]; + + using FragC = typename Operation::CRegisters[PIPE_M][PIPE_N][BATCH_M][BATCH_N]; + using AccumC = FragC[ITER_M][ITER_N]; + + static constexpr int kStepMA = (OP_M * TILE_K) >> 4; + static constexpr int kStepNB = (OP_N * TILE_K) >> 4; + static constexpr int kStepKA = (OP_K) >> 4; + static constexpr int kStepKB = (OP_K) >> 4; + + // Weight 128-block partitioning along OUT (GMMA-M / problem N). + static constexpr int OUTER_M = std::gcd(TILE_OUT, 128); + + // scale_w[2]: V[0], V[1] for the tile's N-straddle. + // pred_W[i]: which V to use for OUT subtile i (OUTER_M chunks). + // act_scales[OP_N/4]: this thread's dense U scales for owned BATCH columns, + // laid out stripe-major: for stripe s, indices [2s]=sa(n0), [2s+1]=sa(n1). + template + __device__ static void scale_batch_to_accum_wa(AccumC& accum_C, + const FragC& frag_C, + const FragW& scale_w, + const FragA& act_scales, + const PredW& pred_W, + int offset_W) + { + PRAGMA_UNROLL + for (int m = 0; m < BATCH_M; ++m) { + // 64-aligned OUT atoms: both M-halves share one weight 128-block. + // Still index pred by offset_W (atom OUT start) for TILE_OUT straddles. + const int wi = offset_W / OUTER_M; + const bool pw = pred_W[wi]; + const float sw0 = pw ? scale_w[1] : scale_w[0]; + const float sw1 = sw0; + + PRAGMA_UNROLL + for (int n = 0; n < BATCH_N; ++n) { + PRAGMA_UNROLL + for (int c = 0, s = 0; c < OP_N; c += 8, ++s) { + const float sa0 = act_scales[2 * s + 0]; + const float sa1 = act_scales[2 * s + 1]; + accum_C[m][n][c / 2 + 0] += sw0 * sa0 * frag_C[m][n][c / 2 + 0]; + accum_C[m][n][c / 2 + 1] += sw0 * sa1 * frag_C[m][n][c / 2 + 1]; + accum_C[m][n][c / 2 + 2] += sw1 * sa0 * frag_C[m][n][c / 2 + 2]; + accum_C[m][n][c / 2 + 3] += sw1 * sa1 * frag_C[m][n][c / 2 + 3]; + } + } + } + } + + __device__ static void warpgroup_wait(int n) + { + if (n == 0) { + cute::warpgroup_wait<0>(); + } + else if (n == 1) { + cute::warpgroup_wait<1>(); + } + else if (n == 2) { + cute::warpgroup_wait<2>(); + } + else if (n == 3) { + cute::warpgroup_wait<3>(); + } + else if (n == 4) { + cute::warpgroup_wait<4>(); + } + else if (n == 5) { + cute::warpgroup_wait<5>(); + } + else if (n == 6) { + cute::warpgroup_wait<6>(); + } + else if (n == 7) { + cute::warpgroup_wait<7>(); + } + } + + template + __device__ static void gmma_batch(SmemIterA& iter_A, SmemIterB& iter_B, FragC& frag_C) + { + constexpr int BATCH_K = TILE_K / OP_K; + PRAGMA_UNROLL + for (int k = 0; k < BATCH_K; ++k) { + PRAGMA_UNROLL + for (int m = 0; m < BATCH_M; ++m) { + PRAGMA_UNROLL + for (int n = 0; n < BATCH_N; ++n) { + wgmma(iter_A, iter_B, frag_C[m][n], k == 0); + iter_B += kStepNB; + } + iter_B -= kStepNB * BATCH_N; + iter_A += kStepMA; + } + iter_A -= kStepMA * BATCH_M; + iter_A += kStepKA; + iter_B += kStepKB; + } + iter_A -= kStepKA * BATCH_K; + iter_B -= kStepKB * BATCH_K; + cute::warpgroup_commit_batch(); + } + + template + __device__ static void gmma_pipe(AccumC& accum_C, + SmemIterA& iter_A, + SmemIterB& iter_B, + FragC& frag_C, + const FragW& scale_w, + const FragA& act_scales, + const PredW& pred_W, + int offset_W) + { + cute::warpgroup_arrive(); + PRAGMA_UNROLL + for (int m = 0; m < PIPE_M; ++m) { + PRAGMA_UNROLL + for (int n = 0; n < PIPE_N; ++n) { + gmma_batch(iter_A, iter_B, frag_C[m][n]); + iter_B += kStepNB * BATCH_N; + } + iter_B -= kStepNB * BATCH_N * PIPE_N; + iter_A += kStepMA * BATCH_M; + } + iter_A -= kStepMA * BATCH_M * PIPE_M; + + int i = 0; + PRAGMA_UNROLL + for (int m = 0; m < PIPE_M; ++m) { + PRAGMA_UNROLL + for (int n = 0; n < PIPE_N; ++n, ++i) { + warpgroup_wait(PIPE_M * PIPE_N - i - 1); + int offset = offset_W + m * BATCH_M * OP_M; + scale_batch_to_accum_wa(accum_C[m][n], frag_C[m][n], scale_w, act_scales, pred_W, offset); + } + } + } + + template + __device__ static void apply(SmemIterA& iter_A, + SmemIterB& iter_B, + FragC& frag_C, + AccumC& accum_C, + const FragW& scale_w, + const FragA& act_scales, + const PredW& pred_W) + { + PRAGMA_UNROLL + for (int m = 0; m < ITER_M; ++m) { + PRAGMA_UNROLL + for (int n = 0; n < ITER_N; ++n) { + int offset_W = m * PIPE_M * BATCH_M * OP_M; + gmma_pipe(accum_C[m][n], iter_A, iter_B, frag_C, scale_w, act_scales, pred_W, offset_W); + iter_B += kStepNB * BATCH_N * PIPE_N; + } + iter_B -= kStepNB * BATCH_N * PIPE_N * ITER_N; + iter_A += kStepMA * BATCH_M * PIPE_M; + } + iter_A -= kStepMA * BATCH_M * PIPE_M * ITER_M; + } + + template + __device__ static void foreach_C(Frag& frag, Func&& func) + { + PRAGMA_UNROLL + for (int i_m = 0; i_m < ITER_M; ++i_m) { + PRAGMA_UNROLL + for (int i_n = 0; i_n < ITER_N; ++i_n) { + PRAGMA_UNROLL + for (int p_m = 0; p_m < PIPE_M; ++p_m) { + PRAGMA_UNROLL + for (int p_n = 0; p_n < PIPE_N; ++p_n) { + PRAGMA_UNROLL + for (int b_m = 0; b_m < BATCH_M; ++b_m) { + PRAGMA_UNROLL + for (int b_n = 0; b_n < BATCH_N; ++b_n) { + int m = ((i_m * PIPE_M) + p_m * BATCH_M) + b_m; + int n = ((i_n * PIPE_N) + p_n * BATCH_N) + b_n; + func(frag[i_m][i_n][p_m][p_n][b_m][b_n], m, n); + } + } + } + } + } + } + } +}; + +} // namespace turbomind::gemm diff --git a/src/turbomind/kernels/gemm/iterator_sm70.h b/src/turbomind/kernels/gemm/iterator_sm70.h index fc52463377..f439f3d747 100644 --- a/src/turbomind/kernels/gemm/iterator_sm70.h +++ b/src/turbomind/kernels/gemm/iterator_sm70.h @@ -61,8 +61,8 @@ struct GmemIteratorSm70 { const char* src_data_; - int src_offset_; - int dst_offset_; + int64_t src_offset_; + int dst_offset_; int offset_c_; int offset_s_; @@ -138,7 +138,7 @@ struct GmemIteratorSm70 { } } - const int src_offset = is_indexed ? offsets.x : offsets.x + offsets.y * ld; + const int64_t src_offset = is_indexed ? offsets.x : offsets.x + (int64_t)offsets.y * ld; src_offset_ = src_offset * bitsof / bitsof; @@ -153,12 +153,12 @@ struct GmemIteratorSm70 { for (int s = 0; s < ITER_S; ++s) { const int ss = cta_cs.y + offset_s_ + s * Map::kDeltaS; const int idx = (mat.idxs && pred_(s, 0)) ? __ldg(mat.idxs + ss) : ss; - const auto tmp = data + cs2idx({cta_cs.x, idx}, ld); + const auto tmp = data + cs2idx({cta_cs.x, idx}, ld); src_data_vec_[s] = reinterpret_cast((T*)tmp) + src_offset_; } } else { - auto src_data = data + cs2idx(to_cs(pack(offset)), ld); + auto src_data = data + cs2idx(to_cs(pack(offset)), ld); src_data_ = reinterpret_cast((T*)src_data) + src_offset_; } } diff --git a/src/turbomind/kernels/gemm/iterator_sm80.h b/src/turbomind/kernels/gemm/iterator_sm80.h index a625c68b65..f0480883c0 100644 --- a/src/turbomind/kernels/gemm/iterator_sm80.h +++ b/src/turbomind/kernels/gemm/iterator_sm80.h @@ -43,8 +43,8 @@ struct GmemIteratorSm80 { const char* src_data_; - int src_offset_; - int dst_offset_; + int64_t src_offset_; + int dst_offset_; int offset_c_; int offset_s_; @@ -118,7 +118,7 @@ struct GmemIteratorSm80 { } } - const int src_offset = is_indexed ? offsets.x : offsets.x + offsets.y * ld; + const int64_t src_offset = is_indexed ? offsets.x : offsets.x + (int64_t)offsets.y * ld; src_offset_ = src_offset * bitsof / bitsof; @@ -133,12 +133,12 @@ struct GmemIteratorSm80 { for (int s = 0; s < ITER_S; ++s) { const int ss = cta_cs.y + offset_s_ + s * Map::kDeltaS; const int idx = (mat.idxs && pred_(s, 0)) ? __ldg(mat.idxs + ss) : ss; - const auto tmp = data + cs2idx({cta_cs.x, idx}, ld); + const auto tmp = data + cs2idx({cta_cs.x, idx}, ld); src_data_vec_[s] = reinterpret_cast((T*)tmp) + src_offset_; } } else { - auto src_data = data + cs2idx(to_cs(pack(offset)), ld); + auto src_data = data + cs2idx(to_cs(pack(offset)), ld); src_data_ = reinterpret_cast((T*)src_data) + src_offset_; } diff --git a/src/turbomind/kernels/gemm/kernel.cu b/src/turbomind/kernels/gemm/kernel.cu index b5449ea8cb..ffa3887df5 100644 --- a/src/turbomind/kernels/gemm/kernel.cu +++ b/src/turbomind/kernels/gemm/kernel.cu @@ -153,13 +153,23 @@ std::string Kernel::GetName() const << "_" << desc_.cluster_shape.x << "x" << desc_.cluster_shape.y // << "_" << to_string(desc_.op_class) // << "_" << desc_.mma_tile.x << "x" << desc_.mma_tile.y << "x" << desc_.mma_tile.z; + if (desc_.atom_layout.x) { + ss << "_atom" << desc_.atom_layout.x << "x" << desc_.atom_layout.y << "x" << desc_.atom_layout.z; + } + if (desc_.supports_fused_silu) { + ss << "_fused_silu"; + } if (desc_.group_axis >= 0) { ss << "_" << "mn"[desc_.group_axis] << "group"; } - ss << "_c" << desc_.c_tile.x << "x" << desc_.c_tile.y // - << "_a" << desc_.align.x << "x" << desc_.align.y << "x" << desc_.align.z // - << "_" << desc_.policy_a << desc_.policy_b; + ss << "_c" << desc_.c_tile.x << "x" << desc_.c_tile.y // + << "_a" << desc_.align.x << "x" << desc_.align.y << "x" << desc_.align.z; + if (desc_.algo) { + ss << "_algo" << std::hex << desc_.algo << std::dec // + << "_" << (desc_.raster == kRowMajor ? 'r' : 'c'); + } + ss << "_" << desc_.policy_a << desc_.policy_b; return ss.str(); } @@ -189,6 +199,8 @@ public: const MatrixLayout& Cdesc, void* D, const MatrixLayout& Ddesc, + void* W, + const MatrixLayout& Wdesc, int swizzle, int splits, Workspace& workspace, @@ -209,6 +221,8 @@ public: transpose(Cdesc), D, transpose(Ddesc), + W, + Wdesc, swizzle, splits, workspace, @@ -217,6 +231,10 @@ public: bool is_feasible(const GemmDesc& desc) const noexcept override { + // The fused-SiLU gate pairing is defined along n and does not transpose. + if ((int)desc.epilogue & (int)Epilogue::kGatedSilu) { + return false; + } return kernel_->is_feasible(desc); } diff --git a/src/turbomind/kernels/gemm/kernel.h b/src/turbomind/kernels/gemm/kernel.h index 1c2349eb96..0f101893b1 100644 --- a/src/turbomind/kernels/gemm/kernel.h +++ b/src/turbomind/kernels/gemm/kernel.h @@ -41,6 +41,8 @@ class Kernel { const MatrixLayout& Cdesc, void* D, const MatrixLayout& Ddesc, + void* W, + const MatrixLayout& Wdesc, int swizzle, int splits, Workspace& workspace, diff --git a/src/turbomind/kernels/gemm/kernel/sm70_884_16.cu b/src/turbomind/kernels/gemm/kernel/sm70_884_16.cu index 5ac5d77cf1..fb9e2ada24 100644 --- a/src/turbomind/kernels/gemm/kernel/sm70_884_16.cu +++ b/src/turbomind/kernels/gemm/kernel/sm70_884_16.cu @@ -1,7 +1,7 @@ // Copyright (c) OpenMMLab. All rights reserved. #include "src/turbomind/kernels/gemm/arch/config_sm70_s884.h" -#include "src/turbomind/kernels/gemm/registry.h" +#include "src/turbomind/kernels/gemm/registrar.h" #include "src/turbomind/kernels/gemm/types.h" namespace turbomind::gemm { @@ -11,24 +11,25 @@ using namespace cache_policy; using S = cache_policy::Stream; using D = cache_policy::Default; -void Registry::sm70_884_16() -{ +namespace { +Registrar reg([](Collector& c, int /*arch*/) { if constexpr (1) { // clang-format off using C = Config_F16; - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); // clang-format on } -} +}); +} // namespace } // namespace turbomind::gemm diff --git a/src/turbomind/kernels/gemm/kernel/sm70_884_4.cu b/src/turbomind/kernels/gemm/kernel/sm70_884_4.cu index 690aa331b6..458d2fe73b 100644 --- a/src/turbomind/kernels/gemm/kernel/sm70_884_4.cu +++ b/src/turbomind/kernels/gemm/kernel/sm70_884_4.cu @@ -1,7 +1,7 @@ // Copyright (c) OpenMMLab. All rights reserved. #include "src/turbomind/kernels/gemm/arch/config_sm70_s884.h" -#include "src/turbomind/kernels/gemm/registry.h" +#include "src/turbomind/kernels/gemm/registrar.h" #include "src/turbomind/kernels/gemm/types.h" namespace turbomind::gemm { @@ -11,106 +11,107 @@ using namespace cache_policy; using S = cache_policy::Stream; using D = cache_policy::Default; -void Registry::sm70_884_4() -{ +namespace { +Registrar reg([](Collector& c, int /*arch*/) { if constexpr (1) { // clang-format off using C = Config_U4_d; - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); // clang-format on } if constexpr (1) { // clang-format off using C = Config_U4_g; - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); // clang-format on } if constexpr (1) { // clang-format off using C = Config_U4_d; - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); // clang-format on } if constexpr (1) { // clang-format off using C = Config_U4_g; - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); // clang-format on } if constexpr (1) { // clang-format off using C = Config_MXF4; - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); // clang-format on } -} +}); +} // namespace } // namespace turbomind::gemm diff --git a/src/turbomind/kernels/gemm/kernel/sm70_884_8.cu b/src/turbomind/kernels/gemm/kernel/sm70_884_8.cu index 8f9bc9ea87..57b2b2024f 100644 --- a/src/turbomind/kernels/gemm/kernel/sm70_884_8.cu +++ b/src/turbomind/kernels/gemm/kernel/sm70_884_8.cu @@ -1,7 +1,7 @@ // Copyright (c) OpenMMLab. All rights reserved. #include "src/turbomind/kernels/gemm/arch/config_sm70_s884.h" -#include "src/turbomind/kernels/gemm/registry.h" +#include "src/turbomind/kernels/gemm/registrar.h" #include "src/turbomind/kernels/gemm/types.h" namespace turbomind::gemm { @@ -11,18 +11,19 @@ using namespace cache_policy; using S = cache_policy::Stream; using D = cache_policy::Default; -void Registry::sm70_884_8() -{ +namespace { +Registrar reg([](Collector& c, int /*arch*/) { if constexpr (1) { // clang-format off using C = Config_E4M3; - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); // clang-format on } -} +}); +} // namespace } // namespace turbomind::gemm diff --git a/src/turbomind/kernels/gemm/kernel/sm75_16816_16.cu b/src/turbomind/kernels/gemm/kernel/sm75_16816_16.cu index 8bb65038c9..aabc2c9db1 100644 --- a/src/turbomind/kernels/gemm/kernel/sm75_16816_16.cu +++ b/src/turbomind/kernels/gemm/kernel/sm75_16816_16.cu @@ -2,7 +2,7 @@ #include "src/turbomind/kernels/gemm/arch.h" #include "src/turbomind/kernels/gemm/arch/config_sm75_s16816.h" -#include "src/turbomind/kernels/gemm/registry.h" +#include "src/turbomind/kernels/gemm/registrar.h" #include "src/turbomind/kernels/gemm/types.h" namespace turbomind::gemm { @@ -12,23 +12,24 @@ using namespace cache_policy; using S = cache_policy::Stream; using D = cache_policy::Default; -void Registry::sm75_16816_16() -{ +namespace { +Registrar reg([](Collector& c, int /*arch*/) { if constexpr (1) { // clang-format off using C = Config_F16; - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); // clang-format on } +}); } } // namespace turbomind::gemm diff --git a/src/turbomind/kernels/gemm/kernel/sm75_16816_4.cu b/src/turbomind/kernels/gemm/kernel/sm75_16816_4.cu index 8c443fc37e..805fc7e244 100644 --- a/src/turbomind/kernels/gemm/kernel/sm75_16816_4.cu +++ b/src/turbomind/kernels/gemm/kernel/sm75_16816_4.cu @@ -1,7 +1,7 @@ // Copyright (c) OpenMMLab. All rights reserved. #include "src/turbomind/kernels/gemm/arch/config_sm75_s16816.h" -#include "src/turbomind/kernels/gemm/registry.h" +#include "src/turbomind/kernels/gemm/registrar.h" #include "src/turbomind/kernels/gemm/types.h" #include @@ -12,41 +12,41 @@ using namespace cache_policy; using S = cache_policy::Stream; using D = cache_policy::Default; -void Registry::sm75_16816_4() -{ - auto register_u4_d = [this](auto group_size_tag) { +namespace { +Registrar reg([](Collector& c, int /*arch*/) { + auto register_u4_d = [&c](auto group_size_tag) { constexpr int kGroupSize = decltype(group_size_tag)::value; // clang-format off using C = Config_U4_d; - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); // clang-format on }; register_u4_d(std::integral_constant{}); register_u4_d(std::integral_constant{}); - auto register_u4_g = [this](auto group_size_tag) { + auto register_u4_g = [&c](auto group_size_tag) { constexpr int kGroupSize = decltype(group_size_tag)::value; // clang-format off using C = Config_U4_g; - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); // clang-format on }; @@ -56,14 +56,15 @@ void Registry::sm75_16816_4() if constexpr (1) { // clang-format off using C = Config_MXF4; - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); // clang-format on } +}); } } // namespace turbomind::gemm diff --git a/src/turbomind/kernels/gemm/kernel/sm75_16816_8.cu b/src/turbomind/kernels/gemm/kernel/sm75_16816_8.cu index 147e33317f..7b157d3c20 100644 --- a/src/turbomind/kernels/gemm/kernel/sm75_16816_8.cu +++ b/src/turbomind/kernels/gemm/kernel/sm75_16816_8.cu @@ -2,7 +2,7 @@ #include "src/turbomind/kernels/gemm/arch.h" #include "src/turbomind/kernels/gemm/arch/config_sm75_s16816.h" -#include "src/turbomind/kernels/gemm/registry.h" +#include "src/turbomind/kernels/gemm/registrar.h" #include "src/turbomind/kernels/gemm/types.h" namespace turbomind::gemm { @@ -12,21 +12,22 @@ using namespace cache_policy; using S = cache_policy::Stream; using D = cache_policy::Default; -void Registry::sm75_16816_8() -{ +namespace { +Registrar reg([](Collector& c, int /*arch*/) { if constexpr (1) { // clang-format off using Cg = Config_E4M3; - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); // clang-format on } +}); } } // namespace turbomind::gemm diff --git a/src/turbomind/kernels/gemm/kernel/sm80_16816_16.cu b/src/turbomind/kernels/gemm/kernel/sm80_16816_16.cu index 89d61c17f5..ad46754c44 100644 --- a/src/turbomind/kernels/gemm/kernel/sm80_16816_16.cu +++ b/src/turbomind/kernels/gemm/kernel/sm80_16816_16.cu @@ -2,7 +2,7 @@ #include "src/turbomind/kernels/gemm/arch.h" #include "src/turbomind/kernels/gemm/arch/config_sm80_s16816.h" -#include "src/turbomind/kernels/gemm/registry.h" +#include "src/turbomind/kernels/gemm/registrar.h" #include "src/turbomind/kernels/gemm/types.h" namespace turbomind::gemm { @@ -12,49 +12,50 @@ using namespace cache_policy; using S = cache_policy::Stream; using D = cache_policy::Default; -void Registry::sm80_16816_16() -{ +namespace { +Registrar reg([](Collector& c, int /*arch*/) { if constexpr (1) { // clang-format off using C = Config_F16_g; - Add>(); - Add>(); // 10 - Add>(); - Add>(); // 6 - Add>(); - Add>(); - Add>(); // 2 - Add>(); - Add>(); // * - Add>(); - Add>(); // 4 - Add>(); - Add>(); - Add>(); // 10 - Add>(); + c.add>(); + c.add>(); // 10 + c.add>(); + c.add>(); // 6 + c.add>(); + c.add>(); + c.add>(); // 2 + c.add>(); + c.add>(); // * + c.add>(); + c.add>(); // 4 + c.add>(); + c.add>(); + c.add>(); // 10 + c.add>(); // clang-format on } if constexpr (1) { // clang-format off using C = Config_F16_g; - Add>(); - Add>(); // 10 - Add>(); - Add>(); // 6 - Add>(); - Add>(); - Add>(); // 2 - Add>(); - Add>(); // * - Add>(); - Add>(); // 4 - Add>(); - Add>(); - Add>(); // 10 - Add>(); + c.add>(); + c.add>(); // 10 + c.add>(); + c.add>(); // 6 + c.add>(); + c.add>(); + c.add>(); // 2 + c.add>(); + c.add>(); // * + c.add>(); + c.add>(); // 4 + c.add>(); + c.add>(); + c.add>(); // 10 + c.add>(); // clang-format on } +}); } } // namespace turbomind::gemm diff --git a/src/turbomind/kernels/gemm/kernel/sm80_16816_4.cu b/src/turbomind/kernels/gemm/kernel/sm80_16816_4.cu index 190e7beabb..34b6b1c2a9 100644 --- a/src/turbomind/kernels/gemm/kernel/sm80_16816_4.cu +++ b/src/turbomind/kernels/gemm/kernel/sm80_16816_4.cu @@ -2,7 +2,7 @@ #include "src/turbomind/kernels/gemm/arch.h" #include "src/turbomind/kernels/gemm/arch/config_sm80_s16816.h" -#include "src/turbomind/kernels/gemm/registry.h" +#include "src/turbomind/kernels/gemm/registrar.h" #include "src/turbomind/kernels/gemm/types.h" #include @@ -13,68 +13,68 @@ using namespace cache_policy; using S = cache_policy::Stream; using D = cache_policy::Default; -void Registry::sm80_16816_4() -{ - auto register_u4_d = [this](auto group_size_tag) { +namespace { +Registrar reg([](Collector& c, int /*arch*/) { + auto register_u4_d = [&c](auto group_size_tag) { constexpr int kGroupSize = decltype(group_size_tag)::value; // clang-format off using C = Config_U4_d; - // Add>(); // 0/0 - Add>(); // 30/3 - Add>(); // --/20 - Add>(); // --/13 - Add>(); // 21/13 - Add>(); // 6/6 - - Add>(); // --/3 - Add>(); // 13/13 - Add>(); // 14/10 - Add>(); // 2/2 - - Add>(); // --/21 - Add>(); // 27/13 - Add>(); // 8/5 - Add>(); // 7/5 - Add>(); // 6/7 - Add>(); - - Add>(); // 1/1 - Add>(); // 1/1 - Add>(); // 4/4 - Add>(); - - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); + // c.add>(); // 0/0 + c.add>(); // 30/3 + c.add>(); // --/20 + c.add>(); // --/13 + c.add>(); // 21/13 + c.add>(); // 6/6 + + c.add>(); // --/3 + c.add>(); // 13/13 + c.add>(); // 14/10 + c.add>(); // 2/2 + + c.add>(); // --/21 + c.add>(); // 27/13 + c.add>(); // 8/5 + c.add>(); // 7/5 + c.add>(); // 6/7 + c.add>(); + + c.add>(); // 1/1 + c.add>(); // 1/1 + c.add>(); // 4/4 + c.add>(); + + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); // clang-format on }; register_u4_d(std::integral_constant{}); register_u4_d(std::integral_constant{}); - auto register_u4_g = [this](auto group_size_tag) { + auto register_u4_g = [&c](auto group_size_tag) { constexpr int kGroupSize = decltype(group_size_tag)::value; // clang-format off using C = Config_U4_g; - Add>(); // 10 + 5 + 4 + 10 + 10, 37 - Add>(); // 1 + 6 + 4 + 4 + 2, 3 - Add>(); // 7 + 4 + 6 + 2 + 4, 26 - Add>(); // 18 - Add>(); // 2 - Add>(); // 1 + 2 + 2 + 2 + 2, 2 - Add>(); // 9 - Add>(); // 22 - Add>(); // 8 - Add>(); // 1 + 13 + 9 + 13 + 7, 7 - Add>(); // 12 + 2 + 6 + 2 + 8, 42 + c.add>(); // 10 + 5 + 4 + 10 + 10, 37 + c.add>(); // 1 + 6 + 4 + 4 + 2, 3 + c.add>(); // 7 + 4 + 6 + 2 + 4, 26 + c.add>(); // 18 + c.add>(); // 2 + c.add>(); // 1 + 2 + 2 + 2 + 2, 2 + c.add>(); // 9 + c.add>(); // 22 + c.add>(); // 8 + c.add>(); // 1 + 13 + 9 + 13 + 7, 7 + c.add>(); // 12 + 2 + 6 + 2 + 8, 42 // clang-format on }; @@ -84,26 +84,27 @@ void Registry::sm80_16816_4() if constexpr (1) { // clang-format off using Cd = Config_MXF4; - // Add>(); + // c.add>(); using Cg = Config_MXF4; - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); using C8 = Config_MXF4; - Add>(); - Add>(); - Add>(); - Add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); // clang-format on } +}); } } // namespace turbomind::gemm diff --git a/src/turbomind/kernels/gemm/kernel/sm80_16816_8.cu b/src/turbomind/kernels/gemm/kernel/sm80_16816_8.cu index 1ea37e7ffd..ca013143cd 100644 --- a/src/turbomind/kernels/gemm/kernel/sm80_16816_8.cu +++ b/src/turbomind/kernels/gemm/kernel/sm80_16816_8.cu @@ -2,7 +2,7 @@ #include "src/turbomind/kernels/gemm/arch.h" #include "src/turbomind/kernels/gemm/arch/config_sm80_s16816.h" -#include "src/turbomind/kernels/gemm/registry.h" +#include "src/turbomind/kernels/gemm/registrar.h" #include "src/turbomind/kernels/gemm/types.h" namespace turbomind::gemm { @@ -12,30 +12,31 @@ using namespace cache_policy; using S = cache_policy::Stream; using D = cache_policy::Default; -void Registry::sm80_16816_8() -{ +namespace { +Registrar reg([](Collector& c, int /*arch*/) { if constexpr (1) { // clang-format off using Cd = Config_E4M3; - // Add>(); + // c.add>(); using Cg = Config_E4M3; - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); using C8 = Config_E4M3; - Add>(); - Add>(); - Add>(); + c.add>(); + c.add>(); + c.add>(); // clang-format on } +}); } } // namespace turbomind::gemm diff --git a/src/turbomind/kernels/gemm/kernel/sm90_16816_16.cu b/src/turbomind/kernels/gemm/kernel/sm90_16816_16.cu deleted file mode 100644 index 2252fc80f5..0000000000 --- a/src/turbomind/kernels/gemm/kernel/sm90_16816_16.cu +++ /dev/null @@ -1,60 +0,0 @@ -// Copyright (c) OpenMMLab. All rights reserved. - -#include "src/turbomind/kernels/gemm/arch.h" -#include "src/turbomind/kernels/gemm/arch/config_sm80_s16816.h" -#include "src/turbomind/kernels/gemm/registry.h" -#include "src/turbomind/kernels/gemm/types.h" - -namespace turbomind::gemm { - -using namespace sm80_s16816; -using namespace cache_policy; -using S = cache_policy::Stream; -using D = cache_policy::Default; - -void Registry::sm90_16816_16() -{ - if constexpr (1) { - // clang-format off - using C = Config_F16_g; - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - // clang-format on - } - - if constexpr (1) { - // clang-format off - using C = Config_F16_g; - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - // clang-format on - } -} - -} // namespace turbomind::gemm diff --git a/src/turbomind/kernels/gemm/kernel/sm90_16816_4.cu b/src/turbomind/kernels/gemm/kernel/sm90_16816_4.cu index 0981875da0..d64630c46b 100644 --- a/src/turbomind/kernels/gemm/kernel/sm90_16816_4.cu +++ b/src/turbomind/kernels/gemm/kernel/sm90_16816_4.cu @@ -2,7 +2,7 @@ #include "src/turbomind/kernels/gemm/arch.h" #include "src/turbomind/kernels/gemm/arch/config_sm80_s16816.h" -#include "src/turbomind/kernels/gemm/registry.h" +#include "src/turbomind/kernels/gemm/registrar.h" #include "src/turbomind/kernels/gemm/types.h" #include @@ -13,68 +13,68 @@ using namespace cache_policy; using S = cache_policy::Stream; using D = cache_policy::Default; -void Registry::sm90_16816_4() -{ - auto register_u4_d = [this](auto group_size_tag) { +namespace { +Registrar reg([](Collector& c, int /*arch*/) { + auto register_u4_d = [&c](auto group_size_tag) { constexpr int kGroupSize = decltype(group_size_tag)::value; // clang-format off using C = Config_U4_d; - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - - Add>(); - Add>(); - Add>(); - Add>(); - - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - - Add>(); - Add>(); - Add>(); - Add>(); - - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + + c.add>(); + c.add>(); + c.add>(); + c.add>(); + + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + + c.add>(); + c.add>(); + c.add>(); + c.add>(); + + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); // clang-format on }; register_u4_d(std::integral_constant{}); register_u4_d(std::integral_constant{}); - auto register_u4_g = [this](auto group_size_tag) { + auto register_u4_g = [&c](auto group_size_tag) { constexpr int kGroupSize = decltype(group_size_tag)::value; // clang-format off using C = Config_U4_g; - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); // clang-format on }; @@ -84,26 +84,27 @@ void Registry::sm90_16816_4() if constexpr (1) { // clang-format off using Cd = Config_MXF4; - // Add>(); + // c.add>(); using Cg = Config_MXF4; - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); using C8 = Config_MXF4; - Add>(); - Add>(); - Add>(); - Add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); // clang-format on } +}); } } // namespace turbomind::gemm diff --git a/src/turbomind/kernels/gemm/kernel/sm90_16816_8.cu b/src/turbomind/kernels/gemm/kernel/sm90_16816_8.cu index a7f285a772..ffa9e009e9 100644 --- a/src/turbomind/kernels/gemm/kernel/sm90_16816_8.cu +++ b/src/turbomind/kernels/gemm/kernel/sm90_16816_8.cu @@ -2,7 +2,7 @@ #include "src/turbomind/kernels/gemm/arch.h" #include "src/turbomind/kernels/gemm/arch/config_sm80_s16816.h" -#include "src/turbomind/kernels/gemm/registry.h" +#include "src/turbomind/kernels/gemm/registrar.h" #include "src/turbomind/kernels/gemm/types.h" namespace turbomind::gemm { @@ -12,30 +12,31 @@ using namespace cache_policy; using S = cache_policy::Stream; using D = cache_policy::Default; -void Registry::sm90_16816_8() -{ +namespace { +Registrar reg([](Collector& c, int /*arch*/) { if constexpr (1) { // clang-format off using Cd = Config_E4M3; - // Add>(); + // c.add>(); using Cg = Config_E4M3; - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); - Add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); + c.add>(); using C8 = Config_E4M3; - Add>(); - Add>(); - Add>(); + c.add>(); + c.add>(); + c.add>(); // clang-format on } +}); } } // namespace turbomind::gemm diff --git a/src/turbomind/kernels/gemm/kernel/sm90_64n16_16.cu b/src/turbomind/kernels/gemm/kernel/sm90_64n16_16.cu new file mode 100644 index 0000000000..53c3c6cb75 --- /dev/null +++ b/src/turbomind/kernels/gemm/kernel/sm90_64n16_16.cu @@ -0,0 +1,191 @@ +// Copyright (c) OpenMMLab. All rights reserved. + +#include + +// We need modifiable TMA, which is added in 12.3 +#if (__CUDACC_VER_MAJOR__ > 12 || (__CUDACC_VER_MAJOR__ >= 12 && __CUDACC_VER_MINOR__ >= 3)) + +#include "src/turbomind/kernels/gemm/arch.h" +#include "src/turbomind/kernels/gemm/gemm_universal_sm90_bf16.h" +#include "src/turbomind/kernels/gemm/kernel_impl_sm90_bf16.h" +#include "src/turbomind/kernels/gemm/sm90_bf16_traits.h" +#include "src/turbomind/kernels/gemm/types.h" + +#include "src/turbomind/kernels/gemm/registrar.h" + +namespace turbomind::gemm { + +namespace { + +// Registers one SM90 BF16 GMMA kernel: cluster (1,1), no multicast. Grouped-ness follows +// striding: dense (kFlat) is ungrouped, MoE (kIndexed / kBlocked) is grouped. Tile_ carries +// CuTe AtomLayoutMNK (2x1 along OUT, 1x2 along BATCH). +template +void add_kernel(Collector& c) +{ + constexpr bool grouped = striding != Striding::kFlat; + c.add(std::make_unique< + KernelImplSm90Bf16>>()); +} + +Registrar reg([](Collector& c, int /*arch*/) { + // Catalog pruned per full-suite scan tmp/sm90_bf16_scan5; refs refreshed per + // tmp/sm90_bf16_scan8 (2026-07-26, H200, TP/EP 1/2/4/8, swizzle 0-3). `refs: N` = + // dispatch records (tuned selections) the kernel accumulated across the scan; + // `// unused` entries had zero refs and are kept visible for re-enabling. + // Only cluster (1,1) was ever selected; (2,1) / (1,2) variants were dropped entirely. + + // --- Dense (kFlat), row raster --- + // N128 tiles never selected. + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + add_kernel(c); // refs: 274 + add_kernel(c); // refs: 1267 + add_kernel(c); // refs: 21 + add_kernel(c); // refs: 235 + add_kernel(c); // refs: 268 + add_kernel(c); // refs: 295 + add_kernel(c); // refs: 1160 + + // --- Dense (kFlat), col raster: never selected --- + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + + // --- Dense N128 AtomLayout 1x2 (either raster): never selected --- + // add_kernel(c); // unused + // add_kernel(c); // unused + + // --- MoE gate_up (kIndexed), col raster --- + add_kernel(c); // refs: 1436 + // add_kernel(c); // unused + // add_kernel(c); // unused + add_kernel(c); // refs: 846 + add_kernel(c); // refs: 183 + add_kernel(c); // refs: 665 + // add_kernel(c); // unused + // add_kernel(c); // unused + add_kernel(c); // refs: 2207 + add_kernel(c); // refs: 399 + add_kernel(c); // refs: 1362 + add_kernel(c); // refs: 1035 + add_kernel(c); // refs: 495 + add_kernel(c); // refs: 803 + add_kernel(c); // refs: 604 + add_kernel(c); // refs: 2420 + + // --- MoE gate_up (kIndexed), row raster --- + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + add_kernel(c); // refs: 1200 + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + + // --- MoE gate_up N128 AtomLayout 1x2 / N256 AtomLayout 2x1 (either raster): never selected --- + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + + // --- MoE down (kBlocked), col raster --- + add_kernel(c); // refs: 2637 + add_kernel(c); // refs: 1419 + add_kernel(c); // refs: 1375 + // add_kernel(c); // unused + add_kernel(c); // refs: 740 + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + + // --- MoE down (kBlocked), row raster: never selected --- + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + + // --- MoE down N128 AtomLayout 1x2 (either raster): never selected --- + // add_kernel(c); // unused + // add_kernel(c); // unused + + // --- Weight L2 evict-first hint (l2_hint_w=1, desc policy_b=1), 1-CTA tiles only --- + // On the 2-CTA small tiles (8/16/32x128, 16x256) the hint measured 5-9% worse on H200. + // Only indexed 256x128 was selected; everything else pruned per the scan. + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + add_kernel(c); // refs: 1543 + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused + // add_kernel(c); // unused +}); +} // namespace + +} // namespace turbomind::gemm + +#else + +// CUDA too old for modifiable TMA: no SM90 GMMA kernels from this TU. + +#endif diff --git a/src/turbomind/kernels/gemm/kernel/sm90_64n32_8.cu b/src/turbomind/kernels/gemm/kernel/sm90_64n32_8.cu index 3299293ad5..12ab0aedf3 100644 --- a/src/turbomind/kernels/gemm/kernel/sm90_64n32_8.cu +++ b/src/turbomind/kernels/gemm/kernel/sm90_64n32_8.cu @@ -1,46 +1,184 @@ #include -#include "src/turbomind/kernels/gemm/registry.h" - // We need modifiable TMA, which is added in 12.3 #if (__CUDACC_VER_MAJOR__ > 12 || (__CUDACC_VER_MAJOR__ >= 12 && __CUDACC_VER_MINOR__ >= 3)) #include "src/turbomind/kernels/gemm/arch.h" +#include "src/turbomind/kernels/gemm/gemm_universal_sm90_fp8_wa.h" #include "src/turbomind/kernels/gemm/gemm_universal_sm90_v3.h" -#include "src/turbomind/kernels/gemm/gemm_universal_sm90_v5.h" #include "src/turbomind/kernels/gemm/kernel_impl_sm90.h" +#include "src/turbomind/kernels/gemm/sm90_fp8_wa_traits.h" +#include "src/turbomind/kernels/gemm/sm90_v3_traits.h" #include "src/turbomind/kernels/gemm/types.h" +#include "src/turbomind/kernels/gemm/registrar.h" + namespace turbomind::gemm { -void Registry::sm90_64n32_8() +namespace { + +// Registers one SM90 v3 (FP8 act-as-A) GMMA kernel. Grouped-ness follows striding: +// dense (kFlat) is ungrouped, MoE (kIndexed / kBlocked) is grouped. +template +void add_v3(Collector& c) { - Add(std::make_unique>>()); - Add(std::make_unique>>()); - Add(std::make_unique>>()); + constexpr bool grouped = striding != Striding::kFlat; + c.add(std::make_unique>>()); +} - Add(std::make_unique>>()); - Add(std::make_unique>>()); - Add(std::make_unique>>()); +// Registers one SM90 FP8 weight-as-A GMMA kernel. TILE_M < 64 dense tiles have no +// (2,1) variant: the U multicast box (kBoxU*4/mcA) is not 128B-aligned. +template +void add_wa(Collector& c) +{ + constexpr bool grouped = striding != Striding::kFlat; + c.add( + std::make_unique>>()); +} - Add(std::make_unique>>()); - Add(std::make_unique>>()); - Add(std::make_unique>>()); +Registrar reg([](Collector& c, int /*arch*/) { + // Catalog pruned per full-suite scans tmp/sm90_fp8wa_scan1 + tmp/sm90_v3_scan1 + // (2026-07-26, H200, TP/EP 1/2/4/8, swizzle 0-3; both FP8 types exercise the same + // kernel pool, refs combined). `refs: N` = dispatch records (tuned selections); + // `// unused` entries had zero refs and are kept visible for re-enabling. + // --- Dense (v3 act-as-A), row raster --- + add_v3(c); // refs: 2319 + add_v3(c); // refs: 1020 + add_v3(c); // refs: 1813 + add_v3(c); // refs: 4648 + add_v3(c); // refs: 86 + add_v3(c); // refs: 1331 + add_v3(c); // refs: 2133 + add_v3(c); // refs: 76 + add_v3(c); // refs: 323 - Add(std::make_unique>>()); - Add(std::make_unique>>()); - Add(std::make_unique>>()); -} + // --- Dense grouped (kBlocked), col raster --- + add_v3(c); // refs: 1247 + add_v3(c); // refs: 777 + add_v3(c); // refs: 99 -} // namespace turbomind::gemm + // --- Indexed (gate/up), col raster: N192 non-fused; N256 supports both epilogues --- + add_v3(c); // refs: 173 + add_v3(c); // refs: 99 + add_v3(c); // refs: 37 + add_v3(c); // refs: 6903 + add_v3(c); // refs: 49 + add_v3(c); // refs: 196 + add_v3(c); // refs: 3824 + add_v3(c); // refs: 67 + add_v3(c); // refs: 16 -#else + // --- Weight-as-A FP8, dense (kFlat), row raster --- + // OUT=128: plain bf16 epilogue. TILE_M<64: no (2,1) (TMA 128B alignment). + add_wa(c); // refs: 5602 + add_wa(c); // refs: 266 + add_wa(c); // refs: 1745 + add_wa(c); // refs: 66 + add_wa(c); // refs: 1516 + add_wa(c); // refs: 145 + add_wa(c); // refs: 1403 + add_wa(c); // refs: 46 + add_wa(c); // refs: 211 + // OUT=256: plain bf16 and fused SiLU->FP8 epilogues. + add_wa(c); // refs: 1556 + add_wa(c); // refs: 111 + add_wa(c); // refs: 329 + add_wa(c); // refs: 25 + add_wa(c); // refs: 221 + // add_wa(c); // unused + add_wa(c); // refs: 13 + // add_wa(c); // unused + // add_wa(c); // unused -namespace turbomind::gemm { + // --- Weight-as-A FP8, blocked, col raster --- + add_wa(c); // refs: 37 + // add_wa(c); // unused + // add_wa(c); // unused + add_wa(c); // refs: 36 + // add_wa(c); // unused + // add_wa(c); // unused + add_wa(c); // refs: 63 + // add_wa(c); // unused + // add_wa(c); // unused + add_wa(c); // refs: 10 + // add_wa(c); // unused + // add_wa(c); // unused + add_wa(c); // refs: 1377 + add_wa(c); // refs: 2 + // add_wa(c); // unused + add_wa(c); // refs: 439 + // add_wa(c); // unused + // add_wa(c); // unused + add_wa(c); // refs: 87 + // add_wa(c); // unused + // add_wa(c); // unused + // add_wa(c); // unused + // add_wa(c); // unused + // add_wa(c); // unused + + // --- Weight-as-A FP8, indexed (gate/up), col raster --- + add_wa(c); // refs: 570 + // add_wa(c); // unused + add_wa(c); // refs: 20 + add_wa(c); // refs: 157 + // add_wa(c); // unused + // add_wa(c); // unused + add_wa(c); // refs: 96 + // add_wa(c); // unused + // add_wa(c); // unused + add_wa(c); // refs: 36 + // add_wa(c); // unused + // add_wa(c); // unused + // Indexed OUT=256 kernels also select the epilogue at runtime. + add_wa(c); // refs: 4553 + add_wa(c); // refs: 32 + add_wa(c); // refs: 4 + add_wa(c); // refs: 1253 + add_wa(c); // refs: 7 + // add_wa(c); // unused + add_wa(c); // refs: 340 + // add_wa(c); // unused + // add_wa(c); // unused + // add_wa(c); // unused + // add_wa(c); // unused + // add_wa(c); // unused -void Registry::sm90_64n32_8() {} + // --- 2 math WGs (per-WG BATCH=64): 128x128 plain, 128x256 fused, 64x256_n2 --- + add_wa(c); + add_wa(c); + add_wa(c); + add_wa(c); + add_wa(c); + add_wa(c); + add_wa(c); + add_wa(c); + add_wa(c); + add_wa(c); + add_wa(c); + add_wa(c); + add_wa(c); + add_wa(c); + add_wa(c); + add_wa(c); + add_wa(c); + add_wa(c); + add_wa(c); + add_wa(c); + add_wa(c); + add_wa(c); + add_wa(c); + add_wa(c); + add_wa(c); + add_wa(c); + add_wa(c); +}); +} // namespace } // namespace turbomind::gemm +#else + +// CUDA too old for modifiable TMA: no SM90 GMMA kernels from this TU. + #endif diff --git a/src/turbomind/kernels/gemm/kernel_impl.h b/src/turbomind/kernels/gemm/kernel_impl.h index d9a702d118..841bae06e0 100644 --- a/src/turbomind/kernels/gemm/kernel_impl.h +++ b/src/turbomind/kernels/gemm/kernel_impl.h @@ -123,11 +123,15 @@ class KernelImpl: public Kernel { const MatrixLayout& Cdesc, void* D, const MatrixLayout& Ddesc, + void* W, + const MatrixLayout& Wdesc, int swizzle, int splits, Workspace& workspace, cudaStream_t stream) override { + (void)W; + (void)Wdesc; MatrixLayout Adesc = _Adesc; const int m = Ddesc.rows; diff --git a/src/turbomind/kernels/gemm/kernel_impl_sm90.h b/src/turbomind/kernels/gemm/kernel_impl_sm90.h index ee91f21689..569d7d1030 100644 --- a/src/turbomind/kernels/gemm/kernel_impl_sm90.h +++ b/src/turbomind/kernels/gemm/kernel_impl_sm90.h @@ -2,6 +2,10 @@ #pragma once +#include +#include +#include + #include "cute/util/debug.hpp" #include "src/turbomind/core/check.h" #include "src/turbomind/kernels/core/common.h" @@ -32,7 +36,7 @@ namespace turbomind::gemm { -extern __shared__ char smem_buf[]; +extern __shared__ __align__(1024) char smem_buf[]; template __global__ void __launch_bounds__(Kernel::CTA_SIZE, 1) gemm_kernel_name(const __grid_constant__ CUtensorMap tm_a, @@ -45,6 +49,8 @@ __global__ void __launch_bounds__(Kernel::CTA_SIZE, 1) gemm_kernel_name(const __ const MatrixParam param_U, const MatrixParam param_V, const MatrixParam param_C, + const MatrixParam param_W, + bool fuse_silu, typename Kernel::Scheduler sched, void* tensormap_buf) { @@ -62,6 +68,8 @@ __global__ void __launch_bounds__(Kernel::CTA_SIZE, 1) gemm_kernel_name(const __ param_U, param_V, param_C, + param_W, + fuse_silu, sched, (CUtensorMap*)tensormap_buf, smem_buf); @@ -79,19 +87,34 @@ class KernelImplSm90: public Kernel { static constexpr auto is_grouped_gemm = Gemm::is_grouped_gemm; + struct AlgoBits { + uint32_t family : 8; + uint32_t math_wgs : 8; + uint32_t : 16; + + uint32_t u32() const + { + static_assert(sizeof(AlgoBits) == sizeof(uint32_t)); + uint32_t v; + std::memcpy(&v, this, sizeof(v)); + return v; + } + }; + KernelImplSm90() { desc_.order_a = kRowMajor; // m, k desc_.order_b = kColMajor; // k, n desc_.order_c = kRowMajor; - desc_.type_a = data_type_v; - desc_.type_b = data_type_v; - desc_.type_c = data_type_v; + desc_.type_a = data_type_v; + desc_.type_b = data_type_v; + desc_.type_c = data_type_v; + desc_.supports_fused_silu = Gemm::kSupportsFusedSilu; - desc_.striding_a = {is_grouped_gemm ? Striding::kBlocked : Striding::kFlat}; // IterA::kMode; - desc_.striding_b = {is_grouped_gemm ? Striding::kBlocked : Striding::kFlat}; // IterB::kMode; - desc_.striding_c = {is_grouped_gemm ? Striding::kBlocked : Striding::kFlat}; // Gemm::Epilogue::kMode; + desc_.striding_a = Gemm::kStridingA; + desc_.striding_b = Gemm::kStridingB; + desc_.striding_c = Gemm::kStridingC; desc_.pack_a = {}; // OpA::kPack; desc_.pack_b = {}; // OpB::kPack; @@ -113,7 +136,13 @@ class KernelImplSm90: public Kernel { desc_.policy_a = 0; // (int)IterA::Policy::kEvictPolicy; desc_.policy_b = 0; // (int)IterB::Policy::kEvictPolicy; desc_.c_tile = {TILE_M, TILE_N}; // {Gemm::Epilogue::TM, Gemm::Epilogue::TN}; - desc_.op_class = OpClass::kGMMA_s64n16; + desc_.op_class = OpClass::kGMMA_q64n32; + desc_.raster = Gemm::kRasterOrder; + + AlgoBits algo{}; + algo.family = Gemm::kAlgoFamily; + algo.math_wgs = Gemm::WARPGROUPS; + desc_.algo = algo.u32(); desc_.cluster_shape = {Gemm::Cluster::M, Gemm::Cluster::N}; @@ -160,6 +189,8 @@ class KernelImplSm90: public Kernel { const MatrixLayout& Cdesc, void* D, const MatrixLayout& Ddesc, + void* W, + const MatrixLayout& Wdesc, int swizzle, int splits, Workspace& workspace, @@ -169,9 +200,11 @@ class KernelImplSm90: public Kernel { MatrixLayout Adesc = _Adesc; - [[maybe_unused]] const int m = Ddesc.rows; - [[maybe_unused]] const int n = Ddesc.cols; - [[maybe_unused]] const int k = Adesc.cols; + const int m = Ddesc.rows; + const int n = Ddesc.cols; + const int k = Adesc.cols; + const int num_groups = is_grouped_gemm ? std::max(Adesc.num, 1) : Adesc.num; + const bool fuse_silu = ((int)operation.epilogue & (int)Epilogue::kGatedSilu) != 0; TM_CHECK_GE(cdiv(k, TILE_K), 2) << "The kernel requires at least 2 k-tiles to work"; @@ -189,7 +222,7 @@ class KernelImplSm90: public Kernel { auto sched = [&] { const int2 tiles = get_tiled_shape(m, n, TILE_M, TILE_N); - const int4 shape{m, n, k, Adesc.num}; + const int4 shape{m, n, k, num_groups}; swizzle = Sched::get_log_tile(tiles, 1 << swizzle); @@ -198,7 +231,7 @@ class KernelImplSm90: public Kernel { sched.next_cluster_id_ = TM_CHECK_NOTNULL(workspace.flags); - sched.offsets_ = Adesc.offsets; + sched.offsets_ = nullptr; return sched; }(); @@ -214,8 +247,11 @@ class KernelImplSm90: public Kernel { TM_CUDA_CHECK(cudaMemsetAsync(workspace.flags, 0, sizeof(int), stream)); } - // std::cout << "A: " << Adesc << "\n"; - auto tm_a = make_2d_tma_desc((void*)A, Adesc, {kTileM / kMulticastA, TILE_K}, CU_TENSOR_MAP_SWIZZLE_128B); + // Indexed-A: gather activations in-kernel; host A TMA template unused. + auto tm_a = make_2d_tma_desc(Gemm::kStridingA == Striding::kIndexed ? nullptr : (void*)A, + Adesc, + {kTileM / kMulticastA, TILE_K}, + CU_TENSOR_MAP_SWIZZLE_128B); // std::cout << "B: " << Bdesc << "\n"; auto tm_b = make_2d_tma_desc(Gemm::is_grouped_gemm ? nullptr : (void*)B, @@ -224,11 +260,31 @@ class KernelImplSm90: public Kernel { CU_TENSOR_MAP_SWIZZLE_128B); // std::cout << "C: " << Cdesc << "\n"; - using LayoutC = typename Gemm::LayoutC; - auto tm_c = make_2d_tma_desc((void*)C, Cdesc, {LayoutC::S0, LayoutC::C0}, get_tma_swizzle(Gemm::kSwizzleC)); + auto make_tm_c = [&](auto fused_silu) { + constexpr bool kFuseSilu = decltype(fused_silu)::value; + using Output = typename Gemm::template Output; + using LayoutC = typename Output::LayoutC; + + MatrixLayout Cdesc_tma = Cdesc; + if constexpr (kFuseSilu) { + TM_CHECK_EQ(Cdesc_tma.cols % 2, 0); + Cdesc_tma.cols /= 2; + } + return make_2d_tma_desc( + (void*)C, Cdesc_tma, {LayoutC::S0, LayoutC::C0}, get_tma_swizzle(Output::kSwizzleC)); + }; + + CUtensorMap tm_c; + if constexpr (Gemm::kSupportsFusedSilu) { + tm_c = fuse_silu ? make_tm_c(std::true_type{}) : make_tm_c(std::false_type{}); + } + else { + tm_c = make_tm_c(std::false_type{}); + } CUtensorMap tm_u{}; - if (U) { + // Indexed-A also gathers U by idxs; no U TMA template. + if (U && Gemm::kStridingA != Striding::kIndexed) { // std::cout << "U: " << Udesc << "\n"; tm_u = make_2d_tma_desc((void*)U, Udesc, {Gemm::kBoxU / kMulticastU, 1}, CU_TENSOR_MAP_SWIZZLE_NONE); } @@ -242,6 +298,32 @@ class KernelImplSm90: public Kernel { // tm_v = make_2d_tma_desc((void*)V, Vdesc, {box_v.y, box_v.x}, CU_TENSOR_MAP_SWIZZLE_NONE); } + const auto param_A = to_param((void*)A, Adesc); + const auto param_B = to_param((void*)B, Bdesc); + const auto param_U = to_param((void*)U, Udesc); + const auto param_V = to_param((void*)V, Vdesc); + const auto param_C = to_param((void*)D, Ddesc); + const auto param_W = to_param((void*)W, Wdesc); + + // Grouped: prepare_moe_tma_descs rebases per-expert A/B/U/C before GEMM launch. + if constexpr (is_grouped_gemm) { + sched.offsets_ = Gemm::PrepareTmaDescs(tm_a, + tm_b, + tm_u, + tm_c, + param_A, + param_B, + param_U, + param_C, + fuse_silu, + (CUtensorMap*)workspace.tensormaps, + num_groups, + m, + n, + stream); + TM_CUDA_CHECK(cudaGetLastError()); + } + const int sm_count = sm_count_; static constexpr int cluster_size = Gemm::kClusterSize; @@ -252,7 +334,7 @@ class KernelImplSm90: public Kernel { cudaLaunchConfig_t config{}; config.gridDim = grid; config.blockDim = block; - config.dynamicSmemBytes = info_.dynamic_smem_size; + config.dynamicSmemBytes = Gemm::GetSmemSize(fuse_silu); config.stream = stream; auto func = gemm_kernel_name; @@ -289,11 +371,13 @@ class KernelImplSm90: public Kernel { tm_c, tm_u, tm_v, - to_param((void*)A, Adesc), - to_param((void*)B, Bdesc), - to_param((void*)U, Udesc), - to_param((void*)V, Vdesc), - to_param((void*)D, Ddesc), + param_A, + param_B, + param_U, + param_V, + param_C, + param_W, + fuse_silu, sched, workspace.tensormaps); TM_CUDA_CHECK(ec); @@ -331,14 +415,14 @@ class KernelImplSm90: public Kernel { bool is_feasible(const GemmDesc& desc) const noexcept override { - if (desc.striding_a != desc_.striding_a) { - return false; - } - if (desc.striding_b != desc_.striding_b) { - return false; - } - if (desc.striding_c != desc_.striding_c) { - return false; + const bool fuse_silu = ((int)desc.epilogue & (int)Epilogue::kGatedSilu) != 0; + if (fuse_silu) { + if (!Gemm::kSupportsFusedSilu || desc.type_c != kFloat8_e4m3) { + return false; + } + GemmDesc canonical = desc; + canonical.type_c = desc_.type_c; + return Kernel::is_feasible(canonical); } return Kernel::is_feasible(desc); } diff --git a/src/turbomind/kernels/gemm/kernel_impl_sm90_bf16.h b/src/turbomind/kernels/gemm/kernel_impl_sm90_bf16.h new file mode 100644 index 0000000000..8834d2cdcf --- /dev/null +++ b/src/turbomind/kernels/gemm/kernel_impl_sm90_bf16.h @@ -0,0 +1,379 @@ +// Copyright (c) OpenMMLab. All rights reserved. + +#pragma once + +#include + +#include "src/turbomind/core/check.h" +#include "src/turbomind/kernels/core/common.h" +#include "src/turbomind/kernels/core/data_type.h" +#include "src/turbomind/kernels/gemm/cta_map.h" +#include "src/turbomind/kernels/gemm/desc.h" +#include "src/turbomind/kernels/gemm/kernel.h" +#include "src/turbomind/kernels/gemm/matrix_ptr.h" +#include "src/turbomind/kernels/gemm/types.h" +#include "src/turbomind/kernels/gemm/utils.h" + +#include "src/turbomind/kernels/gemm/tma.h" + +#include "src/turbomind/utils/cuda_utils.h" + +namespace turbomind::gemm { + +extern __shared__ __align__(1024) char smem_buf[]; + +// Distinct symbol from gemm_kernel_sm90 (FP8) to avoid ODR clash. +template +__global__ void __launch_bounds__(Kernel::CTA_SIZE, 1) gemm_kernel_sm90_bf16(const __grid_constant__ CUtensorMap tm_a, + const __grid_constant__ CUtensorMap tm_b, + const __grid_constant__ CUtensorMap tm_c, + const __grid_constant__ CUtensorMap tm_u, + const __grid_constant__ CUtensorMap tm_v, + const MatrixParam param_A, + const MatrixParam param_B, + const MatrixParam param_U, + const MatrixParam param_V, + const MatrixParam param_C, + bool fuse_silu, + typename Kernel::Scheduler sched, + void* tensormap_buf) +{ + +#if __CUDA_ARCH__ + if constexpr (Kernel::Arch::is_compatible(__CUDA_ARCH__)) { + Kernel kernel; + kernel(tm_a, + tm_b, + tm_c, + tm_u, + tm_v, + param_A, + param_B, + param_U, + param_V, + param_C, + fuse_silu, + sched, + (CUtensorMap*)tensormap_buf, + smem_buf); + } +#endif +} + +template +class KernelImplSm90Bf16: public Kernel { +public: + // import frequently used constants + static constexpr int TILE_M = Gemm::TILE_M; + static constexpr int TILE_N = Gemm::TILE_N; + static constexpr int TILE_K = Gemm::TILE_K; + + static constexpr auto is_grouped_gemm = Gemm::is_grouped_gemm; + + struct AlgoBits { + uint32_t family : 8; + uint32_t math_wgs : 8; + uint32_t : 16; + + uint32_t u32() const + { + static_assert(sizeof(AlgoBits) == sizeof(uint32_t)); + uint32_t v; + std::memcpy(&v, this, sizeof(v)); + return v; + } + }; + + KernelImplSm90Bf16() + { + // After SM90 BF16 prepare: B is (K, N) ColMajor view of physical (N, K). + desc_.order_a = kRowMajor; // A: (M, K) + desc_.order_b = kColMajor; // B: (K, N) col-major == (N, K) row-major storage + desc_.order_c = kRowMajor; // C: (M, N) + + desc_.type_a = data_type_v; + desc_.type_b = data_type_v; + desc_.type_c = data_type_v; + + desc_.striding_a = Gemm::kStridingA; + desc_.striding_b = Gemm::kStridingB; + desc_.striding_c = Gemm::kStridingC; + + desc_.pack_a = {}; + desc_.pack_b = {}; + desc_.pack_u = {}; + desc_.pack_v = {}; + + // BF16 path: no FP8 scales / quant metadata + desc_.quant_a = {}; + desc_.quant_b = {}; + + desc_.cta_tile = {TILE_M, TILE_N, TILE_K}; + desc_.mma_tile = {1, 1, 1}; + desc_.atom_layout = {cute::size<0>(typename Gemm::AtomLayoutMNK{}), + cute::size<1>(typename Gemm::AtomLayoutMNK{}), + cute::size<2>(typename Gemm::AtomLayoutMNK{})}; + desc_.supports_fused_silu = Gemm::kSupportsFusedSilu; + + info_.chunk_size_k = Gemm::TILE_K; + + desc_.align.x = 1; + desc_.align.y = 1; + desc_.align.z = 1; + + // desc policies mirror the actual operand cache hints: A (activations) / B (weights) + desc_.policy_a = 0; + desc_.policy_b = Gemm::kL2HintW; + desc_.c_tile = {TILE_M, TILE_N}; + desc_.op_class = OpClass::kGMMA_h64n16; + desc_.raster = Gemm::kRasterOrder; + + AlgoBits algo{}; + algo.family = 1; + algo.math_wgs = Gemm::WARPGROUPS; + desc_.algo = algo.u32(); + + desc_.cluster_shape = {Gemm::Cluster::M, Gemm::Cluster::N}; + + info_.dynamic_smem_size = Gemm::kSmemSize; + + desc_.stages = Gemm::Stages; + desc_.split_k = 1; + desc_.group_axis = is_grouped_gemm ? 0 : -1; + + desc_.arch = Gemm::Arch::value; + + auto func = gemm_kernel_sm90_bf16; + + cudaFuncGetAttributes(&info_.attr, func); + + if (info_.dynamic_smem_size > (48 << 10)) { + cudaFuncSetAttribute(func, cudaFuncAttributeMaxDynamicSharedMemorySize, info_.dynamic_smem_size); + } + + cudaFuncSetAttribute(func, cudaFuncAttributeNonPortableClusterSizeAllowed, 16); + + cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &info_.max_active_ctas, func, Gemm::CTA_SIZE, info_.dynamic_smem_size); + + sm_count_ = getSMCount(); + + info_.name = GetName(); + } + + int Launch(const Operation& operation, + float alpha, + const void* A, + const MatrixLayout& _Adesc, + const void* U, + const MatrixLayout& Udesc, + const void* B, + const MatrixLayout& _Bdesc, + const void* V, + const MatrixLayout& _Vdesc, + float beta, + const void* C, + const MatrixLayout& Cdesc, + void* D, + const MatrixLayout& Ddesc, + void* W, + const MatrixLayout& Wdesc, + int swizzle, + int splits, + Workspace& workspace, + cudaStream_t stream) override + { + (void)W; + (void)Wdesc; + using Sched = typename Gemm::Scheduler; + + MatrixLayout Adesc = _Adesc; + + const int m = Ddesc.rows; + const int n = Ddesc.cols; + const int k = Adesc.cols; + const int num_groups = std::max(Adesc.num, 1); + const bool fuse_silu = ((int)operation.epilogue & (int)Epilogue::kGatedSilu) != 0; + + TM_CHECK_GE(cdiv(k, TILE_K), 2) << "The kernel requires at least 2 k-tiles to work"; + + auto transpose = [](MatrixLayout x) { + std::swap(x.rows, x.cols); + x.order = gemm::transpose(x.order); + return x; + }; + + // (K, N) ColMajor → (N, K) RowMajor for TMA (K-contiguous, SW128) + MatrixLayout Bdesc = transpose(_Bdesc); + MatrixLayout Vdesc = transpose(_Vdesc); + + auto sched = [&] { + const int2 tiles = get_tiled_shape(m, n, TILE_M, TILE_N); + const int4 shape{m, n, k, num_groups}; + + swizzle = Sched::get_log_tile(tiles, 1 << swizzle); + + Sched sched{}; + sched.init(shape, swizzle, {TILE_M, TILE_N, TILE_K}); + + sched.next_cluster_id_ = TM_CHECK_NOTNULL(workspace.flags); + + sched.offsets_ = nullptr; + + return sched; + }(); + + constexpr int kMulticastA = Gemm::kMulticastA; + constexpr int kMulticastB = Gemm::kMulticastB; + + constexpr int kTileM = Gemm::TILE_M; + constexpr int kTileN = Gemm::TILE_N; + + if (Gemm::Scheduler::is_dynamic) { + TM_CUDA_CHECK(cudaMemsetAsync(workspace.flags, 0, sizeof(int), stream)); + } + + // Indexed A: no static full-A TMA; Flat/Blocked use static TMA templates. + // Grouped: the prepare kernel rebases TMA maps and materializes scheduler offsets. + auto tm_a = make_2d_tma_desc(Gemm::kStridingA == Striding::kIndexed ? nullptr : (void*)A, + Adesc, + {kTileM / kMulticastA, TILE_K}, + CU_TENSOR_MAP_SWIZZLE_128B); + + // Grouped B: nullptr template + prepare-kernel rebase. + auto tm_b = make_2d_tma_desc( + is_grouped_gemm ? nullptr : (void*)B, Bdesc, {kTileN / kMulticastB, TILE_K}, CU_TENSOR_MAP_SWIZZLE_128B); + + using LayoutC = typename Gemm::LayoutC; + // Fused SiLU: scheduler walks full weight N; C TMA / output are half-width. + // LlamaLinear: Ddesc.cols = weight.output_dim (full), ld = output.stride (half). + MatrixLayout Cdesc_tma = Cdesc; + if (fuse_silu) { + TM_CHECK_EQ(Cdesc_tma.cols % 2, 0); + Cdesc_tma.cols /= 2; + } + auto tm_c = make_2d_tma_desc((void*)C, Cdesc_tma, {LayoutC::S0, LayoutC::C0}, get_tma_swizzle(Gemm::kSwizzleC)); + + // BF16: no scale tensors + CUtensorMap tm_u{}; + CUtensorMap tm_v{}; + + const auto param_A = to_param((void*)A, Adesc); + const auto param_B = to_param((void*)B, Bdesc); + const auto param_U = to_param((void*)U, Udesc); + const auto param_V = to_param((void*)V, Vdesc); + const auto param_C = to_param((void*)D, Ddesc); + + if constexpr (is_grouped_gemm) { + sched.offsets_ = Gemm::PrepareTmaDescs(tm_a, + tm_b, + tm_c, + param_A, + param_B, + param_C, + (CUtensorMap*)workspace.tensormaps, + num_groups, + m, + n, + stream); + TM_CUDA_CHECK(cudaGetLastError()); + } + + const int sm_count = sm_count_; + + static constexpr int cluster_size = Gemm::kClusterSize; + + // Persistent scheduler: fill all co-resident CTAs (small tiles reach + // max_active_ctas > 1 per SM); clamped to the true co-resident count below. + auto grid = sm_count * info_.max_active_ctas / cluster_size * cluster_size; + const auto block = Gemm::CTA_SIZE; + + cudaLaunchConfig_t config{}; + config.gridDim = grid; + config.blockDim = block; + config.dynamicSmemBytes = info_.dynamic_smem_size; + config.stream = stream; + + auto func = gemm_kernel_sm90_bf16; + + [[maybe_unused]] static bool _ = [&] { + int max_cluster_size = 0; + cudaOccupancyMaxPotentialClusterSize(&max_cluster_size, func, &config); + return false; + }(); + + cudaLaunchAttribute attrs[1]; + + attrs[0].id = cudaLaunchAttributeClusterDimension; + attrs[0].val.clusterDim.x = cluster_size; + attrs[0].val.clusterDim.y = 1; + attrs[0].val.clusterDim.z = 1; + + config.attrs = attrs; + config.numAttrs = std::size(attrs); + + int max_active_cluster{}; + cudaOccupancyMaxActiveClusters(&max_active_cluster, func, &config); + config.gridDim = std::min(config.gridDim.x, max_active_cluster * cluster_size); + + auto ec = cudaLaunchKernelEx(&config, + func, + tm_a, + tm_b, + tm_c, + tm_u, + tm_v, + param_A, + param_B, + param_U, + param_V, + param_C, + fuse_silu, + sched, + workspace.tensormaps); + TM_CUDA_CHECK(ec); + + return 0; + } + + std::array GetWorkspaceSize(int tiles, int splits) const + { + static constexpr bool kSerial = true; + + size_t barriers_size = sizeof(int) * tiles; + size_t partials_size = sizeof(float) * TILE_M * TILE_N * tiles; + + if constexpr (!kSerial) { + barriers_size *= splits; + partials_size *= splits; + } + + return {barriers_size, partials_size}; + } + + int GetMaxSplits(const int4& shape, int swizzle, size_t bsize, size_t psize) const override + { + return 1; + } + + int GetMaxSwizzle(const int4& shape) const override + { + using Map = typename Gemm::Scheduler; + const auto tiles = get_tiled_shape(shape.x, shape.y, TILE_M, TILE_N); + return Map::get_log_tile(tiles, 1 << 10); + } + + bool is_feasible(const GemmDesc& desc) const noexcept override + { + const bool want_fused = ((int)desc.epilogue & (int)Epilogue::kGatedSilu) != 0; + if (want_fused && !Gemm::kSupportsFusedSilu) { + return false; + } + return Kernel::is_feasible(desc); + } + +private: + int sm_count_ = 0; +}; + +} // namespace turbomind::gemm diff --git a/src/turbomind/kernels/gemm/prepare_moe_tma_descs_sm90_fp8.h b/src/turbomind/kernels/gemm/prepare_moe_tma_descs_sm90_fp8.h new file mode 100644 index 0000000000..a0422b9d04 --- /dev/null +++ b/src/turbomind/kernels/gemm/prepare_moe_tma_descs_sm90_fp8.h @@ -0,0 +1,193 @@ +#pragma once + +#include +#include + +#include "src/turbomind/kernels/core/array.h" +#include "src/turbomind/kernels/core/common.h" +#include "src/turbomind/kernels/core/math.h" +#include "src/turbomind/kernels/core/smem.h" +#include "src/turbomind/kernels/gemm/matrix_ptr.h" +#include "src/turbomind/kernels/gemm/sm90_utils.h" +#include "src/turbomind/kernels/gemm/types.h" +#include "src/turbomind/kernels/gemm/utils.h" + +namespace turbomind::gemm { +namespace detail { + +// Device TMA map helpers for FP8 MoE prepare (copy → replace addr/dim → publish). +__device__ __forceinline__ void copy_tma_desc_fp8(CUtensorMap* dst, const CUtensorMap* src, int lane) +{ + constexpr int kWords = (int)(sizeof(CUtensorMap) / sizeof(uint2)); + if (lane < kWords) { + ((uint2*)dst)[lane] = ((const uint2*)src)[lane]; + } +} + +__device__ __forceinline__ void replace_tma_addr_dim(CUtensorMap* desc, void* global_addr, int dim_idx, int dim) +{ + uint32_t uint_ptr = cast_smem_ptr_to_uint(desc); + // clang-format off + asm volatile("tensormap.replace.tile.global_address.shared::cta.b1024.b64 [%0], %1;" ::"r"(uint_ptr), "l"(global_addr)); + if (dim_idx == 0) { + asm volatile("tensormap.replace.tile.global_dim.shared::cta.b1024.b32 [%0], 0, %1;" ::"r"(uint_ptr), "r"(dim)); + } + else { + asm volatile("tensormap.replace.tile.global_dim.shared::cta.b1024.b32 [%0], 1, %1;" ::"r"(uint_ptr), "r"(dim)); + } + // clang-format on +} + +__device__ __forceinline__ void publish_tma_desc_fp8(CUtensorMap* gmem_desc, CUtensorMap* smem_desc) +{ + uint32_t uint_ptr = cast_smem_ptr_to_uint(smem_desc); + // clang-format off + asm volatile("tensormap.cp_fenceproxy.global.shared::cta.tensormap::generic.release.gpu.sync.aligned [%0], [%1], 128;" :: "l"(gmem_desc), "r"(uint_ptr)); + // clang-format on +} + +template +__device__ __forceinline__ void rebase_publish_tma_descs_fp8(CUtensorMap* gmem_out, + CUtensorMap* smem_desc, + Array templates, + Array global_addrs, + Array dims, + Array dim_idxs, + int stride_desc_idx, + uint64_t stride_bytes, + int lane) +{ + PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + copy_tma_desc_fp8(&smem_desc[i], templates[i], lane); + } + __syncwarp(); + if (lane == 0) { + PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + replace_tma_addr_dim(&smem_desc[i], global_addrs[i], dim_idxs[i], dims[i]); + } + replace_tma_global_stride(&smem_desc[stride_desc_idx], stride_bytes); + } + __syncwarp(); + PRAGMA_UNROLL + for (int i = 0; i < N; ++i) { + publish_tma_desc_fp8(&gmem_out[i], &smem_desc[i]); + } + __syncwarp(); +} + +} // namespace detail + +// Per-group MoE TMA prepare for FP8 GMMA. +// Blocked-A kernels accept Blocked or Flat input; Indexed-A kernels accept Indexed, Blocked, or Flat input. +// Blocked-A: workspace [A, B, U, C]. Indexed-A: gather A/U in GEMM; workspace [B, C]. +// Weight B uses either a direct Flat pointer or a StridedPtr table selected by MatrixLayout.ld == 0. +// Output C is bf16 (non-fused) or fp8_e4m3 (fused SiLU + post-quant). +template +__global__ void __launch_bounds__(32, 1) prepare_moe_tma_descs_sm90_fp8(const __grid_constant__ CUtensorMap tm_a, + const __grid_constant__ CUtensorMap tm_b, + const __grid_constant__ CUtensorMap tm_u, + const __grid_constant__ CUtensorMap tm_c, + MatrixParam param_A, + MatrixParam param_B, + MatrixParam param_U, + MatrixParam param_C, + bool fuse_silu, + CUtensorMap* out, + int* offsets, + int M_total, + int N) +{ + const int g = (int)blockIdx.x; + const int lane = (int)threadIdx.x & 31; + + using Ta = __nv_fp8_e4m3; + using Tb = __nv_fp8_e4m3; + using Tu = float; + + const int m0 = param_A.offsets ? __ldg(param_A.offsets + g) : 0; + const int m1 = param_A.offsets ? __ldg(param_A.offsets + g + 1) : M_total; + const int M = m1 - m0; + + if (lane == 0) { + offsets[g] = m0; + if (g + 1 == gridDim.x) { + offsets[g + 1] = m1; + } + } + + if constexpr (kStridingA == Striding::kBlocked) { + constexpr int kNum = 4; + + __shared__ __align__(128) CUtensorMap smem_desc[kNum]; + + const int beg_u = m0 / kAlignmentU * kAlignmentU; + const int end_u = round_up(m1, kAlignmentU); + + CUtensorMap* gmem_out = out + g * kNum; + + Array templates; + templates[0] = &tm_a; + templates[1] = &tm_b; + templates[2] = &tm_u; + templates[3] = &tm_c; + + Array addrs; + addrs[0] = resolve(param_A, g).ptr.ptr; + const auto b = resolve(param_B, g); + addrs[1] = b.ptr.ptr; + addrs[2] = (Tu*)param_U.ptr + beg_u; + addrs[3] = fuse_silu ? resolve<__nv_fp8_e4m3, Striding::kBlocked>(param_C, g).ptr.ptr : + resolve(param_C, g).ptr.ptr; + + Array dims; + dims[0] = M; + dims[1] = N; + dims[2] = end_u - beg_u; + dims[3] = M; + + Array dim_idxs; + dim_idxs[0] = 1; // A + dim_idxs[1] = 1; // B + dim_idxs[2] = 0; // U (scale extent along dim0) + dim_idxs[3] = 1; // C + + detail::rebase_publish_tma_descs_fp8( + gmem_out, smem_desc, templates, addrs, dims, dim_idxs, 1, (uint64_t)b.ptr.stride * sizeof(Tb), lane); + } + else { + // Indexed-A: gather activations + U scales; prepare weight B + output C only. + constexpr int kNum = 2; + + __shared__ __align__(128) CUtensorMap smem_desc[kNum]; + + CUtensorMap* gmem_out = out + g * kNum; + + Array templates; + templates[0] = &tm_b; + templates[1] = &tm_c; + + Array addrs; + const auto b = resolve(param_B, g); + addrs[0] = b.ptr.ptr; + addrs[1] = fuse_silu ? resolve<__nv_fp8_e4m3, Striding::kBlocked>(param_C, g).ptr.ptr : + resolve(param_C, g).ptr.ptr; + + Array dims; + dims[0] = N; + dims[1] = M; + + Array dim_idxs; + dim_idxs[0] = 1; // B + dim_idxs[1] = 1; // C + + (void)tm_a; + (void)tm_u; + (void)param_U; + detail::rebase_publish_tma_descs_fp8( + gmem_out, smem_desc, templates, addrs, dims, dim_idxs, 0, (uint64_t)b.ptr.stride * sizeof(Tb), lane); + } +} + +} // namespace turbomind::gemm diff --git a/src/turbomind/kernels/gemm/registrar.h b/src/turbomind/kernels/gemm/registrar.h new file mode 100644 index 0000000000..f37f86e986 --- /dev/null +++ b/src/turbomind/kernels/gemm/registrar.h @@ -0,0 +1,53 @@ +// Copyright (c) OpenMMLab. All rights reserved. + +#pragma once + +#include +#include +#include +#include + +#include "src/turbomind/kernels/gemm/kernel.h" +#include "src/turbomind/kernels/gemm/kernel_impl.h" + +namespace turbomind::gemm { + +class Collector { +public: + // Matches Registry::Add(): Config has nested ::Kernel + template + void add() + { + kernels_.emplace_back(std::make_unique>()); + } + + void add(std::unique_ptr kernel) + { + kernels_.emplace_back(std::move(kernel)); + } + + std::vector> release() + { + return std::move(kernels_); + } + +private: + std::vector> kernels_; +}; + +using RegisterFn = std::function; + +inline std::vector& gKernelFactories() +{ + static std::vector v; + return v; +} + +struct Registrar { + explicit Registrar(RegisterFn fn) + { + gKernelFactories().push_back(std::move(fn)); + } +}; + +} // namespace turbomind::gemm diff --git a/src/turbomind/kernels/gemm/registry.cu b/src/turbomind/kernels/gemm/registry.cu index 55937a4db0..4454ffa91c 100644 --- a/src/turbomind/kernels/gemm/registry.cu +++ b/src/turbomind/kernels/gemm/registry.cu @@ -1,6 +1,10 @@ // Copyright (c) OpenMMLab. All rights reserved. +#include +#include + #include "src/turbomind/kernels/gemm/arch.h" +#include "src/turbomind/kernels/gemm/registrar.h" #include "src/turbomind/kernels/gemm/registry.h" namespace turbomind::gemm { @@ -8,32 +12,13 @@ namespace turbomind::gemm { Registry::Registry(std::shared_ptr device_prop): device_prop_{std::move(device_prop)}, arch_{device_prop_->major * 100 + device_prop_->minor * 10} { -#if defined(GEMM2_ARCH_90_ENABLED) - sm90_16816_4(); - sm90_16816_8(); - sm90_16816_16(); - sm90_64n32_8(); -#endif - - sm80_16816_4(); - sm80_16816_8(); - sm80_16816_16(); - - sm75_16816_4(); - sm75_16816_8(); - sm75_16816_16(); - - sm70_884_4(); - sm70_884_8(); - sm70_884_16(); - - cublas_float(); - -#if defined(ENABLE_CUBLAS_GROUPED) - if (Sm100::is_compatible(arch_)) { - sm100_cublas_grouped_float(); + for (auto& register_fn : gKernelFactories()) { + Collector collector; + register_fn(collector, arch_); + for (auto& k : collector.release()) { + Add(std::move(k)); + } } -#endif } bool Registry::Add(std::unique_ptr kernel) @@ -44,13 +29,13 @@ bool Registry::Add(std::unique_ptr kernel) is_valid = false; } - // if (is_valid) { - // std::cout << "register: " << kernel->name() // - // << ", shared: " << (kernel->smem_size() >> 10) << " KB" // - // << ", regs: " << kernel->info().attr.numRegs // - // << ", local: " << (float)kernel->info().attr.localSizeBytes << " bytes" // - // << ", max_active_ctas: " << kernel->info().max_active_ctas << " \n"; - // } + if (is_valid && std::getenv("TM_GEMM_DEBUG_OCCUPANCY")) { + std::cout << "register: " << kernel->name() // + << ", shared: " << (kernel->smem_size() >> 10) << " KB" // + << ", regs: " << kernel->info().attr.numRegs // + << ", local: " << (float)kernel->info().attr.localSizeBytes << " bytes" // + << ", max_active_ctas: " << kernel->info().max_active_ctas << " \n"; + } if ((int)device_prop_->sharedMemPerBlockOptin < kernel->smem_size()) { is_valid = false; diff --git a/src/turbomind/kernels/gemm/registry.h b/src/turbomind/kernels/gemm/registry.h index f5cac6c86a..27cdef20d8 100644 --- a/src/turbomind/kernels/gemm/registry.h +++ b/src/turbomind/kernels/gemm/registry.h @@ -2,8 +2,10 @@ #pragma once -#include "src/turbomind/kernels/gemm/kernel_impl.h" #include +#include + +#include "src/turbomind/kernels/gemm/kernel.h" namespace turbomind::gemm { @@ -11,13 +13,6 @@ class Registry { public: explicit Registry(std::shared_ptr device_prop); - /// TODO: remove this - template - [[maybe_unused]] bool Add() - { - return Add(std::make_unique>()); - } - [[nodiscard]] const std::vector& kernels() const { return ptrs_; @@ -26,28 +21,6 @@ class Registry { private: bool Add(std::unique_ptr kernel); - void sm90_16816_4(); - void sm90_16816_8(); - void sm90_16816_16(); - - void sm80_16816_4(); - void sm80_16816_8(); - void sm80_16816_16(); - - void sm75_16816_4(); - void sm75_16816_8(); - void sm75_16816_16(); - - void sm70_884_4(); - void sm70_884_8(); - void sm70_884_16(); - - void sm90_64n32_8(); - - void cublas_float(); - void sm100_cublas_grouped_float(); - -private: std::shared_ptr device_prop_; int arch_; std::vector> kernels_; diff --git a/src/turbomind/kernels/gemm/scaled_gmma_fp8_sm90.h b/src/turbomind/kernels/gemm/scaled_gmma_fp8_sm90.h index 637358b169..06b70cde37 100644 --- a/src/turbomind/kernels/gemm/scaled_gmma_fp8_sm90.h +++ b/src/turbomind/kernels/gemm/scaled_gmma_fp8_sm90.h @@ -12,7 +12,9 @@ namespace turbomind::gemm { -template +// MAX_OP_N caps the WGMMA atom N (largest divisor of TILE_N that is ≤ MAX_OP_N). +// Smaller atoms shrink FragC vs AccumC peak RF (e.g. 192 tile: 64×3 peak≈128 vs 192×1 peak≈192). +template struct ScaledGmmaFP8_TN { static constexpr auto select_gmma_operation() @@ -27,25 +29,25 @@ struct ScaledGmmaFP8_TN { using namespace cute::SM90::GMMA; - if constexpr (N % 256 == 0) { + if constexpr (N % 256 == 0 && 256 <= MAX_OP_N) { return type_c>; } - else if constexpr (N % 224 == 0) { + else if constexpr (N % 224 == 0 && 224 <= MAX_OP_N) { return type_c>; } - else if constexpr (N % 192 == 0) { + else if constexpr (N % 192 == 0 && 192 <= MAX_OP_N) { return type_c>; } - else if constexpr (N % 160 == 0) { + else if constexpr (N % 160 == 0 && 160 <= MAX_OP_N) { return type_c>; } - else if constexpr (N % 128 == 0) { + else if constexpr (N % 128 == 0 && 128 <= MAX_OP_N) { return type_c>; } - else if constexpr (N % 96 == 0) { + else if constexpr (N % 96 == 0 && 96 <= MAX_OP_N) { return type_c>; } - else if constexpr (N % 64 == 0) { + else if constexpr (N % 64 == 0 && 64 <= MAX_OP_N) { return type_c>; } else { @@ -96,18 +98,19 @@ struct ScaledGmmaFP8_TN { scales[1][1] = frag_U[m][1] * frag_V[1]; PRAGMA_UNROLL for (int n = 0; n < BATCH_N; ++n) { + // Step by 8 (GMMA C reg group), not OUTER_N: OP_N may not be a + // multiple of OUTER_N (e.g. OP_N=96, OUTER_N=64), and atom + // offset_V may not be OUTER_N-aligned (second 96-atom at 96), + // so a coarse c0+=OUTER_N chunk can OOB FragC and straddle a + // V-scale boundary at N%128. PRAGMA_UNROLL - for (int c0 = 0; c0 < OP_N; c0 += OUTER_N) { - int i = (offset_V + c0) / OUTER_N; - bool p = pred_V[i]; - PRAGMA_UNROLL - for (int c1 = 0; c1 < OUTER_N; c1 += 8) { - int c = c0 + c1; - accum_C[m][n][c / 2 + 0] += (p ? scales[0][1] : scales[0][0]) * frag_C[m][n][c / 2 + 0]; - accum_C[m][n][c / 2 + 1] += (p ? scales[0][1] : scales[0][0]) * frag_C[m][n][c / 2 + 1]; - accum_C[m][n][c / 2 + 2] += (p ? scales[1][1] : scales[1][0]) * frag_C[m][n][c / 2 + 2]; - accum_C[m][n][c / 2 + 3] += (p ? scales[1][1] : scales[1][0]) * frag_C[m][n][c / 2 + 3]; - } + for (int c = 0; c < OP_N; c += 8) { + const int i = (offset_V + c) / OUTER_N; + const bool p = pred_V[i]; + accum_C[m][n][c / 2 + 0] += (p ? scales[0][1] : scales[0][0]) * frag_C[m][n][c / 2 + 0]; + accum_C[m][n][c / 2 + 1] += (p ? scales[0][1] : scales[0][0]) * frag_C[m][n][c / 2 + 1]; + accum_C[m][n][c / 2 + 2] += (p ? scales[1][1] : scales[1][0]) * frag_C[m][n][c / 2 + 2]; + accum_C[m][n][c / 2 + 3] += (p ? scales[1][1] : scales[1][0]) * frag_C[m][n][c / 2 + 3]; } } } diff --git a/src/turbomind/kernels/gemm/sm90_bf16_traits.h b/src/turbomind/kernels/gemm/sm90_bf16_traits.h new file mode 100644 index 0000000000..8212d8abe5 --- /dev/null +++ b/src/turbomind/kernels/gemm/sm90_bf16_traits.h @@ -0,0 +1,94 @@ +#pragma once + +#include "cute/layout.hpp" +#include "cute/underscore.hpp" + +namespace turbomind::gemm { + +// Tile / pipeline knobs for GemmUniversalSm90_Bf16. One struct = one registered +// kernel shape. WG partition is a CuTe AtomLayoutMNK on the tile: +// Layout> — 2 math WGs along GMMA-M (problem N / OUT) +// Layout> — 2 math WGs along GMMA-N (problem M / BATCH) +// Striding / grouped TMA policy is NOT here — that comes from the kernel's +// (is_grouped, StridingA) template args. +// +// Register budgets are per A-path (TMA vs indexed) and per tile. Pack: 2 math WGs +// → producer + 2*math ≤ 504 (regfile 64K/SM, 128 thr/WG, 8 regs slack). Budgets are +// only tightened where the lower static allocation buys occupancy: +// Small tiles (8–32x128, 16x256): every region ≤ 80 → 2 CTAs/SM (15–33% faster +// at small m; 64x128 loses from co-residency, so it stays at 1 CTA/SM). +// 3 CTAs/SM (56-reg regions, stages 3 for 16/32x128) measured 2–17% WORSE +// (co-residency contention; these kernels are memory-bound) — 2 is the +// sweet spot, 4 CTAs is not worth pursuing. +// All other tiles: (120,192,192) = 1 CTA/SM — deviating from the proven math=192 +// without an occupancy gain only risks spills/codegen regressions. Exceptions +// are measured winners noted per tile below. + +namespace detail { + +template +struct Sm90Bf16TileBase { + static constexpr int TILE_M = TM; + static constexpr int TILE_N = TN; + static constexpr int TILE_K = TK; + using AtomLayoutMNK = AtomLayoutMNK_; + static constexpr int Stages = Stages_; + + static constexpr int kProducerRegsTma = kProdTma; + static constexpr int kMathRegsTma = kMathTma; + static constexpr int kProducerRegsIndexed = kProdIdx; + static constexpr int kMathRegsIndexed = kMathIdx; +}; + +} // namespace detail + +using AtomLayout_1x1 = cute::Layout>; +using AtomLayout_2x1 = cute::Layout>; +using AtomLayout_1x2 = cute::Layout>; + +// N128 × AtomLayout 2x1 (WGs along OUT). 8–32: all region caps ≤ 80 so the +// static allocation allows 2 CTAs/SM (15–33% faster at small m). 64+: 1 CTA/SM +// (64x128 loses 6–11% from 2-CTA co-residency), sum ≤ 504. +using Sm90Bf16Tile_8x128_2x1 = detail::Sm90Bf16TileBase<8, 128, 64, AtomLayout_2x1, 4, 24, 80, 40, 80>; +using Sm90Bf16Tile_16x128_2x1 = detail::Sm90Bf16TileBase<16, 128, 64, AtomLayout_2x1, 4, 24, 80, 40, 80>; +using Sm90Bf16Tile_32x128_2x1 = detail::Sm90Bf16TileBase<32, 128, 64, AtomLayout_2x1, 4, 24, 80, 48, 80>; +using Sm90Bf16Tile_64x128_2x1 = detail::Sm90Bf16TileBase<64, 128, 64, AtomLayout_2x1, 4, 120, 192, 120, 192>; +using Sm90Bf16Tile_96x128_2x1 = detail::Sm90Bf16TileBase<96, 128, 64, AtomLayout_2x1, 4, 120, 192, 120, 192>; +using Sm90Bf16Tile_128x128_2x1 = detail::Sm90Bf16TileBase<128, 128, 64, AtomLayout_2x1, 4, 120, 192, 120, 192>; +// Measured winners: indexed 192x128 (104,200,200) is ~8% faster, 256x128 +// (128,184,184) is ~18% faster than the (120,192,192) baseline. +using Sm90Bf16Tile_192x128_2x1 = detail::Sm90Bf16TileBase<192, 128, 64, AtomLayout_2x1, 4, 120, 192, 104, 200>; +using Sm90Bf16Tile_224x128_2x1 = detail::Sm90Bf16TileBase<224, 128, 64, AtomLayout_2x1, 4, 120, 192, 120, 192>; +using Sm90Bf16Tile_256x128_2x1 = detail::Sm90Bf16TileBase<256, 128, 64, AtomLayout_2x1, 4, 120, 192, 128, 184>; + +// N128 × AtomLayout 1x2 (WGs along BATCH; TILE_M must be legal for 2 WGs) +using Sm90Bf16Tile_64x128_1x2 = detail::Sm90Bf16TileBase<64, 128, 64, AtomLayout_1x2, 4, 120, 192, 120, 192>; +using Sm90Bf16Tile_128x128_1x2 = detail::Sm90Bf16TileBase<128, 128, 64, AtomLayout_1x2, 4, 120, 192, 120, 192>; + +// N256: Stages=3 (SMEM). Fused SiLU: kAtomM == 1 tiles (each WG covers full OUT) +// use the C-ownership R2S/TMA epilogue directly; 8x256_2x1 pairs the [g64|u64] +// gate/up blocks (split across its two WGs) through an f32 smem staging buffer +// (gemm_universal_sm90_bf16.h run_epilogue). 16x256 is 1x2, 8x256 also has a +// single-WG 1x1 form (on 1x2 its per-WG GMMA-N is 4, below the GMMA atom minimum +// of 8); the epilogue STSM atom tiers (gemm_universal_sm90_bf16.h EpiStsmAtoms) +// cover the narrow per-WG GMMA-N (8/16x256: 4 vals/thread). +// 8/16x256: caps ≤ 80 → 2 CTAs/SM; 32x256 is just over the 2-CTA smem budget → ≤ 504. +using Sm90Bf16Tile_8x256_2x1 = detail::Sm90Bf16TileBase<8, 256, 64, AtomLayout_2x1, 3, 24, 80, 40, 80>; +using Sm90Bf16Tile_8x256_1x1 = detail::Sm90Bf16TileBase<8, 256, 64, AtomLayout_1x1, 3, 24, 80, 40, 80>; +using Sm90Bf16Tile_16x256_1x2 = detail::Sm90Bf16TileBase<16, 256, 64, AtomLayout_1x2, 3, 24, 80, 40, 80>; +using Sm90Bf16Tile_64x256_2x1 = detail::Sm90Bf16TileBase<64, 256, 64, AtomLayout_2x1, 3, 120, 192, 120, 192>; +using Sm90Bf16Tile_128x256_2x1 = detail::Sm90Bf16TileBase<128, 256, 64, AtomLayout_2x1, 3, 120, 192, 120, 192>; +using Sm90Bf16Tile_32x256_1x2 = detail::Sm90Bf16TileBase<32, 256, 64, AtomLayout_1x2, 3, 120, 192, 120, 192>; +using Sm90Bf16Tile_64x256_1x2 = detail::Sm90Bf16TileBase<64, 256, 64, AtomLayout_1x2, 3, 120, 192, 120, 192>; +using Sm90Bf16Tile_96x256_1x2 = detail::Sm90Bf16TileBase<96, 256, 64, AtomLayout_1x2, 3, 120, 192, 120, 192>; +using Sm90Bf16Tile_128x256_1x2 = detail::Sm90Bf16TileBase<128, 256, 64, AtomLayout_1x2, 3, 120, 192, 120, 192>; + +} // namespace turbomind::gemm diff --git a/src/turbomind/kernels/gemm/sm90_fp8_wa_traits.h b/src/turbomind/kernels/gemm/sm90_fp8_wa_traits.h new file mode 100644 index 0000000000..191f9d784c --- /dev/null +++ b/src/turbomind/kernels/gemm/sm90_fp8_wa_traits.h @@ -0,0 +1,52 @@ +#pragma once + +namespace turbomind::gemm { + +// Tile / pipeline knobs for GemmUniversalSm90_Fp8Wa (weight-as-A FP8). +// Problem TILE_M = BATCH (act rows) → GMMA-N; TILE_N = OUT (weight) → GMMA-M. +// +// RF policy: dense act scales live on GMMA-N (= TILE_M). Prefer smaller TILE_M +// when math WG RF exceeds setmaxnreg pack (2 WGs → 504). + +namespace detail { + +template +struct Sm90Fp8WaTileBase { + static constexpr int TILE_M = TM; // BATCH + static constexpr int TILE_N = TN; // OUT + static constexpr int TILE_K = 128; + static constexpr int WG_M = WG_M_; // split BATCH across math WGs + static constexpr int WG_N = WG_N_; // split OUT across math WGs + static constexpr int Stages = Stages_; + static constexpr int kMaxOpN = MaxOpN; + + // Small-BATCH tiles need less math RF; 2 math WGs pack to 504 (producer + 2*math). + static constexpr int kProducerRegsTma = 40; + static constexpr int kMathRegsTma = (WG_M_ * WG_N_ == 2) ? 232 : ((TM <= 32) ? 168 : 208); + static constexpr int kProducerRegsIndexed = 88; + static constexpr int kMathRegsIndexed = (WG_M_ * WG_N_ == 2) ? 208 : ((TM <= 32) ? 168 : 208); +}; + +} // namespace detail + +// Small BATCH (RF-friendly / decode-friendly) +using Sm90Fp8WaTile_8x128 = detail::Sm90Fp8WaTileBase<8, 128, 8, 4>; +using Sm90Fp8WaTile_16x128 = detail::Sm90Fp8WaTileBase<16, 128, 16, 4>; +using Sm90Fp8WaTile_32x128 = detail::Sm90Fp8WaTileBase<32, 128, 32, 4>; + +// Mid BATCH +using Sm90Fp8WaTile_64x128 = detail::Sm90Fp8WaTileBase<64, 128, 64, 4>; + +// OUT=256 (gate/up width); Stages=3 for weight SMEM +using Sm90Fp8WaTile_8x256 = detail::Sm90Fp8WaTileBase<8, 256, 8, 3>; +using Sm90Fp8WaTile_16x256 = detail::Sm90Fp8WaTileBase<16, 256, 16, 3>; +using Sm90Fp8WaTile_32x256 = detail::Sm90Fp8WaTileBase<32, 256, 32, 3>; +using Sm90Fp8WaTile_64x256 = detail::Sm90Fp8WaTileBase<64, 256, 64, 3>; + +// 2 math WGs (per-WG BATCH=64). 128x* split BATCH (WG_M=2); 64x256_n2 splits OUT +// (WG_N=2, gate/up staged through smem in the fused epilogue). +using Sm90Fp8WaTile_128x128 = detail::Sm90Fp8WaTileBase<128, 128, 64, 3, 2, 1>; +using Sm90Fp8WaTile_128x256 = detail::Sm90Fp8WaTileBase<128, 256, 64, 3, 2, 1>; +using Sm90Fp8WaTile_64x256_n2 = detail::Sm90Fp8WaTileBase<64, 256, 64, 3, 1, 2>; + +} // namespace turbomind::gemm diff --git a/src/turbomind/kernels/gemm/sm90_utils.h b/src/turbomind/kernels/gemm/sm90_utils.h index 1b576569e5..07dc8fe29e 100644 --- a/src/turbomind/kernels/gemm/sm90_utils.h +++ b/src/turbomind/kernels/gemm/sm90_utils.h @@ -2,6 +2,7 @@ #pragma once +#include "cute/arch/copy_sm90_desc.hpp" #include "cute/arch/mma_sm90_gmma.hpp" #include "cute/atom/mma_traits.hpp" #include "cute/atom/mma_traits_sm90_gmma.hpp" @@ -16,6 +17,24 @@ namespace turbomind::gemm { namespace GMMA = cute::SM90::GMMA; +namespace detail { + +__device__ __forceinline__ void replace_tma_global_stride(CUtensorMap* desc, uint64_t stride_bytes) +{ + uint32_t uint_ptr = cast_smem_ptr_to_uint(desc); + // CUDA < 12.5 omits the four zero alignment bits; CUDA >= 12.5 takes bytes. +#if ((__CUDACC_VER_MAJOR__ > 12) || ((__CUDACC_VER_MAJOR__ == 12) && (__CUDACC_VER_MINOR__ >= 5))) + const uint64_t encoded_stride = stride_bytes; +#else + const uint64_t encoded_stride = stride_bytes >> 4; +#endif + asm volatile("tensormap.replace.tile.global_stride.shared::cta.b1024.b64 [%0], 0, %1;" + : + : "r"(uint_ptr), "l"(encoded_stride)); +} + +} // namespace detail + inline __device__ cute::GmmaDescriptor make_smem_desc(void* smem_ptr, int layout_type) { auto uint_ptr = cast_smem_ptr_to_uint(smem_ptr); diff --git a/src/turbomind/kernels/gemm/sm90_v3_traits.h b/src/turbomind/kernels/gemm/sm90_v3_traits.h new file mode 100644 index 0000000000..6c56340828 --- /dev/null +++ b/src/turbomind/kernels/gemm/sm90_v3_traits.h @@ -0,0 +1,65 @@ +#pragma once + +namespace turbomind::gemm { + +// Tile / pipeline knobs for GemmUniversalSm90_v3. One struct = one registered +// kernel shape. Striding / grouped TMA policy is NOT here — that comes from +// the kernel's (is_grouped, StridingA) template args. +// +// Register budgets are per A-path (TMA vs indexed). Kernel selects via +// kIndexedGather. setmaxnreg (PTX): +// producer: .dec → imm ≤ current max (release down to imm) +// math: .inc → imm ≥ current max (request up to imm) +// imm ∈ [24,256], multiple of 8 +// Pack: 2 math WGs → producer + 2*math = 504; 1 math WG → producer + math ≤ 512. + +struct Sm90V3Tile_128x192 { + static constexpr int TILE_M = 128; + static constexpr int TILE_N = 192; + static constexpr int TILE_K = 128; + static constexpr int WG_M = 2; + static constexpr int WG_N = 1; + static constexpr int Stages = 4; + static constexpr int kMaxOpN = 192; + + static constexpr int kProducerRegsTma = 40; + static constexpr int kMathRegsTma = 232; // (504-40)/2 + static constexpr int kProducerRegsIndexed = 88; + static constexpr int kMathRegsIndexed = 208; // (504-88)/2 +}; + +struct Sm90V3Tile_128x256 { + static constexpr int TILE_M = 128; + static constexpr int TILE_N = 256; + static constexpr int TILE_K = 128; + static constexpr int WG_M = 2; + static constexpr int WG_N = 1; + static constexpr int Stages = 3; // SMEM: N=256 needs Stages≤3 + static constexpr int kMaxOpN = 128; + + static constexpr int kProducerRegsTma = 40; + static constexpr int kMathRegsTma = 232; + static constexpr int kProducerRegsIndexed = 88; + static constexpr int kMathRegsIndexed = 208; +}; + +// CTA 64×256: single math WG (WG_TILE = CTA tile). Stages=4 fits (smaller A/C). +struct Sm90V3Tile_64x256 { + static constexpr int TILE_M = 64; + static constexpr int TILE_N = 256; + static constexpr int TILE_K = 128; + static constexpr int WG_M = 1; + static constexpr int WG_N = 1; + static constexpr int Stages = 4; + static constexpr int kMaxOpN = 128; + + // Math RF model (kMaxOpN=128 → atom 64×128, ITER_N=2): + // AccumC 128 + FragC 64 peak ≈ 192, +U/V/ctrl ≈ 208. + // Producer .dec must be ≤ ptxas temporal (~192); 256 is illegal. + static constexpr int kProducerRegsTma = 40; + static constexpr int kMathRegsTma = 208; + static constexpr int kProducerRegsIndexed = 88; + static constexpr int kMathRegsIndexed = 208; +}; + +} // namespace turbomind::gemm diff --git a/src/turbomind/kernels/gemm/test/gemm_bench.cu b/src/turbomind/kernels/gemm/test/gemm_bench.cu deleted file mode 100644 index f6e2a624de..0000000000 --- a/src/turbomind/kernels/gemm/test/gemm_bench.cu +++ /dev/null @@ -1,87 +0,0 @@ -// Copyright (c) OpenMMLab. All rights reserved. - -#include "nvbench/main.cuh" -#include "src/turbomind/kernels/gemm/operand.h" -#include "src/turbomind/kernels/gemm/test/models.h" -#include "src/turbomind/kernels/gemm/test/testbed.h" -#include -#include -#include -#include - -void gemm_bench(nvbench::state& state) -{ - const auto idx = state.get_int64("idx"); - - const auto bs = state.get_int64("bs"); - const auto tp = state.get_int64("tp"); - - const auto expert_num = state.get_int64("e_num"); - const auto exp_per_tok = state.get_int64("e_tok"); - - auto [output_dims, input_dims] = config[idx]; - - constexpr int group_size = 128; - - if (idx % 4 == 0 || idx % 4 == 2) { - if (output_dims % tp) - return; - output_dims /= tp; - } - else { - if (input_dims % tp) - return; - input_dims /= tp; - } - - if (input_dims % group_size) - return; - - using turbomind::gemm::get_test; - - { - int m = bs; - int n = output_dims; - int k = input_dims; - if (get_test().kBatchDim == 1) { - std::swap(m, n); - } - std::cerr << "m" << m << "n" << n << "k" << k << "\n"; - - get_test().Initialize(m, n, k, group_size, expert_num, exp_per_tok, state.get_cuda_stream()); - } - - state.add_element_count(get_test().get_element_count()); - - // state.collect_dram_throughput(); - // state.collect_l2_hit_rates(); - - if constexpr (1) { - state.add_global_memory_reads(get_test().get_global_memory_reads()); - get_test().Run(); - state.exec(nvbench::exec_tag::sync, [&](nvbench::launch&) { // - get_test().Run(); - }); - } - else { - state.add_global_memory_reads(get_test().get_ref_global_memory_reads()); - state.exec(nvbench::exec_tag::sync, [&](nvbench::launch&) { // - get_test().RunCublas(); - }); - } - - get_test().ctx_.reset(); -} - -NVBENCH_BENCH(gemm_bench) - .add_int64_axis("idx", nvbench::range(0, (int)config.size() - 1)) - .add_int64_power_of_two_axis("bs", nvbench::range(0, 14)) - .add_int64_axis("tp", {1, 2, 4}) - .add_int64_axis("e_num", {0}) - .add_int64_axis("e_tok", {1}); - -int main(int argc, char* argv[]) -{ - NVBENCH_MAIN_BODY(argc, argv); - return 0; -} diff --git a/src/turbomind/kernels/gemm/test/models.h b/src/turbomind/kernels/gemm/test/models.h deleted file mode 100644 index b4fca0fbe9..0000000000 --- a/src/turbomind/kernels/gemm/test/models.h +++ /dev/null @@ -1,29 +0,0 @@ -// Copyright (c) OpenMMLab. All rights reserved. - -#pragma once - -#include -#include -#include -#include - -static const std::vector> config{ - {11008 * 2, 4096}, {4096, 11008}, {12288, 4096}, {4096, 4096}, // llama2-7b - {14336 * 2, 4096}, {4096, 14336}, {6144, 4096}, {4096, 4096}, // llama3-8b / internlm2.5-7b - {16384 * 2, 6144}, {6144, 16384}, {8192, 6144}, {6144, 6144}, // internlm2-20b - {13696 * 2, 4096}, {4096, 13696}, {4608, 4096}, {4096, 4096}, // glm4-9b - {18944 * 2, 3584}, {3584, 18944}, {4608, 3584}, {3584, 3584}, // qwen2-7b - {20480 * 2, 7168}, {7168, 20480}, {9216, 7168}, {7168, 7168}, // yi-34b - {28672 * 2, 8192}, {8192, 28672}, {10240, 8192}, {8192, 8192}, // llama2-70b / llama3-70b - {29696 * 2, 8192}, {8192, 29696}, {10240, 8192}, {8192, 8192}, // qwen2-72b-instruct-awq - {14336 * 2, 4096}, {4096, 14336}, {6144, 4096}, {4096, 4096}, // mixtral-8x7b, E8e2 - {16384 * 2, 6144}, {6144, 16384}, {0, 0}, {0, 0}, // mixtral-8x22b, E8e2 - {1536 * 2, 5120}, {5120, 1536}, {0, 0}, {0, 0}, // deepseek-v2, E160e6 - {1536 * 2, 2048}, {2048, 1536}, {0, 0}, {0, 0}, // deepseek-v2-lite, E64e6 - {2560 * 2, 3840}, {3840, 2560}, {0, 0}, {0, 0}, // qwen2-a14b, E64e8 - {6400 * 2, 4096}, {4096, 6400}, {0, 0}, {0, 0}, // phi-3.5-MoE, E16e2 -}; - -// static const std::map> moe_config{{32, {8, 2}}, {33, {8, 2}}}; - -// {29568 * 2, 8192}, {8192, 29568}, {10240, 8192}, {8192, 8192}, // qwen2-72b diff --git a/src/turbomind/kernels/gemm/test/test_gemm_v2.cc b/src/turbomind/kernels/gemm/test/test_gemm_v2.cc deleted file mode 100644 index 93f3308c25..0000000000 --- a/src/turbomind/kernels/gemm/test/test_gemm_v2.cc +++ /dev/null @@ -1,101 +0,0 @@ - - -#include "src/turbomind/core/context.h" -#include "src/turbomind/core/data_type.h" - -#include "testbed_v3.h" - -using namespace turbomind; - -struct TestParameter: Testbed_v3::Parameter { - TestParameter(DataType dtype, DataType wtype, DataType itype, int group_size = 128): Testbed_v3::Parameter{} - { - data_type = dtype; - weight_type = wtype; - input_type = itype; - - this->group_size = group_size; - } -}; - -int main() -{ - auto stream = core::Stream::create(); - - core::ContextGuard ctx{stream, core::Allocator{kCPU}, core::Allocator{stream, false}}; - // clang-format off - // TestParameter p{kHalf, kUint4 , kHalf, 128}; - // TestParameter p{kHalf, kFloat4_e2m1, kHalf, 32}; - // TestParameter p{kHalf, kFloat8_e4m3, kHalf, 128}; - // TestParameter p{kHalf, kHalf , kHalf}; - - // TestParameter p{kBfloat16, kBfloat16 , kBfloat16}; - // TestParameter p{kBfloat16, kFloat8_e4m3, kFloat8_e4m3, 128}; - TestParameter p{kBfloat16, kFloat8_e4m3, kBfloat16 , 128}; - // TestParameter p{kBfloat16, kFloat4_e2m1, kBfloat16 , 32}; - // clang-format on - - // p.input_dim = 512; - // p.output_dim = 1024; - // p.max_batch_size = 256; - - // p.input_dim = 1024; - // p.output_dim = 1024; - // p.max_batch_size = 1024; - - // p.input_dim = 12288; - // p.output_dim = 16384; - // p.max_batch_size = 8192; - - // p.expert_num = 1; - // p.experts_per_token = 1; - - // p.input_dim = 2880; - // p.output_dim = 2880; - // p.max_batch_size = 64; - - // p.input_dim = 7168; - // p.output_dim = 4096; - // p.max_batch_size = 16384; - // p.expert_num = 256; - // p.experts_per_token = 8; - - // Qwen3-MoE - p.expert_num = 128; - p.experts_per_token = 8; - // 30B - // p.input_dim = 2048; - // p.output_dim = 768 * 2; - // 235B - // p.input_dim = 4096; - // p.output_dim = 1536 * 2; - // 480B - p.input_dim = 6144; - p.output_dim = 2560 * 2; - - p.max_batch_size = 256; - - // p.input_dim = 16384; - // p.output_dim = 16384; - // p.max_batch_size = 16384; - - // p.input_dim = 2880; - // p.output_dim = 5760; - // p.max_batch_size = 16384; - // p.expert_num = 32; - // p.experts_per_token = 4; - - // p.input_dim = 128; - // p.output_dim = 32; - // p.max_batch_size = 1; - - Testbed_v3 test{p}; - - test.GetReference(); - test.Run(); - test.Compare(); - - cudaDeviceSynchronize(); - - return 0; -} diff --git a/src/turbomind/kernels/gemm/test/testbed_v3.h b/src/turbomind/kernels/gemm/test/testbed_v3.h deleted file mode 100644 index d633baa18b..0000000000 --- a/src/turbomind/kernels/gemm/test/testbed_v3.h +++ /dev/null @@ -1,520 +0,0 @@ - -#include - -#include "src/turbomind/core/allocator.h" -#include "src/turbomind/core/core.h" - -#include "src/turbomind/core/tensor.h" -#include "src/turbomind/kernels/gemm/moe_utils_v2.h" -#include "src/turbomind/kernels/gemm/test/reference.h" -#include "src/turbomind/kernels/gemm/test/test_utils.h" -#include "src/turbomind/kernels/gemm/types.h" -#include "src/turbomind/kernels/quantization.h" -#include "src/turbomind/utils/cuda_utils.h" - -#include "src/turbomind/models/linear_weight.h" -#include "src/turbomind/models/llama/LlamaLinear.h" - -#include "src/turbomind/kernels/gpt_kernels.h" - -namespace turbomind { - -using std::vector; -using std::unique_ptr; - -using Linear = LlamaLinear; - -using namespace gemm; - -struct Parameter { - int input_dim; - int output_dim; - - DataType data_type; - DataType weight_type; - DataType input_type; - - int group_size; - - int max_batch_size; - - int expert_num; - int experts_per_token; - bool combine_experts; -}; - -/// TODO: add a generic copy / casting for non-sub-byte Tensor -static Tensor CopyTransposed(const Tensor& src, Tensor out = {}) -{ - if (out) { - TM_CHECK(out.shapes(0, 1) == src.shapes(1, 0)) << src << " vs " << out; - TM_CHECK_EQ(out.dtype(), src.dtype()); - } - else { - out = {{src.shape(1), src.shape(0)}, src.dtype(), src.device()}; - } - - auto invoke = [&](auto t) { - using T = decltype(t); - TM_SCOPE_CALL(invokeTransposeAxis01( - (T*)out.raw_data(), (T*)src.raw_data(), src.shape(0), src.shape(1), 1, core::Context::stream().handle())); - }; - - const int bits = byte_size(src.dtype(), 8); - if (bits == 8) { - invoke(uint8_t{}); - } - else if (bits == 16) { - invoke(uint16_t{}); - } - else if (bits == 32) { - invoke(int{}); - } - else { - TM_CHECK(0) << "Not implemented. bits = " << bits; - } - - return out; -} - -/// Link individual expert weights into a batched block view for fused MoE. -static void LinkExperts(std::function experts, int n, LinearWeight& d) -{ - const auto& e0 = *experts(0); - - e0.copy_metadata_to(d); - - d.k_desc.num = d.q_desc.num = n; - - if (e0.bias()) { - d.bias() = Tensor{{n, e0.output_dim}, e0.bias().dtype(), kDEVICE}; - } - - std::vector> weights; - std::vector> scales; - - for (int i = 0; i < n; ++i) { - auto& e = *experts(i); - weights.emplace_back(e.weight().raw_data(), e.k_desc.ld); - if (e.scales()) { - scales.emplace_back(e.scales().raw_data(), e.q_desc.ld); - } - if (e.bias()) { - Copy(e.bias(), d.bias().slice(i, 1).squeeze(0)); - } - } - - auto stream = core::Context::stream().handle(); - - if (d.weight_format.dtype == kFloat8_e4m3 && d.input_dtype() == kFloat8_e4m3) { - auto make_blocked_ptr = [&](const auto& ptrs) { - return std::shared_ptr{gemm::MakeBlockedPtrs(ptrs, stream), [](auto p) { cudaFree(p); }}; - }; - d.weight() = Tensor{make_blocked_ptr(weights), {n}, e0.weight().dtype(), kDEVICE}; - d.scales() = Tensor{make_blocked_ptr(scales), {n}, e0.scales().dtype(), kDEVICE}; - d.k_desc.offsets = d.q_desc.offsets = (int*)1; - } - else { - auto make_strided_ptr = [&](const auto& ptrs) { - return std::shared_ptr{gemm::MakeStridedPtrs(ptrs, stream), [](auto p) { cudaFree(p); }}; - }; - d.weight() = Tensor{make_strided_ptr(weights), {n}, d.weight_format.dtype, kDEVICE}; - if (e0.scales()) { - d.scales() = Tensor{make_strided_ptr(scales), {n}, e0.scales().dtype(), kDEVICE}; - } - d.k_desc.ld = d.q_desc.ld = 0; - } -} - -struct Testbed_v3: Parameter { - - Testbed_v3(const Parameter& param): Parameter{param}, stream_{core::Context::stream().handle()}, linear_{} - { - rng_.set_stream(stream_); - ref_.set_stream(stream_); - - if (auto str = std::getenv("TM_GEMM_IMPORT")) { - import_file_ = str; - std::ifstream ifs(import_file_, std::ios::binary); - auto n = linear_.Import(ifs); - std::cout << "Records imported: " << n << "\n"; - } - if (auto str = std::getenv("TM_GEMM_TUNE"); str && import_file_.empty()) { - tuning_ = true; - std::cout << "Enable tuning\n"; - } - if (auto str = std::getenv("TM_GEMM_EXPORT"); str && import_file_.empty()) { - export_file_ = str; - } - - cudaGetDeviceProperties(&prop_, 0); - - w_original_ = std::make_unique(); - w_quant_ = std::make_unique(); - w_dequant_ = std::make_unique(); - - for (int i = 0; i < expert_num; ++i) { - e_original_.push_back(std::make_unique()); - e_quant_.push_back(std::make_unique()); - e_dequant_.push_back(std::make_unique()); - } - - GenerateWeight(); - GenerateInput(); - - if (expert_num) { - LinkExperts([&](int i) { return e_original_[i].get(); }, expert_num, *w_original_); - LinkExperts([&](int i) { return e_quant_[i].get(); }, expert_num, *w_quant_); - LinkExperts([&](int i) { return e_dequant_[i].get(); }, expert_num, *w_dequant_); - Route(); - } - } - - ~Testbed_v3() - { - if (!export_file_.empty()) { - std::cerr << "export file: " << export_file_ << "\n"; - std::ofstream ofs(export_file_, std::ios::binary); - if (ofs.is_open()) { - auto n = linear_.Export(ofs); - std::cout << "Records exported: " << n << "\n"; - } - } - } - - void GenerateInput() - { - x_original_ = Tensor{{max_batch_size, input_dim}, data_type, kDEVICE}; - rng_.NormalFloat(x_original_, 1., 1.); - - if (input_type == data_type) { - x_quant_ = empty_like(x_original_); - x_dequant_ = empty_like(x_original_); - Copy(x_original_, x_quant_); - Copy(x_original_, x_dequant_); - } - else if (input_type == kFloat8_e4m3) { - QuantizeSymm(x_quant_, x_scale_, x_original_, stream_); - DequantizeSymm(x_dequant_, x_quant_, x_scale_, stream_); - } - else { - TM_CHECK(0) << "Not implemented for input type " << to_string(input_type); - } - } - - void Route() - { - const int bsz = max_batch_size; - - std::mt19937 g{}; - - /// TODO: Control the distribution - auto expert_ids = SampleUniform(bsz, expert_num, experts_per_token, g); - - std::uniform_real_distribution dist(1e-3, 1.f); - - Buffer_ tmp(experts_per_token, kCPU); - Buffer_ scales(bsz * experts_per_token, kCPU); - - for (int i = 0; i < bsz; ++i) { - float sum{}; - for (auto& x : tmp) { - x = dist(g); - sum += x; - } - for (int e = 0; e < experts_per_token; ++e) { - scales[e * bsz + i] = tmp[e] / sum; - } - } - - vector count(expert_num); - vector> f2i(expert_num); - for (int i = 0; i < (int)expert_ids.size(); ++i) { - ++count[expert_ids[i]]; - f2i[expert_ids[i]].push_back(i); - } - - Buffer_ offsets(expert_num + 1, kCPU); - offsets[0] = 0; - for (int i = 0; i < expert_num; ++i) { - offsets[i + 1] = offsets[i] + count[i]; - } - - for (const auto& x : count) { - std::cout << x << " "; - } - std::cout << "\n"; - - Buffer_ f2n(expert_ids.size(), kCPU); - Buffer_ en2f(expert_ids.size(), kCPU); - for (int e = 0, i = 0; e < expert_num; ++e) { - for (auto x : f2i[e]) { - f2n[i] = x / experts_per_token; - int en = x % experts_per_token * bsz + x / experts_per_token; - en2f[en] = i; - ++i; - } - } - - f2n_ = {f2n.size(), kDEVICE}; - Copy(f2n, f2n_); - - en2f_ = {en2f.size(), kDEVICE}; - Copy(en2f, en2f_); - - scales_ = {scales.size(), kDEVICE}; - Copy(scales, scales_); - - offsets_ = {offsets.size(), kDEVICE}; - Copy(offsets, offsets_); - h_offsets_ = offsets; - - core::Context::stream().Sync(); - } - - void GenerateWeight() - { - if (expert_num) { - for (int i = 0; i < expert_num; ++i) { - GenerateWeight(*e_original_[i], *e_quant_[i], *e_dequant_[i]); - } - } - else { - GenerateWeight(*w_original_, *w_quant_, *w_dequant_); - } - } - - // - quantize weight - // - dequantize weight - void GenerateWeight(LinearWeight& original, LinearWeight& quant, LinearWeight& dequant) - { - auto make_cfg = [&](DataType wt) -> core::LinearConfig { - core::LinearConfig cfg; - cfg.input_dim = input_dim; - cfg.output_dim = output_dim; - cfg.data_type = data_type; - cfg.format = ResolveLinearWeightFormat(data_type, wt, group_size, 1); - cfg.has_bias = false; - return cfg; - }; - - new (&original) LinearWeight(make_cfg(data_type)); - original.param("weight").alloc({(size_t)input_dim, (size_t)output_dim}, data_type); - rng_.NormalFloat(original.weight(), 1., .1); - - new (&quant) LinearWeight(make_cfg(weight_type)); - quant.param("weight").alloc({(size_t)input_dim, (size_t)output_dim}, weight_type); - new (&dequant) LinearWeight(make_cfg(data_type)); - dequant.param("weight").alloc({(size_t)input_dim, (size_t)output_dim}, data_type); - - Buffer_ rbits; - // rbits = {original.weight().size(), kDEVICE}; - // rng_.RandomBytes(Tensor{rbits}); - - /// Weights are allocated in MN-major, but some quantization requires K-major tensor - - if (weight_type == data_type) { - Copy(original.weight(), quant.weight()); - Copy(original.weight(), dequant.weight()); - } - else if (weight_type == kFloat8_e4m3) { - QuantizeSymmBlock(quant.weight(), quant.scales(), original.weight(), stream_); - DequantizeSymmBlock(dequant.weight(), quant.weight(), quant.scales(), stream_); - } - else if (weight_type == kUint4) { - /// Weights are allocated in (M,N), quantization needs K-major tensor - QuantizeGroupwise(quant.weight().t(), - quant.scales().t(), - quant.zeros().t(), - dequant.weight().t(), - original.weight().t(), - {}, - group_size); - } - else if (weight_type == kFloat4_e2m1) { - QuantizeGroupwise(quant.weight().t(), // - quant.scales().t(), - {}, - dequant.weight().t(), - original.weight().t(), - rbits, - group_size); - } - else { - TM_CHECK(0); - } - - original.prepare(); - quant.prepare(); - dequant.prepare(); - } - - void GetReference() - { - if (expert_num) { - GetReference(x_original_, e_original_, d_original_); - GetReference(x_dequant_, e_dequant_, d_dequant_); - } - else { - GetReference(x_original_, w_original_, d_original_); - GetReference(x_dequant_, w_dequant_, d_dequant_); - } - } - - void GetReference(const Tensor& x, const unique_ptr& dense, Ref d_) - { - auto& d = d_.get(); - if (!d) { - d = Tensor{{x.shape(0), dense->output_dim}, x.dtype(), x.device()}; - } - /// TODO: refactor reference API - const MatrixLayout desc_A{x.dtype(), kRowMajor, (int)x.shape(0), (int)x.shape(1), (int)x.stride(0)}; // m,k - const MatrixLayout desc_D{d.dtype(), kRowMajor, (int)d.shape(0), (int)d.shape(1), (int)d.stride(0)}; // m,n - ref_.gemm(x.raw_data(), desc_A, dense->weight.raw_data(), dense->k_desc, d.raw_data(), desc_D); - } - - void GetReference(const Tensor& x, const vector>& experts, Ref d_) - { - Tensor xe{{x.shape(0) * experts_per_token, input_dim}, data_type, kDEVICE}; - Tensor de{{x.shape(0) * experts_per_token, output_dim}, data_type, kDEVICE}; - - TM_SCOPE_CALL(invokeMoeDispatch(xe, x, f2n_.data(), x.shape(0) * experts_per_token, nullptr, stream_)); - - for (int i = 0; i < expert_num; ++i) { - const int base = h_offsets_[i], size = h_offsets_[i + 1] - base; - GetReference(xe.slice(base, size), experts[i], de.slice(base, size)); - } - - auto& d = d_.get(); - if (combine_experts) { - d = Tensor{{x.shape(0), output_dim}, data_type, kDEVICE}; - TM_SCOPE_CALL(invokeMoeCombine(d, // - de, - {}, - scales_.data(), - en2f_.data(), - nullptr, - nullptr, - experts_per_token, - 1., - 0., - stream_)); - } - else { - d = de; - } - } - - void Run() - { - if (tuning_) { - linear_.set_measure(true); - } - if (expert_num) { - Tensor de; - TM_SCOPE_CALL(linear_.Forward(x_original_, *w_quant_, f2n_, offsets_, de)); - if (combine_experts) { - d_quant_ = Tensor{{x_original_.shape(0), output_dim}, data_type, kDEVICE}; - TM_SCOPE_CALL(invokeMoeCombine(d_quant_, - de, - {}, - scales_.data(), - en2f_.data(), - nullptr, - nullptr, - experts_per_token, - 1., - 0., - stream_)); - } - else { - d_quant_ = de; - } - } - else { - TM_SCOPE_CALL(linear_.Forward(x_original_, *w_quant_, d_quant_)); - } - if (tuning_) { - linear_.set_measure(false); - } - } - - void Run(const Tensor& x, const vector>& experts) {} - - void Compare() - { - // Buffer_ h(16 * 16, kCPU); - // Buffer_ x(linear_.buf, 16 * 16, kDEVICE); - // Copy(x, h); - - // auto y = empty_like(w_dequant_->weight, kCPU); - // Copy(w_dequant_->weight, y); - - // clang-format off - printf("%20s", ""); FC_Header(); - if (!expert_num) { - printf("%20s", "w_dequant v w_origi"); FC_Print(FastCompare(w_dequant_->weight, w_original_->weight, stream_)); - } - printf("%20s", "quant vs dequant"); FC_Print(FastCompare(d_quant_, d_dequant_, stream_)); - printf("%20s", "quant vs original"); FC_Print(FastCompare(d_quant_, d_original_, stream_)); - printf("%20s", "dequant vs original"); FC_Print(FastCompare(d_dequant_, d_original_, stream_)); - // clang-format on - - // for (int m = 0; m < 16; ++m) { - // for (int k = 0; k < 16; ++k) { - // printf("%5.1f", h[m * 16 + k]); - // } - // printf("\n"); - // } - - // printf("\n"); - - // for (int m = 0; m < 16; ++m) { - // for (int k = 0; k < 16; ++k) { - // printf("%5.1f", (float)y.data()[k * output_dim + m]); - // } - // printf("\n"); - // } - } - - cudaStream_t stream_; - - cudaDeviceProp prop_; - - Linear linear_; - - // ! weights are non-movable - unique_ptr w_original_; - unique_ptr w_quant_; - unique_ptr w_dequant_; - - Tensor x_original_; - Tensor x_quant_, x_scale_; - Tensor x_dequant_; - - Tensor d_original_; // x_original * w_original - Tensor d_quant_; // x_original * w_quant, quant for X done by `Linear` - Tensor d_dequant_; // x_dequant * w_dequant - - vector> e_original_; - vector> e_quant_; - vector> e_dequant_; - - Buffer_ f2n_; - Buffer_ en2f_; - - Buffer_ offsets_; - Buffer_ scales_; - - Buffer_ h_offsets_; - - bool tuning_{}; - - std::string import_file_; - std::string export_file_; - - RNG rng_; - Reference ref_; -}; - -} // namespace turbomind diff --git a/src/turbomind/models/ffn_weight.cc b/src/turbomind/models/ffn_weight.cc index 9ecb66923d..06e3639575 100644 --- a/src/turbomind/models/ffn_weight.cc +++ b/src/turbomind/models/ffn_weight.cc @@ -2,8 +2,10 @@ #include "src/turbomind/models/ffn_weight.h" +#include "src/turbomind/core/data_type.h" #include "src/turbomind/core/registry.h" #include "src/turbomind/kernels/gemm/types.h" +#include "src/turbomind/utils/cuda_utils.h" namespace turbomind { @@ -26,6 +28,10 @@ void FfnWeight::prepare() auto* fused = static_cast(w1w3.get()); if (is_fused_silu) { fused->epilogue = gemm::Epilogue::kGatedSilu; + // SM90 FP8 fused SiLU quantizes to e4m3 + dynamic group-128 scales in-kernel. + if (fused->weight_format.dtype == kFloat8_e4m3 && getSMVersion() == 90) { + fused->set_fp8_fused_silu_output(); + } } } diff --git a/src/turbomind/models/linear_weight.cc b/src/turbomind/models/linear_weight.cc index 2c5eb73758..6b11fa1362 100644 --- a/src/turbomind/models/linear_weight.cc +++ b/src/turbomind/models/linear_weight.cc @@ -124,6 +124,11 @@ void LinearWeight::prepare() auto stream = core::Context::stream().handle(); + // TM_GEMM_WEIGHT_PACK=0: keep load-time storage; no transpose / tiled pack. + if (gemm::WeightPackEnv() == 0) { + return; + } + if (weight_format.dtype == kFloat8_e4m3 && input_dtype() == kFloat8_e4m3) { // FP8 native path: transpose weight and scales for native kernels. auto process = [&](Tensor& x, MatrixLayout& d, auto dtype) { @@ -153,6 +158,35 @@ void LinearWeight::prepare() auto [conv_w, conv_s] = GetConverters(data_type, weight_format.dtype, input_dtype(), is_grouped_, getSMVersion()); + // SM90 dense + grouped BF16/FP16: no pack converter (GetConverters returns {}), + // but GMMA+TMA need physical (N, K) with K contiguous (same as FP8 native + // prepare). Plain (K, N) row-major cannot use SW128 TMA boxes along K. + // Grouped experts prepare as 2D before MoeWeight::LinkLinearExperts; SM100 + // grouped stays (K, N) for cuBLAS (sm gate excludes >= 100). + if (!conv_w && getSMVersion() >= 90 && getSMVersion() < 100 + && (weight_format.dtype == kBfloat16 || weight_format.dtype == kHalf)) { + Tensor trans{{weight.shape(1), weight.shape(0)}, weight.dtype(), kDEVICE}; + if (weight.dtype() == kBfloat16) { + invokeTransposeAxis01((nv_bfloat16*)trans.raw_data(), + (nv_bfloat16*)weight.raw_data(), + weight.shape(0), + weight.shape(1), + 1, + stream); + } + else { + invokeTransposeAxis01( + (half*)trans.raw_data(), (half*)weight.raw_data(), weight.shape(0), weight.shape(1), 1, stream); + } + weight = std::move(trans); + k_desc.type = weight.dtype(); + k_desc.order = gemm::kColMajor; + k_desc.rows = (int)weight.shape(1); + k_desc.cols = (int)weight.shape(0); + k_desc.ld = (int)weight.stride(0); + k_desc.pack = {}; + } + if (conv_w) { const auto order_w = conv_w->order; const bool is_A = get_operand_tag(conv_w->pack) == OPERAND_A; diff --git a/src/turbomind/models/linear_weight.h b/src/turbomind/models/linear_weight.h index dd4680d061..a2fb834eb4 100644 --- a/src/turbomind/models/linear_weight.h +++ b/src/turbomind/models/linear_weight.h @@ -58,6 +58,14 @@ class LinearWeight: public core::Module { is_grouped_ = grouped; } + /// SM90 FP8 fused-SiLU: output is e4m3 with dynamic group-128 scales. + void set_fp8_fused_silu_output() + { + output_format.dtype = kFloat8_e4m3; + output_format.block_sizes = {128, 1}; + output_format.scales.dtype = kFloat; + } + explicit operator bool() const noexcept { return static_cast(weight); diff --git a/src/turbomind/models/llama/LlamaFfnLayer.cc b/src/turbomind/models/llama/LlamaFfnLayer.cc index 7efc02568b..406eb546b4 100644 --- a/src/turbomind/models/llama/LlamaFfnLayer.cc +++ b/src/turbomind/models/llama/LlamaFfnLayer.cc @@ -45,13 +45,21 @@ void LlamaFfnLayer::forward(ForwardParam param) auto* fused = mlp.w1w3.get(); bool use_fused = fused && fused->weight; + Tensor inter_scales; + Tensor unused_scales; + if (use_fused) { Tensor mix; - TM_SCOPE_CALL(linear_.Forward(param.input, *fused, mix)); - - gating = mix.slice({0, 0}, {(int)token_num, inter_size}); - if (!mlp.is_fused_silu) { - inter = mix.slice({0, inter_size}, {(ssize_t)token_num, inter_size}); + if (mlp.is_fused_silu && fused->output_dtype() == kFloat8_e4m3) { + TM_SCOPE_CALL(linear_.Forward(param.input, unused_scales, *fused, mix, inter_scales)); + gating = mix; // FP8 fused output is already half-N (inter_size) + } + else { + TM_SCOPE_CALL(linear_.Forward(param.input, *fused, mix)); + gating = mix.slice({0, 0}, {(int)token_num, inter_size}); + if (!mlp.is_fused_silu) { + inter = mix.slice({0, inter_size}, {(ssize_t)token_num, inter_size}); + } } } else { @@ -73,7 +81,12 @@ void LlamaFfnLayer::forward(ForwardParam param) { // w2(x) NvtxScope scope("w2"); - TM_SCOPE_CALL(linear_.Forward(gating, *mlp.w2, param.output)); + if (inter_scales) { + TM_SCOPE_CALL(linear_.Forward(gating, inter_scales, *mlp.w2, param.output, unused_scales)); + } + else { + TM_SCOPE_CALL(linear_.Forward(gating, *mlp.w2, param.output)); + } } } diff --git a/src/turbomind/models/llama/LlamaLinear.cu b/src/turbomind/models/llama/LlamaLinear.cu index f0a3126c7e..1935220f18 100644 --- a/src/turbomind/models/llama/LlamaLinear.cu +++ b/src/turbomind/models/llama/LlamaLinear.cu @@ -7,6 +7,7 @@ #include "src/turbomind/core/data_type.h" #include "src/turbomind/core/scope.h" +#include "src/turbomind/kernels/core/math.h" #include "src/turbomind/kernels/gemm/gemm.h" #include "src/turbomind/kernels/gemm/moe_utils_v2.h" #include "src/turbomind/kernels/gemm/types.h" @@ -63,15 +64,20 @@ struct LlamaLinear::Impl { return {B, desc_B, V, desc_V}; } - std::tuple - GetOperandA(const LinearWeight& weight, const Tensor& input, Buffer_ indices, const Buffer_& offsets) + std::tuple GetOperandA(const LinearWeight& weight, + const Tensor& input, + const Tensor& input_scales, + Buffer_ indices, + const Buffer_& offsets) { auto st = core::Context::stream().handle(); Tensor A; Tensor U; - const int m = indices ? indices.size() : input.shape(0); + // Size-0 / null Buffer must not count as indexed MoE (w2 / down path). + const bool has_indices = static_cast(indices) && indices.size() > 0; + const int m = has_indices ? indices.size() : input.shape(0); // Currently, FP8 only; INT8 may be added later if (input.dtype() != weight.input_dtype()) { @@ -79,11 +85,16 @@ struct LlamaLinear::Impl { } else { A = input; + if (weight.input_dtype() == kFloat8_e4m3) { + TM_CHECK(input_scales) << "FP8 input requires input_scales companion (dynamic group scales)"; + U = input_scales; + } } // SM100+ grouped bf16/fp16: use chunk() weights so Activation() runs separately. + // FP8 SM90 gathers A/U in-kernel (GemmUniversalSm90_v3 indexed); do not host-dispatch. const bool is_cublas_grouped = offsets && getSMVersion() == 100 && weight.weight_format.dtype == kBfloat16; - if (indices && (A.dtype() == kFloat8_e4m3 || is_cublas_grouped)) { + if (has_indices && is_cublas_grouped) { const int k = A.shape(1); Tensor A_e = {{m, k}, A.dtype(), kDEVICE}; const int* num_valid_tokens = offsets ? offsets.data() + offsets.size() - 1 : nullptr; @@ -95,6 +106,7 @@ struct LlamaLinear::Impl { } A = A_e; indices = {}; // indices already applied + // has_indices no longer used below } MatrixLayout desc_A{A.dtype(), gemm::Order::kRowMajor, m, (int)A.shape(1), (int)A.stride(0)}; @@ -106,18 +118,32 @@ struct LlamaLinear::Impl { desc_A.num = desc_U.num = weight.k_desc.num; desc_A.offsets = desc_U.offsets = const_cast(offsets.data()); } - if (indices) { + // Re-check after optional dispatch clears indices. + if (indices && indices.size() > 0) { desc_A.idxs = desc_U.idxs = const_cast(indices.data()); } return {A, desc_A, U, desc_U}; } - void Forward(Tensor& output, - const Tensor& input, // + void AllocOutputScales(Tensor& scales, int m, int out_dim) + { + constexpr int group_size = 128; + constexpr int alignment = 16 / sizeof(float); // match QuantizeSymm + const int s_dim = cdiv(out_dim, group_size); + const int aligned_m = round_up(m, alignment); + if (!scales || scales.shape(0) != s_dim || scales.shape(1) != m || scales.stride(0) != aligned_m) { + scales = Tensor_{{{s_dim, m}, {aligned_m, 1}}, kDEVICE}; + } + } + + void Forward(const Tensor& input, + const Tensor& input_scales, const LinearWeight& weight, const Buffer_& indices, - const Buffer_& offsets) + const Buffer_& offsets, + Tensor& output, + Tensor& output_scales) { TM_FUNCTION_SCOPE(); using namespace gemm; @@ -129,7 +155,7 @@ struct LlamaLinear::Impl { op.quant_b = MakeQuantDesc(weight.weight_format); op.batch_dim = 0; - auto&& [A, desc_A, U, desc_U] = GetOperandA(weight, input, indices, offsets); + auto&& [A, desc_A, U, desc_U] = GetOperandA(weight, input, input_scales, indices, offsets); auto&& [B, desc_B, V, desc_V] = GetOperandB(weight); Tensor& D = output; @@ -138,8 +164,6 @@ struct LlamaLinear::Impl { D = Tensor{{desc_A.rows, dim}, weight.output_dtype(), kDEVICE}; } - // std::cout << "D: " << D << " " << desc_B.num << "\n"; - MatrixLayout desc_D{ output.dtype(), kRowMajor, @@ -153,6 +177,19 @@ struct LlamaLinear::Impl { desc_D.offsets = const_cast(offsets.data()); } + MatrixLayout desc_W{}; + void* W_ptr = nullptr; + if (weight.output_dtype() == kFloat8_e4m3) { + const int out_cols = weight.epilogue == Epilogue::kGatedSilu ? weight.output_dim / 2 : weight.output_dim; + AllocOutputScales(output_scales, desc_A.rows, out_cols); + desc_W = {output_scales.dtype(), + kColMajor, + (int)output_scales.shape(1), + (int)output_scales.shape(0), + (int)output_scales.stride(0)}; + W_ptr = output_scales.raw_data(); + } + auto ec = gemm_.Run(op, 1.f, A.raw_data(), @@ -168,6 +205,8 @@ struct LlamaLinear::Impl { desc_D, D.raw_data(), desc_D, + W_ptr, + desc_W, workspace_, core::Context::stream().handle()); @@ -203,7 +242,34 @@ void LlamaLinear::Forward(const Tensor& input, // output.get() = output.get().view({-1, output.get().shape(-1)}); } - impl_->Forward(output.get(), in, weight, indices, offsets); + Tensor in_s, out_s; + impl_->Forward(in, in_s, weight, indices, offsets, output.get(), out_s); +} + +void LlamaLinear::Forward(const Tensor& input, + const Tensor& input_scales, + const LinearWeight& weight, + const Buffer_& indices, + const Buffer_& offsets, + Ref output, + Ref output_scales) +{ + Tensor in = input.view({-1, input.shape(-1)}); + + if (output.get()) { + output.get() = output.get().view({-1, output.get().shape(-1)}); + } + + impl_->Forward(in, input_scales, weight, indices, offsets, output.get(), output_scales.get()); +} + +void LlamaLinear::Forward(const Tensor& input, + const Tensor& input_scales, + const LinearWeight& weight, + Ref output, + Ref output_scales) +{ + Forward(input, input_scales, weight, {}, {}, output, output_scales); } void LlamaLinear::set_measure(bool measure) diff --git a/src/turbomind/models/llama/LlamaLinear.h b/src/turbomind/models/llama/LlamaLinear.h index ee3315a73d..4b79b11fef 100644 --- a/src/turbomind/models/llama/LlamaLinear.h +++ b/src/turbomind/models/llama/LlamaLinear.h @@ -1,5 +1,4 @@ // Copyright (c) OpenMMLab. All rights reserved. - #pragma once #include @@ -24,6 +23,23 @@ class LlamaLinear { const Buffer_& offsets, Ref output); + /// Forward with optional dynamic act-scale companions. + /// ``input_scales``: when ``input`` is already FP8, used as GEMM U (skip QuantizeSymm). + /// ``output_scales``: when ``weight.output_dtype()`` is FP8, filled with group-128 scales (W). + void Forward(const Tensor& input, + const Tensor& input_scales, + const LinearWeight& weight, + const Buffer_& indices, + const Buffer_& offsets, + Ref output, + Ref output_scales); + + void Forward(const Tensor& input, + const Tensor& input_scales, + const LinearWeight& weight, + Ref output, + Ref output_scales); + void set_measure(bool measure); [[maybe_unused]] int Export(std::ostream& os); diff --git a/src/turbomind/models/llama/moe_ffn_layer.cc b/src/turbomind/models/llama/moe_ffn_layer.cc index 18a44874f2..7231d61afe 100644 --- a/src/turbomind/models/llama/moe_ffn_layer.cc +++ b/src/turbomind/models/llama/moe_ffn_layer.cc @@ -44,9 +44,8 @@ Tensor_ MoeFfnLayerImpl::Gate(const Tensor& input, const LinearWeight& ga { TM_FUNCTION_SCOPE(); - auto& w = gate.weight; - TM_CHECK_EQ(input.shape(1), w.shape(0)); - Tensor_ logits{{input.shape(0), w.shape(1)}, kDEVICE}; + TM_CHECK_EQ(input.shape(1), gate.input_dim); + Tensor_ logits{{input.shape(0), gate.output_dim}, kDEVICE}; TM_SCOPE_CALL(linear_.Forward(input, gate, logits)); ApplyBias(logits, gate.bias, core::Context::stream().handle()); TM_CUDA_CHECK(cudaGetLastError()); @@ -258,14 +257,24 @@ void MoeFfnDefaultImpl::Forward(MoeFfnLayer::ForwardParam& p) if (block.w1w3) { Tensor inter; - TM_SCOPE_CALL(linear_.Forward(p.input, *block.w1w3, indices, offsets, inter)); - - if (!block.is_fused_silu) { - Activation(inter, block.w1w3->bias, f2E_, block.act_type, num_valid_tokens, st); - TM_CUDA_CHECK(cudaGetLastError()); + Tensor inter_scales; + Tensor unused_in_scales; + if (block.is_fused_silu && block.w1w3->output_dtype() == kFloat8_e4m3) { + TM_SCOPE_CALL( + linear_.Forward(p.input, unused_in_scales, *block.w1w3, indices, offsets, inter, inter_scales)); + TM_SCOPE_CALL(linear_.Forward( + inter.slice({0, 0}, {-1, inter_size}), inter_scales, *block.w2, {}, offsets, temp_, unused_in_scales)); } + else { + TM_SCOPE_CALL(linear_.Forward(p.input, *block.w1w3, indices, offsets, inter)); - TM_SCOPE_CALL(linear_.Forward(inter.slice({0, 0}, {-1, inter_size}), *block.w2, {}, offsets, temp_)); + if (!block.is_fused_silu) { + Activation(inter, block.w1w3->bias, f2E_, block.act_type, num_valid_tokens, st); + TM_CUDA_CHECK(cudaGetLastError()); + } + + TM_SCOPE_CALL(linear_.Forward(inter.slice({0, 0}, {-1, inter_size}), *block.w2, {}, offsets, temp_)); + } } else { // Separate w1/w3 path diff --git a/src/turbomind/models/moe_weight.cc b/src/turbomind/models/moe_weight.cc index 3c34b2a183..c12a65ecbd 100644 --- a/src/turbomind/models/moe_weight.cc +++ b/src/turbomind/models/moe_weight.cc @@ -56,24 +56,16 @@ static void LinkLinearExperts(std::function experts, int n, auto stream = core::Context::stream().handle(); - if (d.weight_format.dtype == kFloat8_e4m3 && d.input_dtype() == kFloat8_e4m3) { - auto make_blocked_ptr = [&](const auto& ptrs) { - return std::shared_ptr{gemm::MakeBlockedPtrs(ptrs, stream), [](auto p) { cudaFree(p); }}; - }; - d.weight = Tensor{make_blocked_ptr(weights), {n}, e0.weight.dtype(), kDEVICE}; - d.scales = Tensor{make_blocked_ptr(scales), {n}, e0.scales.dtype(), kDEVICE}; - d.k_desc.offsets = d.q_desc.offsets = (int*)1; - } - else { - auto make_strided_ptr = [&](const auto& ptrs) { - return std::shared_ptr{gemm::MakeStridedPtrs(ptrs, stream), [](auto p) { cudaFree(p); }}; - }; - d.weight = Tensor{make_strided_ptr(weights), {n}, d.weight_format.dtype, kDEVICE}; - if (e0.scales) { - d.scales = Tensor{make_strided_ptr(scales), {n}, e0.scales.dtype(), kDEVICE}; - } - d.k_desc.ld = d.q_desc.ld = 0; + auto make_strided_ptr = [&](const auto& ptrs) { + return std::shared_ptr{gemm::MakeStridedPtrs(ptrs, stream), [](auto p) { cudaFree(p); }}; + }; + d.weight = Tensor{make_strided_ptr(weights), {n}, d.weight_format.dtype, kDEVICE}; + if (e0.scales) { + d.scales = Tensor{make_strided_ptr(scales), {n}, e0.scales.dtype(), kDEVICE}; } + // MatrixLayout.ld == 0 identifies the StridedPtr expert table. + d.k_desc.ld = d.q_desc.ld = 0; + d.k_desc.offsets = d.q_desc.offsets = nullptr; } FfnWeight* MoeWeight::expert(int i) const diff --git a/src/turbomind/python/CMakeLists.txt b/src/turbomind/python/CMakeLists.txt index d4ad8ca101..31cbd85981 100644 --- a/src/turbomind/python/CMakeLists.txt +++ b/src/turbomind/python/CMakeLists.txt @@ -14,6 +14,7 @@ endif() pybind11_add_module(${PROJECT_NAME} bind.cpp + linear_bind.cpp ../kernels/linear_attn/python_bind.cpp) target_link_libraries(${PROJECT_NAME} PRIVATE turbomind xgrammar) diff --git a/src/turbomind/python/bind.cpp b/src/turbomind/python/bind.cpp index dbbd3760a3..512c6936d3 100644 --- a/src/turbomind/python/bind.cpp +++ b/src/turbomind/python/bind.cpp @@ -51,6 +51,10 @@ namespace turbomind::linear_attn::delta_rule { void bind_delta_rule(pybind11::module_& m); } +namespace turbomind::python_linear { +void bind_linear(pybind11::module_& m); +} + using ft::core::Tensor; // prepare to bind container @@ -144,6 +148,11 @@ DLManagedTensor* TritonTensorToDLManagedTensor(Tensor& tensor) data_type.code = DLDataTypeCode::kDLBfloat; data_type.bits = 16; break; + case data_type_v: + // Export as opaque uint8; consumers reinterpret (torch float8 via uint8 view). + data_type.code = DLDataTypeCode::kDLUInt; + data_type.bits = 8; + break; default: break; } @@ -561,15 +570,15 @@ PYBIND11_MODULE(_turbomind, m) // DataFormat descriptors py::class_(m, "QuantParamDesc") - .def_readonly("dtype", &turbomind::QuantParamDesc::dtype) - .def_readonly("transposed", &turbomind::QuantParamDesc::transposed) + .def_readwrite("dtype", &turbomind::QuantParamDesc::dtype) + .def_readwrite("transposed", &turbomind::QuantParamDesc::transposed) .def("present", &turbomind::QuantParamDesc::present); py::class_(m, "DataFormat") - .def_readonly("dtype", &turbomind::DataFormat::dtype) - .def_readonly("block_sizes", &turbomind::DataFormat::block_sizes) - .def_readonly("scales", &turbomind::DataFormat::scales) - .def_readonly("zeros", &turbomind::DataFormat::zeros) + .def_readwrite("dtype", &turbomind::DataFormat::dtype) + .def_readwrite("block_sizes", &turbomind::DataFormat::block_sizes) + .def_readwrite("scales", &turbomind::DataFormat::scales) + .def_readwrite("zeros", &turbomind::DataFormat::zeros) .def("is_quantized", &turbomind::DataFormat::is_quantized) .def("rank", &turbomind::DataFormat::rank); @@ -641,10 +650,12 @@ PYBIND11_MODULE(_turbomind, m) }); }, "stream"_a = 0) - .def("__dlpack_device__", [](const Tensor& self) { - auto device = getDLDevice(self); - return std::tuple(int(device.device_type), device.device_id); - }); + .def("__dlpack_device__", + [](const Tensor& self) { + auto device = getDLDevice(self); + return std::tuple(int(device.device_type), device.device_id); + }) + .def("t", [](const Tensor& self) { return std::make_shared(self.t()); }); m.def( "from_dlpack", [](py::object obj) { @@ -735,18 +746,35 @@ PYBIND11_MODULE(_turbomind, m) "grammar"_a); // Python context manager wrapper for ContextGuard. - // Stores copies of Stream + Allocator; constructs the real guard + // Stores copies of Stream + Allocator(s); constructs the real guard // in-place on __enter__ and destroys it on __exit__. struct PyContextGuard { ft::core::Stream stream; - ft::core::Allocator alloc; + ft::core::Allocator host; + ft::core::Allocator device; + bool has_device{false}; std::unique_ptr guard; - PyContextGuard(ft::core::Stream s, ft::core::Allocator a): stream(std::move(s)), alloc(std::move(a)) {} + // existing TurboMind path: stream + single allocator + PyContextGuard(ft::core::Stream s, ft::core::Allocator a): + stream(std::move(s)), host(std::move(a)), device{}, has_device(false) + { + } + + // harness / standalone path: stream + host + device (test_gemm_v2 shape) + PyContextGuard(ft::core::Stream s, ft::core::Allocator h, ft::core::Allocator d): + stream(std::move(s)), host(std::move(h)), device(std::move(d)), has_device(true) + { + } void enter() { - guard = std::make_unique(stream, alloc); + if (has_device) { + guard = std::make_unique(stream, host, device); + } + else { + guard = std::make_unique(stream, host); + } } void exit() { @@ -760,7 +788,22 @@ PYBIND11_MODULE(_turbomind, m) g.enter(); return g; }) - .def("__exit__", [](PyContextGuard& g, py::object, py::object, py::object) { g.exit(); }); + .def("__exit__", [](PyContextGuard& g, py::object, py::object, py::object) { g.exit(); }) + .def_property_readonly( + "stream_ptr", + [](const PyContextGuard& g) { return reinterpret_cast(g.stream.handle()); }, + "Underlying cudaStream_t as an integer (for torch.cuda.ExternalStream)."); + + m.def( + "create_device_context", + []() { + auto stream = ft::core::Stream::create(); + return std::make_unique( + stream, ft::core::Allocator{ft::kCPU}, ft::core::Allocator{stream, false}); + }, + "Create a ContextGuard with stream + host + device allocators.\n\n" + "Objects that use core::Context::stream() in their constructor or destructor " + "(notably LlamaLinear) must be destroyed before this context exits."); // Param — lightweight handle to a Module parameter slot py::class_(m, "Param") @@ -772,6 +815,8 @@ PYBIND11_MODULE(_turbomind, m) "shape"_a, "dtype"_a) .def("get", [](ft::core::Param& p) { return std::make_shared(p.get()); }) + .def( + "set", [](ft::core::Param& p, std::shared_ptr t) { p.set(t ? *t : Tensor{}); }, "tensor"_a) .def("__bool__", [](ft::core::Param& p) { return static_cast(p); }); // Module class — navigation and allocation interface @@ -831,8 +876,7 @@ PYBIND11_MODULE(_turbomind, m) py::return_value_policy::reference, "config"_a); - // LinearWeight — specific interface for weight loading - py::class_(m, "LinearWeight"); + // LinearWeight is fully bound in bind_linear() (after Module is registered). // transformer model using ft::TurboMind; @@ -900,4 +944,5 @@ PYBIND11_MODULE(_turbomind, m) .def("model_tp_rank", &TurboMind::GetModelTpRank, "index"_a); turbomind::linear_attn::delta_rule::bind_delta_rule(m); + turbomind::python_linear::bind_linear(m); } diff --git a/src/turbomind/python/linear_bind.cpp b/src/turbomind/python/linear_bind.cpp new file mode 100644 index 0000000000..6c82fa7234 --- /dev/null +++ b/src/turbomind/python/linear_bind.cpp @@ -0,0 +1,376 @@ +// Copyright (c) OpenMMLab. All rights reserved. + +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +#include "src/turbomind/core/buffer.h" +#include "src/turbomind/core/context.h" +#include "src/turbomind/core/data_type.h" +#include "src/turbomind/core/tensor.h" +#include "src/turbomind/kernels/gemm/convert.h" +#include "src/turbomind/kernels/gemm/moe_utils_v2.h" +#include "src/turbomind/kernels/gemm/types.h" +#include "src/turbomind/kernels/quantization.h" +#include "src/turbomind/models/linear_weight.h" +#include "src/turbomind/models/llama/LlamaLinear.h" + +namespace py = pybind11; + +namespace turbomind::python_linear { +namespace { + +core::Tensor TensorFromShared(const std::shared_ptr& p, const char* name) +{ + if (!p) { + throw py::value_error(std::string(name) + " is null"); + } + return *p; +} + +core::Tensor TensorOrEmpty(const std::shared_ptr& p) +{ + return p ? *p : core::Tensor{}; +} + +cudaStream_t DefaultStream() +{ + return core::Context::stream().handle(); +} + +std::shared_ptr OwnCudaPtr(void* ptr) +{ + return std::shared_ptr{ptr, [](void* p) { + if (p) { + cudaFree(p); + } + }}; +} + +std::vector> PtrPairsFromPy(const std::vector>& ptrs) +{ + std::vector> out; + out.reserve(ptrs.size()); + for (const auto& [p, ld] : ptrs) { + out.emplace_back(reinterpret_cast(p), ld); + } + return out; +} + +std::shared_ptr MakePtrTensor(void* raw, ssize_t n, DataType dtype, cudaStream_t /*stream*/) +{ + auto data = OwnCudaPtr(raw); + return std::make_shared(std::move(data), core::Layout{n}, dtype, kDEVICE); +} + +} // namespace + +void bind_linear(py::module_& m) +{ + py::class_(m, "MatrixLayout") + .def(py::init<>()) + .def_readwrite("type", &gemm::MatrixLayout::type) + .def_readwrite("rows", &gemm::MatrixLayout::rows) + .def_readwrite("cols", &gemm::MatrixLayout::cols) + .def_readwrite("ld", &gemm::MatrixLayout::ld) + .def_readwrite("num", &gemm::MatrixLayout::num) + .def_property( + "offsets", + [](const gemm::MatrixLayout& d) { return reinterpret_cast(d.offsets); }, + [](gemm::MatrixLayout& d, std::uintptr_t p) { d.offsets = reinterpret_cast(p); }); + + py::enum_(m, "Epilogue") + .value("kNone", gemm::Epilogue::kNone) + .value("kChannelCombination", gemm::Epilogue::kChannelCombination) + .value("kGatedSilu", gemm::Epilogue::kGatedSilu); + + py::class_(m, "LinearWeight") + .def(py::init()) + .def("prepare", &LinearWeight::prepare) + .def("copy_metadata_to", &LinearWeight::copy_metadata_to, py::arg("dst")) + .def("set_grouped", &LinearWeight::set_grouped, py::arg("grouped")) + .def("set_fp8_fused_silu_output", &LinearWeight::set_fp8_fused_silu_output) + .def_readwrite("input_dim", &LinearWeight::input_dim) + .def_readwrite("output_dim", &LinearWeight::output_dim) + .def_readwrite("data_type", &LinearWeight::data_type) + .def_readwrite("weight_format", &LinearWeight::weight_format) + .def_readwrite("input_format", &LinearWeight::input_format) + .def_readwrite("output_format", &LinearWeight::output_format) + .def_readwrite("epilogue", &LinearWeight::epilogue) + .def_readwrite("k_desc", &LinearWeight::k_desc) + .def_readwrite("q_desc", &LinearWeight::q_desc); + + py::class_>(m, "LlamaLinear") + .def(py::init<>()) + .def( + "forward_dense", + [](LlamaLinear& self, + std::shared_ptr input, + LinearWeight& weight, + std::shared_ptr output, + std::shared_ptr input_scales, + std::shared_ptr output_scales) { + core::Tensor out = TensorOrEmpty(output); + core::Tensor in_s = TensorOrEmpty(input_scales); + core::Tensor out_s = TensorOrEmpty(output_scales); + if (weight.output_dtype() == kFloat8_e4m3 || (input_scales && *input_scales)) { + self.Forward(TensorFromShared(input, "input"), in_s, weight, out, out_s); + } + else { + self.Forward(TensorFromShared(input, "input"), weight, out); + } + return std::make_tuple(std::make_shared(out), std::make_shared(out_s)); + }, + py::arg("input"), + py::arg("weight"), + py::arg("output") = py::none(), + py::arg("input_scales") = py::none(), + py::arg("output_scales") = py::none(), + py::call_guard()) + .def( + "forward_moe", + [](LlamaLinear& self, + std::shared_ptr input, + LinearWeight& weight, + std::shared_ptr indices_t, + std::shared_ptr offsets_t, + std::shared_ptr output, + std::shared_ptr input_scales, + std::shared_ptr output_scales) { + // Null-check while GIL is held (no call_guard); release only for Forward. + // indices=None → empty Buffer (MoE down / w2: expert-packed A, offsets only). + core::Tensor in = TensorFromShared(input, "input"); + core::Tensor offsets_tensor = TensorFromShared(offsets_t, "offsets"); + Buffer_ indices{}; + if (indices_t && *indices_t) { + core::Tensor indices_tensor = *indices_t; + indices = + Buffer_((int*)indices_tensor.raw_data(), indices_tensor.size(), indices_tensor.device()); + } + auto offsets = + Buffer_((int*)offsets_tensor.raw_data(), offsets_tensor.size(), offsets_tensor.device()); + core::Tensor out = TensorOrEmpty(output); + core::Tensor in_s = TensorOrEmpty(input_scales); + core::Tensor out_s = TensorOrEmpty(output_scales); + { + py::gil_scoped_release release; + if (weight.output_dtype() == kFloat8_e4m3 || (input_scales && *input_scales)) { + self.Forward(in, in_s, weight, indices, offsets, out, out_s); + } + else { + self.Forward(in, weight, indices, offsets, out); + } + } + return std::make_tuple(std::make_shared(out), std::make_shared(out_s)); + }, + py::arg("input"), + py::arg("weight"), + py::arg("indices") = py::none(), + py::arg("offsets"), + py::arg("output") = py::none(), + py::arg("input_scales") = py::none(), + py::arg("output_scales") = py::none()) + .def("set_measure", &LlamaLinear::set_measure) + .def( + "import_records", + [](LlamaLinear& self, const std::string& path) { + std::ifstream ifs(path, std::ios::binary); + return self.Import(ifs); + }, + py::arg("path")) + .def( + "export_records", + [](LlamaLinear& self, const std::string& path) { + std::ofstream ofs(path, std::ios::binary); + return self.Export(ofs); + }, + py::arg("path")); + + // --- Quantization helpers (stream defaults to Context::stream) --- + + m.def( + "QuantizeSymm", + [](std::shared_ptr out, std::shared_ptr scale, std::shared_ptr src) { + // Null-check / snapshot under GIL; release only for CUDA; return C++ + // tuple so pybind converts after GIL is restored (no py::make_tuple). + core::Tensor o = TensorOrEmpty(out); + core::Tensor s = TensorOrEmpty(scale); + core::Tensor src_t = TensorFromShared(src, "src"); + { + py::gil_scoped_release release; + QuantizeSymm(o, s, src_t, DefaultStream()); + } + return std::make_tuple(std::make_shared(o), std::make_shared(s)); + }, + py::arg("out") = py::none(), + py::arg("scale") = py::none(), + py::arg("src")); + + m.def( + "DequantizeSymm", + [](std::shared_ptr out, std::shared_ptr src, std::shared_ptr scale) { + core::Tensor o = TensorOrEmpty(out); + core::Tensor src_t = TensorFromShared(src, "src"); + core::Tensor scale_t = TensorFromShared(scale, "scale"); + { + py::gil_scoped_release release; + DequantizeSymm(o, src_t, scale_t, DefaultStream()); + } + return std::make_shared(o); + }, + py::arg("out") = py::none(), + py::arg("src"), + py::arg("scale")); + + m.def( + "QuantizeSymmBlock", + [](std::shared_ptr out, std::shared_ptr scale, std::shared_ptr src) { + core::Tensor o = TensorOrEmpty(out); + core::Tensor s = TensorOrEmpty(scale); + core::Tensor src_t = TensorFromShared(src, "src"); + { + py::gil_scoped_release release; + QuantizeSymmBlock(o, s, src_t, DefaultStream()); + } + return std::make_tuple(std::make_shared(o), std::make_shared(s)); + }, + py::arg("out") = py::none(), + py::arg("scale") = py::none(), + py::arg("src")); + + m.def( + "DequantizeSymmBlock", + [](std::shared_ptr out, std::shared_ptr src, std::shared_ptr scale) { + core::Tensor o = TensorOrEmpty(out); + core::Tensor src_t = TensorFromShared(src, "src"); + core::Tensor scale_t = TensorFromShared(scale, "scale"); + { + py::gil_scoped_release release; + DequantizeSymmBlock(o, src_t, scale_t, DefaultStream()); + } + return std::make_shared(o); + }, + py::arg("out") = py::none(), + py::arg("src"), + py::arg("scale")); + + m.def( + "QuantizeGroupwise", + [](std::shared_ptr quant, + std::shared_ptr scales, + std::shared_ptr zeros, + std::shared_ptr dequant, + std::shared_ptr src, + std::shared_ptr rbits, + int group_size) { + core::Tensor quant_t = TensorFromShared(quant, "quant"); + core::Tensor scales_t = TensorFromShared(scales, "scales"); + core::Tensor zeros_t = TensorOrEmpty(zeros); + core::Tensor dequant_t = TensorFromShared(dequant, "dequant"); + core::Tensor src_t = TensorFromShared(src, "src"); + Buffer_ r; + if (rbits && *rbits) { + r = Buffer_((unsigned*)rbits->raw_data(), rbits->size(), rbits->device()); + } + { + py::gil_scoped_release release; + QuantizeGroupwise(quant_t, scales_t, zeros_t, dequant_t, src_t, r, group_size); + } + }, + py::arg("quant"), + py::arg("scales"), + py::arg("zeros") = py::none(), + py::arg("dequant"), + py::arg("src"), + py::arg("rbits") = py::none(), + py::arg("group_size")); + + // --- MoE pointer builders (owned by returned Tensor; freed with cudaFree) --- + + m.def( + "MakeStridedPtrs", + [](const std::vector>& ptrs, DataType dtype) { + auto pairs = PtrPairsFromPy(ptrs); + void* raw = gemm::MakeStridedPtrs(pairs, DefaultStream()); + return MakePtrTensor(raw, (ssize_t)ptrs.size(), dtype, DefaultStream()); + }, + py::arg("ptrs"), + py::arg("dtype"), + py::call_guard()); + + // --- MoE dispatch / combine --- + + m.def( + "invokeMoeDispatch", + [](std::shared_ptr out, + std::shared_ptr src, + std::shared_ptr f2n, + int expert_per_token) { + core::Tensor o = TensorOrEmpty(out); + core::Tensor idx = TensorFromShared(f2n, "f2n"); + invokeMoeDispatch(o, + TensorFromShared(src, "src"), + (const int*)idx.raw_data(), + expert_per_token, + nullptr, + DefaultStream()); + return std::make_shared(o); + }, + py::arg("out") = py::none(), + py::arg("src"), + py::arg("f2n"), + py::arg("expert_per_token"), + py::call_guard()); + + m.def( + "invokeMoeCombine", + [](std::shared_ptr out, + std::shared_ptr src, + std::shared_ptr bias, + std::shared_ptr scales, + std::shared_ptr en2f, + std::shared_ptr f2E, + std::shared_ptr dst_scales, + int experts_per_token, + float bscale, + float dst_scale) { + core::Tensor o = TensorOrEmpty(out); + const float* scales_ptr = (scales && *scales) ? scales->data() : nullptr; + const int* en2f_ptr = (en2f && *en2f) ? (const int*)en2f->raw_data() : nullptr; + const int* f2E_ptr = (f2E && *f2E) ? (const int*)f2E->raw_data() : nullptr; + const float* dst_scales_ptr = (dst_scales && *dst_scales) ? dst_scales->data() : nullptr; + invokeMoeCombine(o, + TensorFromShared(src, "src"), + TensorOrEmpty(bias), + scales_ptr, + en2f_ptr, + f2E_ptr, + dst_scales_ptr, + experts_per_token, + bscale, + dst_scale, + DefaultStream()); + return std::make_shared(o); + }, + py::arg("out") = py::none(), + py::arg("src"), + py::arg("bias") = py::none(), + py::arg("scales") = py::none(), + py::arg("en2f") = py::none(), + py::arg("f2E") = py::none(), + py::arg("dst_scales") = py::none(), + py::arg("experts_per_token"), + py::arg("bscale") = 1.f, + py::arg("dst_scale") = 0.f, + py::call_guard()); +} + +} // namespace turbomind::python_linear diff --git a/tests/turbomind/linear/__init__.py b/tests/turbomind/linear/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/turbomind/linear/bench_linear.py b/tests/turbomind/linear/bench_linear.py new file mode 100644 index 0000000000..36be9db9fa --- /dev/null +++ b/tests/turbomind/linear/bench_linear.py @@ -0,0 +1,4 @@ +from .benchmark import main + +if __name__ == '__main__': + raise SystemExit(main()) diff --git a/tests/turbomind/linear/benchmark.py b/tests/turbomind/linear/benchmark.py new file mode 100644 index 0000000000..fef0735396 --- /dev/null +++ b/tests/turbomind/linear/benchmark.py @@ -0,0 +1,196 @@ +from __future__ import annotations + +import argparse +import os +from dataclasses import dataclass + +import torch + +from .cases import VALID_SUITES, expand_suite +from .fixture import LinearFixture + + +@dataclass(frozen=True) +class BenchmarkRequest: + validate_outputs: bool + print_diffs: bool + l2_flush: bool + warmup: int + iters: int + tune: bool + import_path: str | None + export_path: str | None + + +class L2CacheFlusher: + def __init__(self, device: torch.device): + props = torch.cuda.get_device_properties(device) + # allocate ~2x L2 to flush; mirror linear_attn.benchmark.L2CacheFlusher if present + nbytes = int(getattr(props, 'L2_cache_size', 4 * 1024 * 1024) * 2) + self._buf = torch.empty(nbytes, dtype=torch.uint8, device=device) + + def __call__(self) -> None: + self._buf.fill_(0) + + +def resolve_tune_paths(args) -> tuple[bool, str | None, str | None]: + tune = bool(args.tune) or bool(os.environ.get('TM_GEMM_TUNE')) + import_path = args.import_path or os.environ.get('TM_GEMM_IMPORT') + export_path = args.export_path or os.environ.get('TM_GEMM_EXPORT') + if import_path: + tune = False # match testbed_v3: import disables tuning + return tune, import_path, export_path + + +def make_parser() -> argparse.ArgumentParser: + p = argparse.ArgumentParser(prog='bench_linear') + p.add_argument('--suite', choices=VALID_SUITES, default='smoke') + p.add_argument('--case', type=str, default='', help='comma-separated LinearCase names') + p.add_argument('--type', dest='type_names', type=str, default='', help='comma-separated TypeSpec names') + p.add_argument('--shape', dest='shape_names', type=str, default='', help='comma-separated ShapeSpec names') + p.add_argument('--batch', type=str, default='', help='comma-separated batch sizes') + p.add_argument('--tp', type=str, default='1', help='comma-separated TP sizes') + p.add_argument('--ep', type=str, default='1', help='comma-separated EP sizes') + p.add_argument('--warmup', type=int, default=10) + p.add_argument('--iters', type=int, default=50) + p.add_argument('--tune', action='store_true') + p.add_argument('--import', dest='import_path', default=None) + p.add_argument('--export', dest='export_path', default=None) + p.add_argument('--print-diffs', action='store_true') + p.add_argument('--no-validate', action='store_true') + p.add_argument('--no-l2-flush', action='store_true') + return p + + +def parse_int_list(s: str) -> tuple[int, ...] | None: + if not s.strip(): + return None + return tuple(int(x) for x in s.split(',') if x.strip()) + + +def parse_name_list(s: str) -> tuple[str, ...] | None: + if not s.strip(): + return None + return tuple(x for x in s.split(',') if x.strip()) + + +def flop_count(m: int, n: int, k: int, expert_tokens: int | None = None) -> int: + # dense: 2*m*n*k; MoE packed: 2 * token_count * n * k + if expert_tokens is None: + return 2 * m * n * k + return 2 * expert_tokens * n * k + + +def main(argv: list[str] | None = None) -> int: + args = make_parser().parse_args(argv) + case_names = parse_name_list(args.case) + type_names = parse_name_list(args.type_names) + shape_names = parse_name_list(args.shape_names) + batches = parse_int_list(args.batch) + tps = parse_int_list(args.tp) + eps = parse_int_list(args.ep) + runs = expand_suite( + args.suite, + case_names, + batches, + type_names, + shape_names, + tps=tps, + eps=eps, + ) + tune, import_path, export_path = resolve_tune_paths(args) + device = torch.device('cuda') + flusher = None if args.no_l2_flush else L2CacheFlusher(device) + + # Group runs by concrete local shape so weights build once per batch sweep. + by_case: dict[tuple[str, int, int], list] = {} + for run in runs: + key = (run.case.name, run.case.tp, run.case.ep) + by_case.setdefault(key, []).append(run) + + for case_runs in by_case.values(): + fx = LinearFixture(case_runs[0].case, device=device) + try: + assert fx.linear is not None + if import_path: + fx.linear.import_records(import_path) + for run in case_runs: + fx.prepare_batch(run.batch_size) + row = { + 'case': run.case.name, + 'type': run.case.type_name, + 'shape': run.case.shape_name, + 'm': run.batch_size, + 'n': run.case.output_dim, + 'k': run.case.input_dim, + 'data_type': run.case.data_type, + 'weight_type': run.case.weight_type, + 'input_type': run.case.input_type, + 'expert_num': run.case.expert_num, + 'experts_per_token': run.case.experts_per_token, + 'tp': run.case.tp, + 'ep': run.case.ep, + 'max_tp': run.case.max_tp, + 'max_ep': run.case.max_ep, + 'fuse_silu': run.case.fuse_silu, + } + if not args.no_validate or args.print_diffs: + fx.run_reference() + fx.run_linear() + metrics = fx.compare() + if args.print_diffs or not args.no_validate: + row.update({f'{a}.{b}': c for a, d in metrics.items() for b, c in d.items()}) + if not args.no_validate: + fx.check_tolerances(metrics) + if tune or args.iters > 0: + assert fx.linear is not None + with fx.on_tm_stream() as stream: + # Tune before warmup/timed so measure overhead is not in TFLOPS. + # One forward is enough: GEMM measure policy runs its own internal iters. + if tune: + fx.linear.set_measure(True) + fx.run_linear_forward() + fx.sync_tm() + fx.linear.set_measure(False) + fx.release_forward_result() + if args.iters > 0: + for _ in range(args.warmup): + fx.run_linear_forward() + fx.release_forward_result() + fx.sync_tm() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + elapsed_ms = 0.0 + # L2 flush is on the TM stream, before the timing window. + for _ in range(args.iters): + if flusher is not None: + flusher() + start.record(stream) + fx.run_linear_forward() + end.record(stream) + fx.release_forward_result() + end.synchronize() + elapsed_ms += start.elapsed_time(end) + ms = elapsed_ms / args.iters + row['latency_ms'] = ms + tokens = None + if run.case.expert_num: + tokens = run.batch_size * run.case.experts_per_token + flops = flop_count(run.batch_size, run.case.output_dim, run.case.input_dim, tokens) + row['tflops'] = (flops / (ms * 1e-3)) / 1e12 + if args.iters == 0: + row['latency_ms'] = 0.0 + row['tflops'] = 0.0 + print(row) + assert fx.linear is not None + if export_path: + # Per-case export so a full-suite scan does not overwrite previous records; + # the dispatch cache is per-LinearFixture (per-Linear / per-Gemm). + suffix = f'{case_runs[0].case.name}__tp{case_runs[0].case.tp}__ep{case_runs[0].case.ep}' + fx.linear.export_records(f'{export_path}.{suffix}') + finally: + fx.close() + # Release PyTorch cache between case groups; the MoE cases can be large and + # the full-suite scan is run on a shared GPU. + torch.cuda.empty_cache() + return 0 diff --git a/tests/turbomind/linear/cases.py b/tests/turbomind/linear/cases.py new file mode 100644 index 0000000000..db2bb07dce --- /dev/null +++ b/tests/turbomind/linear/cases.py @@ -0,0 +1,565 @@ +from __future__ import annotations + +from dataclasses import dataclass, replace +from typing import cast + +from .model_specs import ( + SUPPORTED_TP_SIZES, + ModelSpec, + MoeSpec, + TpAxis, + WeightSpec, + load_model_specs, +) + +DTYPE_ALIASES = { + 'bf16': 'bf16', + 'fp16': 'fp16', + 'fp8_e4m3': 'fp8_e4m3', + 'uint4': 'uint4', + 'fp4_e2m1': 'fp4_e2m1', +} + +VALID_SUITES = ('smoke', 'quick', 'full', 'custom') + +# Powers-of-two cover aligned CTA tiles; odd / tile-residue Ms hit epilogue tails. +SMOKE_BATCHES = (1, 3, 16) +QUICK_BATCHES = (1, 3, 16, 17, 64, 65, 256, 1024, 4096) +FULL_BATCHES = ( + 1, + 2, + 3, + 4, + 5, + 7, + 8, + 15, + 16, + 17, + 31, + 32, + 33, + 63, + 64, + 65, + 127, + 128, + 129, + 255, + 256, + 257, + 511, + 512, + 513, + 1023, + 1024, + 1025, + 2047, + 2048, + 2049, + 4095, + 4096, + 4097, + 8191, + 8192, + 8233, + 16384, + 16385, + 32768, +) + + +@dataclass(frozen=True, kw_only=True) +class TypeSpec: + """Dtype / quant axis (mirrors testbed_v3 TestParameter ctor fields).""" + + name: str + data_type: str + weight_type: str + input_type: str + group_size: int + + def __post_init__(self) -> None: + for field, value in ( + ('data_type', self.data_type), + ('weight_type', self.weight_type), + ('input_type', self.input_type), + ): + if value not in DTYPE_ALIASES: + raise ValueError(f'unsupported_{field}_{value}') + if self.group_size < 0: + raise ValueError('group_size_must_be_non_negative') + + +@dataclass(frozen=True, kw_only=True) +class ShapeSpec: + """Problem geometry axis (dims + MoE layout).""" + + name: str + input_dim: int + output_dim: int + expert_num: int + experts_per_token: int + combine_experts: bool + tp_axis: TpAxis + max_tp: int + max_ep: int + # True: token-major x + f2n indices (w1/gate). False: expert-packed x, offsets only (w2/down). + moe_indexed: bool = False + + def __post_init__(self) -> None: + if self.input_dim <= 0 or self.output_dim <= 0: + raise ValueError('dims_must_be_positive') + if self.expert_num < 0: + raise ValueError('expert_num_must_be_non_negative') + if self.expert_num == 0 and self.experts_per_token != 0: + raise ValueError('dense_shape_requires_experts_per_token_0') + if self.expert_num > 0 and self.experts_per_token <= 0: + raise ValueError('moe_shape_requires_positive_experts_per_token') + if self.expert_num == 0 and self.moe_indexed: + raise ValueError('dense_shape_cannot_be_moe_indexed') + if self.tp_axis not in ('input', 'output'): + raise ValueError(f'unsupported_tp_axis_{self.tp_axis}') + if self.max_tp <= 0 or self.max_ep <= 0: + raise ValueError('parallel_caps_must_be_positive') + tp_dim = self.input_dim if self.tp_axis == 'input' else self.output_dim + if tp_dim % self.max_tp: + raise ValueError('tp_dimension_not_divisible_by_max_tp') + if self.expert_num == 0: + if self.max_ep != 1: + raise ValueError('dense_shape_requires_max_ep_1') + else: + if self.expert_num % self.max_ep: + raise ValueError('expert_num_not_divisible_by_max_ep') + local_experts = self.expert_num // self.max_ep + if local_experts <= 1: + raise ValueError('expert_parallel_requires_multiple_local_experts') + if local_experts < self.experts_per_token: + raise ValueError('local_experts_less_than_experts_per_token') + inter_size = self.output_dim // 2 if self.tp_axis == 'output' else self.input_dim + if inter_size % self.max_tp or inter_size // self.max_tp < 256: + raise ValueError('expert_inter_size_per_tp_less_than_256') + + +@dataclass(frozen=True, kw_only=True) +class LinearCase: + """One cell of the type × shape product.""" + + name: str + input_dim: int + output_dim: int + data_type: str + weight_type: str + input_type: str + group_size: int + expert_num: int + experts_per_token: int + combine_experts: bool + moe_indexed: bool + type_name: str + shape_name: str + tp_axis: TpAxis + max_tp: int + max_ep: int + tp: int = 1 + ep: int = 1 + # Optional SM90 gate/up block-pack + kGatedSilu epilogue. + fuse_silu: bool = False + + def __post_init__(self) -> None: + for field, value in ( + ('data_type', self.data_type), + ('weight_type', self.weight_type), + ('input_type', self.input_type), + ): + if value not in DTYPE_ALIASES: + raise ValueError(f'unsupported_{field}_{value}') + if self.expert_num < 0: + raise ValueError('expert_num_must_be_non_negative') + if self.expert_num == 0 and self.experts_per_token != 0: + raise ValueError('dense_case_requires_experts_per_token_0') + if self.expert_num > 0 and self.experts_per_token <= 0: + raise ValueError('moe_case_requires_positive_experts_per_token') + if self.expert_num == 0 and self.moe_indexed: + raise ValueError('dense_case_cannot_be_moe_indexed') + if self.tp_axis not in ('input', 'output'): + raise ValueError(f'unsupported_tp_axis_{self.tp_axis}') + if self.max_tp <= 0 or self.max_ep <= 0 or self.tp <= 0 or self.ep <= 0: + raise ValueError('parallel_sizes_must_be_positive') + if self.tp > self.max_tp or self.ep > self.max_ep: + raise ValueError('parallel_size_exceeds_case_cap') + if self.tp > 1 and self.ep > 1: + raise ValueError('combined_tp_ep_case_not_supported') + if self.fuse_silu and not self.shape_name.endswith('_gate_up'): + raise ValueError('fuse_silu_requires_gate_up_shape') + if self.fuse_silu and self.expert_num > 0 and not self.moe_indexed: + raise ValueError('moe_fuse_silu_requires_indexed_gate_up') + if self.fuse_silu and self.weight_type not in ('fp8_e4m3', 'bf16'): + raise ValueError('fuse_silu_requires_fp8_or_bf16_weights') + + +@dataclass(frozen=True, kw_only=True) +class RunCase: + case: LinearCase + batch_size: int + + def __post_init__(self) -> None: + if self.batch_size <= 0: + raise ValueError('batch_size_must_be_positive') + + +def _compose(type_spec: TypeSpec, shape: ShapeSpec, *, fuse_silu: bool = False) -> LinearCase: + name = f'{shape.name}__{type_spec.name}' + if fuse_silu: + name = f'{name}__fuse_silu' + return LinearCase( + name=name, + input_dim=shape.input_dim, + output_dim=shape.output_dim, + data_type=type_spec.data_type, + weight_type=type_spec.weight_type, + input_type=type_spec.input_type, + group_size=type_spec.group_size, + expert_num=shape.expert_num, + experts_per_token=shape.experts_per_token, + combine_experts=shape.combine_experts, + moe_indexed=shape.moe_indexed, + type_name=type_spec.name, + shape_name=shape.name, + tp_axis=shape.tp_axis, + max_tp=shape.max_tp, + max_ep=shape.max_ep, + fuse_silu=fuse_silu, + ) + + +# --- Types: commented combos from test_gemm_v2 / testbed_v3 --- +# QuantizeGroupwise supports: fp16×uint4, fp16×fp4, bf16×fp4 (not bf16×uint4). +# +# TypeSpec.name matches GemmDesc type substring (desc.h to_string): +# {Ta}[{Qa}]_{Tb}[{Qb}]_{Tc} with dtype tokens f16/bf16/e4m3/e2m1/u4 +# and glued QuantDesc (k128/b128/k32). Order is act → weight → output. + +TYPE_BF16 = TypeSpec(name='bf16_bf16_bf16', data_type='bf16', weight_type='bf16', input_type='bf16', group_size=0) +TYPE_BF16_E4M3B128_BF16 = TypeSpec( + name='bf16_e4m3b128_bf16', + data_type='bf16', + weight_type='fp8_e4m3', + input_type='bf16', + group_size=128, +) +TYPE_E4M3K128_E4M3B128_BF16 = TypeSpec( + name='e4m3k128_e4m3b128_bf16', + data_type='bf16', + weight_type='fp8_e4m3', + input_type='fp8_e4m3', + group_size=128, +) +TYPE_BF16_E2M1K32_BF16 = TypeSpec( + name='bf16_e2m1k32_bf16', + data_type='bf16', + weight_type='fp4_e2m1', + input_type='bf16', + group_size=32, +) +TYPE_F16 = TypeSpec(name='f16_f16_f16', data_type='fp16', weight_type='fp16', input_type='fp16', group_size=0) +TYPE_F16_U4K128_F16 = TypeSpec( + name='f16_u4k128_f16', + data_type='fp16', + weight_type='uint4', + input_type='fp16', + group_size=128, +) +TYPE_F16_E2M1K32_F16 = TypeSpec( + name='f16_e2m1k32_f16', + data_type='fp16', + weight_type='fp4_e2m1', + input_type='fp16', + group_size=32, +) +TYPE_F16_E4M3B128_F16 = TypeSpec( + name='f16_e4m3b128_f16', + data_type='fp16', + weight_type='fp8_e4m3', + input_type='fp16', + group_size=128, +) + +ALL_TYPES: tuple[TypeSpec, ...] = ( + TYPE_BF16, + TYPE_BF16_E4M3B128_BF16, + TYPE_E4M3K128_E4M3B128_BF16, + TYPE_BF16_E2M1K32_BF16, + TYPE_F16, + TYPE_F16_U4K128_F16, + TYPE_F16_E2M1K32_F16, + TYPE_F16_E4M3B128_F16, +) + +SMOKE_TYPES: tuple[TypeSpec, ...] = (TYPE_BF16_E4M3B128_BF16,) + + +# --- Shapes: curated model weight configurations loaded from YAML --- + + +def _moe_max_tp(model: ModelSpec, weight: WeightSpec) -> int: + if weight.kind == 'moe_gate_up': + inter_size = weight.output_dim // 2 + else: + inter_size = weight.input_dim + + for tp in SUPPORTED_TP_SIZES: + if tp <= model.max_tp and inter_size % tp == 0 and inter_size // tp >= 256: + return tp + + raise ValueError( + f'no_expert_safe_tp_{model.name}_{weight.name}' + f'_inter_size_{inter_size}_max_tp_{model.max_tp}' + ) + + +def _make_shape(model: ModelSpec, weight: WeightSpec) -> ShapeSpec: + is_moe = weight.kind in ('moe_gate_up', 'moe_down') + if weight.kind == 'dense': + tp_axis = cast(TpAxis, weight.tp_axis) + max_tp = model.max_tp + elif weight.kind == 'replicated': + tp_axis = 'output' + max_tp = 1 + else: + tp_axis = 'output' if weight.kind == 'moe_gate_up' else 'input' + max_tp = _moe_max_tp(model, weight) + + if is_moe: + moe = cast(MoeSpec, model.moe) + expert_num = moe.expert_num + experts_per_token = moe.experts_per_token + else: + expert_num = 0 + experts_per_token = 0 + + return ShapeSpec( + name=f'{model.name}_{weight.name}', + input_dim=weight.input_dim, + output_dim=weight.output_dim, + expert_num=expert_num, + experts_per_token=experts_per_token, + combine_experts=is_moe, + tp_axis=tp_axis, + max_tp=max_tp, + max_ep=model.max_ep if is_moe else 1, + moe_indexed=weight.kind == 'moe_gate_up', + ) + + +MODEL_SPECS: tuple[ModelSpec, ...] = load_model_specs() +ALL_SHAPES: tuple[ShapeSpec, ...] = tuple( + _make_shape(model, weight) + for model in MODEL_SPECS + for weight in model.weights +) +_SHAPE_BY_NAME = {shape.name: shape for shape in ALL_SHAPES} +SMOKE_SHAPES: tuple[ShapeSpec, ...] = (_SHAPE_BY_NAME['llama2_7b_o'],) + + +def is_supported(type_spec: TypeSpec, shape: ShapeSpec) -> bool: + """Drop type×shape cells with no TurboMind kernel / quant path. + + - uint4 MoE: gemm dispatch has no SM90 f16×u4 grouped kernel today + (``No feasible kernel ... sm90_f16_u4k128_..._ibb_...``). + - fp16×fp4: SM90 MXF4 configs are registered for bfloat16 only + (``sm90_16816_4.cu``); half×e2m1 dense hits ``..._f16_e2m1k32_f16_...`` miss. + - fp16×fp8: SM90 E4M3 configs use ``bfloat16_t`` as Tc only + (``sm90_16816_8.cu``); half Tc hits ``..._e4m3..._f16_tnt_...`` miss. + - fp16 MoE: SM90 gemm dispatch has no fp16 indexed/blocked grouped kernel + (``No feasible kernel ... sm90_f16_f16_f16_tnt_ibb_...``); the GMMA + grouped kernels are bf16-only. + """ + if type_spec.weight_type == 'uint4' and shape.expert_num > 0: + return False + if type_spec.data_type == 'fp16' and type_spec.weight_type == 'fp4_e2m1': + return False + if type_spec.data_type == 'fp16' and type_spec.weight_type == 'fp8_e4m3': + return False + if type_spec.data_type == 'fp16' and type_spec.weight_type == 'fp16' and shape.expert_num > 0: + return False + return True + + +def supports_fuse_silu(type_spec: TypeSpec, shape: ShapeSpec) -> bool: + """FP8/BF16 SM90 gate_up can optionally use block-pack + kGatedSilu. + + BF16 pairs 64-wide gate/up blocks; FP8 pairs 128-wide blocks. + """ + if not shape.name.endswith('_gate_up'): + return False + if shape.expert_num > 0 and not shape.moe_indexed: + return False + if type_spec.weight_type not in ('fp8_e4m3', 'bf16'): + return False + block = 128 if type_spec.weight_type == 'fp8_e4m3' else 64 + return shape.output_dim % (2 * block) == 0 + + +def expand_cases( + types: tuple[TypeSpec, ...] | None = None, + shapes: tuple[ShapeSpec, ...] | None = None, +) -> tuple[LinearCase, ...]: + ts = types if types is not None else ALL_TYPES + ss = shapes if shapes is not None else ALL_SHAPES + cases: list[LinearCase] = [] + for t in ts: + for s in ss: + if not is_supported(t, s): + continue + cases.append(_compose(t, s, fuse_silu=False)) + if supports_fuse_silu(t, s): + cases.append(_compose(t, s, fuse_silu=True)) + return tuple(cases) + + +ALL_CASES: tuple[LinearCase, ...] = expand_cases() + + +def case_by_name() -> dict[str, LinearCase]: + return {c.name: c for c in ALL_CASES} + + +def select_types(suite: str, type_names: tuple[str, ...] | None) -> tuple[TypeSpec, ...]: + by_name = {t.name: t for t in ALL_TYPES} + if type_names: + missing = [n for n in type_names if n not in by_name] + if missing: + raise ValueError(f'unknown_types_{missing}') + return tuple(by_name[n] for n in type_names) + if suite == 'smoke': + return SMOKE_TYPES + if suite in ('quick', 'full', 'custom'): + return ALL_TYPES + raise ValueError(f'unknown_suite_{suite}') + + +def select_shapes(suite: str, shape_names: tuple[str, ...] | None) -> tuple[ShapeSpec, ...]: + by_name = {s.name: s for s in ALL_SHAPES} + if shape_names: + missing = [n for n in shape_names if n not in by_name] + if missing: + raise ValueError(f'unknown_shapes_{missing}') + return tuple(by_name[n] for n in shape_names) + if suite == 'smoke': + return SMOKE_SHAPES + if suite in ('quick', 'full', 'custom'): + return ALL_SHAPES + raise ValueError(f'unknown_suite_{suite}') + + +def select_cases( + suite: str, + case_names: tuple[str, ...] | None = None, + type_names: tuple[str, ...] | None = None, + shape_names: tuple[str, ...] | None = None, +) -> tuple[LinearCase, ...]: + if case_names: + by_name = case_by_name() + missing = [n for n in case_names if n not in by_name] + if missing: + raise ValueError(f'unknown_cases_{missing}') + return tuple(by_name[n] for n in case_names) + return expand_cases(select_types(suite, type_names), select_shapes(suite, shape_names)) + + +def select_batches(suite: str, batches: tuple[int, ...] | None) -> tuple[int, ...]: + if batches is not None: + if not batches: + raise ValueError('batches_must_be_non_empty') + return batches + if suite == 'smoke': + return SMOKE_BATCHES + if suite == 'quick': + return QUICK_BATCHES + if suite == 'full': + return FULL_BATCHES + raise ValueError('custom_suite_requires_explicit_batches') + + +def _select_parallel_sizes(name: str, sizes: tuple[int, ...] | None) -> tuple[int, ...]: + values = sizes if sizes is not None else (1,) + if not values: + raise ValueError(f'{name}_sizes_must_be_non_empty') + if any(value <= 0 for value in values): + raise ValueError(f'{name}_sizes_must_be_positive') + return tuple(dict.fromkeys(values)) + + +def _parallel_variants( + case: LinearCase, + tps: tuple[int, ...], + eps: tuple[int, ...], +) -> tuple[tuple[int, int], ...]: + variants = [(tp, 1) for tp in tps] + if case.expert_num > 0: + variants.extend((1, ep) for ep in eps) + return tuple(dict.fromkeys(variants)) + + +def _slice_case(case: LinearCase, tp: int, ep: int) -> LinearCase | None: + if tp > case.max_tp or ep > case.max_ep: + return None + + input_dim = case.input_dim + output_dim = case.output_dim + tp_dim = input_dim if case.tp_axis == 'input' else output_dim + if tp_dim % tp: + return None + if case.tp_axis == 'input': + input_dim //= tp + else: + output_dim //= tp + + expert_num = case.expert_num + if expert_num: + if expert_num % ep: + return None + expert_num //= ep + + if case.group_size and input_dim % case.group_size: + return None + if case.fuse_silu: + block = 128 if case.weight_type == 'fp8_e4m3' else 64 + if output_dim % (2 * block): + return None + + return replace( + case, + input_dim=input_dim, + output_dim=output_dim, + expert_num=expert_num, + tp=tp, + ep=ep, + ) + + +def expand_suite( + suite: str, + case_names: tuple[str, ...] | None = None, + batches: tuple[int, ...] | None = None, + type_names: tuple[str, ...] | None = None, + shape_names: tuple[str, ...] | None = None, + tps: tuple[int, ...] | None = None, + eps: tuple[int, ...] | None = None, +) -> tuple[RunCase, ...]: + if suite not in VALID_SUITES: + raise ValueError(f'unknown_suite_{suite}') + selected_tps = _select_parallel_sizes('tp', tps) + selected_eps = _select_parallel_sizes('ep', eps) + selected_batches = select_batches(suite, batches) + runs: list[RunCase] = [] + for case in select_cases(suite, case_names, type_names, shape_names): + for tp, ep in _parallel_variants(case, selected_tps, selected_eps): + sliced = _slice_case(case, tp, ep) + if sliced is None: + continue + runs.extend(RunCase(case=sliced, batch_size=m) for m in selected_batches) + if not runs: + raise ValueError('no_supported_runs') + return tuple(runs) diff --git a/tests/turbomind/linear/fixture.py b/tests/turbomind/linear/fixture.py new file mode 100644 index 0000000000..44b9e0b050 --- /dev/null +++ b/tests/turbomind/linear/fixture.py @@ -0,0 +1,507 @@ +from __future__ import annotations + +import math +import random +from collections.abc import Iterator +from contextlib import contextmanager + +import torch + +from .cases import LinearCase +from .linear import ( + Linear, + Weight, + activation_needs_quantize, + dequantize_symm, + device_context, + link_experts, + quantize_symm, +) +from .reference import ( + apply_block_fused_silu, + block_pack_w1w3, + compare_tensors, + dense_gemm, + fused_silu_block, + moe_reference, + quantize_symm_row_fp8, +) + +# quant_vs_dequant: LlamaLinear vs torch on matched dequant oracle (incl. act +# quant round-trip when Linear quantizes activations). Gate on abs only. +# Weights are scaled by 0.1/sqrt(K) so |C| stays O(0.1) and bf16/fp16 ULP +# (~1e-3 at that magnitude) remains well under identity abs gates. +TOLERANCES = { + ('bf16', 'bf16'): {'quant_vs_dequant': {'max_abs': 1e-2, 'mean_abs': 1e-3}}, + ('fp16', 'fp16'): {'quant_vs_dequant': {'max_abs': 1e-2, 'mean_abs': 1e-3}}, + ('bf16', 'fp8_e4m3'): {'quant_vs_dequant': {'max_abs': 0.25, 'mean_abs': 0.05}}, + ('fp16', 'fp8_e4m3'): {'quant_vs_dequant': {'max_abs': 0.25, 'mean_abs': 0.05}}, + ('fp16', 'uint4'): {'quant_vs_dequant': {'max_abs': 0.25, 'mean_abs': 0.05}}, + ('bf16', 'fp4_e2m1'): {'quant_vs_dequant': {'max_abs': 0.25, 'mean_abs': 0.05}}, + ('fp16', 'fp4_e2m1'): {'quant_vs_dequant': {'max_abs': 0.25, 'mean_abs': 0.05}}, +} + + +def _weight_fill_scale(input_dim: int) -> float: + return 0.1 / math.sqrt(max(input_dim, 1)) + +_TORCH_DTYPE = { + 'bf16': torch.bfloat16, + 'fp16': torch.float16, +} + + +def sample_moe_routing( + batch_size: int, + expert_num: int, + experts_per_token: int, + device: torch.device, + seed: int = 5489, +) -> dict[str, torch.Tensor]: + """Port of testbed_v3 Route() layouts (f2n / en2f / offsets / scales). + + Sampling uses Python random.sample (not libstdc++ std::sample), but the + buffer shapes and index algebra match Route() exactly. + """ + # std::mt19937 default seed is 5489u + rng = random.Random(seed) + expert_ids: list[int] = [] + for _ in range(batch_size): + expert_ids.extend(rng.sample(range(expert_num), experts_per_token)) + + scales_cpu = torch.empty(batch_size * experts_per_token, dtype=torch.float32) + for i in range(batch_size): + tmp = [rng.uniform(1e-3, 1.0) for _ in range(experts_per_token)] + s = sum(tmp) + for e in range(experts_per_token): + scales_cpu[e * batch_size + i] = tmp[e] / s + + count = [0] * expert_num + f2i: list[list[int]] = [[] for _ in range(expert_num)] + for i, eid in enumerate(expert_ids): + count[eid] += 1 + f2i[eid].append(i) + + offsets_cpu = torch.empty(expert_num + 1, dtype=torch.int32) + offsets_cpu[0] = 0 + for i in range(expert_num): + offsets_cpu[i + 1] = offsets_cpu[i] + count[i] + + token_slots = len(expert_ids) + f2n_cpu = torch.empty(token_slots, dtype=torch.int32) + en2f_cpu = torch.empty(token_slots, dtype=torch.int32) + i = 0 + for e in range(expert_num): + for x in f2i[e]: + f2n_cpu[i] = x // experts_per_token + en = x % experts_per_token * batch_size + x // experts_per_token + en2f_cpu[en] = i + i += 1 + + # Keep routing copies blocking so the active stream can consume them immediately. + return { + 'f2n': f2n_cpu.to(device=device, non_blocking=False), + 'en2f': en2f_cpu.to(device=device, non_blocking=False), + 'offsets': offsets_cpu.to(device=device, non_blocking=False), + 'scales': scales_cpu.to(device=device, non_blocking=False), + } + + +class _CudaStreamBoundary: + + def __init__(self, stream: torch.cuda.Stream, device: torch.device): + self.stream = stream + self.device = device + + @contextmanager + def enter(self) -> Iterator[torch.cuda.Stream]: + caller = torch.cuda.current_stream(self.device) + if caller == self.stream: + yield self.stream + return + + self.stream.wait_stream(caller) + try: + with torch.cuda.stream(self.stream): + yield self.stream + finally: + caller.wait_stream(self.stream) + + def synchronize(self) -> None: + self.stream.synchronize() + + def assert_active(self) -> None: + assert torch.cuda.current_stream(self.device) == self.stream + + def clone_to_caller(self, tensor: torch.Tensor) -> torch.Tensor: + caller = torch.cuda.current_stream(self.device) + caller.wait_stream(self.stream) + out = tensor.detach().clone() + self.stream.wait_stream(caller) + return out + + +class LinearFixture: + def __init__(self, case: LinearCase, device: torch.device | None = None): + self.case = case + self.device = device or torch.device('cuda') + self._ctx = device_context() + self._ctx.__enter__() + self._stream_boundary: _CudaStreamBoundary | None = _CudaStreamBoundary( + torch.cuda.ExternalStream(self._ctx.stream_ptr, device=self.device), + self.device, + ) + self.linear: Linear | None = None + self.w_original: Weight | None = None + self.w_quant: Weight | None = None + self.w_dequant: Weight | None = None + self.e_original: list[Weight] = [] + self.e_quant: list[Weight] = [] + self.e_dequant: list[Weight] = [] + self.w_original_torch: torch.Tensor | list[torch.Tensor] | None = None + self.w_dequant_torch: torch.Tensor | list[torch.Tensor] | None = None + self.x_original: torch.Tensor | None = None + self.x_dequant: torch.Tensor | None = None + self.f2n: torch.Tensor | None = None + self.en2f: torch.Tensor | None = None + self.offsets: torch.Tensor | None = None + self.scales: torch.Tensor | None = None + self.d_original: torch.Tensor | None = None + self.d_dequant: torch.Tensor | None = None + self.d_quant: torch.Tensor | None = None + try: + with self.on_tm_stream(): + self.linear = Linear() + if case.expert_num > 0: + self._build_moe_weights() + else: + self._build_dense_weights() + except Exception: + self.close() + raise + + def close(self) -> None: + # Join streams, then drop TM-backed torch tensors while the Context mempool + # is still alive (from_dlpack / Forward outputs alias pool storage). + if self._stream_boundary is not None: + with self.on_tm_stream(): + pass + self.sync_tm() + torch.cuda.current_stream(self.device).synchronize() + self.d_quant = None + self.x_dequant = None + self.d_original = None + self.d_dequant = None + self.x_original = None + self.f2n = None + self.en2f = None + self.offsets = None + self.scales = None + self.w_original_torch = None + self.w_dequant_torch = None + # Destroy stream users before leaving the device Context. + # Fused views alias expert storage — drop fused first, then experts. + self.linear = None + self.w_original = None + self.w_quant = None + self.w_dequant = None + self.e_original = [] + self.e_quant = [] + self.e_dequant = [] + self._stream_boundary = None + ctx = getattr(self, '_ctx', None) + if ctx is not None: + self._ctx = None + ctx.__exit__(None, None, None) + + @contextmanager + def on_tm_stream(self) -> Iterator[torch.cuda.Stream]: + assert self._stream_boundary is not None + with self._stream_boundary.enter() as stream: + yield stream + + def sync_tm(self) -> None: + """Host wait until Context::stream() is idle (e.g. before del / set_measure).""" + assert self._stream_boundary is not None + self._stream_boundary.synchronize() + + def _torch_dtype(self) -> torch.dtype: + try: + return _TORCH_DTYPE[self.case.data_type] + except KeyError as e: + raise NotImplementedError(f'torch_dtype_for_{self.case.data_type}') from e + + def _clone_tm_weight(self, w: Weight) -> torch.Tensor: + """Clone TM weight storage into caller-stream-owned torch storage.""" + assert self._stream_boundary is not None + return self._stream_boundary.clone_to_caller(w.weight_tensor()) + + def _fill_weight_triple( + self, + w_original: Weight, + w_quant: Weight, + w_dequant: Weight, + w: torch.Tensor, + ) -> torch.Tensor: + """Quantize into w_quant / w_dequant; return torch dequant view (pre- + prepare).""" + c = self.case + src = w.contiguous() + with self.on_tm_stream(): + sp = self._ctx.stream_ptr + w_original.copy_weight_from(src, stream_ptr=sp) + if c.weight_type == c.data_type: + w_quant.copy_weight_from(src, stream_ptr=sp) + w_dequant.copy_weight_from(src, stream_ptr=sp) + elif c.weight_type == 'fp8_e4m3': + w_quant.quantize_symm_block_from(w_original) + w_dequant.dequantize_symm_block_from(w_quant) + elif c.weight_type in ('uint4', 'fp4_e2m1'): + w_quant.quantize_groupwise_into(w_original, w_dequant) + else: + raise NotImplementedError(f'weight_type_{c.weight_type}') + if c.weight_type == c.data_type: + return w.clone() + return self._clone_tm_weight(w_dequant) + + def _make_random_weight(self) -> torch.Tensor: + """Random weight; fuse_silu uses the dtype-specific gate/up block + layout.""" + c = self.case + dtype = self._torch_dtype() + scale = _weight_fill_scale(c.input_dim) + if c.fuse_silu: + inter = c.output_dim // 2 + w1 = torch.randn(c.input_dim, inter, device=self.device, dtype=dtype) * scale + w3 = torch.randn(c.input_dim, inter, device=self.device, dtype=dtype) * scale + return block_pack_w1w3(w1, w3, fused_silu_block(c.weight_type)) + return torch.randn(c.input_dim, c.output_dim, device=self.device, dtype=dtype) * scale + + def _apply_fuse_silu_epilogue(self, weight: Weight) -> None: + from .linear import _tm + + weight.set_epilogue(_tm().Epilogue.kGatedSilu) + + def _allocate_weight_triple(self) -> tuple[Weight, Weight, Weight]: + c = self.case + with self.on_tm_stream(): + return ( + Weight(c.input_dim, c.output_dim, c.data_type, c.data_type, 0), + Weight(c.input_dim, c.output_dim, c.data_type, c.weight_type, c.group_size), + Weight(c.input_dim, c.output_dim, c.data_type, c.data_type, 0), + ) + + def _make_weight_triple(self) -> tuple[Weight, Weight, Weight, torch.Tensor, torch.Tensor]: + c = self.case + w_original, w_quant, w_dequant = self._allocate_weight_triple() + w = self._make_random_weight() + w_deq_torch = self._fill_weight_triple(w_original, w_quant, w_dequant, w) + # Clone dequant before prepare(): groupwise/fp8 prepare may repack storage. + with self.on_tm_stream(): + w_original.prepare() + w_quant.prepare() + w_dequant.prepare() + if c.fuse_silu: + self._apply_fuse_silu_epilogue(w_quant) + return w_original, w_quant, w_dequant, w, w_deq_torch + + def _build_dense_weights(self) -> None: + w_o, w_q, w_d, w_torch, w_deq = self._make_weight_triple() + self.w_original = w_o + self.w_quant = w_q + self.w_dequant = w_d + self.w_original_torch = w_torch + self.w_dequant_torch = w_deq + self.sync_tm() + + def _build_moe_weights(self) -> None: + """Build per-expert quant Weights and torch reference tensors. + + Prepare is deferred until every expert has been quantized: preparing + quantized weights can free/transpose scale storage and corrupt peers. + """ + c = self.case + orig_torch: list[torch.Tensor] = [] + deq_torch: list[torch.Tensor] = [] + for _ in range(c.expert_num): + w_original, w_quant, w_dequant = self._allocate_weight_triple() + w = self._make_random_weight() + w_deq = self._fill_weight_triple(w_original, w_quant, w_dequant, w) + orig_torch.append(w) + deq_torch.append(w_deq) + self.e_quant.append(w_quant) + self.sync_tm() + del w_original, w_dequant + with self.on_tm_stream(): + for w_quant in self.e_quant: + # Production MoE sets this via FfnWeight::prepare (is_expert_). + # Required so GetConverters packs bf16/fp16 B for Config_F16_g. + w_quant.set_grouped(True) + w_quant.prepare() + if c.fuse_silu: + self._apply_fuse_silu_epilogue(w_quant) + self.sync_tm() + fused = link_experts(self.e_quant) + if c.fuse_silu: + self._apply_fuse_silu_epilogue(fused) + self.w_quant = fused + self.sync_tm() + self.w_original_torch = orig_torch + self.w_dequant_torch = deq_torch + + def _prepare_x_dequant(self) -> None: + """Match LlamaLinear GetOperandA: QuantizeSymm then DequantizeSymm. + + On SM90, FP8 weights derive input_format=fp8 even when the caller passes + bf16 activations. The dequant oracle must use the same act quant round-trip + or quant_vs_dequant is dominated by act error (~0.1 mean abs), not GEMM. + """ + assert self.x_original is not None + assert self.w_quant is not None + if not activation_needs_quantize(self.w_quant): + self.x_dequant = None + return + from .linear import _tm + + with self.on_tm_stream(): + x_tm = _tm().from_dlpack_with_strides(self.x_original) + x_q, x_s = quantize_symm(x_tm) + x_d = dequantize_symm(x_q, x_s) + x_d_torch = torch.from_dlpack(x_d).to(dtype=self.x_original.dtype) + # Clone so x_dequant never aliases Context-pool storage (same-dtype .to is a no-op). + assert self._stream_boundary is not None + self.x_dequant = self._stream_boundary.clone_to_caller(x_d_torch) + + def prepare_batch(self, batch_size: int) -> None: + c = self.case + dtype = self._torch_dtype() + x_tokens = torch.randn(batch_size, c.input_dim, device=self.device, dtype=dtype) + # testbed_v3 still feeds data_type activations into Forward; input_type + # only changes how the dequant reference is formed. We always match + # LlamaLinear's act quant via activation_needs_quantize(). + if c.input_type not in (c.data_type, 'fp8_e4m3'): + raise NotImplementedError(f'input_type_{c.input_type}') + if c.expert_num > 0: + route = sample_moe_routing(batch_size, c.expert_num, c.experts_per_token, self.device) + self.f2n = route['f2n'] + self.en2f = route['en2f'] + self.offsets = route['offsets'] + self.scales = route['scales'] + if c.moe_indexed: + # w1/gate: token-major x + f2n indices. + self.x_original = x_tokens + else: + # w2/down: expert-packed x, offsets only (matches moe_ffn_layer). + self.x_original = x_tokens[self.f2n.long()].contiguous() + else: + self.f2n = None + self.en2f = None + self.offsets = None + self.scales = None + self.x_original = x_tokens + self._prepare_x_dequant() + + def run_reference(self) -> None: + assert self.x_original is not None + assert self.w_original_torch is not None + assert self.w_dequant_torch is not None + c = self.case + if c.expert_num > 0: + assert self.offsets is not None + assert self.scales is not None and self.en2f is not None + assert isinstance(self.w_original_torch, list) + assert isinstance(self.w_dequant_torch, list) + x_deq = self.x_dequant if self.x_dequant is not None else self.x_original + f2n = self.f2n if c.moe_indexed else None + # Validate the grouped GEMM (packed expert-major outputs). + self.d_original = moe_reference( + self.x_original, + self.w_original_torch, + f2n, + self.offsets, + self.scales, + self.en2f, + c.experts_per_token, + False, + ) + self.d_dequant = moe_reference( + x_deq, + self.w_dequant_torch, + f2n, + self.offsets, + self.scales, + self.en2f, + c.experts_per_token, + False, + ) + else: + assert isinstance(self.w_original_torch, torch.Tensor) + assert isinstance(self.w_dequant_torch, torch.Tensor) + self.d_original = dense_gemm(self.x_original, self.w_original_torch) + x_deq = self.x_dequant if self.x_dequant is not None else self.x_original + self.d_dequant = dense_gemm(x_deq, self.w_dequant_torch) + if c.fuse_silu: + block = fused_silu_block(c.weight_type) + self.d_original = apply_block_fused_silu(self.d_original, block) + self.d_dequant = apply_block_fused_silu(self.d_dequant, block) + # SM90 FP8 fused path quantizes after SiLU; reference matches QuantizeSymm. + if c.weight_type == 'fp8_e4m3': + _, _, self.d_original = quantize_symm_row_fp8(self.d_original) + _, _, self.d_dequant = quantize_symm_row_fp8(self.d_dequant) + + def _forward_linear(self) -> tuple[torch.Tensor, torch.Tensor | None]: + assert self._stream_boundary is not None + self._stream_boundary.assert_active() + assert self.linear is not None + assert self.x_original is not None + assert self.w_quant is not None + c = self.case + if c.expert_num > 0: + assert self.offsets is not None + indices = self.f2n if c.moe_indexed else None + return self.linear.forward_moe( + self.x_original, self.w_quant, indices, self.offsets) + return self.linear.forward_dense(self.x_original, self.w_quant) + + def run_linear(self) -> None: + with self.on_tm_stream(): + out, out_scales = self._forward_linear() + # FP8 fused path is already dequantized to bf16 in Linear._pack_forward_result. + assert self._stream_boundary is not None + self.d_quant = self._stream_boundary.clone_to_caller(out) + del out, out_scales + with self.on_tm_stream(): + self.release_forward_result() + + def release_forward_result(self) -> None: + assert self._stream_boundary is not None + self._stream_boundary.assert_active() + assert self.linear is not None + self.linear.release_forward_result() + + def run_linear_forward(self) -> None: + """Launch one forward while on_tm_stream() owns the active stream.""" + self._forward_linear() + + def compare(self) -> dict[str, dict[str, float]]: + assert self.d_quant is not None + assert self.d_dequant is not None + assert self.d_original is not None + return { + 'quant_vs_dequant': compare_tensors(self.d_quant, self.d_dequant), + 'quant_vs_original': compare_tensors(self.d_quant, self.d_original), + 'dequant_vs_original': compare_tensors(self.d_dequant, self.d_original), + } + + def check_tolerances(self, metrics: dict[str, dict[str, float]]) -> None: + key = (self.case.data_type, self.case.weight_type) + gates = TOLERANCES.get(key) + if gates is None: + raise ValueError(f'no_tolerances_for_{key}') + for pair, limits in gates.items(): + for metric_name, limit in limits.items(): + value = metrics[pair][metric_name] + if not math.isfinite(value): + raise AssertionError(f'{pair}.{metric_name}={value} (non-finite)') + if value > limit: + raise AssertionError(f'{pair}.{metric_name}={value} > {limit}') diff --git a/tests/turbomind/linear/kernel_usage_report.py b/tests/turbomind/linear/kernel_usage_report.py new file mode 100755 index 0000000000..bb2d96ac29 --- /dev/null +++ b/tests/turbomind/linear/kernel_usage_report.py @@ -0,0 +1,73 @@ +#!/usr/bin/env python3 +"""Aggregate SM90 BF16 kernel usage from a bench_linear log. + +`DispatchCache::Impl::Summary` (src/turbomind/kernels/gemm/dispatch_cache.cu) prints one block of +``kernel_name: count`` lines every time `Gemm::Export` is called. This script sums those counts +across all per-case exports emitted by a full-suite run and reports used/unused SM90 BF16 kernels. +""" + +from __future__ import annotations + +import argparse +import json +import re +import sys +from pathlib import Path + +DEFAULT_PATTERN = re.compile(r'^(sm90_bf16_bf16_bf16_.+):\s+(\d+)$') + + +def parse_log(path: Path, pattern: re.Pattern[str]) -> dict[str, int]: + counts: dict[str, int] = {} + for line in path.read_text().splitlines(): + m = pattern.match(line.strip()) + if not m: + continue + name = m.group(1) + counts[name] = counts.get(name, 0) + int(m.group(2)) + return counts + + +def main(argv: list[str] | None = None) -> int: + p = argparse.ArgumentParser(description=__doc__) + p.add_argument('log', type=Path, help='bench_linear log file') + p.add_argument('--pattern', default=r'^(sm90_bf16_bf16_bf16_.+):\s+(\d+)$', + help='regex with two groups (name, count)') + p.add_argument('--json', type=Path, default=None, help='write aggregated counts to JSON') + p.add_argument('--unused-only', action='store_true', help='only print unused kernels') + args = p.parse_args(argv) + + pattern = re.compile(args.pattern) + counts = parse_log(args.log, pattern) + if not counts: + print('No SM90 BF16 kernel usage lines found in log.', file=sys.stderr) + return 1 + + used = {name: c for name, c in counts.items() if c > 0} + unused = {name: c for name, c in counts.items() if c == 0} + + if not args.unused_only: + print(f'SM90 BF16 kernels found in log: {len(counts)}') + print(f' used : {len(used)}') + print(f' unused : {len(unused)}') + print(f' total dispatch records across all case exports: {sum(counts.values())}') + print() + print('Used kernels (sorted by count):') + for name, c in sorted(used.items(), key=lambda kv: (-kv[1], kv[0])): + print(f' {c:6d} {name}') + print() + + print('Unused kernels (candidates for pruning):') + for name in sorted(unused): + print(f' 0 {name}') + + if args.json: + args.json.write_text(json.dumps({ + 'used': used, + 'unused': unused, + }, indent=2, sort_keys=True)) + return 0 + + +if __name__ == '__main__': + raise SystemExit(main()) diff --git a/tests/turbomind/linear/linear.py b/tests/turbomind/linear/linear.py new file mode 100644 index 0000000000..131e70ae87 --- /dev/null +++ b/tests/turbomind/linear/linear.py @@ -0,0 +1,415 @@ +from __future__ import annotations + +import ctypes +from collections.abc import Sequence + +import torch + +_TM = None + + +def _tm(): + global _TM + if _TM is None: + import _turbomind as tm + _TM = tm + return _TM + + +def is_available() -> bool: + try: + _tm() + return True + except Exception: + return False + + +_DTYPE_TO_TM = { + 'bf16': 'TYPE_BF16', + 'fp16': 'TYPE_FP16', + 'fp8_e4m3': 'TYPE_FP8_E4M3', + 'uint4': 'TYPE_UINT4', + 'fp4_e2m1': 'TYPE_FP4_E2M1', +} + + +def to_tm_dtype(name: str): + tm = _tm() + key = _DTYPE_TO_TM[name] + return getattr(tm.DataType, key) + + +def _resolve_block_sizes(weight_type: str, group_size: int) -> tuple[int, int]: + # Match ResolveLinearWeightFormat rules in data_format.cc. + if weight_type == 'fp8_e4m3': + return 128, 128 + if weight_type in ('uint4', 'fp4_e2m1'): + return (group_size or 1), 1 + return 1, 1 + + +def quantize_symm_block(src, out=None, scale=None): + """Thin wrapper over tm.QuantizeSymmBlock (Context stream). + + Returns (out, scale) TM tensors. Caller must Param.set them onto Weight when scale was newly allocated (typical for + empty scales). + """ + return _tm().QuantizeSymmBlock(out=out, scale=scale, src=src) + + +def dequantize_symm_block(src, scale, out=None): + """Thin wrapper over tm.DequantizeSymmBlock (Context stream).""" + return _tm().DequantizeSymmBlock(out=out, src=src, scale=scale) + + +def quantize_symm(src, out=None, scale=None): + """Thin wrapper over tm.QuantizeSymm (row/group act quant, Context + stream).""" + return _tm().QuantizeSymm(out=out, scale=scale, src=src) + + +def dequantize_symm(src, scale, out=None): + """Thin wrapper over tm.DequantizeSymm (Context stream).""" + return _tm().DequantizeSymm(out=out, src=src, scale=scale) + + +def quantize_groupwise( + *, + quant, + scales, + dequant, + src, + group_size: int, + zeros=None, + rbits=None, +) -> None: + """Thin wrapper over tm.QuantizeGroupwise (Context stream). + + Callers must pass K-major views (``.t()``) matching testbed_v3 GenerateWeight. + """ + _tm().QuantizeGroupwise( + quant=quant, + scales=scales, + zeros=zeros, + dequant=dequant, + src=src, + rbits=rbits, + group_size=group_size, + ) + + +def activation_needs_quantize(weight: Weight) -> bool: + """True when LlamaLinear will QuantizeSymm activations for this weight.""" + return weight._impl.input_format.dtype != weight._impl.data_type + + +def tensor_data_ptr(tm_tensor) -> int: + """Device pointer from a TurboMind Tensor `.data` capsule.""" + cap = tm_tensor.data + fn = ctypes.pythonapi.PyCapsule_GetPointer + fn.restype = ctypes.c_void_p + fn.argtypes = [ctypes.py_object, ctypes.c_char_p] + ptr = fn(cap, None) + if not ptr: + raise RuntimeError('null_tensor_data_ptr') + return int(ptr) + + +def make_strided_ptrs(ptrs: Sequence[tuple[int, int]], dtype): + """Build owned StridedPtr tensor via tm.MakeStridedPtrs (Context + stream).""" + return _tm().MakeStridedPtrs(list(ptrs), dtype) + + +def invoke_moe_dispatch(src: torch.Tensor, f2n: torch.Tensor, experts_per_token: int, out: torch.Tensor | None = None): + tm = _tm() + y = tm.invokeMoeDispatch( + out=None if out is None else tm.from_dlpack_with_strides(out), + src=tm.from_dlpack_with_strides(src), + f2n=tm.from_dlpack_with_strides(f2n), + expert_per_token=experts_per_token, + ) + return torch.from_dlpack(y) if out is None else out + + +def invoke_moe_combine( + out: torch.Tensor, + src: torch.Tensor, + scales: torch.Tensor, + en2f: torch.Tensor, + experts_per_token: int, + *, + bias: torch.Tensor | None = None, + f2E: torch.Tensor | None = None, + dst_scales: torch.Tensor | None = None, + bscale: float = 1.0, + dst_scale: float = 0.0, +) -> torch.Tensor: + """Mirror testbed_v3 Run() invokeMoeCombine (out must be preallocated).""" + tm = _tm() + tm.invokeMoeCombine( + out=tm.from_dlpack_with_strides(out), + src=tm.from_dlpack_with_strides(src), + bias=None if bias is None else tm.from_dlpack_with_strides(bias), + scales=tm.from_dlpack_with_strides(scales), + en2f=tm.from_dlpack_with_strides(en2f), + f2E=None if f2E is None else tm.from_dlpack_with_strides(f2E), + dst_scales=None if dst_scales is None else tm.from_dlpack_with_strides(dst_scales), + experts_per_token=experts_per_token, + bscale=bscale, + dst_scale=dst_scale, + ) + return out + + +def link_experts(experts: Sequence[Weight]) -> Weight: + """Port of testbed_v3 / moe_weight LinkExperts into a fused Weight view. + + Experts must already be prepare()'d and kept alive for as long as the fused view is used (pointers alias expert + storage). Call on the Context stream. + """ + if not experts: + raise ValueError('link_experts_requires_non_empty') + e0 = experts[0] + fused = Weight( + e0._input_dim, + e0._output_dim, + e0._data_type, + e0._weight_type, + e0._group_size, + has_bias=e0._has_bias, + ) + e0._impl.copy_metadata_to(fused._impl) + + n = len(experts) + fused._impl.k_desc.num = n + fused._impl.q_desc.num = n + + weights: list[tuple[int, int]] = [] + scales: list[tuple[int, int]] = [] + for e in experts: + weights.append((tensor_data_ptr(e.param_tensor('weight')), int(e._impl.k_desc.ld))) + scales_t = e.param_tensor('scales') + if scales_t: + scales.append((tensor_data_ptr(scales_t), int(e._impl.q_desc.ld))) + + fused.set_param( + 'weight', + make_strided_ptrs(weights, fused._impl.weight_format.dtype), + ) + if scales: + fused.set_param('scales', make_strided_ptrs(scales, e0.param_tensor('scales').type)) + fused._impl.k_desc.ld = 0 + fused._impl.q_desc.ld = 0 + fused._impl.k_desc.offsets = 0 + fused._impl.q_desc.offsets = 0 + return fused + + +class Weight: + """Adapter over tm.LinearWeight. + + Instances use the TurboMind Context stream; destroy them before exiting the enclosing device_context(). + """ + + def __init__( + self, + input_dim: int, + output_dim: int, + data_type: str, + weight_type: str, + group_size: int = 0, + has_bias: bool = False, + ) -> None: + tm = _tm() + dt = to_tm_dtype(data_type) + wt = to_tm_dtype(weight_type) + block_in, block_out = _resolve_block_sizes(weight_type, group_size) + cfg = tm.LinearConfig() + cfg.input_dim = input_dim + cfg.output_dim = output_dim + cfg.data_type = dt + cfg.format = tm.ResolveLinearWeightFormat(dt, wt, block_in, block_out) + cfg.has_bias = has_bias + self._input_dim = input_dim + self._output_dim = output_dim + self._data_type = data_type + self._weight_type = weight_type + self._group_size = group_size + self._has_bias = has_bias + self._impl = tm.LinearWeight(cfg) + self._impl.param('weight').alloc([input_dim, output_dim], wt) + # Groupwise formats need scales (/zeros) in MN-major (K/g, N). + if weight_type in ('uint4', 'fp4_e2m1'): + if group_size <= 0: + raise ValueError('groupwise_weight_requires_positive_group_size') + if input_dim % group_size != 0: + raise ValueError(f'input_dim_{input_dim}_not_divisible_by_group_size_{group_size}') + scale_shape = [input_dim // group_size, output_dim] + # uint4: f16/bf16 scales+zeros; fp4: ue8m0 scales (uint8). + if weight_type == 'fp4_e2m1': + scale_dtype = tm.DataType.TYPE_UINT8 + else: + scale_dtype = dt + self._impl.param('scales').alloc(scale_shape, scale_dtype) + if weight_type == 'uint4': + self._impl.param('zeros').alloc(scale_shape, scale_dtype) + + def set_grouped(self, grouped: bool) -> None: + """Mark expert weights for grouped-GEMM layout conversion in prepare(). + + Mirrors FfnWeight::prepare for is_expert_: without this, bf16/fp16 MoE + weights stay flat and SM90 Config_F16_g (ibb + packed B) cannot match. + """ + self._impl.set_grouped(grouped) + + def set_epilogue(self, epilogue) -> None: + """Set GEMM epilogue (e.g. tm.Epilogue.kGatedSilu for fused SiLU).""" + tm = _tm() + self._impl.epilogue = epilogue + # Mirror FfnWeight::prepare: SM90 FP8 fused SiLU writes FP8 + group scales. + major, _ = torch.cuda.get_device_capability() + if (epilogue == tm.Epilogue.kGatedSilu and self._weight_type == 'fp8_e4m3' and major == 9): + self._impl.set_fp8_fused_silu_output() + + def prepare(self) -> None: + self._impl.prepare() + + def weight_tensor(self) -> torch.Tensor: + return torch.from_dlpack(self._impl.param('weight').get()) + + def copy_weight_from(self, src: torch.Tensor, *, stream_ptr: int) -> None: + """Copy contiguous ``src`` into this weight on ``stream_ptr`` (Context::stream). + + Must not use Tensor.copy_from: that path is default-stream cudaMemcpy and + cannot participate in the harness's scoped Context-stream ownership. + + ``src`` must be materialized and contiguous before entering the Context-stream + boundary. This method will not create caller-owned storage inside that scope. + """ + if not src.is_contiguous(): + raise ValueError('copy_weight_from_requires_contiguous_src') + tm = _tm() + src_tm = tm.from_dlpack(src) + tm.generic_copy_on_stream(src_tm, self.param_tensor('weight'), stream_ptr) + + def param_tensor(self, name: str): + return self._impl.param(name).get() + + def set_param(self, name: str, tensor) -> None: + self._impl.param(name).set(tensor) + + def quantize_symm_block_from(self, src_weight: Weight) -> None: + out, scale = quantize_symm_block( + src_weight.param_tensor('weight'), + out=self.param_tensor('weight'), + scale=None, + ) + self.set_param('weight', out) + self.set_param('scales', scale) + + def dequantize_symm_block_from(self, quant_weight: Weight) -> None: + out = dequantize_symm_block( + quant_weight.param_tensor('weight'), + quant_weight.param_tensor('scales'), + out=self.param_tensor('weight'), + ) + self.set_param('weight', out) + + def quantize_groupwise_into(self, src_weight: Weight, dequant_weight: Weight) -> None: + """Port of testbed_v3 GenerateWeight uint4/fp4 path (K-major via + .t()).""" + zeros = self.param_tensor('zeros') if self._weight_type == 'uint4' else None + quantize_groupwise( + quant=self.param_tensor('weight').t(), + scales=self.param_tensor('scales').t(), + zeros=None if zeros is None else zeros.t(), + dequant=dequant_weight.param_tensor('weight').t(), + src=src_weight.param_tensor('weight').t(), + group_size=self._group_size, + ) + + +class Linear: + """Adapter over tm.LlamaLinear. + + Instances use the TurboMind Context stream; destroy them before exiting the enclosing device_context(). + """ + + def __init__(self) -> None: + self._impl = _tm().LlamaLinear() + # TM tensors backing the latest forward dlpack views (Context pool). + self._forward_keep_alive: tuple | None = None + + def forward_dense( + self, + x: torch.Tensor, + weight: Weight, + out: torch.Tensor | None = None, + input_scales: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + tm = _tm() + xin = tm.from_dlpack_with_strides(x) + oin = None if out is None else tm.from_dlpack_with_strides(out) + iscales = None if input_scales is None else tm.from_dlpack_with_strides(input_scales) + y, ys = self._impl.forward_dense(xin, weight._impl, oin, iscales, None) + return self._pack_forward_result(y, ys) + + def forward_moe( + self, + x: torch.Tensor, + weight: Weight, + f2n: torch.Tensor | None, + offsets: torch.Tensor, + out: torch.Tensor | None = None, + input_scales: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + tm = _tm() + y, ys = self._impl.forward_moe( + tm.from_dlpack_with_strides(x), + weight._impl, + None if f2n is None else tm.from_dlpack_with_strides(f2n), + tm.from_dlpack_with_strides(offsets), + None if out is None else tm.from_dlpack_with_strides(out), + None if input_scales is None else tm.from_dlpack_with_strides(input_scales), + None, + ) + return self._pack_forward_result(y, ys) + + def _pack_forward_result(self, y, ys) -> tuple[torch.Tensor, torch.Tensor | None]: + """Return torch views of Context-stream results (FP8 dequantized to + bf16). + + Views alias TM pool storage. Callers must transfer them through LinearFixture's stream boundary before releasing + the backing TM tensors. + """ + if ys: + y_bf16 = dequantize_symm(y, ys) + self._forward_keep_alive = (y, ys, y_bf16) + y_t = torch.from_dlpack(y_bf16) + if y_t.dtype != torch.bfloat16: + raise TypeError(f'expected bf16 dequant output, got {y_t.dtype}') + return y_t, torch.from_dlpack(ys) + self._forward_keep_alive = (y, ys) + return torch.from_dlpack(y), None + + def release_forward_result(self) -> None: + self._forward_keep_alive = None + + def set_measure(self, on: bool) -> None: + self._impl.set_measure(on) + + def import_records(self, path: str) -> int: + return int(self._impl.import_records(path)) + + def export_records(self, path: str) -> int: + return int(self._impl.export_records(path)) + + +def device_context(): + """Enter/exit a TurboMind device Context (stream). + + Objects that use the Context stream — notably Linear/LlamaLinear and Weight/LinearWeight — must be destroyed before + exiting this context. + """ + return _tm().create_device_context() diff --git a/tests/turbomind/linear/model_specs.py b/tests/turbomind/linear/model_specs.py new file mode 100644 index 0000000000..01a94bd8c1 --- /dev/null +++ b/tests/turbomind/linear/model_specs.py @@ -0,0 +1,112 @@ +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Literal + +import yaml + +TpAxis = Literal['input', 'output'] +WeightKind = Literal['dense', 'replicated', 'moe_gate_up', 'moe_down'] + +SUPPORTED_TP_SIZES = (8, 4, 2, 1) +WEIGHT_KINDS = ('dense', 'replicated', 'moe_gate_up', 'moe_down') +MODEL_SPECS_PATH = Path(__file__).with_name('models.yaml') + + +@dataclass(frozen=True, kw_only=True) +class MoeSpec: + expert_num: int + experts_per_token: int + + def __post_init__(self) -> None: + if self.expert_num <= 0: + raise ValueError('expert_num_must_be_positive') + if self.experts_per_token <= 0: + raise ValueError('experts_per_token_must_be_positive') + if self.experts_per_token > self.expert_num: + raise ValueError('experts_per_token_exceeds_expert_num') + + +@dataclass(frozen=True, kw_only=True) +class WeightSpec: + name: str + kind: WeightKind + output_dim: int + input_dim: int + tp_axis: TpAxis | None = None + + def __post_init__(self) -> None: + if not self.name: + raise ValueError('weight_name_must_be_non_empty') + if self.kind not in WEIGHT_KINDS: + raise ValueError(f'unsupported_weight_kind_{self.kind}') + if self.output_dim <= 0 or self.input_dim <= 0: + raise ValueError('weight_dims_must_be_positive') + if self.kind == 'dense': + if self.tp_axis not in ('input', 'output'): + raise ValueError('dense_weight_requires_tp_axis') + elif self.tp_axis is not None: + raise ValueError(f'{self.kind}_weight_forbids_tp_axis') + if self.kind == 'moe_gate_up' and self.output_dim % 2: + raise ValueError('moe_gate_up_output_dim_must_be_even') + + +@dataclass(frozen=True, kw_only=True) +class ModelSpec: + name: str + max_tp: int + max_ep: int + weights: tuple[WeightSpec, ...] + moe: MoeSpec | None = None + + def __post_init__(self) -> None: + if not self.name: + raise ValueError('model_name_must_be_non_empty') + if self.max_tp not in SUPPORTED_TP_SIZES: + raise ValueError(f'unsupported_model_max_tp_{self.max_tp}') + if self.max_ep <= 0: + raise ValueError('model_max_ep_must_be_positive') + if not self.weights: + raise ValueError('model_weights_must_be_non_empty') + + names = tuple(weight.name for weight in self.weights) + if len(names) != len(set(names)): + raise ValueError(f'duplicate_weight_names_{self.name}') + + has_moe_weights = any(weight.kind in ('moe_gate_up', 'moe_down') for weight in self.weights) + if has_moe_weights != (self.moe is not None): + raise ValueError(f'moe_metadata_and_weights_disagree_{self.name}') + if self.moe is None: + if self.max_ep != 1: + raise ValueError(f'dense_model_requires_max_ep_1_{self.name}') + return + + if self.moe.expert_num % self.max_ep: + raise ValueError(f'expert_num_not_divisible_by_max_ep_{self.name}') + local_experts = self.moe.expert_num // self.max_ep + if local_experts <= 1: + raise ValueError(f'expert_parallel_requires_multiple_local_experts_{self.name}') + if local_experts < self.moe.experts_per_token: + raise ValueError(f'local_experts_less_than_experts_per_token_{self.name}') + + +def _make_weight(config: dict[str, Any]) -> WeightSpec: + return WeightSpec(**config) + + +def _make_model(config: dict[str, Any]) -> ModelSpec: + values = dict(config) + moe_config = values.get('moe') + values['moe'] = MoeSpec(**moe_config) if moe_config is not None else None + values['weights'] = tuple(_make_weight(weight) for weight in values['weights']) + return ModelSpec(**values) + + +def load_model_specs(path: Path = MODEL_SPECS_PATH) -> tuple[ModelSpec, ...]: + config = yaml.safe_load(path.read_text()) + models = tuple(_make_model(model) for model in config['models']) + names = tuple(model.name for model in models) + if len(names) != len(set(names)): + raise ValueError('duplicate_model_names') + return models diff --git a/tests/turbomind/linear/models.yaml b/tests/turbomind/linear/models.yaml new file mode 100644 index 0000000000..b63e693646 --- /dev/null +++ b/tests/turbomind/linear/models.yaml @@ -0,0 +1,298 @@ +models: + - name: llama2_7b + max_tp: 2 + max_ep: 1 + weights: + - {name: gate_up, kind: dense, output_dim: 22016, input_dim: 4096, tp_axis: output} + - {name: down, kind: dense, output_dim: 4096, input_dim: 11008, tp_axis: input} + - {name: qkv, kind: dense, output_dim: 12288, input_dim: 4096, tp_axis: output} + - {name: o, kind: dense, output_dim: 4096, input_dim: 4096, tp_axis: input} + - name: llama3_8b + max_tp: 2 + max_ep: 1 + weights: + - {name: gate_up, kind: dense, output_dim: 28672, input_dim: 4096, tp_axis: output} + - {name: down, kind: dense, output_dim: 4096, input_dim: 14336, tp_axis: input} + - {name: qkv, kind: dense, output_dim: 6144, input_dim: 4096, tp_axis: output} + - {name: o, kind: dense, output_dim: 4096, input_dim: 4096, tp_axis: input} + - name: internlm2_20b + max_tp: 4 + max_ep: 1 + weights: + - {name: gate_up, kind: dense, output_dim: 32768, input_dim: 6144, tp_axis: output} + - {name: down, kind: dense, output_dim: 6144, input_dim: 16384, tp_axis: input} + - {name: qkv, kind: dense, output_dim: 8192, input_dim: 6144, tp_axis: output} + - {name: o, kind: dense, output_dim: 6144, input_dim: 6144, tp_axis: input} + - name: glm4_9b + max_tp: 2 + max_ep: 1 + weights: + - {name: gate_up, kind: dense, output_dim: 27392, input_dim: 4096, tp_axis: output} + - {name: down, kind: dense, output_dim: 4096, input_dim: 13696, tp_axis: input} + - {name: qkv, kind: dense, output_dim: 4608, input_dim: 4096, tp_axis: output} + - {name: o, kind: dense, output_dim: 4096, input_dim: 4096, tp_axis: input} + - name: qwen2_7b + max_tp: 2 + max_ep: 1 + weights: + - {name: gate_up, kind: dense, output_dim: 37888, input_dim: 3584, tp_axis: output} + - {name: down, kind: dense, output_dim: 3584, input_dim: 18944, tp_axis: input} + - {name: qkv, kind: dense, output_dim: 4608, input_dim: 3584, tp_axis: output} + - {name: o, kind: dense, output_dim: 3584, input_dim: 3584, tp_axis: input} + - name: yi_34b + max_tp: 4 + max_ep: 1 + weights: + - {name: gate_up, kind: dense, output_dim: 40960, input_dim: 7168, tp_axis: output} + - {name: down, kind: dense, output_dim: 7168, input_dim: 20480, tp_axis: input} + - {name: qkv, kind: dense, output_dim: 9216, input_dim: 7168, tp_axis: output} + - {name: o, kind: dense, output_dim: 7168, input_dim: 7168, tp_axis: input} + - name: llama_70b + max_tp: 8 + max_ep: 1 + weights: + - {name: gate_up, kind: dense, output_dim: 57344, input_dim: 8192, tp_axis: output} + - {name: down, kind: dense, output_dim: 8192, input_dim: 28672, tp_axis: input} + - {name: qkv, kind: dense, output_dim: 10240, input_dim: 8192, tp_axis: output} + - {name: o, kind: dense, output_dim: 8192, input_dim: 8192, tp_axis: input} + - name: qwen2_72b + max_tp: 8 + max_ep: 1 + weights: + - {name: gate_up, kind: dense, output_dim: 59392, input_dim: 8192, tp_axis: output} + - {name: down, kind: dense, output_dim: 8192, input_dim: 29696, tp_axis: input} + - {name: qkv, kind: dense, output_dim: 10240, input_dim: 8192, tp_axis: output} + - {name: o, kind: dense, output_dim: 8192, input_dim: 8192, tp_axis: input} + - name: mixtral_8x7b + max_tp: 4 + max_ep: 4 + moe: + expert_num: 8 + experts_per_token: 2 + weights: + - {name: gate_up, kind: moe_gate_up, output_dim: 28672, input_dim: 4096} + - {name: down, kind: moe_down, output_dim: 4096, input_dim: 14336} + - {name: qkv, kind: dense, output_dim: 6144, input_dim: 4096, tp_axis: output} + - {name: o, kind: dense, output_dim: 4096, input_dim: 4096, tp_axis: input} + - name: mixtral_8x22b + max_tp: 8 + max_ep: 4 + moe: + expert_num: 8 + experts_per_token: 2 + weights: + - {name: gate_up, kind: moe_gate_up, output_dim: 32768, input_dim: 6144} + - {name: down, kind: moe_down, output_dim: 6144, input_dim: 16384} + - name: deepseek_v2 + max_tp: 4 + max_ep: 20 + moe: + expert_num: 160 + experts_per_token: 6 + weights: + - {name: gate_up, kind: moe_gate_up, output_dim: 3072, input_dim: 5120} + - {name: down, kind: moe_down, output_dim: 5120, input_dim: 1536} + - name: deepseek_v2_lite + max_tp: 2 + max_ep: 8 + moe: + expert_num: 64 + experts_per_token: 6 + weights: + - {name: gate_up, kind: moe_gate_up, output_dim: 3072, input_dim: 2048} + - {name: down, kind: moe_down, output_dim: 2048, input_dim: 1536} + - name: qwen2_a14b + max_tp: 4 + max_ep: 8 + moe: + expert_num: 64 + experts_per_token: 8 + weights: + - {name: gate_up, kind: moe_gate_up, output_dim: 5120, input_dim: 3840} + - {name: down, kind: moe_down, output_dim: 3840, input_dim: 2560} + - name: phi35_moe + max_tp: 4 + max_ep: 8 + moe: + expert_num: 16 + experts_per_token: 2 + weights: + - {name: gate_up, kind: moe_gate_up, output_dim: 12800, input_dim: 4096} + - {name: down, kind: moe_down, output_dim: 4096, input_dim: 6400} + - name: qwen3_30b_a3b + max_tp: 4 + max_ep: 16 + moe: + expert_num: 128 + experts_per_token: 8 + weights: + - {name: gate_up, kind: moe_gate_up, output_dim: 1536, input_dim: 2048} + - {name: down, kind: moe_down, output_dim: 2048, input_dim: 768} + - {name: qkv, kind: dense, output_dim: 5120, input_dim: 2048, tp_axis: output} + - {name: o, kind: dense, output_dim: 2048, input_dim: 4096, tp_axis: input} + - name: qwen3_235b_a22b + max_tp: 8 + max_ep: 16 + moe: + expert_num: 128 + experts_per_token: 8 + weights: + - {name: gate_up, kind: moe_gate_up, output_dim: 3072, input_dim: 4096} + - {name: down, kind: moe_down, output_dim: 4096, input_dim: 1536} + - {name: qkv, kind: dense, output_dim: 9216, input_dim: 4096, tp_axis: output} + - {name: o, kind: dense, output_dim: 4096, input_dim: 8192, tp_axis: input} + - name: qwen35_4b + max_tp: 2 + max_ep: 1 + weights: + - {name: gate_up, kind: dense, output_dim: 18432, input_dim: 2560, tp_axis: output} + - {name: down, kind: dense, output_dim: 2560, input_dim: 9216, tp_axis: input} + - {name: qkv, kind: dense, output_dim: 10240, input_dim: 2560, tp_axis: output} + - {name: o, kind: dense, output_dim: 2560, input_dim: 4096, tp_axis: input} + - {name: lin_in_proj_all, kind: dense, output_dim: 12352, input_dim: 2560, tp_axis: output} + - {name: lin_o, kind: dense, output_dim: 2560, input_dim: 4096, tp_axis: input} + - name: qwen35_9b + max_tp: 2 + max_ep: 1 + weights: + - {name: gate_up, kind: dense, output_dim: 24576, input_dim: 4096, tp_axis: output} + - {name: down, kind: dense, output_dim: 4096, input_dim: 12288, tp_axis: input} + - {name: qkv, kind: dense, output_dim: 10240, input_dim: 4096, tp_axis: output} + - {name: o, kind: dense, output_dim: 4096, input_dim: 4096, tp_axis: input} + - {name: lin_in_proj_all, kind: dense, output_dim: 12352, input_dim: 4096, tp_axis: output} + - {name: lin_o, kind: dense, output_dim: 4096, input_dim: 4096, tp_axis: input} + - name: qwen35_27b + max_tp: 4 + max_ep: 1 + weights: + - {name: gate_up, kind: dense, output_dim: 34816, input_dim: 5120, tp_axis: output} + - {name: down, kind: dense, output_dim: 5120, input_dim: 17408, tp_axis: input} + - {name: qkv, kind: dense, output_dim: 14336, input_dim: 5120, tp_axis: output} + - {name: o, kind: dense, output_dim: 5120, input_dim: 6144, tp_axis: input} + - {name: lin_in_proj_all, kind: dense, output_dim: 16480, input_dim: 5120, tp_axis: output} + - {name: lin_o, kind: dense, output_dim: 5120, input_dim: 6144, tp_axis: input} + - name: qwen35_35b_a3b + max_tp: 4 + max_ep: 32 + moe: + expert_num: 256 + experts_per_token: 8 + weights: + - {name: gate_up, kind: moe_gate_up, output_dim: 1024, input_dim: 2048} + - {name: down, kind: moe_down, output_dim: 2048, input_dim: 512} + - {name: shared_gate_up, kind: dense, output_dim: 1024, input_dim: 2048, tp_axis: output} + - {name: shared_down, kind: dense, output_dim: 2048, input_dim: 512, tp_axis: input} + - {name: qkv, kind: dense, output_dim: 9216, input_dim: 2048, tp_axis: output} + - {name: o, kind: dense, output_dim: 2048, input_dim: 4096, tp_axis: input} + - {name: lin_in_proj_all, kind: dense, output_dim: 12352, input_dim: 2048, tp_axis: output} + - {name: lin_o, kind: dense, output_dim: 2048, input_dim: 4096, tp_axis: input} + - name: qwen35_122b_a10b + max_tp: 8 + max_ep: 32 + moe: + expert_num: 256 + experts_per_token: 8 + weights: + - {name: gate_up, kind: moe_gate_up, output_dim: 2048, input_dim: 3072} + - {name: down, kind: moe_down, output_dim: 3072, input_dim: 1024} + - {name: shared_gate_up, kind: dense, output_dim: 2048, input_dim: 3072, tp_axis: output} + - {name: shared_down, kind: dense, output_dim: 3072, input_dim: 1024, tp_axis: input} + - {name: qkv, kind: dense, output_dim: 17408, input_dim: 3072, tp_axis: output} + - {name: o, kind: dense, output_dim: 3072, input_dim: 8192, tp_axis: input} + - {name: lin_in_proj_all, kind: dense, output_dim: 20608, input_dim: 3072, tp_axis: output} + - {name: lin_o, kind: dense, output_dim: 3072, input_dim: 8192, tp_axis: input} + - name: qwen35_397b_a17b + max_tp: 8 + max_ep: 32 + moe: + expert_num: 512 + experts_per_token: 10 + weights: + - {name: gate_up, kind: moe_gate_up, output_dim: 2048, input_dim: 4096} + - {name: down, kind: moe_down, output_dim: 4096, input_dim: 1024} + - {name: shared_gate_up, kind: dense, output_dim: 2048, input_dim: 4096, tp_axis: output} + - {name: shared_down, kind: dense, output_dim: 4096, input_dim: 1024, tp_axis: input} + - {name: qkv, kind: dense, output_dim: 17408, input_dim: 4096, tp_axis: output} + - {name: o, kind: dense, output_dim: 4096, input_dim: 8192, tp_axis: input} + - {name: lin_in_proj_all, kind: dense, output_dim: 20608, input_dim: 4096, tp_axis: output} + - {name: lin_o, kind: dense, output_dim: 4096, input_dim: 8192, tp_axis: input} + - name: glm52 + max_tp: 8 + max_ep: 32 + moe: + expert_num: 256 + experts_per_token: 8 + weights: + - {name: dense_gate_up, kind: dense, output_dim: 24576, input_dim: 6144, tp_axis: output} + - {name: dense_down, kind: dense, output_dim: 6144, input_dim: 12288, tp_axis: input} + - {name: gate_up, kind: moe_gate_up, output_dim: 4096, input_dim: 6144} + - {name: down, kind: moe_down, output_dim: 6144, input_dim: 2048} + - {name: shared_gate_up, kind: dense, output_dim: 4096, input_dim: 6144, tp_axis: output} + - {name: shared_down, kind: dense, output_dim: 6144, input_dim: 2048, tp_axis: input} + - {name: mla_q_a, kind: replicated, output_dim: 2048, input_dim: 6144} + - {name: mla_q_b, kind: dense, output_dim: 36864, input_dim: 2048, tp_axis: output} + - {name: mla_kv_a, kind: replicated, output_dim: 576, input_dim: 6144} + - {name: mla_o, kind: dense, output_dim: 6144, input_dim: 36864, tp_axis: input} + - {name: index_q, kind: replicated, output_dim: 4096, input_dim: 2048} + - {name: index_k, kind: replicated, output_dim: 128, input_dim: 6144} + - {name: index_weight, kind: replicated, output_dim: 32, input_dim: 6144} + - name: kimi_k25 + max_tp: 8 + max_ep: 48 + moe: + expert_num: 384 + experts_per_token: 8 + weights: + - {name: dense_gate_up, kind: dense, output_dim: 36864, input_dim: 7168, tp_axis: output} + - {name: dense_down, kind: dense, output_dim: 7168, input_dim: 18432, tp_axis: input} + - {name: gate_up, kind: moe_gate_up, output_dim: 4096, input_dim: 7168} + - {name: down, kind: moe_down, output_dim: 7168, input_dim: 2048} + - {name: shared_gate_up, kind: dense, output_dim: 4096, input_dim: 7168, tp_axis: output} + - {name: shared_down, kind: dense, output_dim: 7168, input_dim: 2048, tp_axis: input} + - {name: mla_q_a, kind: replicated, output_dim: 1536, input_dim: 7168} + - {name: mla_q_b, kind: dense, output_dim: 36864, input_dim: 1536, tp_axis: output} + - {name: mla_kv_a, kind: replicated, output_dim: 576, input_dim: 7168} + - {name: mla_o, kind: dense, output_dim: 7168, input_dim: 36864, tp_axis: input} + - name: deepseek_v4_pro + max_tp: 8 + max_ep: 64 + moe: + expert_num: 384 + experts_per_token: 6 + weights: + - {name: gate_up, kind: moe_gate_up, output_dim: 6144, input_dim: 7168} + - {name: down, kind: moe_down, output_dim: 7168, input_dim: 3072} + - {name: shared_gate_up, kind: dense, output_dim: 6144, input_dim: 7168, tp_axis: output} + - {name: shared_down, kind: dense, output_dim: 7168, input_dim: 3072, tp_axis: input} + - {name: attn_q_a, kind: replicated, output_dim: 1536, input_dim: 7168} + - {name: attn_q_b, kind: dense, output_dim: 65536, input_dim: 1536, tp_axis: output} + - {name: attn_kv, kind: replicated, output_dim: 512, input_dim: 7168} + - {name: attn_o_b, kind: dense, output_dim: 7168, input_dim: 16384, tp_axis: input} + - {name: index_q, kind: dense, output_dim: 8192, input_dim: 1536, tp_axis: output} + - {name: index_weight, kind: dense, output_dim: 64, input_dim: 7168, tp_axis: output} + - name: deepseek_v4_flash + max_tp: 8 + max_ep: 32 + moe: + expert_num: 256 + experts_per_token: 6 + weights: + - {name: gate_up, kind: moe_gate_up, output_dim: 4096, input_dim: 4096} + - {name: down, kind: moe_down, output_dim: 4096, input_dim: 2048} + - {name: shared_gate_up, kind: dense, output_dim: 4096, input_dim: 4096, tp_axis: output} + - {name: shared_down, kind: dense, output_dim: 4096, input_dim: 2048, tp_axis: input} + - {name: attn_q_a, kind: replicated, output_dim: 1024, input_dim: 4096} + - {name: attn_q_b, kind: dense, output_dim: 32768, input_dim: 1024, tp_axis: output} + - {name: attn_kv, kind: replicated, output_dim: 512, input_dim: 4096} + - {name: attn_o_b, kind: dense, output_dim: 4096, input_dim: 8192, tp_axis: input} + - {name: index_q, kind: dense, output_dim: 8192, input_dim: 1024, tp_axis: output} + - {name: index_weight, kind: dense, output_dim: 64, input_dim: 4096, tp_axis: output} + - name: qwen3x_35b_a3b + max_tp: 4 + max_ep: 320 + moe: + expert_num: 2560 + experts_per_token: 8 + weights: + - {name: gate_up, kind: moe_gate_up, output_dim: 1024, input_dim: 2048} + - {name: down, kind: moe_down, output_dim: 2048, input_dim: 512} diff --git a/tests/turbomind/linear/reference.py b/tests/turbomind/linear/reference.py new file mode 100644 index 0000000000..c864c93ff8 --- /dev/null +++ b/tests/turbomind/linear/reference.py @@ -0,0 +1,141 @@ +from __future__ import annotations + +import torch + +_SM90_FP8_FUSED_SILU_BLOCK = 128 +_SM90_BF16_FUSED_SILU_BLOCK = 64 + + +def fused_silu_block(weight_type: str) -> int: + return { + 'fp8_e4m3': _SM90_FP8_FUSED_SILU_BLOCK, + 'bf16': _SM90_BF16_FUSED_SILU_BLOCK, + }[weight_type] + + +def compare_tensors(actual: torch.Tensor, expected: torch.Tensor) -> dict[str, float]: + if actual.shape != expected.shape: + raise ValueError(f'shape_mismatch_{tuple(actual.shape)}_{tuple(expected.shape)}') + a = actual.detach().float() + e = expected.detach().float() + abs_diff = (a - e).abs() + denom = e.abs().clamp_min(1e-6) + rel = abs_diff / denom + return { + 'max_abs': float(abs_diff.max().item()) if abs_diff.numel() else 0.0, + 'mean_abs': float(abs_diff.mean().item()) if abs_diff.numel() else 0.0, + 'max_rel': float(rel.max().item()) if rel.numel() else 0.0, + 'mean_rel': float(rel.mean().item()) if rel.numel() else 0.0, + } + + +def dense_gemm(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + # x: [M, K], weight: [K, N] (MN-major storage as in testbed_v3) + return x @ weight + + +def block_pack_w1w3(w1: torch.Tensor, w3: torch.Tensor, block: int) -> torch.Tensor: + """Pack [g0:block|u0:block|…] along the output dim (matches builder + `_block_pack_w1w3`).""" + n = w1.shape[-1] + if n % block != 0: + raise ValueError(f'inter_{n}_not_divisible_by_block_{block}') + shape = w1.shape[:-1] + w1b = w1.reshape(*shape, n // block, block) + w3b = w3.reshape(*shape, n // block, block) + return torch.stack([w1b, w3b], dim=-2).reshape(*shape, n * 2).contiguous() + + +def apply_block_fused_silu(c: torch.Tensor, block: int) -> torch.Tensor: + """Apply silu(gate)*up on block-packed GEMM output ``[M, 2*inter]`` → ``[M, + inter]``.""" + m, n = c.shape + if n % (2 * block) != 0: + raise ValueError(f'packed_n_{n}_not_aligned_to_{2 * block}') + inter = n // 2 + out = torch.empty((m, inter), device=c.device, dtype=torch.float32) + cf = c.float() + for b in range(inter // block): + g0 = b * 2 * block + u0 = g0 + block + o0 = b * block + gate = cf[:, g0:g0 + block] + up = cf[:, u0:u0 + block] + out[:, o0:o0 + block] = gate * torch.sigmoid(gate) * up + return out.to(dtype=c.dtype) + + +def quantize_symm_row_fp8( + src: torch.Tensor, + group_size: int = _SM90_FP8_FUSED_SILU_BLOCK, + qmax: float = 448.0): + """Match QuantizeSymm: per-row group-absmax → e4m3 + float scales [ceil(N/gs), M]. + + Returns (fp8 e4m3 tensor, scales float32, dequant bf16). + """ + m, n = src.shape + xf = src.float() + s_dim = (n + group_size - 1) // group_size + scales = torch.empty((s_dim, m), device=src.device, dtype=torch.float32) + out_f = torch.empty_like(xf) + for g in range(s_dim): + c0 = g * group_size + c1 = min(c0 + group_size, n) + block = xf[:, c0:c1] + amax = block.abs().amax(dim=1).clamp_min(1e-8) + scales[g] = amax / qmax + inv = qmax / amax + out_f[:, c0:c1] = block * inv[:, None] + # Round-trip through CUDA fp8 to match kernel cast. + fp8 = out_f.to(torch.float8_e4m3fn) + dequant = (fp8.float() * scales.t().repeat_interleave(group_size, dim=1)[:, :n]) + return fp8, scales, dequant.to(dtype=torch.bfloat16) + + +def moe_reference( + x: torch.Tensor, + expert_weights: list[torch.Tensor], + f2n: torch.Tensor | None, + offsets: torch.Tensor, + scales: torch.Tensor, + en2f: torch.Tensor, + experts_per_token: int, + combine_experts: bool, +) -> torch.Tensor: + """Mirror testbed_v3 GetReference MoE path using gather/GEMM/scatter. + + Layouts match testbed_v3 Route()/GetReference(): + - f2n: [tokens * experts_per_token] packed expert-major token indices. + If None, ``x`` is already expert-packed (MoE down / w2). + - offsets: [expert_num + 1] exclusive prefix counts + - scales: [experts_per_token * batch] with index e * batch + tok + - en2f: [experts_per_token * batch] maps (e, tok) -> packed slot + - expert_weights[e]: [K, N] + """ + n = expert_weights[0].shape[1] + device = x.device + dtype = x.dtype + if f2n is None: + xe = x + token_count = int(x.shape[0]) + m = token_count // experts_per_token + else: + m, k = x.shape + token_count = int(f2n.numel()) + xe = x[f2n.long()] + de = torch.empty((token_count, n), device=device, dtype=dtype) + h_offsets = offsets.detach().cpu().tolist() + for e, w in enumerate(expert_weights): + base = int(h_offsets[e]) + end = int(h_offsets[e + 1]) + if end > base: + de[base:end] = dense_gemm(xe[base:end], w) + if not combine_experts: + return de + # Match invokeMoeCombine / MoeReduceKernel: accumulate in float, then cast. + out_f = torch.zeros((m, n), device=device, dtype=torch.float32) + for e in range(experts_per_token): + for tok in range(m): + packed = int(en2f[e * m + tok].item()) + out_f[tok] += de[packed].float() * float(scales[e * m + tok].item()) + return out_f.to(dtype) diff --git a/tests/turbomind/linear/regtune_cmp_report.py b/tests/turbomind/linear/regtune_cmp_report.py new file mode 100644 index 0000000000..8bb9904736 --- /dev/null +++ b/tests/turbomind/linear/regtune_cmp_report.py @@ -0,0 +1,82 @@ +#!/usr/bin/env python3 +"""Compare two tune-verbose sweep dirs: per-problem best measured time, new vs old. + +Usage: regtune_cmp_report.py OLD_DIR NEW_DIR [--threshold 0.005] +Problem key = (case, m, n, k[, fuse]) taken from the [tune] line prefix. +""" + +from __future__ import annotations + +import argparse +import re +from pathlib import Path + +TUNE_RE = re.compile(r'^\[tune\] (\S+?) (\S+) swizzle=(\d+) splits=(\d+) measured=([0-9.eE+-]+)') + + +def best_per_problem(d: Path) -> dict: + best = {} + for log in sorted(d.glob('*.log')): + case = log.stem + for line in log.read_text(errors='ignore').splitlines(): + m = TUNE_RE.match(line) + if not m: + continue + prob, kern, sw, sp, t = m.group(1), m.group(2), m.group(3), m.group(4), float(m.group(5)) + if t <= 0: + continue + key = (case, prob) + cur = best.get(key) + if cur is None or t < cur[0]: + best[key] = (t, kern, sw, sp) + return best + + +def main() -> int: + ap = argparse.ArgumentParser() + ap.add_argument('old', type=Path) + ap.add_argument('new', type=Path) + ap.add_argument('--threshold', type=float, default=0.005, help='relative change to report') + args = ap.parse_args() + + old, new = best_per_problem(args.old), best_per_problem(args.new) + common = sorted(set(old) & set(new)) + only_new = set(new) - set(old) + only_old = set(old) - set(new) + if only_new or only_old: + print(f'problems only in new: {len(only_new)}, only in old: {len(only_old)}') + + ratios = [] + improved, regressed = [], [] + for key in common: + t_old, t_new = old[key][0], new[key][0] + r = t_new / t_old + ratios.append(r) + if r < 1 - args.threshold: + improved.append((r, key, old[key], new[key])) + elif r > 1 + args.threshold: + regressed.append((r, key, old[key], new[key])) + + n = len(ratios) + mean = sum(ratios) / n + geo = 1.0 + for r in ratios: + geo *= r + geo **= 1.0 / n + print(f'problems: {n} mean new/old={mean:.4f} geomean={geo:.4f}') + print(f'improved >{args.threshold:.1%}: {len(improved)} ({len(improved)/n:.1%}) ' + f'regressed: {len(regressed)} ({len(regressed)/n:.1%})') + + print('\n== top improvements ==') + for r, key, o, nn in sorted(improved)[:15]: + print(f'{r:.3f} {key[0]} {key[1]} {o[0]:.4f}->{nn[0]:.4f}ms ' + f"old={o[1].split('tnt_')[-1]} sw{o[2]} new={nn[1].split('tnt_')[-1]} sw{nn[2]}") + print('\n== top regressions ==') + for r, key, o, nn in sorted(regressed, reverse=True)[:15]: + print(f'{r:.3f} {key[0]} {key[1]} {o[0]:.4f}->{nn[0]:.4f}ms ' + f"old={o[1].split('tnt_')[-1]} sw{o[2]} new={nn[1].split('tnt_')[-1]} sw{nn[2]}") + return 0 + + +if __name__ == '__main__': + raise SystemExit(main()) diff --git a/tests/turbomind/linear/run_sched_cmp.py b/tests/turbomind/linear/run_sched_cmp.py new file mode 100644 index 0000000000..4cf23730aa --- /dev/null +++ b/tests/turbomind/linear/run_sched_cmp.py @@ -0,0 +1,97 @@ +#!/usr/bin/env python3 +"""Row-major vs col-major tile-scheduler comparison for MoE gate_up/down (tp=1, +ep=1). + +Runs bench_linear with real tuning (--iters > 0 skips the pre-tune forward, so the +GEMM Measure path actually measures) and dumps per-spec measurements via +TM_GEMM_TUNE_VERBOSE=1. Kernel names encode the scheduler raster order in the +trailing policy digits: `_00` = col-major scheduler, `_01` = row-major scheduler. +""" + +from __future__ import annotations + +import argparse +import os +import subprocess +import sys +from pathlib import Path + +CASES = [ + # gate_up -> indexed (ibb) kernels + 'deepseek_v2_gate_up__bf16_bf16_bf16', + 'deepseek_v2_gate_up__bf16_bf16_bf16__fuse_silu', + 'qwen3_30b_a3b_gate_up__bf16_bf16_bf16', + 'mixtral_8x7b_gate_up__bf16_bf16_bf16', + # down -> blocked (bbb) kernels + 'deepseek_v2_down__bf16_bf16_bf16', + 'qwen3_30b_a3b_down__bf16_bf16_bf16', + 'mixtral_8x7b_down__bf16_bf16_bf16', +] + +BATCHES = '1,3,16,17,64,65,256,1024,4096' + + +def discover_cases(kind: str = 'moe') -> list[str]: + """All tp=1, ep=1 cases; kind='moe': gate_up/down with experts; + kind='dense': expert_num==0.""" + from .cases import TYPE_BF16, expand_suite + + names = set() + for r in expand_suite('full', None, None, (TYPE_BF16.name,)): + c = r.case + if c.tp != 1 or c.ep != 1: + continue + if kind == 'moe' and c.expert_num > 0 and ('gate_up' in c.name or 'down' in c.name): + names.add(c.name) + elif kind == 'dense' and c.expert_num == 0: + names.add(c.name) + return sorted(names) + + +def main() -> int: + p = argparse.ArgumentParser(description=__doc__) + p.add_argument('--outdir', type=Path, default=Path('tmp/sched_cmp')) + p.add_argument('--gpu', type=str, default='0') + p.add_argument('--iters', type=int, default=20) + p.add_argument('--cases', type=str, default=None, help='comma list; default: all MoE gate_up/down tp1ep1') + p.add_argument('--kind', choices=['moe', 'dense'], default='moe') + p.add_argument('--batches', type=str, default=BATCHES) + args = p.parse_args() + + case_names = args.cases.split(',') if args.cases else discover_cases(args.kind) + print(f'cases: {len(case_names)}', file=sys.stderr) + + args.outdir.mkdir(parents=True, exist_ok=True) + + env = os.environ.copy() + env['CUDA_VISIBLE_DEVICES'] = args.gpu + env['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True' + env['TM_GEMM_TUNE'] = 'swizzle=[0,1,2,3]' + env['TM_GEMM_TUNE_VERBOSE'] = '1' + + for name in case_names: + log = args.outdir / f'{name}.log' + cmd = [ + sys.executable, + 'tests/turbomind/linear/bench_linear.py', + '--suite', 'full', + '--case', name, + '--type', 'bf16_bf16_bf16', + '--batch', args.batches, + '--tune', + '--iters', str(args.iters), + '--warmup', '5', + '--no-validate', + '--export', str(args.outdir / 'records'), + ] + print(f'=== {name} ===', file=sys.stderr) + with log.open('w') as f: + rc = subprocess.call(cmd, stdout=f, stderr=subprocess.STDOUT, env=env, + cwd=str(Path(__file__).resolve().parents[3])) + print(f' exit={rc} log={log}', file=sys.stderr) + + return 0 + + +if __name__ == '__main__': + raise SystemExit(main()) diff --git a/tests/turbomind/linear/run_sm90_scan.py b/tests/turbomind/linear/run_sm90_scan.py new file mode 100755 index 0000000000..bc17601a65 --- /dev/null +++ b/tests/turbomind/linear/run_sm90_scan.py @@ -0,0 +1,183 @@ +#!/usr/bin/env python3 +"""Chunked full SM90 kernel-usage scan. + +The full benchmark run in one process accumulates GPU memory across the large/MoE +model cases and eventually OOMs on the testbed. This driver splits the full suite +into smaller per-name chunks, runs each chunk in a fresh Python process, +and appends the per-case Summary blocks to a single log for aggregation. + +Chunks are distributed across one worker per GPU (--gpus); each worker processes +its chunks sequentially in fresh subprocesses and appends to its own +``bench.gpu.log``. The per-GPU logs are concatenated into ``bench.log`` when +all workers finish, so downstream aggregation is unchanged: + + python tests/turbomind/linear/kernel_usage_report.py tmp/sm90_bf16_scan5/bench.log + +Use --type to select the TypeSpec to scan, e.g. bf16_bf16_bf16 (SM90 BF16 kernels), +bf16_e4m3b128_bf16 (FP8 weight-as-A) or e4m3k128_e4m3b128_bf16 (FP8 v3 act-as-A). +""" + +from __future__ import annotations + +import argparse +import os +import queue +import subprocess +import sys +import threading +from pathlib import Path + +from .cases import expand_suite + + +def chunk_names(names: list[str], size: int) -> list[list[str]]: + return [names[i : i + size] for i in range(0, len(names), size)] + + +def parse_int_list(value: str) -> tuple[int, ...]: + return tuple(int(x) for x in value.split(',') if x.strip()) + + +def run_chunk( + case_names: list[str], + type_name: str, + outdir: Path, + log: Path, + env: dict[str, str], + python: str, + tps: tuple[int, ...], + eps: tuple[int, ...], +) -> int: + cmd = [ + python, + 'tests/turbomind/linear/bench_linear.py', + '--suite', + 'full', + '--case', + ','.join(case_names), + '--type', + type_name, + '--tp', + ','.join(map(str, tps)), + '--ep', + ','.join(map(str, eps)), + '--tune', + '--iters', + '0', + '--no-validate', + '--export', + str(outdir / 'records'), + ] + with log.open('ab') as f: + # Append a small marker so we know where each chunk starts. + f.write(f'\n# CHUNK: {case_names[0]} .. {case_names[-1]}\n'.encode()) + proc = subprocess.Popen( + cmd, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + env=env, + cwd=str(Path(__file__).resolve().parents[3]), + ) + assert proc.stdout is not None + for line in proc.stdout: + f.write(line) + proc.wait() + return proc.returncode + + +def worker( + tag: str, + gpu: str, + chunks: queue.Queue[list[str]], + type_name: str, + outdir: Path, + base_env: dict[str, str], + python: str, + tps: tuple[int, ...], + eps: tuple[int, ...], + failures: list[list[str]], + lock: threading.Lock, +) -> None: + env = base_env.copy() + env['CUDA_VISIBLE_DEVICES'] = gpu + log = outdir / f'bench.gpu{gpu}.log' + while True: + try: + chunk = chunks.get_nowait() + except queue.Empty: + return + with lock: + print(f'[gpu{gpu}] start ({len(chunk)} cases): {chunk[0]} .. {chunk[-1]}', flush=True) + rc = run_chunk(chunk, type_name, outdir, log, env, python, tps, eps) + with lock: + print(f'[gpu{gpu}] done rc={rc}: {chunk[0]} .. {chunk[-1]}', flush=True) + if rc != 0: + failures.append(chunk) + + +def main(argv: list[str] | None = None) -> int: + p = argparse.ArgumentParser(description=__doc__) + p.add_argument('--outdir', type=Path, default=Path('tmp/sm90_scan')) + p.add_argument('--type', type=str, default='bf16_bf16_bf16', help='TypeSpec name to scan') + p.add_argument('--chunk', type=int, default=6, help='case names per subprocess') + p.add_argument('--gpus', type=str, default='0', help='comma-separated CUDA_VISIBLE_DEVICES values') + p.add_argument('--tp', type=str, default='1,2,4,8', help='comma-separated TP sizes') + p.add_argument('--ep', type=str, default='1,2,4,8', help='comma-separated EP sizes') + p.add_argument('--python', type=str, default=sys.executable) + args = p.parse_args(argv) + + tps = parse_int_list(args.tp) + eps = parse_int_list(args.ep) + gpus = [g.strip() for g in args.gpus.split(',') if g.strip()] + + args.outdir.mkdir(parents=True, exist_ok=True) + + runs = expand_suite('full', None, None, (args.type,), tps=tps, eps=eps) + names = sorted({r.case.name for r in runs}) + print(f'Total {args.type} case names: {len(names)}, runs: {len(runs)}', file=sys.stderr) + + base_env = os.environ.copy() + base_env['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True' + base_env['TM_GEMM_TUNE'] = 'swizzle=[0,1,2,3]' + + chunks: queue.Queue[list[str]] = queue.Queue() + for chunk in chunk_names(names, args.chunk): + chunks.put(chunk) + + failures: list[list[str]] = [] + lock = threading.Lock() + threads = [ + threading.Thread( + target=worker, + args=(f'w{i}', gpu, chunks, args.type, args.outdir, base_env, args.python, tps, eps, failures, lock), + ) + for i, gpu in enumerate(gpus) + ] + for t in threads: + t.start() + for t in threads: + t.join() + + # Merge per-GPU logs into a single log for the aggregation tool. + merged = args.outdir / 'bench.log' + with merged.open('wb') as out: + for gpu in gpus: + part = args.outdir / f'bench.gpu{gpu}.log' + if part.exists(): + out.write(part.read_bytes()) + + if failures: + print(f'\n{len(failures)} chunk(s) failed:', file=sys.stderr) + for chunk in failures: + print(f' {chunk[0]} .. {chunk[-1]}', file=sys.stderr) + + print(f'\nDone. Log: {merged}', file=sys.stderr) + print( + f'Aggregate usage: python tests/turbomind/linear/kernel_usage_report.py {merged}', + file=sys.stderr, + ) + return 1 if failures else 0 + + +if __name__ == '__main__': + raise SystemExit(main()) diff --git a/tests/turbomind/linear/sched_cmp_report.py b/tests/turbomind/linear/sched_cmp_report.py new file mode 100644 index 0000000000..b59c190ba4 --- /dev/null +++ b/tests/turbomind/linear/sched_cmp_report.py @@ -0,0 +1,107 @@ +#!/usr/bin/env python3 +"""Parse [tune] logs from run_sched_cmp.py and compare row-major vs col-major +tile schedulers for the same kernel tile. + +Kernel-name suffix: `_00` = col-major scheduler, `_01` = row-major scheduler. +""" + +from __future__ import annotations + +import argparse +import ast +import re +import sys +from collections import defaultdict +from pathlib import Path + +TUNE_RE = re.compile( + r'^\[tune\] (?P\S+) (?P\S+) swizzle=(?P\d+) ' + r'splits=(?P\d+) measured=(?P[0-9.eE+-]+)$' +) +ROW_RE = re.compile(r"\{'case':[^}]*\}") + + +def strip_order(kernel: str) -> tuple[str, str]: + """Split kernel name into (name without trailing policy digits, + scheduler).""" + base, _, pol = kernel.rpartition('_') + if pol in ('00', '01'): + return base, ('row' if pol == '01' else 'col') + return kernel, '?' + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument('logs', nargs='+', type=Path) + ap.add_argument('--detail', action='store_true', help='print per-tile row/col table') + args = ap.parse_args() + + # problem desc -> kernel base -> scheduler -> best (min over swizzle) ms + best: dict[str, dict[str, dict[str, float]]] = defaultdict(lambda: defaultdict(dict)) + # problem desc -> overall winner (kernel, ms) from first sorted [tune] line + winner: dict[str, tuple[str, float]] = {} + # benchmark rows: (case, m) -> tflops + rows: dict[tuple[str, int], float] = {} + + for log in args.logs: + case = log.name[: -len('.log')] + for line in log.read_text().splitlines(): + m = TUNE_RE.match(line) + if m: + desc = m.group('desc') + base, order = strip_order(m.group('kernel')) + ms = float(m.group('ms')) + cur = best[desc][base].get(order) + if cur is None or ms < cur: + best[desc][base][order] = ms + if desc not in winner: + winner[desc] = (m.group('kernel'), ms) + elif "{'case':" in line: + mrow = ROW_RE.search(line) + if mrow: + d = ast.literal_eval(mrow.group(0)) + if d.get('tflops'): + rows[(case, d['m'])] = d['tflops'] + + n_row_wins = n_col_wins = 0 + print(f"{'problem':<46} {'winner-sched':<12} {'winner-tile':<14} {'row-best':>9} {'col-best':>9} {'row/col':>8}") + for desc in sorted(best, key=lambda d: (d.split('_')[3], [int(x) for x in re.findall(r'\d+', d.split('_')[-2])])): + per_tile = best[desc] + w_kernel, w_ms = winner[desc] + _, w_order = strip_order(w_kernel) + if w_order == 'row': + n_row_wins += 1 + elif w_order == 'col': + n_col_wins += 1 + # same-tile row vs col comparison, best-matched pair + row_best = min((v['row'] for v in per_tile.values() if 'row' in v), default=None) + col_best = min((v['col'] for v in per_tile.values() if 'col' in v), default=None) + w_tile = re.search(r'_(\d+x\d+x\d+)_', w_kernel) + ratio = (row_best / col_best) if (row_best and col_best) else float('nan') + row_str = f"{row_best if row_best else float('nan'):>9.4f}" + col_str = f"{col_best if col_best else float('nan'):>9.4f}" + print(f"{desc:<46} {w_order:<12} {(w_tile.group(1) if w_tile else ''):<14} " + f'{row_str} {col_str} {ratio:>8.3f}') + if args.detail: + for base in sorted(per_tile): + v = per_tile[base] + r = v.get('row') + c = v.get('col') + tile = re.search(r'_(\d+x\d+x\d+)_', base).group(1) + r_s = f'{r:.4f}' if r else '-' + c_s = f'{c:.4f}' if c else '-' + mark = '' + if r and c: + mark = f'row/col={r / c:.3f}' + print(f' {tile:<12} row={r_s:<9} col={c_s:<9} {mark}') + + print(f'\noverall winner scheduler: row={n_row_wins} col={n_col_wins} (of {len(best)} problems)') + + print('\nbenchmark rows (tuned winner tflops):') + for (case, m), tf in sorted(rows.items()): + print(f' {case:<52} m={m:<6} {tf:8.2f} TFLOPS') + return 0 + + +if __name__ == '__main__': + sys.exit(main()) diff --git a/tests/turbomind/linear/test_linear.py b/tests/turbomind/linear/test_linear.py new file mode 100644 index 0000000000..d7c8139bfd --- /dev/null +++ b/tests/turbomind/linear/test_linear.py @@ -0,0 +1,25 @@ +from __future__ import annotations + +import pytest +import torch + +from . import linear as linear_mod +from .cases import expand_suite +from .fixture import LinearFixture + +cuda_required = pytest.mark.skipif(not torch.cuda.is_available(), reason='CUDA required') +tm_required = pytest.mark.skipif(not linear_mod.is_available(), reason='_turbomind required') + + +@cuda_required +@tm_required +@pytest.mark.parametrize('run', expand_suite('smoke'), ids=lambda r: f'{r.case.name}_m{r.batch_size}') +def test_smoke_linear_correctness(run): + fx = LinearFixture(run.case) + try: + fx.prepare_batch(run.batch_size) + fx.run_reference() + fx.run_linear() + fx.check_tolerances(fx.compare()) + finally: + fx.close() diff --git a/tests/turbomind/linear/validate_records.py b/tests/turbomind/linear/validate_records.py new file mode 100644 index 0000000000..ad98f13b7e --- /dev/null +++ b/tests/turbomind/linear/validate_records.py @@ -0,0 +1,91 @@ +#!/usr/bin/env python3 +"""Validate tuned dispatch records against the PyTorch reference, sharded +across GPUs. + +For every case with an exported records dir, run bench_linear with --import (disables +tuning) and --iters 0 (runs reference + compare + check_tolerances). A case fails if +the process exits non-zero or the log contains a Traceback. + +Usage: validate_records.py --records-root DIR --outdir DIR [--gpus 0,1,..] [--kind dense|moe] +""" + +from __future__ import annotations + +import argparse +import os +import subprocess +import sys +from pathlib import Path + +BATCHES = ('1,2,3,4,5,7,8,15,16,17,31,32,33,63,64,65,96,127,128,129,192,224,255,' + '256,257,511,512,513,1023,1024,1025,2047,2048,2049,4095,4096,4097,8191,8192,8233,16384') + + +def main() -> int: + ap = argparse.ArgumentParser() + ap.add_argument('--records-root', type=Path, required=True) + ap.add_argument('--outdir', type=Path, required=True) + ap.add_argument('--gpus', type=str, default='0') + ap.add_argument('--cases', type=str, default=None) + args = ap.parse_args() + + records = sorted(args.records_root.glob('records.*__tp1__ep1')) + if args.cases: + want = set(args.cases.split(',')) + records = [r for r in records if r.name[len('records.'):-len('__tp1__ep1')] in want] + if not records: + print('no records found', file=sys.stderr) + return 1 + + gpus = args.gpus.split(',') + args.outdir.mkdir(parents=True, exist_ok=True) + + # Deterministic modulo split of the case list across GPUs. + for i, gpu in enumerate(gpus): + shard = [str(r) for j, r in enumerate(records) if j % len(gpus) == i] + (args.outdir / f'_shard{gpu}.txt').write_text('\n'.join(shard)) + + running = [] + for gpu in gpus: + env = os.environ.copy() + env['CUDA_VISIBLE_DEVICES'] = gpu + env['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True' + script = args.outdir / f'_run{gpu}.sh' + script.write_text(f"""#!/bin/bash +while read -r rec || [ -n "$rec" ]; do + case_name="${{rec##*records.}}"; case_name="${{case_name%__tp1__ep1}}" + [ -z "$case_name" ] && continue + {sys.executable} tests/turbomind/linear/bench_linear.py --suite full \\ + --case "$case_name" --type bf16_bf16_bf16 --batch {BATCHES} \\ + --import "$rec" --iters 0 > "{args.outdir}/${{case_name}}.log" 2>&1 + echo "$case_name rc=$?" >> "{args.outdir}/_rc.txt" +done < "{args.outdir}/_shard{gpu}.txt" +""") + script.chmod(0o755) + running.append(subprocess.Popen(['bash', str(script)], env=env, + cwd=str(Path(__file__).resolve().parents[3]))) + rcs = [p.wait() for p in running] + + # Report + rc_file = args.outdir / '_rc.txt' + fails = [] + done = 0 + if rc_file.exists(): + for line in rc_file.read_text().splitlines(): + name, rc = line.rsplit(' rc=', 1) + done += 1 + if rc != '0': + fails.append((name, rc)) + tracebacks = [p.stem for p in args.outdir.glob('*.log') + if 'Traceback' in p.read_text(errors='ignore')] + print(f'validated: {done}/{len(records)} worker_rcs={rcs}') + print(f'nonzero exits: {len(fails)} logs with Traceback: {len(tracebacks)}') + for name, rc in fails: + print(f' FAIL rc={rc}: {name}') + for name in tracebacks: + print(f' TRACEBACK: {name}') + return 0 if not fails and not tracebacks and done == len(records) else 1 + + +if __name__ == '__main__': + raise SystemExit(main())